parakeet.ts 1.0.0

package/dist/mlx/conformer.js

@@ -0,0 +1,315 @@
+import { MxArray } from '@mlx-node/core';
+import { Module, Linear, LayerNorm, BatchNorm, Conv1d, Conv2d, silu, relu, glu, } from './nn.js';
+import { MultiHeadAttention, RelPositionMultiHeadAttention, RelPositionMultiHeadLocalAttention, RelPositionalEncoding, LocalRelPositionalEncoding, } from './attention.js';
+function s(...dims) {
+    return BigInt64Array.from(dims.map(BigInt));
+}
+// ---------------------------------------------------------------------------
+// Feed-forward block
+// ---------------------------------------------------------------------------
+class FeedForward extends Module {
+    linear1;
+    linear2;
+    constructor(dModel, dFf, useBias) {
+        super();
+        this.linear1 = new Linear(dModel, dFf, useBias);
+        this.linear2 = new Linear(dFf, dModel, useBias);
+    }
+    forward(x) {
+        return this.linear2.forward(silu(this.linear1.forward(x)));
+    }
+    loadWeights(weights, prefix) {
+        this.linear1.loadWeights(weights, `${prefix}.linear1`);
+        this.linear2.loadWeights(weights, `${prefix}.linear2`);
+    }
+}
+// ---------------------------------------------------------------------------
+// Convolution block
+// ---------------------------------------------------------------------------
+class Convolution extends Module {
+    padding;
+    pointwiseConv1;
+    depthwiseConv;
+    batchNorm;
+    pointwiseConv2;
+    constructor(args) {
+        super();
+        const useBias = args.useBias ?? true;
+        this.padding = Math.floor((args.convKernelSize - 1) / 2);
+        this.pointwiseConv1 = new Conv1d(args.dModel, args.dModel * 2, 1, 1, 0, 1, useBias);
+        this.depthwiseConv = new Conv1d(args.dModel, args.dModel, args.convKernelSize, 1, 0, args.dModel, useBias);
+        this.batchNorm = new BatchNorm(args.dModel);
+        this.pointwiseConv2 = new Conv1d(args.dModel, args.dModel, 1, 1, 0, 1, useBias);
+    }
+    forward(x, cache) {
+        x = this.pointwiseConv1.forward(x);
+        x = glu(x, 2); // split along last axis, gate with sigmoid
+        if (cache !== null) {
+            x = cache.updateAndFetchConv(x, this.padding);
+        }
+        else {
+            x = x.pad(new Int32Array([0, 0, this.padding, this.padding, 0, 0]), 0.0);
+        }
+        x = this.depthwiseConv.forward(x);
+        x = this.batchNorm.forward(x);
+        x = silu(x);
+        x = this.pointwiseConv2.forward(x);
+        return x;
+    }
+    loadWeights(weights, prefix) {
+        this.pointwiseConv1.loadWeights(weights, `${prefix}.pointwise_conv1`);
+        this.depthwiseConv.loadWeights(weights, `${prefix}.depthwise_conv`);
+        this.batchNorm.loadWeights(weights, `${prefix}.batch_norm`);
+        this.pointwiseConv2.loadWeights(weights, `${prefix}.pointwise_conv2`);
+    }
+}
+export class ConformerBlock extends Module {
+    normFF1;
+    ff1;
+    normSelfAtt;
+    selfAttn;
+    normConv;
+    conv;
+    normFF2;
+    ff2;
+    normOut;
+    args;
+    constructor(args) {
+        super();
+        this.args = args;
+        const useBias = args.useBias ?? true;
+        const ffDim = args.dModel * args.ffExpansionFactor;
+        this.normFF1 = new LayerNorm(args.dModel);
+        this.ff1 = new FeedForward(args.dModel, ffDim, useBias);
+        this.normSelfAtt = new LayerNorm(args.dModel);
+        this.selfAttn = this.buildAttention(args.selfAttentionModel, args.attContextSize ?? null);
+        this.normConv = new LayerNorm(args.dModel);
+        this.conv = new Convolution(args);
+        this.normFF2 = new LayerNorm(args.dModel);
+        this.ff2 = new FeedForward(args.dModel, ffDim, useBias);
+        this.normOut = new LayerNorm(args.dModel);
+    }
+    buildAttention(name, contextSize) {
+        const useBias = this.args.useBias ?? true;
+        if (name === 'rel_pos') {
+            return new RelPositionMultiHeadAttention(this.args.nHeads, this.args.dModel, useBias);
+        }
+        else if (name === 'rel_pos_local_attn') {
+            return new RelPositionMultiHeadLocalAttention(this.args.nHeads, this.args.dModel, useBias, contextSize ?? [256, 256]);
+        }
+        else {
+            return new MultiHeadAttention(this.args.nHeads, this.args.dModel, true);
+        }
+    }
+    setAttentionModel(name, contextSize = [256, 256]) {
+        const newAttn = this.buildAttention(name, contextSize);
+        // Copy weights from old attention if possible
+        // (In a real implementation we'd need to transfer parameters)
+        this.selfAttn = newAttn;
+    }
+    forward(x, posEmb, mask, cache) {
+        // FF1
+        x = x.add(this.ff1.forward(this.normFF1.forward(x)).mulScalar(0.5));
+        // Self-attention
+        const xNorm = this.normSelfAtt.forward(x);
+        x = x.add(this.selfAttn.forward(xNorm, xNorm, xNorm, posEmb, mask, cache));
+        // Convolution
+        x = x.add(this.conv.forward(this.normConv.forward(x), cache));
+        // FF2
+        x = x.add(this.ff2.forward(this.normFF2.forward(x)).mulScalar(0.5));
+        return this.normOut.forward(x);
+    }
+    loadWeights(weights, prefix) {
+        this.normFF1.loadWeights(weights, `${prefix}.norm_feed_forward1`);
+        this.ff1.loadWeights(weights, `${prefix}.feed_forward1`);
+        this.normSelfAtt.loadWeights(weights, `${prefix}.norm_self_att`);
+        this.selfAttn.loadWeights(weights, `${prefix}.self_attn`);
+        this.normConv.loadWeights(weights, `${prefix}.norm_conv`);
+        this.conv.loadWeights(weights, `${prefix}.conv`);
+        this.normFF2.loadWeights(weights, `${prefix}.norm_feed_forward2`);
+        this.ff2.loadWeights(weights, `${prefix}.feed_forward2`);
+        this.normOut.loadWeights(weights, `${prefix}.norm_out`);
+    }
+}
+// ---------------------------------------------------------------------------
+// DW-striding subsampling (Conv2D-based)
+// ---------------------------------------------------------------------------
+class DwStridingSubsampling extends Module {
+    samplingNum;
+    stride = 2;
+    kernelSize = 3;
+    padding;
+    convLayers; // null = ReLU placeholder
+    out;
+    constructor(args) {
+        super();
+        this.padding = Math.floor((this.kernelSize - 1) / 2);
+        this.samplingNum = Math.round(Math.log2(args.subsamplingFactor));
+        const convChannels = args.subsamplingConvChannels;
+        // Compute final frequency dimension
+        let finalFreqDim = args.featIn;
+        for (let i = 0; i < this.samplingNum; i++) {
+            finalFreqDim =
+                Math.floor((finalFreqDim + 2 * this.padding - this.kernelSize) / this.stride) + 1;
+        }
+        // Build convolution layers
+        this.convLayers = [];
+        let inCh = 1;
+        // First layer: standard conv2d
+        this.convLayers.push(new Conv2d(inCh, convChannels, this.kernelSize, this.stride, this.padding, 1, true));
+        this.convLayers.push(null); // ReLU
+        inCh = convChannels;
+        for (let i = 1; i < this.samplingNum; i++) {
+            // Depthwise
+            this.convLayers.push(new Conv2d(inCh, inCh, this.kernelSize, this.stride, this.padding, inCh, true));
+            // Pointwise
+            this.convLayers.push(new Conv2d(inCh, convChannels, 1, 1, 0, 1, true));
+            this.convLayers.push(null); // ReLU
+        }
+        this.out = new Linear(convChannels * finalFreqDim, args.dModel);
+    }
+    forward(x, lengths) {
+        // x: [batch, seq, mel] → [batch, 1, seq, mel] then to MLX NHWC: [batch, seq, mel, 1]
+        const xShape = x.shape();
+        const batch = Number(xShape[0]);
+        // lengths update
+        let outLengths = lengths;
+        for (let i = 0; i < this.samplingNum; i++) {
+            const pad = this.padding;
+            const k = this.kernelSize;
+            const st = this.stride;
+            // floor((len + 2*pad - k) / stride) + 1
+            outLengths = outLengths
+                .addScalar(2 * pad - k)
+                .divScalar(st)
+                .floor()
+                .addScalar(1);
+        }
+        outLengths = outLengths.astype(3); // 3 = int32
+        // Reshape x: [batch, seq, mel] → [batch, seq, mel, 1] (NHWC with C=1)
+        let cur = x.expandDims(3); // [batch, seq, mel, 1]
+        for (const layer of this.convLayers) {
+            if (layer === null) {
+                cur = relu(cur);
+            }
+            else {
+                cur = layer.forward(cur);
+            }
+        }
+        // cur: [batch, outSeq, outMel, convChannels] (NHWC layout)
+        // Python flattens as [batch, outSeq, convChannels*outMel] (C-major, F-minor):
+        // it transposes NHWC -> NCHW then swapaxes(1,2) -> [B, T', C', F'] before
+        // reshaping. To match that memory ordering (which the Linear weights expect),
+        // transpose channel axis before freq axis here.
+        const cShape = cur.shape();
+        const outSeq = Number(cShape[1]);
+        const outMel = Number(cShape[2]);
+        const ch = Number(cShape[3]);
+        cur = cur
+            .transpose(new Int32Array([0, 1, 3, 2]))
+            .reshape(s(batch, outSeq, ch * outMel));
+        cur = this.out.forward(cur);
+        return [cur, outLengths];
+    }
+    loadWeights(weights, prefix) {
+        let convIdx = 0;
+        for (const layer of this.convLayers) {
+            if (layer !== null) {
+                layer.loadWeights(weights, `${prefix}.conv.${convIdx}`);
+            }
+            convIdx++;
+        }
+        this.out.loadWeights(weights, `${prefix}.out`);
+    }
+}
+// ---------------------------------------------------------------------------
+// Conformer encoder
+// ---------------------------------------------------------------------------
+export class Conformer extends Module {
+    args;
+    posEnc;
+    preEncode;
+    layers;
+    constructor(args) {
+        super();
+        this.args = args;
+        const selfAttModel = args.selfAttentionModel;
+        const ctxSize = args.attContextSize ?? null;
+        if (selfAttModel === 'rel_pos') {
+            this.posEnc = new RelPositionalEncoding(args.dModel, args.posEmbMaxLen, args.xscaling ?? false);
+        }
+        else if (selfAttModel === 'rel_pos_local_attn') {
+            this.posEnc = new LocalRelPositionalEncoding(args.dModel, args.posEmbMaxLen, args.xscaling ?? false, ctxSize ?? [256, 256]);
+        }
+        else {
+            this.posEnc = null;
+        }
+        if (args.subsamplingFactor > 1) {
+            if (args.subsampling === 'dw_striding' && !(args.causalDownsampling ?? false)) {
+                this.preEncode = new DwStridingSubsampling(args);
+            }
+            else {
+                throw new Error('Only dw_striding non-causal subsampling is supported');
+            }
+        }
+        else {
+            this.preEncode = new Linear(args.featIn, args.dModel);
+        }
+        this.layers = Array.from({ length: args.nLayers }, () => new ConformerBlock(args));
+    }
+    setAttentionModel(name, contextSize = [256, 256]) {
+        if (name === 'rel_pos') {
+            this.posEnc = new RelPositionalEncoding(this.args.dModel, this.args.posEmbMaxLen, this.args.xscaling ?? false);
+        }
+        else if (name === 'rel_pos_local_attn') {
+            this.posEnc = new LocalRelPositionalEncoding(this.args.dModel, this.args.posEmbMaxLen, this.args.xscaling ?? false, contextSize);
+        }
+        else {
+            this.posEnc = null;
+        }
+        for (const layer of this.layers) {
+            layer.setAttentionModel(name, contextSize);
+        }
+    }
+    forward(x, // [batch, seq, mel]
+    lengths, cache) {
+        const xShape = x.shape();
+        const batch = Number(xShape[0]);
+        const seq = Number(xShape[1]);
+        if (lengths === null) {
+            const lenData = new Int32Array(batch).fill(seq);
+            lengths = MxArray.fromInt32(lenData, BigInt64Array.from([BigInt(batch)]));
+        }
+        let outLengths;
+        if (this.preEncode instanceof DwStridingSubsampling) {
+            [x, outLengths] = this.preEncode.forward(x, lengths);
+        }
+        else {
+            x = this.preEncode.forward(x);
+            outLengths = lengths;
+        }
+        const effectiveCache = cache ?? new Array(this.layers.length).fill(null);
+        let posEmb = null;
+        if (this.posEnc !== null) {
+            const offset = effectiveCache[0]?.offset ?? 0;
+            [x, posEmb] = this.posEnc.forward(x, offset);
+        }
+        for (let i = 0; i < this.layers.length; i++) {
+            x = this.layers[i].forward(x, posEmb, null, effectiveCache[i]);
+        }
+        return [x, outLengths];
+    }
+    loadWeights(weights, prefix) {
+        if (this.preEncode instanceof DwStridingSubsampling) {
+            this.preEncode.loadWeights(weights, `${prefix}.pre_encode`);
+        }
+        else {
+            this.preEncode.loadWeights(weights, `${prefix}.pre_encode`);
+        }
+        for (let i = 0; i < this.layers.length; i++) {
+            this.layers[i].loadWeights(weights, `${prefix}.layers.${i}`);
+        }
+        // pos_enc has no learnable weights (pe is computed; pos_bias_{u,v} live in attention)
+    }
+}
+//# sourceMappingURL=conformer.js.map

package/dist/mlx/conformer.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"conformer.js","sourceRoot":"","sources":["../../src/mlx/conformer.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,OAAO,EAAE,MAAM,gBAAgB,CAAC;AACzC,OAAO,EACL,MAAM,EAEN,MAAM,EACN,SAAS,EACT,SAAS,EACT,MAAM,EACN,MAAM,EACN,IAAI,EACJ,IAAI,EACJ,GAAG,GACJ,MAAM,SAAS,CAAC;AACjB,OAAO,EACL,kBAAkB,EAClB,6BAA6B,EAC7B,kCAAkC,EAClC,qBAAqB,EACrB,0BAA0B,GAC3B,MAAM,gBAAgB,CAAC;AAGxB,SAAS,CAAC,CAAC,GAAG,IAAc;IAC1B,OAAO,aAAa,CAAC,IAAI,CAAC,IAAI,CAAC,GAAG,CAAC,MAAM,CAAC,CAAC,CAAC;AAC9C,CAAC;AAyBD,8EAA8E;AAC9E,qBAAqB;AACrB,8EAA8E;AAE9E,MAAM,WAAY,SAAQ,MAAM;IAC9B,OAAO,CAAS;IAChB,OAAO,CAAS;IAEhB,YAAY,MAAc,EAAE,GAAW,EAAE,OAAgB;QACvD,KAAK,EAAE,CAAC;QACR,IAAI,CAAC,OAAO,GAAG,IAAI,MAAM,CAAC,MAAM,EAAE,GAAG,EAAE,OAAO,CAAC,CAAC;QAChD,IAAI,CAAC,OAAO,GAAG,IAAI,MAAM,CAAC,GAAG,EAAE,MAAM,EAAE,OAAO,CAAC,CAAC;IAClD,CAAC;IAED,OAAO,CAAC,CAAU;QAChB,OAAO,IAAI,CAAC,OAAO,CAAC,OAAO,CAAC,IAAI,CAAC,IAAI,CAAC,OAAO,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;IAC7D,CAAC;IAED,WAAW,CAAC,OAAkB,EAAE,MAAc;QAC5C,IAAI,CAAC,OAAO,CAAC,WAAW,CAAC,OAAO,EAAE,GAAG,MAAM,UAAU,CAAC,CAAC;QACvD,IAAI,CAAC,OAAO,CAAC,WAAW,CAAC,OAAO,EAAE,GAAG,MAAM,UAAU,CAAC,CAAC;IACzD,CAAC;CACF;AAED,8EAA8E;AAC9E,oBAAoB;AACpB,8EAA8E;AAE9E,MAAM,WAAY,SAAQ,MAAM;IACrB,OAAO,CAAS;IACzB,cAAc,CAAS;IACvB,aAAa,CAAS;IACtB,SAAS,CAAY;IACrB,cAAc,CAAS;IAEvB,YAAY,IAAmB;QAC7B,KAAK,EAAE,CAAC;QACR,MAAM,OAAO,GAAG,IAAI,CAAC,OAAO,IAAI,IAAI,CAAC;QACrC,IAAI,CAAC,OAAO,GAAG,IAAI,CAAC,KAAK,CAAC,CAAC,IAAI,CAAC,cAAc,GAAG,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC;QAEzD,IAAI,CAAC,cAAc,GAAG,IAAI,MAAM,CAAC,IAAI,CAAC,MAAM,EAAE,IAAI,CAAC,MAAM,GAAG,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,OAAO,CAAC,CAAC;QACpF,IAAI,CAAC,aAAa,GAAG,IAAI,MAAM,CAAC,IAAI,CAAC,MAAM,EAAE,IAAI,CAAC,MAAM,EAAE,IAAI,CAAC,cAAc,EAAE,CAAC,EAAE,CAAC,EAAE,IAAI,CAAC,MAAM,EAAE,OAAO,CAAC,CAAC;QAC3G,IAAI,CAAC,SAAS,GAAG,IAAI,SAAS,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC;QAC5C,IAAI,CAAC,cAAc,GAAG,IAAI,MAAM,CAAC,IAAI,CAAC,MAAM,EAAE,IAAI,CAAC,MAAM,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,OAAO,CAAC,CAAC;IAClF,CAAC;IAED,OAAO,CAAC,CAAU,EAAE,KAA4B;QAC9C,CAAC,GAAG,IAAI,CAAC,cAAc,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC;QACnC,CAAC,GAAG,GAAG,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC,2CAA2C;QAE1D,IAAI,KAAK,KAAK,IAAI,EAAE,CAAC;YACnB,CAAC,GAAG,KAAK,CAAC,kBAAkB,CAAC,CAAC,EAAE,IAAI,CAAC,OAAO,CAAC,CAAC;QAChD,CAAC;aAAM,CAAC;YACN,CAAC,GAAG,CAAC,CAAC,GAAG,CAAC,IAAI,UAAU,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,IAAI,CAAC,OAAO,EAAE,IAAI,CAAC,OAAO,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,EAAE,GAAG,CAAC,CAAC;QAC3E,CAAC;QAED,CAAC,GAAG,IAAI,CAAC,aAAa,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC;QAClC,CAAC,GAAG,IAAI,CAAC,SAAS,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC;QAC9B,CAAC,GAAG,IAAI,CAAC,CAAC,CAAC,CAAC;QACZ,CAAC,GAAG,IAAI,CAAC,cAAc,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC;QACnC,OAAO,CAAC,CAAC;IACX,CAAC;IAED,WAAW,CAAC,OAAkB,EAAE,MAAc;QAC5C,IAAI,CAAC,cAAc,CAAC,WAAW,CAAC,OAAO,EAAE,GAAG,MAAM,kBAAkB,CAAC,CAAC;QACtE,IAAI,CAAC,aAAa,CAAC,WAAW,CAAC,OAAO,EAAE,GAAG,MAAM,iBAAiB,CAAC,CAAC;QACpE,IAAI,CAAC,SAAS,CAAC,WAAW,CAAC,OAAO,EAAE,GAAG,MAAM,aAAa,CAAC,CAAC;QAC5D,IAAI,CAAC,cAAc,CAAC,WAAW,CAAC,OAAO,EAAE,GAAG,MAAM,kBAAkB,CAAC,CAAC;IACxE,CAAC;CACF;AAQD,MAAM,OAAO,cAAe,SAAQ,MAAM;IACxC,OAAO,CAAY;IACnB,GAAG,CAAc;IAEjB,WAAW,CAAY;IACvB,QAAQ,CAA0F;IAElG,QAAQ,CAAY;IACpB,IAAI,CAAc;IAElB,OAAO,CAAY;IACnB,GAAG,CAAc;IAEjB,OAAO,CAAY;IAEF,IAAI,CAAgB;IAErC,YAAY,IAAmB;QAC7B,KAAK,EAAE,CAAC;QACR,IAAI,CAAC,IAAI,GAAG,IAAI,CAAC;QACjB,MAAM,OAAO,GAAG,IAAI,CAAC,OAAO,IAAI,IAAI,CAAC;QACrC,MAAM,KAAK,GAAG,IAAI,CAAC,MAAM,GAAG,IAAI,CAAC,iBAAiB,CAAC;QAEnD,IAAI,CAAC,OAAO,GAAG,IAAI,SAAS,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC;QAC1C,IAAI,CAAC,GAAG,GAAG,IAAI,WAAW,CAAC,IAAI,CAAC,MAAM,EAAE,KAAK,EAAE,OAAO,CAAC,CAAC;QAExD,IAAI,CAAC,WAAW,GAAG,IAAI,SAAS,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC;QAC9C,IAAI,CAAC,QAAQ,GAAG,IAAI,CAAC,cAAc,CAAC,IAAI,CAAC,kBAAkB,EAAE,IAAI,CAAC,cAAc,IAAI,IAAI,CAAC,CAAC;QAE1F,IAAI,CAAC,QAAQ,GAAG,IAAI,SAAS,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC;QAC3C,IAAI,CAAC,IAAI,GAAG,IAAI,WAAW,CAAC,IAAI,CAAC,CAAC;QAElC,IAAI,CAAC,OAAO,GAAG,IAAI,SAAS,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC;QAC1C,IAAI,CAAC,GAAG,GAAG,IAAI,WAAW,CAAC,IAAI,CAAC,MAAM,EAAE,KAAK,EAAE,OAAO,CAAC,CAAC;QAExD,IAAI,CAAC,OAAO,GAAG,IAAI,SAAS,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC;IAC5C,CAAC;IAEO,cAAc,CACpB,IAAY,EACZ,WAAoC;QAEpC,MAAM,OAAO,GAAG,IAAI,CAAC,IAAI,CAAC,OAAO,IAAI,IAAI,CAAC;QAC1C,IAAI,IAAI,KAAK,SAAS,EAAE,CAAC;YACvB,OAAO,IAAI,6BAA6B,CAAC,IAAI,CAAC,IAAI,CAAC,MAAM,EAAE,IAAI,CAAC,IAAI,CAAC,MAAM,EAAE,OAAO,CAAC,CAAC;QACxF,CAAC;aAAM,IAAI,IAAI,KAAK,oBAAoB,EAAE,CAAC;YACzC,OAAO,IAAI,kCAAkC,CAC3C,IAAI,CAAC,IAAI,CAAC,MAAM,EAChB,IAAI,CAAC,IAAI,CAAC,MAAM,EAChB,OAAO,EACP,WAAW,IAAI,CAAC,GAAG,EAAE,GAAG,CAAC,CAC1B,CAAC;QACJ,CAAC;aAAM,CAAC;YACN,OAAO,IAAI,kBAAkB,CAAC,IAAI,CAAC,IAAI,CAAC,MAAM,EAAE,IAAI,CAAC,IAAI,CAAC,MAAM,EAAE,IAAI,CAAC,CAAC;QAC1E,CAAC;IACH,CAAC;IAED,iBAAiB,CAAC,IAAoB,EAAE,cAAgC,CAAC,GAAG,EAAE,GAAG,CAAC;QAChF,MAAM,OAAO,GAAG,IAAI,CAAC,cAAc,CAAC,IAAI,EAAE,WAAW,CAAC,CAAC;QACvD,8CAA8C;QAC9C,8DAA8D;QAC9D,IAAI,CAAC,QAAQ,GAAG,OAAO,CAAC;IAC1B,CAAC;IAED,OAAO,CACL,CAAU,EACV,MAAsB,EACtB,IAAoB,EACpB,KAA4B;QAE5B,MAAM;QACN,CAAC,GAAG,CAAC,CAAC,GAAG,CAAC,IAAI,CAAC,GAAG,CAAC,OAAO,CAAC,IAAI,CAAC,OAAO,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC,CAAC,SAAS,CAAC,GAAG,CAAC,CAAC,CAAC;QAEpE,iBAAiB;QACjB,MAAM,KAAK,GAAG,IAAI,CAAC,WAAW,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC;QAC1C,CAAC,GAAG,CAAC,CAAC,GAAG,CAAC,IAAI,CAAC,QAAQ,CAAC,OAAO,CAAC,KAAK,EAAE,KAAK,EAAE,KAAK,EAAE,MAAM,EAAE,IAAI,EAAE,KAAK,CAAC,CAAC,CAAC;QAE3E,cAAc;QACd,CAAC,GAAG,CAAC,CAAC,GAAG,CAAC,IAAI,CAAC,IAAI,CAAC,OAAO,CAAC,IAAI,CAAC,QAAQ,CAAC,OAAO,CAAC,CAAC,CAAC,EAAE,KAAK,CAAC,CAAC,CAAC;QAE9D,MAAM;QACN,CAAC,GAAG,CAAC,CAAC,GAAG,CAAC,IAAI,CAAC,GAAG,CAAC,OAAO,CAAC,IAAI,CAAC,OAAO,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC,CAAC,SAAS,CAAC,GAAG,CAAC,CAAC,CAAC;QAEpE,OAAO,IAAI,CAAC,OAAO,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC;IACjC,CAAC;IAED,WAAW,CAAC,OAAkB,EAAE,MAAc;QAC5C,IAAI,CAAC,OAAO,CAAC,WAAW,CAAC,OAAO,EAAE,GAAG,MAAM,qBAAqB,CAAC,CAAC;QAClE,IAAI,CAAC,GAAG,CAAC,WAAW,CAAC,OAAO,EAAE,GAAG,MAAM,gBAAgB,CAAC,CAAC;QACzD,IAAI,CAAC,WAAW,CAAC,WAAW,CAAC,OAAO,EAAE,GAAG,MAAM,gBAAgB,CAAC,CAAC;QACjE,IAAI,CAAC,QAAQ,CAAC,WAAW,CAAC,OAAO,EAAE,GAAG,MAAM,YAAY,CAAC,CAAC;QAC1D,IAAI,CAAC,QAAQ,CAAC,WAAW,CAAC,OAAO,EAAE,GAAG,MAAM,YAAY,CAAC,CAAC;QAC1D,IAAI,CAAC,IAAI,CAAC,WAAW,CAAC,OAAO,EAAE,GAAG,MAAM,OAAO,CAAC,CAAC;QACjD,IAAI,CAAC,OAAO,CAAC,WAAW,CAAC,OAAO,EAAE,GAAG,MAAM,qBAAqB,CAAC,CAAC;QAClE,IAAI,CAAC,GAAG,CAAC,WAAW,CAAC,OAAO,EAAE,GAAG,MAAM,gBAAgB,CAAC,CAAC;QACzD,IAAI,CAAC,OAAO,CAAC,WAAW,CAAC,OAAO,EAAE,GAAG,MAAM,WAAW,CAAC,CAAC;IAC1D,CAAC;CACF;AAED,8EAA8E;AAC9E,yCAAyC;AACzC,8EAA8E;AAE9E,MAAM,qBAAsB,SAAQ,MAAM;IACvB,WAAW,CAAS;IACpB,MAAM,GAAG,CAAC,CAAC;IACX,UAAU,GAAG,CAAC,CAAC;IACf,OAAO,CAAS;IAEjC,UAAU,CAAuB,CAAC,0BAA0B;IAC5D,GAAG,CAAS;IAEZ,YAAY,IAAmB;QAC7B,KAAK,EAAE,CAAC;QACR,IAAI,CAAC,OAAO,GAAG,IAAI,CAAC,KAAK,CAAC,CAAC,IAAI,CAAC,UAAU,GAAG,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC;QACrD,IAAI,CAAC,WAAW,GAAG,IAAI,CAAC,KAAK,CAAC,IAAI,CAAC,IAAI,CAAC,IAAI,CAAC,iBAAiB,CAAC,CAAC,CAAC;QAEjE,MAAM,YAAY,GAAG,IAAI,CAAC,uBAAuB,CAAC;QAElD,oCAAoC;QACpC,IAAI,YAAY,GAAG,IAAI,CAAC,MAAM,CAAC;QAC/B,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,IAAI,CAAC,WAAW,EAAE,CAAC,EAAE,EAAE,CAAC;YAC1C,YAAY;gBACV,IAAI,CAAC,KAAK,CAAC,CAAC,YAAY,GAAG,CAAC,GAAG,IAAI,CAAC,OAAO,GAAG,IAAI,CAAC,UAAU,CAAC,GAAG,IAAI,CAAC,MAAM,CAAC,GAAG,CAAC,CAAC;QACtF,CAAC;QAED,2BAA2B;QAC3B,IAAI,CAAC,UAAU,GAAG,EAAE,CAAC;QACrB,IAAI,IAAI,GAAG,CAAC,CAAC;QAEb,+BAA+B;QAC/B,IAAI,CAAC,UAAU,CAAC,IAAI,CAAC,IAAI,MAAM,CAAC,IAAI,EAAE,YAAY,EAAE,IAAI,CAAC,UAAU,EAAE,IAAI,CAAC,MAAM,EAAE,IAAI,CAAC,OAAO,EAAE,CAAC,EAAE,IAAI,CAAC,CAAC,CAAC;QAC1G,IAAI,CAAC,UAAU,CAAC,IAAI,CAAC,IAAI,CAAC,CAAC,CAAC,OAAO;QAEnC,IAAI,GAAG,YAAY,CAAC;QACpB,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,IAAI,CAAC,WAAW,EAAE,CAAC,EAAE,EAAE,CAAC;YAC1C,YAAY;YACZ,IAAI,CAAC,UAAU,CAAC,IAAI,CAAC,IAAI,MAAM,CAAC,IAAI,EAAE,IAAI,EAAE,IAAI,CAAC,UAAU,EAAE,IAAI,CAAC,MAAM,EAAE,IAAI,CAAC,OAAO,EAAE,IAAI,EAAE,IAAI,CAAC,CAAC,CAAC;YACrG,YAAY;YACZ,IAAI,CAAC,UAAU,CAAC,IAAI,CAAC,IAAI,MAAM,CAAC,IAAI,EAAE,YAAY,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,IAAI,CAAC,CAAC,CAAC;YACvE,IAAI,CAAC,UAAU,CAAC,IAAI,CAAC,IAAI,CAAC,CAAC,CAAC,OAAO;QACrC,CAAC;QAED,IAAI,CAAC,GAAG,GAAG,IAAI,MAAM,CAAC,YAAY,GAAG,YAAY,EAAE,IAAI,CAAC,MAAM,CAAC,CAAC;IAClE,CAAC;IAED,OAAO,CAAC,CAAU,EAAE,OAAgB;QAClC,qFAAqF;QACrF,MAAM,MAAM,GAAG,CAAC,CAAC,KAAK,EAAE,CAAC;QACzB,MAAM,KAAK,GAAG,MAAM,CAAC,MAAM,CAAC,CAAC,CAAC,CAAC,CAAC;QAEhC,iBAAiB;QACjB,IAAI,UAAU,GAAG,OAAO,CAAC;QACzB,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,IAAI,CAAC,WAAW,EAAE,CAAC,EAAE,EAAE,CAAC;YAC1C,MAAM,GAAG,GAAG,IAAI,CAAC,OAAO,CAAC;YACzB,MAAM,CAAC,GAAG,IAAI,CAAC,UAAU,CAAC;YAC1B,MAAM,EAAE,GAAG,IAAI,CAAC,MAAM,CAAC;YACvB,wCAAwC;YACxC,UAAU,GAAG,UAAU;iBACpB,SAAS,CAAC,CAAC,GAAG,GAAG,GAAG,CAAC,CAAC;iBACtB,SAAS,CAAC,EAAE,CAAC;iBACb,KAAK,EAAE;iBACP,SAAS,CAAC,CAAC,CAAC,CAAC;QAClB,CAAC;QACD,UAAU,GAAG,UAAU,CAAC,MAAM,CAAC,CAAC,CAAC,CAAC,CAAC,YAAY;QAE/C,sEAAsE;QACtE,IAAI,GAAG,GAAG,CAAC,CAAC,UAAU,CAAC,CAAC,CAAC,CAAC,CAAC,uBAAuB;QAElD,KAAK,MAAM,KAAK,IAAI,IAAI,CAAC,UAAU,EAAE,CAAC;YACpC,IAAI,KAAK,KAAK,IAAI,EAAE,CAAC;gBACnB,GAAG,GAAG,IAAI,CAAC,GAAG,CAAC,CAAC;YAClB,CAAC;iBAAM,CAAC;gBACN,GAAG,GAAG,KAAK,CAAC,OAAO,CAAC,GAAG,CAAC,CAAC;YAC3B,CAAC;QACH,CAAC;QAED,2DAA2D;QAC3D,8EAA8E;QAC9E,0EAA0E;QAC1E,8EAA8E;QAC9E,gDAAgD;QAChD,MAAM,MAAM,GAAG,GAAG,CAAC,KAAK,EAAE,CAAC;QAC3B,MAAM,MAAM,GAAG,MAAM,CAAC,MAAM,CAAC,CAAC,CAAC,CAAC,CAAC;QACjC,MAAM,MAAM,GAAG,MAAM,CAAC,MAAM,CAAC,CAAC,CAAC,CAAC,CAAC;QACjC,MAAM,EAAE,GAAG,MAAM,CAAC,MAAM,CAAC,CAAC,CAAC,CAAC,CAAC;QAE7B,GAAG,GAAG,GAAG;aACN,SAAS,CAAC,IAAI,UAAU,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;aACvC,OAAO,CAAC,CAAC,CAAC,KAAK,EAAE,MAAM,EAAE,EAAE,GAAG,MAAM,CAAC,CAAC,CAAC;QAC1C,GAAG,GAAG,IAAI,CAAC,GAAG,CAAC,OAAO,CAAC,GAAG,CAAC,CAAC;QAE5B,OAAO,CAAC,GAAG,EAAE,UAAU,CAAC,CAAC;IAC3B,CAAC;IAED,WAAW,CAAC,OAAkB,EAAE,MAAc;QAC5C,IAAI,OAAO,GAAG,CAAC,CAAC;QAChB,KAAK,MAAM,KAAK,IAAI,IAAI,CAAC,UAAU,EAAE,CAAC;YACpC,IAAI,KAAK,KAAK,IAAI,EAAE,CAAC;gBACnB,KAAK,CAAC,WAAW,CAAC,OAAO,EAAE,GAAG,MAAM,SAAS,OAAO,EAAE,CAAC,CAAC;YAC1D,CAAC;YACD,OAAO,EAAE,CAAC;QACZ,CAAC;QACD,IAAI,CAAC,GAAG,CAAC,WAAW,CAAC,OAAO,EAAE,GAAG,MAAM,MAAM,CAAC,CAAC;IACjD,CAAC;CACF;AAED,8EAA8E;AAC9E,oBAAoB;AACpB,8EAA8E;AAE9E,MAAM,OAAO,SAAU,SAAQ,MAAM;IAC1B,IAAI,CAAgB;IAC7B,MAAM,CAA4D;IAClE,SAAS,CAAiC;IAC1C,MAAM,CAAmB;IAEzB,YAAY,IAAmB;QAC7B,KAAK,EAAE,CAAC;QACR,IAAI,CAAC,IAAI,GAAG,IAAI,CAAC;QAEjB,MAAM,YAAY,GAAG,IAAI,CAAC,kBAAkB,CAAC;QAC7C,MAAM,OAAO,GAAG,IAAI,CAAC,cAAc,IAAI,IAAI,CAAC;QAE5C,IAAI,YAAY,KAAK,SAAS,EAAE,CAAC;YAC/B,IAAI,CAAC,MAAM,GAAG,IAAI,qBAAqB,CAAC,IAAI,CAAC,MAAM,EAAE,IAAI,CAAC,YAAY,EAAE,IAAI,CAAC,QAAQ,IAAI,KAAK,CAAC,CAAC;QAClG,CAAC;aAAM,IAAI,YAAY,KAAK,oBAAoB,EAAE,CAAC;YACjD,IAAI,CAAC,MAAM,GAAG,IAAI,0BAA0B,CAC1C,IAAI,CAAC,MAAM,EACX,IAAI,CAAC,YAAY,EACjB,IAAI,CAAC,QAAQ,IAAI,KAAK,EACtB,OAAO,IAAI,CAAC,GAAG,EAAE,GAAG,CAAC,CACtB,CAAC;QACJ,CAAC;aAAM,CAAC;YACN,IAAI,CAAC,MAAM,GAAG,IAAI,CAAC;QACrB,CAAC;QAED,IAAI,IAAI,CAAC,iBAAiB,GAAG,CAAC,EAAE,CAAC;YAC/B,IAAI,IAAI,CAAC,WAAW,KAAK,aAAa,IAAI,CAAC,CAAC,IAAI,CAAC,kBAAkB,IAAI,KAAK,CAAC,EAAE,CAAC;gBAC9E,IAAI,CAAC,SAAS,GAAG,IAAI,qBAAqB,CAAC,IAAI,CAAC,CAAC;YACnD,CAAC;iBAAM,CAAC;gBACN,MAAM,IAAI,KAAK,CAAC,sDAAsD,CAAC,CAAC;YAC1E,CAAC;QACH,CAAC;aAAM,CAAC;YACN,IAAI,CAAC,SAAS,GAAG,IAAI,MAAM,CAAC,IAAI,CAAC,MAAM,EAAE,IAAI,CAAC,MAAM,CAAC,CAAC;QACxD,CAAC;QAED,IAAI,CAAC,MAAM,GAAG,KAAK,CAAC,IAAI,CAAC,EAAE,MAAM,EAAE,IAAI,CAAC,OAAO,EAAE,EAAE,GAAG,EAAE,CAAC,IAAI,cAAc,CAAC,IAAI,CAAC,CAAC,CAAC;IACrF,CAAC;IAED,iBAAiB,CACf,IAAoB,EACpB,cAAgC,CAAC,GAAG,EAAE,GAAG,CAAC;QAE1C,IAAI,IAAI,KAAK,SAAS,EAAE,CAAC;YACvB,IAAI,CAAC,MAAM,GAAG,IAAI,qBAAqB,CAAC,IAAI,CAAC,IAAI,CAAC,MAAM,EAAE,IAAI,CAAC,IAAI,CAAC,YAAY,EAAE,IAAI,CAAC,IAAI,CAAC,QAAQ,IAAI,KAAK,CAAC,CAAC;QACjH,CAAC;aAAM,IAAI,IAAI,KAAK,oBAAoB,EAAE,CAAC;YACzC,IAAI,CAAC,MAAM,GAAG,IAAI,0BAA0B,CAC1C,IAAI,CAAC,IAAI,CAAC,MAAM,EAChB,IAAI,CAAC,IAAI,CAAC,YAAY,EACtB,IAAI,CAAC,IAAI,CAAC,QAAQ,IAAI,KAAK,EAC3B,WAAW,CACZ,CAAC;QACJ,CAAC;aAAM,CAAC;YACN,IAAI,CAAC,MAAM,GAAG,IAAI,CAAC;QACrB,CAAC;QAED,KAAK,MAAM,KAAK,IAAI,IAAI,CAAC,MAAM,EAAE,CAAC;YAChC,KAAK,CAAC,iBAAiB,CAAC,IAAI,EAAE,WAAW,CAAC,CAAC;QAC7C,CAAC;IACH,CAAC;IAED,OAAO,CACL,CAAU,EAAE,oBAAoB;IAChC,OAAuB,EACvB,KAA0C;QAE1C,MAAM,MAAM,GAAG,CAAC,CAAC,KAAK,EAAE,CAAC;QACzB,MAAM,KAAK,GAAG,MAAM,CAAC,MAAM,CAAC,CAAC,CAAC,CAAC,CAAC;QAChC,MAAM,GAAG,GAAG,MAAM,CAAC,MAAM,CAAC,CAAC,CAAC,CAAC,CAAC;QAE9B,IAAI,OAAO,KAAK,IAAI,EAAE,CAAC;YACrB,MAAM,OAAO,GAAG,IAAI,UAAU,CAAC,KAAK,CAAC,CAAC,IAAI,CAAC,GAAG,CAAC,CAAC;YAChD,OAAO,GAAG,OAAO,CAAC,SAAS,CAAC,OAAO,EAAE,aAAa,CAAC,IAAI,CAAC,CAAC,MAAM,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC,CAAC;QAC5E,CAAC;QAED,IAAI,UAAmB,CAAC;QAExB,IAAI,IAAI,CAAC,SAAS,YAAY,qBAAqB,EAAE,CAAC;YACpD,CAAC,CAAC,EAAE,UAAU,CAAC,GAAG,IAAI,CAAC,SAAS,CAAC,OAAO,CAAC,CAAC,EAAE,OAAO,CAAC,CAAC;QACvD,CAAC;aAAM,CAAC;YACN,CAAC,GAAG,IAAI,CAAC,SAAS,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC;YAC9B,UAAU,GAAG,OAAO,CAAC;QACvB,CAAC;QAED,MAAM,cAAc,GAAG,KAAK,IAAI,IAAI,KAAK,CAAC,IAAI,CAAC,MAAM,CAAC,MAAM,CAAC,CAAC,IAAI,CAAC,IAAI,CAAC,CAAC;QAEzE,IAAI,MAAM,GAAmB,IAAI,CAAC;QAClC,IAAI,IAAI,CAAC,MAAM,KAAK,IAAI,EAAE,CAAC;YACzB,MAAM,MAAM,GAAG,cAAc,CAAC,CAAC,CAAC,EAAE,MAAM,IAAI,CAAC,CAAC;YAC9C,CAAC,CAAC,EAAE,MAAM,CAAC,GAAG,IAAI,CAAC,MAAM,CAAC,OAAO,CAAC,CAAC,EAAE,MAAM,CAAC,CAAC;QAC/C,CAAC;QAED,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,IAAI,CAAC,MAAM,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE,CAAC;YAC5C,CAAC,GAAG,IAAI,CAAC,MAAM,CAAC,CAAC,CAAC,CAAC,OAAO,CAAC,CAAC,EAAE,MAAM,EAAE,IAAI,EAAE,cAAc,CAAC,CAAC,CAAC,CAAC,CAAC;QACjE,CAAC;QAED,OAAO,CAAC,CAAC,EAAE,UAAU,CAAC,CAAC;IACzB,CAAC;IAED,WAAW,CAAC,OAAkB,EAAE,MAAc;QAC5C,IAAI,IAAI,CAAC,SAAS,YAAY,qBAAqB,EAAE,CAAC;YACpD,IAAI,CAAC,SAAS,CAAC,WAAW,CAAC,OAAO,EAAE,GAAG,MAAM,aAAa,CAAC,CAAC;QAC9D,CAAC;aAAM,CAAC;YACL,IAAI,CAAC,SAAoB,CAAC,WAAW,CAAC,OAAO,EAAE,GAAG,MAAM,aAAa,CAAC,CAAC;QAC1E,CAAC;QAED,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,IAAI,CAAC,MAAM,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE,CAAC;YAC5C,IAAI,CAAC,MAAM,CAAC,CAAC,CAAC,CAAC,WAAW,CAAC,OAAO,EAAE,GAAG,MAAM,WAAW,CAAC,EAAE,CAAC,CAAC;QAC/D,CAAC;QAED,sFAAsF;IACxF,CAAC;CACF"}

package/dist/mlx/index.d.ts

@@ -0,0 +1,6 @@
+export { fromLocal, fromPretrained } from './load.js';
+export type { MlxModelOptions, FromPretrainedOptions } from './load.js';
+export { MlxBackend } from './backend.js';
+export { ParakeetModel, StreamingParakeet, consumePcmStream, decodeTDTGreedy, decodeRNNTGreedy, getLogMel, loadAudioRaw, makePreprocessArgs, computeMelFilterbanks, computeMelFilterbanksInterpolated, tokensToSentences, sentencesToResult, decode, } from '../index.js';
+export type { ParakeetBackend, EncoderOutput, TranscribeOptions, StreamOptions, DecoderState, PreprocessArgs, LogMel, AlignedToken, AlignedSentence, AlignedResult, SentenceConfig, } from '../index.js';
+//# sourceMappingURL=index.d.ts.map

package/dist/mlx/index.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../../src/mlx/index.ts"],"names":[],"mappings":"AAKA,OAAO,EAAE,SAAS,EAAE,cAAc,EAAE,MAAM,WAAW,CAAC;AACtD,YAAY,EAAE,eAAe,EAAE,qBAAqB,EAAE,MAAM,WAAW,CAAC;AACxE,OAAO,EAAE,UAAU,EAAE,MAAM,cAAc,CAAC;AAG1C,OAAO,EACL,aAAa,EACb,iBAAiB,EACjB,gBAAgB,EAChB,eAAe,EACf,gBAAgB,EAChB,SAAS,EACT,YAAY,EACZ,kBAAkB,EAClB,qBAAqB,EACrB,iCAAiC,EACjC,iBAAiB,EACjB,iBAAiB,EACjB,MAAM,GACP,MAAM,aAAa,CAAC;AACrB,YAAY,EACV,eAAe,EACf,aAAa,EACb,iBAAiB,EACjB,aAAa,EACb,YAAY,EACZ,cAAc,EACd,MAAM,EACN,YAAY,EACZ,eAAe,EACf,aAAa,EACb,cAAc,GACf,MAAM,aAAa,CAAC"}

package/dist/mlx/index.js

@@ -0,0 +1,9 @@
+// parakeet.ts/mlx — MLX backend (Apple Silicon)
+//
+// Loads a safetensors checkpoint and returns the shared, backend-agnostic
+// `ParakeetModel` — the same class the ONNX backend returns.
+export { fromLocal, fromPretrained } from './load.js';
+export { MlxBackend } from './backend.js';
+// Shared API, re-exported for convenience
+export { ParakeetModel, StreamingParakeet, consumePcmStream, decodeTDTGreedy, decodeRNNTGreedy, getLogMel, loadAudioRaw, makePreprocessArgs, computeMelFilterbanks, computeMelFilterbanksInterpolated, tokensToSentences, sentencesToResult, decode, } from '../index.js';
+//# sourceMappingURL=index.js.map

package/dist/mlx/index.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"index.js","sourceRoot":"","sources":["../../src/mlx/index.ts"],"names":[],"mappings":"AAAA,gDAAgD;AAChD,EAAE;AACF,0EAA0E;AAC1E,6DAA6D;AAE7D,OAAO,EAAE,SAAS,EAAE,cAAc,EAAE,MAAM,WAAW,CAAC;AAEtD,OAAO,EAAE,UAAU,EAAE,MAAM,cAAc,CAAC;AAE1C,0CAA0C;AAC1C,OAAO,EACL,aAAa,EACb,iBAAiB,EACjB,gBAAgB,EAChB,eAAe,EACf,gBAAgB,EAChB,SAAS,EACT,YAAY,EACZ,kBAAkB,EAClB,qBAAqB,EACrB,iCAAiC,EACjC,iBAAiB,EACjB,iBAAiB,EACjB,MAAM,GACP,MAAM,aAAa,CAAC"}

package/dist/mlx/load.d.ts

@@ -0,0 +1,27 @@
+import { ParakeetModel } from '../model.js';
+export interface MlxModelOptions {
+    /**
+     * 'interpolated' (default) matches NVIDIA's reference preprocessor.
+     * 'floor' reproduces parakeet-mlx's filterbank, which collapses 13 of 128 mel
+     * bins to zero — see docs/cuda.md.
+     */
+    filterbank?: 'floor' | 'interpolated';
+}
+/**
+ * Load a safetensors Parakeet checkpoint from a local directory.
+ *
+ * @param dir directory holding `config.json` and `model.safetensors`
+ */
+export declare function fromLocal(dir: string, options?: MlxModelOptions): ParakeetModel;
+export interface FromPretrainedOptions extends MlxModelOptions {
+    cacheDir?: string;
+    onProgress?: (file: string, downloaded: number, total: number) => void;
+}
+/**
+ * Load a Parakeet checkpoint from a HuggingFace Hub repo or a local directory.
+ *
+ * @param hfIdOrPath HuggingFace repo id (e.g. "mlx-community/parakeet-tdt-0.6b-v3")
+ *                   or a local directory path.
+ */
+export declare function fromPretrained(hfIdOrPath: string, options?: FromPretrainedOptions): Promise<ParakeetModel>;
+//# sourceMappingURL=load.d.ts.map

package/dist/mlx/load.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"load.d.ts","sourceRoot":"","sources":["../../src/mlx/load.ts"],"names":[],"mappings":"AAoBA,OAAO,EAAE,aAAa,EAAE,MAAM,aAAa,CAAC;AAG5C,MAAM,WAAW,eAAe;IAC9B;;;;OAIG;IACH,UAAU,CAAC,EAAE,OAAO,GAAG,cAAc,CAAC;CACvC;AAwID;;;;GAIG;AACH,wBAAgB,SAAS,CAAC,GAAG,EAAE,MAAM,EAAE,OAAO,GAAE,eAAoB,GAAG,aAAa,CAQnF;AAED,MAAM,WAAW,qBAAsB,SAAQ,eAAe;IAC5D,QAAQ,CAAC,EAAE,MAAM,CAAC;IAClB,UAAU,CAAC,EAAE,CAAC,IAAI,EAAE,MAAM,EAAE,UAAU,EAAE,MAAM,EAAE,KAAK,EAAE,MAAM,KAAK,IAAI,CAAC;CACxE;AAED;;;;;GAKG;AACH,wBAAsB,cAAc,CAClC,UAAU,EAAE,MAAM,EAClB,OAAO,GAAE,qBAA0B,GAClC,OAAO,CAAC,aAAa,CAAC,CAexB"}

package/dist/mlx/load.js

@@ -0,0 +1,161 @@
+/**
+ * Loading MLX-weight Parakeet checkpoints into the backend-agnostic
+ * `ParakeetModel`.
+ *
+ * This is the single MLX entry point: it builds the Conformer encoder, the
+ * prediction network and the joint network directly from `config.json`, loads
+ * the safetensors weights into them, and wraps them behind `MlxBackend` so the
+ * MLX path drives the same decode loop, tokenizer, alignment and audio
+ * front-end as the ONNX path. MLX and ONNX therefore produce identical output
+ * from identical features.
+ *
+ * TDT and RNN-T checkpoints are supported (this is also all the ONNX export
+ * covers). CTC / hybrid TDT-CTC checkpoints are not.
+ */
+import fs from 'node:fs';
+import path from 'node:path';
+import { Conformer } from './conformer.js';
+import { PredictNetwork, JointNetwork } from './rnnt.js';
+import { MlxBackend } from './backend.js';
+import { loadSafetensors, downloadFromHub } from './utils.js';
+import { ParakeetModel } from '../model.js';
+import { makePreprocessArgs } from '../audio.js';
+/** Build the MLX modules and metadata from a parsed NeMo-style config. */
+function buildFromConfig(config) {
+    const modelDefaults = config['model_defaults'] ?? {};
+    const encoderRaw = config['encoder'];
+    const decoderRaw = config['decoder'];
+    const jointRaw = config['joint'];
+    const decodingRaw = config['decoding'] ?? {};
+    if (!jointRaw || !decoderRaw?.['prednet']) {
+        throw new Error('Only TDT and RNN-T checkpoints are supported by the MLX loader ' +
+            '(config needs decoder.prednet and a joint network). ' +
+            'CTC / hybrid TDT-CTC checkpoints are not supported.');
+    }
+    const attContextSize = encoderRaw['att_context_size'] ?? null;
+    const encoderArgs = {
+        featIn: encoderRaw['feat_in'],
+        nLayers: encoderRaw['n_layers'],
+        dModel: encoderRaw['d_model'],
+        nHeads: encoderRaw['n_heads'],
+        ffExpansionFactor: encoderRaw['ff_expansion_factor'],
+        subsamplingFactor: encoderRaw['subsampling_factor'],
+        selfAttentionModel: encoderRaw['self_attention_model'],
+        subsampling: encoderRaw['subsampling'],
+        convKernelSize: encoderRaw['conv_kernel_size'],
+        subsamplingConvChannels: encoderRaw['subsampling_conv_channels'],
+        posEmbMaxLen: encoderRaw['pos_emb_max_len'],
+        causalDownsampling: encoderRaw['causal_downsampling'] ?? false,
+        useBias: encoderRaw['use_bias'] ?? true,
+        xscaling: encoderRaw['xscaling'] ?? false,
+        subsamplingConvChunkingFactor: encoderRaw['subsampling_conv_chunking_factor'] ?? 1,
+        attContextSize,
+    };
+    const prednetRaw = decoderRaw['prednet'] ?? {};
+    const predictArgs = {
+        blankAsPad: decoderRaw['blank_as_pad'],
+        vocabSize: decoderRaw['vocab_size'],
+        prednet: {
+            predHidden: prednetRaw['pred_hidden'],
+            predRnnLayers: prednetRaw['pred_rnn_layers'],
+            rnnHiddenSize: prednetRaw['rnn_hidden_size'],
+        },
+    };
+    const jointnetRaw = jointRaw['jointnet'] ?? {};
+    const jointArgs = {
+        numClasses: jointRaw['num_classes'],
+        vocabulary: jointRaw['vocabulary'],
+        jointnet: {
+            jointHidden: jointnetRaw['joint_hidden'],
+            activation: jointnetRaw['activation'],
+            encoderHidden: jointnetRaw['encoder_hidden'],
+            predHidden: jointnetRaw['pred_hidden'],
+        },
+        numExtraOutputs: jointRaw['num_extra_outputs'] ?? 0,
+    };
+    const durations = modelDefaults['tdt_durations'] ?? null;
+    const greedy = decodingRaw['greedy'];
+    const maxSymbols = greedy?.['max_symbols'] != null ? Number(greedy['max_symbols']) : null;
+    return {
+        encoder: new Conformer(encoderArgs),
+        predict: new PredictNetwork(predictArgs),
+        joint: new JointNetwork(jointArgs),
+        vocabulary: jointArgs.vocabulary,
+        durations,
+        maxSymbols,
+        encoderDim: encoderArgs.dModel,
+        subsamplingFactor: encoderArgs.subsamplingFactor,
+    };
+}
+/** Assemble a `ParakeetModel` from a config, a weights file, and options. */
+function assemble(config, weightsPath, options) {
+    const built = buildFromConfig(config);
+    const weights = loadSafetensors(weightsPath);
+    built.encoder.loadWeights(weights, 'encoder');
+    built.predict.loadWeights(weights, 'decoder');
+    built.joint.loadWeights(weights, 'joint');
+    const pre = config['preprocessor'];
+    const preprocessor = makePreprocessArgs({
+        sampleRate: pre['sample_rate'],
+        normalize: pre['normalize'],
+        windowSize: pre['window_size'],
+        windowStride: pre['window_stride'],
+        window: pre['window'],
+        features: pre['features'],
+        nFft: pre['n_fft'],
+        dither: pre['dither'],
+        padTo: pre['pad_to'] ?? 0,
+        padValue: pre['pad_value'] ?? 0,
+        preemph: pre['preemph'],
+        magPower: pre['mag_power'] ?? 2.0,
+        filterbank: options.filterbank ?? 'interpolated',
+    });
+    const backend = new MlxBackend({
+        encoder: built.encoder,
+        predict: built.predict,
+        joint: built.joint,
+        encoderDim: built.encoderDim,
+    });
+    return new ParakeetModel({
+        backend,
+        preprocessor,
+        vocabulary: built.vocabulary,
+        durations: built.durations,
+        maxSymbols: built.maxSymbols,
+        subsamplingFactor: built.subsamplingFactor,
+    });
+}
+/**
+ * Load a safetensors Parakeet checkpoint from a local directory.
+ *
+ * @param dir directory holding `config.json` and `model.safetensors`
+ */
+export function fromLocal(dir, options = {}) {
+    const configPath = path.join(dir, 'config.json');
+    const weightsPath = path.join(dir, 'model.safetensors');
+    if (!fs.existsSync(configPath))
+        throw new Error(`config.json not found in ${dir}`);
+    if (!fs.existsSync(weightsPath))
+        throw new Error(`model.safetensors not found in ${dir}`);
+    const config = JSON.parse(fs.readFileSync(configPath, 'utf8'));
+    return assemble(config, weightsPath, options);
+}
+/**
+ * Load a Parakeet checkpoint from a HuggingFace Hub repo or a local directory.
+ *
+ * @param hfIdOrPath HuggingFace repo id (e.g. "mlx-community/parakeet-tdt-0.6b-v3")
+ *                   or a local directory path.
+ */
+export async function fromPretrained(hfIdOrPath, options = {}) {
+    if (fs.existsSync(hfIdOrPath) && fs.statSync(hfIdOrPath).isDirectory()) {
+        return fromLocal(hfIdOrPath, options);
+    }
+    const makeProgress = (file) => options.onProgress
+        ? (downloaded, total) => options.onProgress(file, downloaded, total)
+        : undefined;
+    const configPath = await downloadFromHub(hfIdOrPath, 'config.json', options.cacheDir, makeProgress('config.json'));
+    const weightsPath = await downloadFromHub(hfIdOrPath, 'model.safetensors', options.cacheDir, makeProgress('model.safetensors'));
+    const config = JSON.parse(fs.readFileSync(configPath, 'utf8'));
+    return assemble(config, weightsPath, options);
+}
+//# sourceMappingURL=load.js.map

package/dist/mlx/load.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"load.js","sourceRoot":"","sources":["../../src/mlx/load.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;GAaG;AACH,OAAO,EAAE,MAAM,SAAS,CAAC;AACzB,OAAO,IAAI,MAAM,WAAW,CAAC;AAC7B,OAAO,EAAE,SAAS,EAAiB,MAAM,gBAAgB,CAAC;AAC1D,OAAO,EAAE,cAAc,EAAE,YAAY,EAA0B,MAAM,WAAW,CAAC;AACjF,OAAO,EAAE,UAAU,EAAE,MAAM,cAAc,CAAC;AAC1C,OAAO,EAAE,eAAe,EAAE,eAAe,EAAE,MAAM,YAAY,CAAC;AAC9D,OAAO,EAAE,aAAa,EAAE,MAAM,aAAa,CAAC;AAC5C,OAAO,EAAE,kBAAkB,EAAE,MAAM,aAAa,CAAC;AAsBjD,0EAA0E;AAC1E,SAAS,eAAe,CAAC,MAA+B;IACtD,MAAM,aAAa,GAAI,MAAM,CAAC,gBAAgB,CAA6B,IAAI,EAAE,CAAC;IAClF,MAAM,UAAU,GAAG,MAAM,CAAC,SAAS,CAA4B,CAAC;IAChE,MAAM,UAAU,GAAG,MAAM,CAAC,SAAS,CAA4B,CAAC;IAChE,MAAM,QAAQ,GAAG,MAAM,CAAC,OAAO,CAAwC,CAAC;IACxE,MAAM,WAAW,GAAI,MAAM,CAAC,UAAU,CAA6B,IAAI,EAAE,CAAC;IAE1E,IAAI,CAAC,QAAQ,IAAI,CAAC,UAAU,EAAE,CAAC,SAAS,CAAC,EAAE,CAAC;QAC1C,MAAM,IAAI,KAAK,CACb,iEAAiE;YACjE,sDAAsD;YACtD,qDAAqD,CACtD,CAAC;IACJ,CAAC;IAED,MAAM,cAAc,GAAI,UAAU,CAAC,kBAAkB,CAA6B,IAAI,IAAI,CAAC;IAC3F,MAAM,WAAW,GAAkB;QACjC,MAAM,EAAE,UAAU,CAAC,SAAS,CAAW;QACvC,OAAO,EAAE,UAAU,CAAC,UAAU,CAAW;QACzC,MAAM,EAAE,UAAU,CAAC,SAAS,CAAW;QACvC,MAAM,EAAE,UAAU,CAAC,SAAS,CAAW;QACvC,iBAAiB,EAAE,UAAU,CAAC,qBAAqB,CAAW;QAC9D,iBAAiB,EAAE,UAAU,CAAC,oBAAoB,CAAW;QAC7D,kBAAkB,EAAE,UAAU,CAAC,sBAAsB,CAAW;QAChE,WAAW,EAAE,UAAU,CAAC,aAAa,CAAW;QAChD,cAAc,EAAE,UAAU,CAAC,kBAAkB,CAAW;QACxD,uBAAuB,EAAE,UAAU,CAAC,2BAA2B,CAAW;QAC1E,YAAY,EAAE,UAAU,CAAC,iBAAiB,CAAW;QACrD,kBAAkB,EAAG,UAAU,CAAC,qBAAqB,CAAa,IAAI,KAAK;QAC3E,OAAO,EAAG,UAAU,CAAC,UAAU,CAAa,IAAI,IAAI;QACpD,QAAQ,EAAG,UAAU,CAAC,UAAU,CAAa,IAAI,KAAK;QACtD,6BAA6B,EAAG,UAAU,CAAC,kCAAkC,CAAY,IAAI,CAAC;QAC9F,cAAc;KACf,CAAC;IAEF,MAAM,UAAU,GAAI,UAAU,CAAC,SAAS,CAA6B,IAAI,EAAE,CAAC;IAC5E,MAAM,WAAW,GAAgB;QAC/B,UAAU,EAAE,UAAU,CAAC,cAAc,CAAY;QACjD,SAAS,EAAE,UAAU,CAAC,YAAY,CAAW;QAC7C,OAAO,EAAE;YACP,UAAU,EAAE,UAAU,CAAC,aAAa,CAAW;YAC/C,aAAa,EAAE,UAAU,CAAC,iBAAiB,CAAW;YACtD,aAAa,EAAE,UAAU,CAAC,iBAAiB,CAAuB;SACnE;KACF,CAAC;IAEF,MAAM,WAAW,GAAI,QAAQ,CAAC,UAAU,CAA6B,IAAI,EAAE,CAAC;IAC5E,MAAM,SAAS,GAAc;QAC3B,UAAU,EAAE,QAAQ,CAAC,aAAa,CAAW;QAC7C,UAAU,EAAE,QAAQ,CAAC,YAAY,CAAa;QAC9C,QAAQ,EAAE;YACR,WAAW,EAAE,WAAW,CAAC,cAAc,CAAW;YAClD,UAAU,EAAE,WAAW,CAAC,YAAY,CAAW;YAC/C,aAAa,EAAE,WAAW,CAAC,gBAAgB,CAAW;YACtD,UAAU,EAAE,WAAW,CAAC,aAAa,CAAW;SACjD;QACD,eAAe,EAAG,QAAQ,CAAC,mBAAmB,CAAY,IAAI,CAAC;KAChE,CAAC;IAEF,MAAM,SAAS,GAAI,aAAa,CAAC,eAAe,CAA0B,IAAI,IAAI,CAAC;IACnF,MAAM,MAAM,GAAG,WAAW,CAAC,QAAQ,CAAwC,CAAC;IAC5E,MAAM,UAAU,GAAG,MAAM,EAAE,CAAC,aAAa,CAAC,IAAI,IAAI,CAAC,CAAC,CAAC,MAAM,CAAC,MAAM,CAAC,aAAa,CAAC,CAAC,CAAC,CAAC,CAAC,IAAI,CAAC;IAE1F,OAAO;QACL,OAAO,EAAE,IAAI,SAAS,CAAC,WAAW,CAAC;QACnC,OAAO,EAAE,IAAI,cAAc,CAAC,WAAW,CAAC;QACxC,KAAK,EAAE,IAAI,YAAY,CAAC,SAAS,CAAC;QAClC,UAAU,EAAE,SAAS,CAAC,UAAU;QAChC,SAAS;QACT,UAAU;QACV,UAAU,EAAE,WAAW,CAAC,MAAM;QAC9B,iBAAiB,EAAE,WAAW,CAAC,iBAAiB;KACjD,CAAC;AACJ,CAAC;AAED,6EAA6E;AAC7E,SAAS,QAAQ,CACf,MAA+B,EAC/B,WAAmB,EACnB,OAAwB;IAExB,MAAM,KAAK,GAAG,eAAe,CAAC,MAAM,CAAC,CAAC;IAEtC,MAAM,OAAO,GAAG,eAAe,CAAC,WAAW,CAAC,CAAC;IAC7C,KAAK,CAAC,OAAO,CAAC,WAAW,CAAC,OAAO,EAAE,SAAS,CAAC,CAAC;IAC9C,KAAK,CAAC,OAAO,CAAC,WAAW,CAAC,OAAO,EAAE,SAAS,CAAC,CAAC;IAC9C,KAAK,CAAC,KAAK,CAAC,WAAW,CAAC,OAAO,EAAE,OAAO,CAAC,CAAC;IAE1C,MAAM,GAAG,GAAG,MAAM,CAAC,cAAc,CAA4B,CAAC;IAC9D,MAAM,YAAY,GAAG,kBAAkB,CAAC;QACtC,UAAU,EAAE,GAAG,CAAC,aAAa,CAAW;QACxC,SAAS,EAAE,GAAG,CAAC,WAAW,CAAW;QACrC,UAAU,EAAE,GAAG,CAAC,aAAa,CAAW;QACxC,YAAY,EAAE,GAAG,CAAC,eAAe,CAAW;QAC5C,MAAM,EAAE,GAAG,CAAC,QAAQ,CAAW;QAC/B,QAAQ,EAAE,GAAG,CAAC,UAAU,CAAW;QACnC,IAAI,EAAE,GAAG,CAAC,OAAO,CAAW;QAC5B,MAAM,EAAE,GAAG,CAAC,QAAQ,CAAW;QAC/B,KAAK,EAAG,GAAG,CAAC,QAAQ,CAAY,IAAI,CAAC;QACrC,QAAQ,EAAG,GAAG,CAAC,WAAW,CAAY,IAAI,CAAC;QAC3C,OAAO,EAAE,GAAG,CAAC,SAAS,CAAkB;QACxC,QAAQ,EAAG,GAAG,CAAC,WAAW,CAAY,IAAI,GAAG;QAC7C,UAAU,EAAE,OAAO,CAAC,UAAU,IAAI,cAAc;KACjD,CAAC,CAAC;IAEH,MAAM,OAAO,GAAG,IAAI,UAAU,CAAC;QAC7B,OAAO,EAAE,KAAK,CAAC,OAAO;QACtB,OAAO,EAAE,KAAK,CAAC,OAAO;QACtB,KAAK,EAAE,KAAK,CAAC,KAAK;QAClB,UAAU,EAAE,KAAK,CAAC,UAAU;KAC7B,CAAC,CAAC;IAEH,OAAO,IAAI,aAAa,CAAC;QACvB,OAAO;QACP,YAAY;QACZ,UAAU,EAAE,KAAK,CAAC,UAAU;QAC5B,SAAS,EAAE,KAAK,CAAC,SAAS;QAC1B,UAAU,EAAE,KAAK,CAAC,UAAU;QAC5B,iBAAiB,EAAE,KAAK,CAAC,iBAAiB;KAC3C,CAAC,CAAC;AACL,CAAC;AAED;;;;GAIG;AACH,MAAM,UAAU,SAAS,CAAC,GAAW,EAAE,UAA2B,EAAE;IAClE,MAAM,UAAU,GAAG,IAAI,CAAC,IAAI,CAAC,GAAG,EAAE,aAAa,CAAC,CAAC;IACjD,MAAM,WAAW,GAAG,IAAI,CAAC,IAAI,CAAC,GAAG,EAAE,mBAAmB,CAAC,CAAC;IACxD,IAAI,CAAC,EAAE,CAAC,UAAU,CAAC,UAAU,CAAC;QAAE,MAAM,IAAI,KAAK,CAAC,4BAA4B,GAAG,EAAE,CAAC,CAAC;IACnF,IAAI,CAAC,EAAE,CAAC,UAAU,CAAC,WAAW,CAAC;QAAE,MAAM,IAAI,KAAK,CAAC,kCAAkC,GAAG,EAAE,CAAC,CAAC;IAE1F,MAAM,MAAM,GAAG,IAAI,CAAC,KAAK,CAAC,EAAE,CAAC,YAAY,CAAC,UAAU,EAAE,MAAM,CAAC,CAA4B,CAAC;IAC1F,OAAO,QAAQ,CAAC,MAAM,EAAE,WAAW,EAAE,OAAO,CAAC,CAAC;AAChD,CAAC;AAOD;;;;;GAKG;AACH,MAAM,CAAC,KAAK,UAAU,cAAc,CAClC,UAAkB,EAClB,UAAiC,EAAE;IAEnC,IAAI,EAAE,CAAC,UAAU,CAAC,UAAU,CAAC,IAAI,EAAE,CAAC,QAAQ,CAAC,UAAU,CAAC,CAAC,WAAW,EAAE,EAAE,CAAC;QACvE,OAAO,SAAS,CAAC,UAAU,EAAE,OAAO,CAAC,CAAC;IACxC,CAAC;IAED,MAAM,YAAY,GAAG,CAAC,IAAY,EAAE,EAAE,CACpC,OAAO,CAAC,UAAU;QAChB,CAAC,CAAC,CAAC,UAAkB,EAAE,KAAa,EAAE,EAAE,CAAC,OAAO,CAAC,UAAW,CAAC,IAAI,EAAE,UAAU,EAAE,KAAK,CAAC;QACrF,CAAC,CAAC,SAAS,CAAC;IAEhB,MAAM,UAAU,GAAG,MAAM,eAAe,CAAC,UAAU,EAAE,aAAa,EAAE,OAAO,CAAC,QAAQ,EAAE,YAAY,CAAC,aAAa,CAAC,CAAC,CAAC;IACnH,MAAM,WAAW,GAAG,MAAM,eAAe,CAAC,UAAU,EAAE,mBAAmB,EAAE,OAAO,CAAC,QAAQ,EAAE,YAAY,CAAC,mBAAmB,CAAC,CAAC,CAAC;IAEhI,MAAM,MAAM,GAAG,IAAI,CAAC,KAAK,CAAC,EAAE,CAAC,YAAY,CAAC,UAAU,EAAE,MAAM,CAAC,CAA4B,CAAC;IAC1F,OAAO,QAAQ,CAAC,MAAM,EAAE,WAAW,EAAE,OAAO,CAAC,CAAC;AAChD,CAAC"}