mambacode.js 1.0.0 → 1.0.1

package/dist/model/mamba_model.js

@@ -0,0 +1,289 @@
+/**
+ * mamba_model.ts – Full Mamba language model.
+ */
+import { MambaBlock } from './mamba_block';
+import { createStorageBuffer, createEmptyStorageBuffer, createUniformBuffer, createComputePipeline, createBindGroup, dispatchKernel, readBuffer, uploadBuffer, cdiv, } from '../utils/gpu_utils';
+import { LINEAR_FORWARD_WGSL } from '../kernels/linear_projection';
+import { ACTIVATIONS_WGSL } from '../kernels/activations';
+export class MambaModel {
+    device;
+    config;
+    gpuEmbedding;
+    blocks;
+    gpuFinalNorm;
+    tiedEmbedding;
+    gpuLMHeadBias;
+    _lmHeadPipeline;
+    _rmsnormPipeline;
+    _embedPipeline;
+    _wslaMode = false;
+    constructor(device, config) {
+        this.device = device;
+        this.config = {
+            dState: 16,
+            dConv: 4,
+            expand: 2,
+            eosId: -1,
+            ...config,
+        };
+        const { vocabSize, dModel, numLayers } = this.config;
+        const embedData = new Float32Array(vocabSize * dModel);
+        const std = 1.0 / Math.sqrt(dModel);
+        for (let i = 0; i < embedData.length; i++) {
+            const u1 = Math.random(), u2 = Math.random();
+            embedData[i] = std * Math.sqrt(-2 * Math.log(u1 + 1e-12)) *
+                Math.cos(2 * Math.PI * u2);
+        }
+        this.gpuEmbedding = createStorageBuffer(device, embedData, true);
+        this.blocks = Array.from({ length: numLayers }, () => new MambaBlock(device, {
+            dModel,
+            dState: this.config.dState,
+            dConv: this.config.dConv,
+            expand: this.config.expand,
+        }));
+        const finalNormW = new Float32Array(dModel).fill(1.0);
+        this.gpuFinalNorm = createStorageBuffer(device, finalNormW, true);
+        this.tiedEmbedding = true;
+        this._lmHeadPipeline = createComputePipeline(device, LINEAR_FORWARD_WGSL, 'linear_forward');
+        this._rmsnormPipeline = createComputePipeline(device, ACTIVATIONS_WGSL, 'rmsnorm_forward');
+        this.gpuLMHeadBias = createStorageBuffer(device, new Float32Array(vocabSize), true);
+        this._embedPipeline = createComputePipeline(device, EMBED_LOOKUP_WGSL, 'embed_lookup');
+    }
+    embedTokens(tokenIds, batch, seqLen) {
+        const { dModel } = this.config;
+        const M = batch * seqLen;
+        const idsBuf = createStorageBuffer(this.device, tokenIds instanceof Uint32Array ? tokenIds : new Uint32Array(tokenIds), false);
+        const outBuf = createEmptyStorageBuffer(this.device, M * dModel * 4, true);
+        const params = new Uint32Array([M, dModel]).buffer;
+        const pBuf = createUniformBuffer(this.device, params);
+        const bg = createBindGroup(this.device, this._embedPipeline, [pBuf, idsBuf, this.gpuEmbedding, outBuf]);
+        dispatchKernel(this.device, this._embedPipeline, bg, [cdiv(M, 64), 1, 1]);
+        idsBuf.destroy();
+        pBuf.destroy();
+        return outBuf;
+    }
+    async forward(tokenIds, batch, seqLen) {
+        const { dModel, vocabSize } = this.config;
+        const M = batch * seqLen;
+        let hidden = this.embedTokens(tokenIds, batch, seqLen);
+        const caches = [];
+        for (const block of this.blocks) {
+            const { output, cache } = block.forward(hidden, batch, seqLen);
+            caches.push(cache);
+            hidden.destroy();
+            hidden = output;
+        }
+        const normOut = createEmptyStorageBuffer(this.device, M * dModel * 4, true);
+        const normInv = createEmptyStorageBuffer(this.device, M * 4, false);
+        {
+            const params = new ArrayBuffer(16);
+            new Uint32Array(params, 0, 2).set([M, dModel]);
+            new Float32Array(params, 8, 1).set([1e-6]);
+            const pBuf = createUniformBuffer(this.device, params);
+            const bg = createBindGroup(this.device, this._rmsnormPipeline, [pBuf, hidden, this.gpuFinalNorm, normOut, normInv]);
+            dispatchKernel(this.device, this._rmsnormPipeline, bg, [cdiv(M, 64), 1, 1]);
+        }
+        const gpuLogits = createEmptyStorageBuffer(this.device, M * vocabSize * 4, true);
+        {
+            const params = new Uint32Array([M, dModel, vocabSize]).buffer;
+            const pBuf = createUniformBuffer(this.device, params);
+            const weightBuf = this.tiedEmbedding ? this.gpuEmbedding : this.gpuLMHeadBias;
+            const bg = createBindGroup(this.device, this._lmHeadPipeline, [pBuf, normOut, weightBuf, this.gpuLMHeadBias, gpuLogits]);
+            dispatchKernel(this.device, this._lmHeadPipeline, bg, [cdiv(M, 16), cdiv(vocabSize, 16), 1]);
+        }
+        normOut.destroy();
+        normInv.destroy();
+        const logits = await readBuffer(this.device, gpuLogits, M * vocabSize * 4);
+        return { logits, gpuLogits, caches };
+    }
+    async generate(promptIds, maxNewTokens = 200, samplingOpts = {}) {
+        const { temperature = 1.0, topK = 50, topP = 0.9 } = samplingOpts;
+        const { vocabSize } = this.config;
+        let ids = [...promptIds];
+        for (let step = 0; step < maxNewTokens; step++) {
+            const { logits } = await this.forward(new Uint32Array(ids), 1, ids.length);
+            const lastLogits = logits.slice((ids.length - 1) * vocabSize, ids.length * vocabSize);
+            const nextId = sampleToken(lastLogits, { temperature, topK, topP });
+            ids.push(nextId);
+            if (nextId === this.config.eosId)
+                break;
+        }
+        return ids;
+    }
+    parameters() {
+        const params = [];
+        params.push({
+            buf: this.gpuEmbedding,
+            numel: this.config.vocabSize * this.config.dModel,
+            name: 'embedding',
+        });
+        for (let i = 0; i < this.blocks.length; i++) {
+            for (const p of this.blocks[i].parameters()) {
+                params.push({ ...p, name: `block${i}.${p.name}` });
+            }
+        }
+        params.push({
+            buf: this.gpuFinalNorm,
+            numel: this.config.dModel,
+            name: 'final_norm',
+        });
+        return params;
+    }
+    setWSLAMode(enabled) {
+        for (const block of this.blocks)
+            block.setWSLAMode(enabled);
+        this._wslaMode = enabled;
+    }
+    /**
+     * Serialise all model parameters to an ArrayBuffer.
+     *
+     * Binary format:
+     *   [0..3]   magic  : uint32  = 0x4D424A53 ('MBJS')
+     *   [4..7]   version: uint32  = 1
+     *   [8..11]  nParams: uint32
+     *   [12 .. 12+4*nParams-1]  numel[i]: uint32 for each parameter i
+     *   [12+4*nParams ..]       float32 data for each parameter, concatenated
+     *
+     * Save the returned buffer to a file or IndexedDB and reload it with
+     * `loadWeights()` to resume from a checkpoint.
+     */
+    async exportWeights() {
+        const params = this.parameters();
+        const nParams = params.length;
+        // Read all GPU buffers into CPU Float32Arrays
+        const arrays = await Promise.all(params.map(p => readBuffer(this.device, p.buf, p.numel * 4)));
+        // Calculate total byte size: header + numel table + all float data
+        const headerBytes = 4 + 4 + 4 + nParams * 4; // magic + version + nParams + numel[]
+        const dataBytes = arrays.reduce((acc, a) => acc + a.byteLength, 0);
+        const out = new ArrayBuffer(headerBytes + dataBytes);
+        const view = new DataView(out);
+        let offset = 0;
+        view.setUint32(offset, 0x4D424A53, true);
+        offset += 4; // magic 'MBJS'
+        view.setUint32(offset, 1, true);
+        offset += 4; // version
+        view.setUint32(offset, nParams, true);
+        offset += 4; // nParams
+        for (const p of params) {
+            view.setUint32(offset, p.numel, true);
+            offset += 4;
+        }
+        for (const arr of arrays) {
+            new Float32Array(out, offset, arr.length).set(arr);
+            offset += arr.byteLength;
+        }
+        return out;
+    }
+    /**
+     * Load model parameters from an ArrayBuffer previously produced by
+     * `exportWeights()`.  The parameter count and element counts must match
+     * the current model configuration exactly.
+     *
+     * @throws {Error} if the magic number, version, or parameter layout do
+     *                 not match the current model.
+     */
+    async loadWeights(buffer) {
+        const view = new DataView(buffer);
+        let offset = 0;
+        const magic = view.getUint32(offset, true);
+        offset += 4;
+        if (magic !== 0x4D424A53) {
+            throw new Error('Invalid weight file: bad magic number. ' +
+                'Ensure the file was exported by MambaModel.exportWeights().');
+        }
+        const version = view.getUint32(offset, true);
+        offset += 4;
+        if (version !== 1) {
+            throw new Error(`Unsupported weight file version: ${version}. Expected version 1.`);
+        }
+        const nParams = view.getUint32(offset, true);
+        offset += 4;
+        const params = this.parameters();
+        if (nParams !== params.length) {
+            throw new Error(`Weight file has ${nParams} parameters but this model has ${params.length}. ` +
+                'Ensure the model configuration matches the one used when exporting.');
+        }
+        const numels = [];
+        for (let i = 0; i < nParams; i++) {
+            numels.push(view.getUint32(offset, true));
+            offset += 4;
+        }
+        for (let i = 0; i < nParams; i++) {
+            // i is guaranteed in-bounds: nParams === params.length was verified above
+            const p = params[i];
+            const numel = numels[i];
+            if (numel !== p.numel) {
+                throw new Error(`Parameter ${i} ("${p.name}") size mismatch: ` +
+                    `file has ${numel} elements, model expects ${p.numel}.`);
+            }
+            const slice = new Float32Array(buffer, offset, p.numel);
+            uploadBuffer(this.device, p.buf, slice);
+            offset += p.numel * 4;
+        }
+    }
+}
+const EMBED_LOOKUP_WGSL = /* wgsl */ `
+struct EmbedParams {
+    num_tokens : u32,
+    d_model    : u32,
+};
+
+@group(0) @binding(0) var<uniform>            params  : EmbedParams;
+@group(0) @binding(1) var<storage, read>      ids     : array<u32>;
+@group(0) @binding(2) var<storage, read>      table   : array<f32>;
+@group(0) @binding(3) var<storage, read_write> out    : array<f32>;
+
+@compute @workgroup_size(64, 1, 1)
+fn embed_lookup(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let token_idx = gid.x;
+    if (token_idx >= params.num_tokens) { return; }
+
+    let D   = params.d_model;
+    let tok = ids[token_idx];
+    let src = tok * D;
+    let dst = token_idx * D;
+
+    for (var i: u32 = 0u; i < D; i = i + 1u) {
+        out[dst + i] = table[src + i];
+    }
+}
+`;
+function sampleToken(logits, { temperature = 1.0, topK = 50, topP = 0.9 } = {}) {
+    const n = logits.length;
+    const scaled = new Float32Array(n);
+    for (let i = 0; i < n; i++)
+        scaled[i] = logits[i] / Math.max(temperature, 1e-7);
+    let maxL = -Infinity;
+    for (let i = 0; i < n; i++)
+        if (scaled[i] > maxL)
+            maxL = scaled[i];
+    let sumE = 0;
+    const exps = new Float32Array(n);
+    for (let i = 0; i < n; i++) {
+        exps[i] = Math.exp(scaled[i] - maxL);
+        sumE += exps[i];
+    }
+    const indices = Array.from({ length: n }, (_, i) => i)
+        .sort((a, b) => exps[b] - exps[a]);
+    const topKIndices = indices.slice(0, topK);
+    let cumSum = 0;
+    const nucleus = [];
+    for (const idx of topKIndices) {
+        cumSum += exps[idx] / sumE;
+        nucleus.push(idx);
+        if (cumSum >= topP)
+            break;
+    }
+    let nucleusSum = 0;
+    for (const idx of nucleus)
+        nucleusSum += exps[idx];
+    const threshold = Math.random() * nucleusSum;
+    let acc = 0;
+    for (const idx of nucleus) {
+        acc += exps[idx];
+        if (acc >= threshold)
+            return idx;
+    }
+    return nucleus[nucleus.length - 1];
+}
+//# sourceMappingURL=mamba_model.js.map

package/dist/model/mamba_model.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"mamba_model.js","sourceRoot":"","sources":["../../src/model/mamba_model.ts"],"names":[],"mappings":"AAAA;;GAEG;AAEH,OAAO,EAAE,UAAU,EAA0B,MAAM,eAAe,CAAC;AACnE,OAAO,EACH,mBAAmB,EACnB,wBAAwB,EACxB,mBAAmB,EACnB,qBAAqB,EACrB,eAAe,EACf,cAAc,EACd,UAAU,EACV,YAAY,EACZ,IAAI,GACP,MAAM,oBAAoB,CAAC;AAC5B,OAAO,EAAE,mBAAmB,EAAE,MAAM,8BAA8B,CAAC;AACnE,OAAO,EAAE,gBAAgB,EAAE,MAAS,wBAAwB,CAAC;AAwB7D,MAAM,OAAO,UAAU;IACnB,MAAM,CAAY;IAClB,MAAM,CAA6B;IACnC,YAAY,CAAY;IACxB,MAAM,CAAe;IACrB,YAAY,CAAY;IACxB,aAAa,CAAU;IACvB,aAAa,CAAY;IACjB,eAAe,CAAqB;IACpC,gBAAgB,CAAqB;IACrC,cAAc,CAAqB;IACnC,SAAS,GAAG,KAAK,CAAC;IAE1B,YAAY,MAAiB,EAAE,MAAwB;QACnD,IAAI,CAAC,MAAM,GAAG,MAAM,CAAC;QACrB,IAAI,CAAC,MAAM,GAAG;YACV,MAAM,EAAM,EAAE;YACd,KAAK,EAAO,CAAC;YACb,MAAM,EAAM,CAAC;YACb,KAAK,EAAO,CAAC,CAAC;YACd,GAAG,MAAM;SACkB,CAAC;QAEhC,MAAM,EAAE,SAAS,EAAE,MAAM,EAAE,SAAS,EAAE,GAAG,IAAI,CAAC,MAAM,CAAC;QAErD,MAAM,SAAS,GAAG,IAAI,YAAY,CAAC,SAAS,GAAG,MAAM,CAAC,CAAC;QACvD,MAAM,GAAG,GAAG,GAAG,GAAG,IAAI,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC;QACpC,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,SAAS,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE,CAAC;YACxC,MAAM,EAAE,GAAG,IAAI,CAAC,MAAM,EAAE,EAAE,EAAE,GAAG,IAAI,CAAC,MAAM,EAAE,CAAC;YAC7C,SAAS,CAAC,CAAC,CAAC,GAAG,GAAG,GAAG,IAAI,CAAC,IAAI,CAAC,CAAC,CAAC,GAAG,IAAI,CAAC,GAAG,CAAC,EAAE,GAAG,KAAK,CAAC,CAAC;gBAC1C,IAAI,CAAC,GAAG,CAAC,CAAC,GAAG,IAAI,CAAC,EAAE,GAAG,EAAE,CAAC,CAAC;QAC9C,CAAC;QACD,IAAI,CAAC,YAAY,GAAG,mBAAmB,CAAC,MAAM,EAAE,SAAS,EAAE,IAAI,CAAC,CAAC;QAEjE,IAAI,CAAC,MAAM,GAAG,KAAK,CAAC,IAAI,CAAC,EAAE,MAAM,EAAE,SAAS,EAAE,EAAE,GAAG,EAAE,CACjD,IAAI,UAAU,CAAC,MAAM,EAAE;YACnB,MAAM;YACN,MAAM,EAAI,IAAI,CAAC,MAAM,CAAC,MAAM;YAC5B,KAAK,EAAK,IAAI,CAAC,MAAM,CAAC,KAAK;YAC3B,MAAM,EAAI,IAAI,CAAC,MAAM,CAAC,MAAM;SAC/B,CAAC,CACL,CAAC;QAEF,MAAM,UAAU,GAAG,IAAI,YAAY,CAAC,MAAM,CAAC,CAAC,IAAI,CAAC,GAAG,CAAC,CAAC;QACtD,IAAI,CAAC,YAAY,GAAG,mBAAmB,CAAC,MAAM,EAAE,UAAU,EAAE,IAAI,CAAC,CAAC;QAElE,IAAI,CAAC,aAAa,GAAG,IAAI,CAAC;QAE1B,IAAI,CAAC,eAAe,GAAI,qBAAqB,CAAC,MAAM,EAAE,mBAAmB,EAAE,gBAAgB,CAAC,CAAC;QAC7F,IAAI,CAAC,gBAAgB,GAAG,qBAAqB,CAAC,MAAM,EAAE,gBAAgB,EAAK,iBAAiB,CAAC,CAAC;QAE9F,IAAI,CAAC,aAAa,GAAG,mBAAmB,CAAC,MAAM,EAAE,IAAI,YAAY,CAAC,SAAS,CAAC,EAAE,IAAI,CAAC,CAAC;QAEpF,IAAI,CAAC,cAAc,GAAG,qBAAqB,CAAC,MAAM,EAAE,iBAAiB,EAAE,cAAc,CAAC,CAAC;IAC3F,CAAC;IAED,WAAW,CAAC,QAAgC,EAAE,KAAa,EAAE,MAAc;QACvE,MAAM,EAAE,MAAM,EAAE,GAAG,IAAI,CAAC,MAAM,CAAC;QAC/B,MAAM,CAAC,GAAG,KAAK,GAAG,MAAM,CAAC;QAEzB,MAAM,MAAM,GAAI,mBAAmB,CAAC,IAAI,CAAC,MAAM,EAC3C,QAAQ,YAAY,WAAW,CAAC,CAAC,CAAC,QAAQ,CAAC,CAAC,CAAC,IAAI,WAAW,CAAC,QAAQ,CAAC,EAAE,KAAK,CAAC,CAAC;QACnF,MAAM,MAAM,GAAI,wBAAwB,CAAC,IAAI,CAAC,MAAM,EAAE,CAAC,GAAG,MAAM,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;QAE5E,MAAM,MAAM,GAAI,IAAI,WAAW,CAAC,CAAC,CAAC,EAAE,MAAM,CAAC,CAAC,CAAC,MAAM,CAAC;QACpD,MAAM,IAAI,GAAM,mBAAmB,CAAC,IAAI,CAAC,MAAM,EAAE,MAAM,CAAC,CAAC;QAEzD,MAAM,EAAE,GAAG,eAAe,CAAC,IAAI,CAAC,MAAM,EAAE,IAAI,CAAC,cAAc,EACvD,CAAC,IAAI,EAAE,MAAM,EAAE,IAAI,CAAC,YAAY,EAAE,MAAM,CAAC,CAAC,CAAC;QAC/C,cAAc,CAAC,IAAI,CAAC,MAAM,EAAE,IAAI,CAAC,cAAc,EAAE,EAAE,EAAE,CAAC,IAAI,CAAC,CAAC,EAAE,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;QAE1E,MAAM,CAAC,OAAO,EAAE,CAAC;QACjB,IAAI,CAAC,OAAO,EAAE,CAAC;QACf,OAAO,MAAM,CAAC;IAClB,CAAC;IAED,KAAK,CAAC,OAAO,CAAC,QAAgC,EAAE,KAAa,EAAE,MAAc;QACzE,MAAM,EAAE,MAAM,EAAE,SAAS,EAAE,GAAG,IAAI,CAAC,MAAM,CAAC;QAC1C,MAAM,CAAC,GAAG,KAAK,GAAG,MAAM,CAAC;QAEzB,IAAI,MAAM,GAAG,IAAI,CAAC,WAAW,CAAC,QAAQ,EAAE,KAAK,EAAE,MAAM,CAAC,CAAC;QAEvD,MAAM,MAAM,GAAiB,EAAE,CAAC;QAChC,KAAK,MAAM,KAAK,IAAI,IAAI,CAAC,MAAM,EAAE,CAAC;YAC9B,MAAM,EAAE,MAAM,EAAE,KAAK,EAAE,GAAG,KAAK,CAAC,OAAO,CAAC,MAAM,EAAE,KAAK,EAAE,MAAM,CAAC,CAAC;YAC/D,MAAM,CAAC,IAAI,CAAC,KAAK,CAAC,CAAC;YACnB,MAAM,CAAC,OAAO,EAAE,CAAC;YACjB,MAAM,GAAG,MAAM,CAAC;QACpB,CAAC;QAED,MAAM,OAAO,GAAG,wBAAwB,CAAC,IAAI,CAAC,MAAM,EAAE,CAAC,GAAG,MAAM,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;QAC5E,MAAM,OAAO,GAAG,wBAAwB,CAAC,IAAI,CAAC,MAAM,EAAE,CAAC,GAAG,CAAC,EAAW,KAAK,CAAC,CAAC;QAC7E,CAAC;YACG,MAAM,MAAM,GAAG,IAAI,WAAW,CAAC,EAAE,CAAC,CAAC;YACnC,IAAI,WAAW,CAAC,MAAM,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,MAAM,CAAC,CAAC,CAAC;YAC/C,IAAI,YAAY,CAAC,MAAM,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,GAAG,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC;YAC3C,MAAM,IAAI,GAAG,mBAAmB,CAAC,IAAI,CAAC,MAAM,EAAE,MAAM,CAAC,CAAC;YACtD,MAAM,EAAE,GAAG,eAAe,CAAC,IAAI,CAAC,MAAM,EAAE,IAAI,CAAC,gBAAgB,EACzD,CAAC,IAAI,EAAE,MAAM,EAAE,IAAI,CAAC,YAAY,EAAE,OAAO,EAAE,OAAO,CAAC,CAAC,CAAC;YACzD,cAAc,CAAC,IAAI,CAAC,MAAM,EAAE,IAAI,CAAC,gBAAgB,EAAE,EAAE,EAAE,CAAC,IAAI,CAAC,CAAC,EAAE,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;QAChF,CAAC;QAED,MAAM,SAAS,GAAG,wBAAwB,CAAC,IAAI,CAAC,MAAM,EAAE,CAAC,GAAG,SAAS,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;QACjF,CAAC;YACG,MAAM,MAAM,GAAG,IAAI,WAAW,CAAC,CAAC,CAAC,EAAE,MAAM,EAAE,SAAS,CAAC,CAAC,CAAC,MAAM,CAAC;YAC9D,MAAM,IAAI,GAAK,mBAAmB,CAAC,IAAI,CAAC,MAAM,EAAE,MAAM,CAAC,CAAC;YACxD,MAAM,SAAS,GAAG,IAAI,CAAC,aAAa,CAAC,CAAC,CAAC,IAAI,CAAC,YAAY,CAAC,CAAC,CAAC,IAAI,CAAC,aAAa,CAAC;YAC9E,MAAM,EAAE,GAAG,eAAe,CAAC,IAAI,CAAC,MAAM,EAAE,IAAI,CAAC,eAAe,EACxD,CAAC,IAAI,EAAE,OAAO,EAAE,SAAS,EAAE,IAAI,CAAC,aAAa,EAAE,SAAS,CAAC,CAAC,CAAC;YAC/D,cAAc,CAAC,IAAI,CAAC,MAAM,EAAE,IAAI,CAAC,eAAe,EAAE,EAAE,EAChD,CAAC,IAAI,CAAC,CAAC,EAAE,EAAE,CAAC,EAAE,IAAI,CAAC,SAAS,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;QAC/C,CAAC;QAED,OAAO,CAAC,OAAO,EAAE,CAAC;QAClB,OAAO,CAAC,OAAO,EAAE,CAAC;QAElB,MAAM,MAAM,GAAG,MAAM,UAAU,CAAC,IAAI,CAAC,MAAM,EAAE,SAAS,EAAE,CAAC,GAAG,SAAS,GAAG,CAAC,CAAC,CAAC;QAE3E,OAAO,EAAE,MAAM,EAAE,SAAS,EAAE,MAAM,EAAE,CAAC;IACzC,CAAC;IAED,KAAK,CAAC,QAAQ,CAAC,SAAmB,EAAE,YAAY,GAAG,GAAG,EAAE,eAAgC,EAAE;QACtF,MAAM,EAAE,WAAW,GAAG,GAAG,EAAE,IAAI,GAAG,EAAE,EAAE,IAAI,GAAG,GAAG,EAAE,GAAG,YAAY,CAAC;QAClE,MAAM,EAAE,SAAS,EAAE,GAAG,IAAI,CAAC,MAAM,CAAC;QAElC,IAAI,GAAG,GAAG,CAAC,GAAG,SAAS,CAAC,CAAC;QAEzB,KAAK,IAAI,IAAI,GAAG,CAAC,EAAE,IAAI,GAAG,YAAY,EAAE,IAAI,EAAE,EAAE,CAAC;YAC7C,MAAM,EAAE,MAAM,EAAE,GAAG,MAAM,IAAI,CAAC,OAAO,CACjC,IAAI,WAAW,CAAC,GAAG,CAAC,EAAE,CAAC,EAAE,GAAG,CAAC,MAAM,CACtC,CAAC;YACF,MAAM,UAAU,GAAG,MAAM,CAAC,KAAK,CAAC,CAAC,GAAG,CAAC,MAAM,GAAG,CAAC,CAAC,GAAG,SAAS,EAAE,GAAG,CAAC,MAAM,GAAG,SAAS,CAAC,CAAC;YAEtF,MAAM,MAAM,GAAG,WAAW,CAAC,UAAU,EAAE,EAAE,WAAW,EAAE,IAAI,EAAE,IAAI,EAAE,CAAC,CAAC;YACpE,GAAG,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC;YAEjB,IAAI,MAAM,KAAK,IAAI,CAAC,MAAM,CAAC,KAAK;gBAAE,MAAM;QAC5C,CAAC;QAED,OAAO,GAAG,CAAC;IACf,CAAC;IAED,UAAU;QACN,MAAM,MAAM,GAAiB,EAAE,CAAC;QAEhC,MAAM,CAAC,IAAI,CAAC;YACR,GAAG,EAAI,IAAI,CAAC,YAAY;YACxB,KAAK,EAAE,IAAI,CAAC,MAAM,CAAC,SAAS,GAAG,IAAI,CAAC,MAAM,CAAC,MAAM;YACjD,IAAI,EAAG,WAAW;SACrB,CAAC,CAAC;QAEH,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,IAAI,CAAC,MAAM,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE,CAAC;YAC1C,KAAK,MAAM,CAAC,IAAI,IAAI,CAAC,MAAM,CAAC,CAAC,CAAE,CAAC,UAAU,EAAE,EAAE,CAAC;gBAC3C,MAAM,CAAC,IAAI,CAAC,EAAE,GAAG,CAAC,EAAE,IAAI,EAAE,QAAQ,CAAC,IAAI,CAAC,CAAC,IAAI,EAAE,EAAE,CAAC,CAAC;YACvD,CAAC;QACL,CAAC;QAED,MAAM,CAAC,IAAI,CAAC;YACR,GAAG,EAAI,IAAI,CAAC,YAAY;YACxB,KAAK,EAAE,IAAI,CAAC,MAAM,CAAC,MAAM;YACzB,IAAI,EAAG,YAAY;SACtB,CAAC,CAAC;QAEH,OAAO,MAAM,CAAC;IAClB,CAAC;IAED,WAAW,CAAC,OAAgB;QACxB,KAAK,MAAM,KAAK,IAAI,IAAI,CAAC,MAAM;YAAE,KAAK,CAAC,WAAW,CAAC,OAAO,CAAC,CAAC;QAC5D,IAAI,CAAC,SAAS,GAAG,OAAO,CAAC;IAC7B,CAAC;IAED;;;;;;;;;;;;OAYG;IACH,KAAK,CAAC,aAAa;QACf,MAAM,MAAM,GAAG,IAAI,CAAC,UAAU,EAAE,CAAC;QACjC,MAAM,OAAO,GAAG,MAAM,CAAC,MAAM,CAAC;QAE9B,8CAA8C;QAC9C,MAAM,MAAM,GAAmB,MAAM,OAAO,CAAC,GAAG,CAC5C,MAAM,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,CAAC,UAAU,CAAC,IAAI,CAAC,MAAM,EAAE,CAAC,CAAC,GAAG,EAAE,CAAC,CAAC,KAAK,GAAG,CAAC,CAAC,CAAC,CAC/D,CAAC;QAEF,mEAAmE;QACnE,MAAM,WAAW,GAAG,CAAC,GAAG,CAAC,GAAG,CAAC,GAAG,OAAO,GAAG,CAAC,CAAC,CAAE,sCAAsC;QACpF,MAAM,SAAS,GAAK,MAAM,CAAC,MAAM,CAAC,CAAC,GAAG,EAAE,CAAC,EAAE,EAAE,CAAC,GAAG,GAAG,CAAC,CAAC,UAAU,EAAE,CAAC,CAAC,CAAC;QACrE,MAAM,GAAG,GAAW,IAAI,WAAW,CAAC,WAAW,GAAG,SAAS,CAAC,CAAC;QAC7D,MAAM,IAAI,GAAU,IAAI,QAAQ,CAAC,GAAG,CAAC,CAAC;QAEtC,IAAI,MAAM,GAAG,CAAC,CAAC;QACf,IAAI,CAAC,SAAS,CAAC,MAAM,EAAE,UAAU,EAAE,IAAI,CAAC,CAAC;QAAC,MAAM,IAAI,CAAC,CAAC,CAAE,eAAe;QACvE,IAAI,CAAC,SAAS,CAAC,MAAM,EAAE,CAAC,EAAY,IAAI,CAAC,CAAC;QAAC,MAAM,IAAI,CAAC,CAAC,CAAE,UAAU;QACnE,IAAI,CAAC,SAAS,CAAC,MAAM,EAAE,OAAO,EAAM,IAAI,CAAC,CAAC;QAAC,MAAM,IAAI,CAAC,CAAC,CAAE,UAAU;QAEnE,KAAK,MAAM,CAAC,IAAI,MAAM,EAAE,CAAC;YACrB,IAAI,CAAC,SAAS,CAAC,MAAM,EAAE,CAAC,CAAC,KAAK,EAAE,IAAI,CAAC,CAAC;YACtC,MAAM,IAAI,CAAC,CAAC;QAChB,CAAC;QAED,KAAK,MAAM,GAAG,IAAI,MAAM,EAAE,CAAC;YACvB,IAAI,YAAY,CAAC,GAAG,EAAE,MAAM,EAAE,GAAG,CAAC,MAAM,CAAC,CAAC,GAAG,CAAC,GAAG,CAAC,CAAC;YACnD,MAAM,IAAI,GAAG,CAAC,UAAU,CAAC;QAC7B,CAAC;QAED,OAAO,GAAG,CAAC;IACf,CAAC;IAED;;;;;;;OAOG;IACH,KAAK,CAAC,WAAW,CAAC,MAAmB;QACjC,MAAM,IAAI,GAAM,IAAI,QAAQ,CAAC,MAAM,CAAC,CAAC;QACrC,IAAI,MAAM,GAAM,CAAC,CAAC;QAElB,MAAM,KAAK,GAAK,IAAI,CAAC,SAAS,CAAC,MAAM,EAAE,IAAI,CAAC,CAAC;QAAC,MAAM,IAAI,CAAC,CAAC;QAC1D,IAAI,KAAK,KAAK,UAAU,EAAE,CAAC;YACvB,MAAM,IAAI,KAAK,CACX,yCAAyC;gBACzC,6DAA6D,CAChE,CAAC;QACN,CAAC;QAED,MAAM,OAAO,GAAG,IAAI,CAAC,SAAS,CAAC,MAAM,EAAE,IAAI,CAAC,CAAC;QAAC,MAAM,IAAI,CAAC,CAAC;QAC1D,IAAI,OAAO,KAAK,CAAC,EAAE,CAAC;YAChB,MAAM,IAAI,KAAK,CAAC,oCAAoC,OAAO,uBAAuB,CAAC,CAAC;QACxF,CAAC;QAED,MAAM,OAAO,GAAG,IAAI,CAAC,SAAS,CAAC,MAAM,EAAE,IAAI,CAAC,CAAC;QAAC,MAAM,IAAI,CAAC,CAAC;QAC1D,MAAM,MAAM,GAAI,IAAI,CAAC,UAAU,EAAE,CAAC;QAElC,IAAI,OAAO,KAAK,MAAM,CAAC,MAAM,EAAE,CAAC;YAC5B,MAAM,IAAI,KAAK,CACX,mBAAmB,OAAO,kCAAkC,MAAM,CAAC,MAAM,IAAI;gBAC7E,qEAAqE,CACxE,CAAC;QACN,CAAC;QAED,MAAM,MAAM,GAAa,EAAE,CAAC;QAC5B,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,OAAO,EAAE,CAAC,EAAE,EAAE,CAAC;YAC/B,MAAM,CAAC,IAAI,CAAC,IAAI,CAAC,SAAS,CAAC,MAAM,EAAE,IAAI,CAAC,CAAC,CAAC;YAC1C,MAAM,IAAI,CAAC,CAAC;QAChB,CAAC;QAED,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,OAAO,EAAE,CAAC,EAAE,EAAE,CAAC;YAC/B,0EAA0E;YAC1E,MAAM,CAAC,GAAQ,MAAM,CAAC,CAAC,CAAE,CAAC;YAC1B,MAAM,KAAK,GAAI,MAAM,CAAC,CAAC,CAAE,CAAC;YAC1B,IAAI,KAAK,KAAK,CAAC,CAAC,KAAK,EAAE,CAAC;gBACpB,MAAM,IAAI,KAAK,CACX,aAAa,CAAC,MAAM,CAAC,CAAC,IAAI,oBAAoB;oBAC9C,YAAY,KAAK,4BAA4B,CAAC,CAAC,KAAK,GAAG,CAC1D,CAAC;YACN,CAAC;YAED,MAAM,KAAK,GAAG,IAAI,YAAY,CAAC,MAAM,EAAE,MAAM,EAAE,CAAC,CAAC,KAAK,CAAC,CAAC;YACxD,YAAY,CAAC,IAAI,CAAC,MAAM,EAAE,CAAC,CAAC,GAAG,EAAE,KAAK,CAAC,CAAC;YACxC,MAAM,IAAI,CAAC,CAAC,KAAK,GAAG,CAAC,CAAC;QAC1B,CAAC;IACL,CAAC;CACJ;AAED,MAAM,iBAAiB,GAAW,UAAU,CAAA;;;;;;;;;;;;;;;;;;;;;;;;;CAyB3C,CAAC;AAEF,SAAS,WAAW,CAAC,MAAoB,EAAE,EAAE,WAAW,GAAG,GAAG,EAAE,IAAI,GAAG,EAAE,EAAE,IAAI,GAAG,GAAG,EAAE,GAAG,EAAE;IACxF,MAAM,CAAC,GAAG,MAAM,CAAC,MAAM,CAAC;IAExB,MAAM,MAAM,GAAG,IAAI,YAAY,CAAC,CAAC,CAAC,CAAC;IACnC,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,CAAC,EAAE,CAAC,EAAE;QAAE,MAAM,CAAC,CAAC,CAAC,GAAG,MAAM,CAAC,CAAC,CAAE,GAAG,IAAI,CAAC,GAAG,CAAC,WAAW,EAAE,IAAI,CAAC,CAAC;IAEjF,IAAI,IAAI,GAAG,CAAC,QAAQ,CAAC;IACrB,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,CAAC,EAAE,CAAC,EAAE;QAAE,IAAI,MAAM,CAAC,CAAC,CAAE,GAAG,IAAI;YAAE,IAAI,GAAG,MAAM,CAAC,CAAC,CAAE,CAAC;IACrE,IAAI,IAAI,GAAG,CAAC,CAAC;IACb,MAAM,IAAI,GAAG,IAAI,YAAY,CAAC,CAAC,CAAC,CAAC;IACjC,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC;QAAC,IAAI,CAAC,CAAC,CAAC,GAAG,IAAI,CAAC,GAAG,CAAC,MAAM,CAAC,CAAC,CAAE,GAAG,IAAI,CAAC,CAAC;QAAC,IAAI,IAAI,IAAI,CAAC,CAAC,CAAE,CAAC;IAAC,CAAC;IAExF,MAAM,OAAO,GAAG,KAAK,CAAC,IAAI,CAAC,EAAE,MAAM,EAAE,CAAC,EAAE,EAAE,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC;SACjD,IAAI,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,IAAI,CAAC,CAAC,CAAE,GAAG,IAAI,CAAC,CAAC,CAAE,CAAC,CAAC;IAEzC,MAAM,WAAW,GAAG,OAAO,CAAC,KAAK,CAAC,CAAC,EAAE,IAAI,CAAC,CAAC;IAE3C,IAAI,MAAM,GAAG,CAAC,CAAC;IACf,MAAM,OAAO,GAAa,EAAE,CAAC;IAC7B,KAAK,MAAM,GAAG,IAAI,WAAW,EAAE,CAAC;QAC5B,MAAM,IAAI,IAAI,CAAC,GAAG,CAAE,GAAG,IAAI,CAAC;QAC5B,OAAO,CAAC,IAAI,CAAC,GAAG,CAAC,CAAC;QAClB,IAAI,MAAM,IAAI,IAAI;YAAE,MAAM;IAC9B,CAAC;IAED,IAAI,UAAU,GAAG,CAAC,CAAC;IACnB,KAAK,MAAM,GAAG,IAAI,OAAO;QAAE,UAAU,IAAI,IAAI,CAAC,GAAG,CAAE,CAAC;IACpD,MAAM,SAAS,GAAG,IAAI,CAAC,MAAM,EAAE,GAAG,UAAU,CAAC;IAC7C,IAAI,GAAG,GAAG,CAAC,CAAC;IACZ,KAAK,MAAM,GAAG,IAAI,OAAO,EAAE,CAAC;QACxB,GAAG,IAAI,IAAI,CAAC,GAAG,CAAE,CAAC;QAClB,IAAI,GAAG,IAAI,SAAS;YAAE,OAAO,GAAG,CAAC;IACrC,CAAC;IACD,OAAO,OAAO,CAAC,OAAO,CAAC,MAAM,GAAG,CAAC,CAAE,CAAC;AACxC,CAAC"}

package/dist/tokenizer/bpe.d.ts

@@ -0,0 +1,29 @@
+/**
+ * bpe.ts – Browser-side Byte Pair Encoding (BPE) tokenizer.
+ */
+export interface BPEEncodeOptions {
+    addBos?: boolean;
+    addEos?: boolean;
+}
+export type PadSide = 'right' | 'left';
+export declare class BPETokenizer {
+    vocab: Map<string, number>;
+    idToToken: Map<number, string>;
+    merges: Map<string, number>;
+    bosToken: string;
+    eosToken: string;
+    padToken: string;
+    unkToken: string;
+    bosId: number | null;
+    eosId: number | null;
+    padId: number | null;
+    constructor();
+    load(vocab: string | Record<string, number>, merges: string | string[]): Promise<void>;
+    loadFromObjects(vocabObj: Record<string, number>, mergeArr: string[]): void;
+    encode(text: string, opts?: BPEEncodeOptions): number[];
+    decode(ids: number[]): string;
+    _bpe(word: string): string[];
+    padOrTruncate(ids: number[], maxLen: number, side?: PadSide): number[];
+    get vocabSize(): number;
+}
+//# sourceMappingURL=bpe.d.ts.map

package/dist/tokenizer/bpe.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"bpe.d.ts","sourceRoot":"","sources":["../../src/tokenizer/bpe.ts"],"names":[],"mappings":"AAAA;;GAEG;AAEH,MAAM,WAAW,gBAAgB;IAC/B,MAAM,CAAC,EAAE,OAAO,CAAC;IACjB,MAAM,CAAC,EAAE,OAAO,CAAC;CAClB;AAED,MAAM,MAAM,OAAO,GAAG,OAAO,GAAG,MAAM,CAAC;AA8BvC,qBAAa,YAAY;IACrB,KAAK,EAAE,GAAG,CAAC,MAAM,EAAE,MAAM,CAAC,CAAC;IAC3B,SAAS,EAAE,GAAG,CAAC,MAAM,EAAE,MAAM,CAAC,CAAC;IAC/B,MAAM,EAAE,GAAG,CAAC,MAAM,EAAE,MAAM,CAAC,CAAC;IAC5B,QAAQ,EAAE,MAAM,CAAC;IACjB,QAAQ,EAAE,MAAM,CAAC;IACjB,QAAQ,EAAE,MAAM,CAAC;IACjB,QAAQ,EAAE,MAAM,CAAC;IACjB,KAAK,EAAE,MAAM,GAAG,IAAI,CAAC;IACrB,KAAK,EAAE,MAAM,GAAG,IAAI,CAAC;IACrB,KAAK,EAAE,MAAM,GAAG,IAAI,CAAC;;IAef,IAAI,CAAC,KAAK,EAAE,MAAM,GAAG,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,EAAE,MAAM,EAAE,MAAM,GAAG,MAAM,EAAE,GAAG,OAAO,CAAC,IAAI,CAAC;IA6B5F,eAAe,CAAC,QAAQ,EAAE,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,EAAE,QAAQ,EAAE,MAAM,EAAE,GAAG,IAAI;IAS3E,MAAM,CAAC,IAAI,EAAE,MAAM,EAAE,IAAI,GAAE,gBAAqB,GAAG,MAAM,EAAE;IA4B3D,MAAM,CAAC,GAAG,EAAE,MAAM,EAAE,GAAG,MAAM;IAgB7B,IAAI,CAAC,IAAI,EAAE,MAAM,GAAG,MAAM,EAAE;IA+B5B,aAAa,CAAC,GAAG,EAAE,MAAM,EAAE,EAAE,MAAM,EAAE,MAAM,EAAE,IAAI,GAAE,OAAiB,GAAG,MAAM,EAAE;IAO/E,IAAI,SAAS,IAAI,MAAM,CAA4B;CACtD"}

package/dist/tokenizer/bpe.js

@@ -0,0 +1,164 @@
+/**
+ * bpe.ts – Browser-side Byte Pair Encoding (BPE) tokenizer.
+ */
+function buildByteEncoder() {
+    const enc = new Map();
+    const ranges = [
+        [0x21, 0x7E],
+        [0xA1, 0xAC],
+        [0xAE, 0xFF],
+    ];
+    let n = 0;
+    for (const [lo, hi] of ranges) {
+        for (let b = lo; b <= hi; b++) {
+            enc.set(b, String.fromCodePoint(b));
+        }
+    }
+    for (let b = 0; b < 256; b++) {
+        if (!enc.has(b)) {
+            enc.set(b, String.fromCodePoint(256 + n));
+            n++;
+        }
+    }
+    return enc;
+}
+const BYTE_ENCODER = buildByteEncoder();
+const BYTE_DECODER = new Map([...BYTE_ENCODER].map(([k, v]) => [v, k]));
+const PRE_TOKENIZE_RE = /(?:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+/gu;
+export class BPETokenizer {
+    vocab;
+    idToToken;
+    merges;
+    bosToken;
+    eosToken;
+    padToken;
+    unkToken;
+    bosId;
+    eosId;
+    padId;
+    constructor() {
+        this.vocab = new Map();
+        this.idToToken = new Map();
+        this.merges = new Map();
+        this.bosToken = '<|im_start|>';
+        this.eosToken = '<|im_end|>';
+        this.padToken = '<|endoftext|>';
+        this.unkToken = '<unk>';
+        this.bosId = null;
+        this.eosId = null;
+        this.padId = null;
+    }
+    async load(vocab, merges) {
+        let vocabObj;
+        if (typeof vocab === 'string') {
+            const res = await fetch(vocab);
+            vocabObj = await res.json();
+        }
+        else {
+            vocabObj = vocab;
+        }
+        this.vocab = new Map(Object.entries(vocabObj).map(([k, v]) => [k, Number(v)]));
+        this.idToToken = new Map([...this.vocab].map(([k, v]) => [v, k]));
+        let mergeLines;
+        if (typeof merges === 'string') {
+            const res = await fetch(merges);
+            const txt = await res.text();
+            mergeLines = txt.split('\n').filter(l => l && !l.startsWith('#'));
+        }
+        else {
+            mergeLines = merges;
+        }
+        this.merges = new Map();
+        mergeLines.forEach((line, rank) => {
+            this.merges.set(line.trim(), rank);
+        });
+        this.bosId = this.vocab.get(this.bosToken) ?? null;
+        this.eosId = this.vocab.get(this.eosToken) ?? null;
+        this.padId = this.vocab.get(this.padToken) ?? null;
+    }
+    loadFromObjects(vocabObj, mergeArr) {
+        this.vocab = new Map(Object.entries(vocabObj).map(([k, v]) => [k, Number(v)]));
+        this.idToToken = new Map([...this.vocab].map(([k, v]) => [v, k]));
+        this.merges = new Map(mergeArr.map((m, i) => [m, i]));
+        this.bosId = this.vocab.get(this.bosToken) ?? null;
+        this.eosId = this.vocab.get(this.eosToken) ?? null;
+        this.padId = this.vocab.get(this.padToken) ?? null;
+    }
+    encode(text, opts = {}) {
+        const words = text.match(PRE_TOKENIZE_RE) ?? [];
+        const ids = [];
+        if (opts.addBos && this.bosId !== null)
+            ids.push(this.bosId);
+        for (const word of words) {
+            const bytes = new TextEncoder().encode(word);
+            const byteStr = Array.from(bytes).map(b => BYTE_ENCODER.get(b) ?? '?').join('');
+            const bpeTokens = this._bpe(byteStr);
+            for (const tok of bpeTokens) {
+                const id = this.vocab.get(tok);
+                if (id !== undefined) {
+                    ids.push(id);
+                }
+                else {
+                    for (const ch of tok) {
+                        const cid = this.vocab.get(ch);
+                        if (cid !== undefined)
+                            ids.push(cid);
+                    }
+                }
+            }
+        }
+        if (opts.addEos && this.eosId !== null)
+            ids.push(this.eosId);
+        return ids;
+    }
+    decode(ids) {
+        let byteStr = '';
+        for (const id of ids) {
+            const tok = this.idToToken.get(id);
+            if (tok !== undefined)
+                byteStr += tok;
+        }
+        const bytes = new Uint8Array([...byteStr].map(ch => BYTE_DECODER.get(ch) ?? ch.codePointAt(0) ?? 0));
+        try {
+            return new TextDecoder('utf-8').decode(bytes);
+        }
+        catch {
+            return byteStr;
+        }
+    }
+    _bpe(word) {
+        if (this.vocab.has(word))
+            return [word];
+        let symbols = [...word];
+        while (symbols.length > 1) {
+            let bestRank = Infinity;
+            let bestIdx = -1;
+            for (let i = 0; i < symbols.length - 1; i++) {
+                const pair = symbols[i] + ' ' + symbols[i + 1];
+                const rank = this.merges.get(pair);
+                if (rank !== undefined && rank < bestRank) {
+                    bestRank = rank;
+                    bestIdx = i;
+                }
+            }
+            if (bestIdx === -1)
+                break;
+            const merged = symbols[bestIdx] + symbols[bestIdx + 1];
+            symbols = [
+                ...symbols.slice(0, bestIdx),
+                merged,
+                ...symbols.slice(bestIdx + 2),
+            ];
+        }
+        return symbols;
+    }
+    padOrTruncate(ids, maxLen, side = 'right') {
+        if (ids.length >= maxLen)
+            return ids.slice(0, maxLen);
+        const padId = this.padId ?? 0;
+        const pad = new Array(maxLen - ids.length).fill(padId);
+        return side === 'right' ? [...ids, ...pad] : [...pad, ...ids];
+    }
+    get vocabSize() { return this.vocab.size; }
+}
+//# sourceMappingURL=bpe.js.map

package/dist/tokenizer/bpe.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"bpe.js","sourceRoot":"","sources":["../../src/tokenizer/bpe.ts"],"names":[],"mappings":"AAAA;;GAEG;AASH,SAAS,gBAAgB;IACrB,MAAM,GAAG,GAAG,IAAI,GAAG,EAAkB,CAAC;IACtC,MAAM,MAAM,GAAuB;QAC/B,CAAC,IAAI,EAAE,IAAI,CAAC;QACZ,CAAC,IAAI,EAAE,IAAI,CAAC;QACZ,CAAC,IAAI,EAAE,IAAI,CAAC;KACf,CAAC;IACF,IAAI,CAAC,GAAG,CAAC,CAAC;IACV,KAAK,MAAM,CAAC,EAAE,EAAE,EAAE,CAAC,IAAI,MAAM,EAAE,CAAC;QAC5B,KAAK,IAAI,CAAC,GAAG,EAAE,EAAE,CAAC,IAAI,EAAE,EAAE,CAAC,EAAE,EAAE,CAAC;YAC5B,GAAG,CAAC,GAAG,CAAC,CAAC,EAAE,MAAM,CAAC,aAAa,CAAC,CAAC,CAAC,CAAC,CAAC;QACxC,CAAC;IACL,CAAC;IACD,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,GAAG,EAAE,CAAC,EAAE,EAAE,CAAC;QAC3B,IAAI,CAAC,GAAG,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,CAAC;YACd,GAAG,CAAC,GAAG,CAAC,CAAC,EAAE,MAAM,CAAC,aAAa,CAAC,GAAG,GAAG,CAAC,CAAC,CAAC,CAAC;YAC1C,CAAC,EAAE,CAAC;QACR,CAAC;IACL,CAAC;IACD,OAAO,GAAG,CAAC;AACf,CAAC;AAED,MAAM,YAAY,GAAG,gBAAgB,EAAE,CAAC;AACxC,MAAM,YAAY,GAAG,IAAI,GAAG,CAAC,CAAC,GAAG,YAAY,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC,CAAC;AAExE,MAAM,eAAe,GACjB,sHAAsH,CAAC;AAE3H,MAAM,OAAO,YAAY;IACrB,KAAK,CAAsB;IAC3B,SAAS,CAAsB;IAC/B,MAAM,CAAsB;IAC5B,QAAQ,CAAS;IACjB,QAAQ,CAAS;IACjB,QAAQ,CAAS;IACjB,QAAQ,CAAS;IACjB,KAAK,CAAgB;IACrB,KAAK,CAAgB;IACrB,KAAK,CAAgB;IAErB;QACI,IAAI,CAAC,KAAK,GAAQ,IAAI,GAAG,EAAE,CAAC;QAC5B,IAAI,CAAC,SAAS,GAAI,IAAI,GAAG,EAAE,CAAC;QAC5B,IAAI,CAAC,MAAM,GAAO,IAAI,GAAG,EAAE,CAAC;QAC5B,IAAI,CAAC,QAAQ,GAAK,cAAc,CAAC;QACjC,IAAI,CAAC,QAAQ,GAAK,YAAY,CAAC;QAC/B,IAAI,CAAC,QAAQ,GAAK,eAAe,CAAC;QAClC,IAAI,CAAC,QAAQ,GAAK,OAAO,CAAC;QAC1B,IAAI,CAAC,KAAK,GAAQ,IAAI,CAAC;QACvB,IAAI,CAAC,KAAK,GAAQ,IAAI,CAAC;QACvB,IAAI,CAAC,KAAK,GAAQ,IAAI,CAAC;IAC3B,CAAC;IAED,KAAK,CAAC,IAAI,CAAC,KAAsC,EAAE,MAAyB;QACxE,IAAI,QAAgC,CAAC;QACrC,IAAI,OAAO,KAAK,KAAK,QAAQ,EAAE,CAAC;YAC5B,MAAM,GAAG,GAAG,MAAM,KAAK,CAAC,KAAK,CAAC,CAAC;YAC/B,QAAQ,GAAG,MAAM,GAAG,CAAC,IAAI,EAA4B,CAAC;QAC1D,CAAC;aAAM,CAAC;YACJ,QAAQ,GAAG,KAAK,CAAC;QACrB,CAAC;QACD,IAAI,CAAC,KAAK,GAAO,IAAI,GAAG,CAAC,MAAM,CAAC,OAAO,CAAC,QAAQ,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC,EAAE,MAAM,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;QACnF,IAAI,CAAC,SAAS,GAAG,IAAI,GAAG,CAAC,CAAC,GAAG,IAAI,CAAC,KAAK,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC,CAAC;QAElE,IAAI,UAAoB,CAAC;QACzB,IAAI,OAAO,MAAM,KAAK,QAAQ,EAAE,CAAC;YAC7B,MAAM,GAAG,GAAG,MAAM,KAAK,CAAC,MAAM,CAAC,CAAC;YAChC,MAAM,GAAG,GAAG,MAAM,GAAG,CAAC,IAAI,EAAE,CAAC;YAC7B,UAAU,GAAG,GAAG,CAAC,KAAK,CAAC,IAAI,CAAC,CAAC,MAAM,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC,UAAU,CAAC,GAAG,CAAC,CAAC,CAAC;QACtE,CAAC;aAAM,CAAC;YACJ,UAAU,GAAG,MAAM,CAAC;QACxB,CAAC;QACD,IAAI,CAAC,MAAM,GAAG,IAAI,GAAG,EAAE,CAAC;QACxB,UAAU,CAAC,OAAO,CAAC,CAAC,IAAI,EAAE,IAAI,EAAE,EAAE;YAC9B,IAAI,CAAC,MAAM,CAAC,GAAG,CAAC,IAAI,CAAC,IAAI,EAAE,EAAE,IAAI,CAAC,CAAC;QACvC,CAAC,CAAC,CAAC;QAEH,IAAI,CAAC,KAAK,GAAG,IAAI,CAAC,KAAK,CAAC,GAAG,CAAC,IAAI,CAAC,QAAQ,CAAC,IAAI,IAAI,CAAC;QACnD,IAAI,CAAC,KAAK,GAAG,IAAI,CAAC,KAAK,CAAC,GAAG,CAAC,IAAI,CAAC,QAAQ,CAAC,IAAI,IAAI,CAAC;QACnD,IAAI,CAAC,KAAK,GAAG,IAAI,CAAC,KAAK,CAAC,GAAG,CAAC,IAAI,CAAC,QAAQ,CAAC,IAAI,IAAI,CAAC;IACvD,CAAC;IAED,eAAe,CAAC,QAAgC,EAAE,QAAkB;QAChE,IAAI,CAAC,KAAK,GAAO,IAAI,GAAG,CAAC,MAAM,CAAC,OAAO,CAAC,QAAQ,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC,EAAE,MAAM,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;QACnF,IAAI,CAAC,SAAS,GAAG,IAAI,GAAG,CAAC,CAAC,GAAG,IAAI,CAAC,KAAK,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC,CAAC;QAClE,IAAI,CAAC,MAAM,GAAM,IAAI,GAAG,CAAC,QAAQ,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC,CAAC;QACzD,IAAI,CAAC,KAAK,GAAG,IAAI,CAAC,KAAK,CAAC,GAAG,CAAC,IAAI,CAAC,QAAQ,CAAC,IAAI,IAAI,CAAC;QACnD,IAAI,CAAC,KAAK,GAAG,IAAI,CAAC,KAAK,CAAC,GAAG,CAAC,IAAI,CAAC,QAAQ,CAAC,IAAI,IAAI,CAAC;QACnD,IAAI,CAAC,KAAK,GAAG,IAAI,CAAC,KAAK,CAAC,GAAG,CAAC,IAAI,CAAC,QAAQ,CAAC,IAAI,IAAI,CAAC;IACvD,CAAC;IAED,MAAM,CAAC,IAAY,EAAE,OAAyB,EAAE;QAC5C,MAAM,KAAK,GAAG,IAAI,CAAC,KAAK,CAAC,eAAe,CAAC,IAAI,EAAE,CAAC;QAChD,MAAM,GAAG,GAAe,EAAE,CAAC;QAE3B,IAAI,IAAI,CAAC,MAAM,IAAI,IAAI,CAAC,KAAK,KAAK,IAAI;YAAE,GAAG,CAAC,IAAI,CAAC,IAAI,CAAC,KAAK,CAAC,CAAC;QAE7D,KAAK,MAAM,IAAI,IAAI,KAAK,EAAE,CAAC;YACvB,MAAM,KAAK,GAAM,IAAI,WAAW,EAAE,CAAC,MAAM,CAAC,IAAI,CAAC,CAAC;YAChD,MAAM,OAAO,GAAI,KAAK,CAAC,IAAI,CAAC,KAAK,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,CAAC,YAAY,CAAC,GAAG,CAAC,CAAC,CAAC,IAAI,GAAG,CAAC,CAAC,IAAI,CAAC,EAAE,CAAC,CAAC;YACjF,MAAM,SAAS,GAAG,IAAI,CAAC,IAAI,CAAC,OAAO,CAAC,CAAC;YAErC,KAAK,MAAM,GAAG,IAAI,SAAS,EAAE,CAAC;gBAC1B,MAAM,EAAE,GAAG,IAAI,CAAC,KAAK,CAAC,GAAG,CAAC,GAAG,CAAC,CAAC;gBAC/B,IAAI,EAAE,KAAK,SAAS,EAAE,CAAC;oBACnB,GAAG,CAAC,IAAI,CAAC,EAAE,CAAC,CAAC;gBACjB,CAAC;qBAAM,CAAC;oBACJ,KAAK,MAAM,EAAE,IAAI,GAAG,EAAE,CAAC;wBACnB,MAAM,GAAG,GAAG,IAAI,CAAC,KAAK,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC;wBAC/B,IAAI,GAAG,KAAK,SAAS;4BAAE,GAAG,CAAC,IAAI,CAAC,GAAG,CAAC,CAAC;oBACzC,CAAC;gBACL,CAAC;YACL,CAAC;QACL,CAAC;QAED,IAAI,IAAI,CAAC,MAAM,IAAI,IAAI,CAAC,KAAK,KAAK,IAAI;YAAE,GAAG,CAAC,IAAI,CAAC,IAAI,CAAC,KAAK,CAAC,CAAC;QAC7D,OAAO,GAAG,CAAC;IACf,CAAC;IAED,MAAM,CAAC,GAAa;QAChB,IAAI,OAAO,GAAG,EAAE,CAAC;QACjB,KAAK,MAAM,EAAE,IAAI,GAAG,EAAE,CAAC;YACnB,MAAM,GAAG,GAAG,IAAI,CAAC,SAAS,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC;YACnC,IAAI,GAAG,KAAK,SAAS;gBAAE,OAAO,IAAI,GAAG,CAAC;QAC1C,CAAC;QACD,MAAM,KAAK,GAAG,IAAI,UAAU,CACxB,CAAC,GAAG,OAAO,CAAC,CAAC,GAAG,CAAC,EAAE,CAAC,EAAE,CAAC,YAAY,CAAC,GAAG,CAAC,EAAE,CAAC,IAAI,EAAE,CAAC,WAAW,CAAC,CAAC,CAAC,IAAI,CAAC,CAAC,CACzE,CAAC;QACF,IAAI,CAAC;YACD,OAAO,IAAI,WAAW,CAAC,OAAO,CAAC,CAAC,MAAM,CAAC,KAAK,CAAC,CAAC;QAClD,CAAC;QAAC,MAAM,CAAC;YACL,OAAO,OAAO,CAAC;QACnB,CAAC;IACL,CAAC;IAED,IAAI,CAAC,IAAY;QACb,IAAI,IAAI,CAAC,KAAK,CAAC,GAAG,CAAC,IAAI,CAAC;YAAE,OAAO,CAAC,IAAI,CAAC,CAAC;QAExC,IAAI,OAAO,GAAG,CAAC,GAAG,IAAI,CAAC,CAAC;QAExB,OAAO,OAAO,CAAC,MAAM,GAAG,CAAC,EAAE,CAAC;YACxB,IAAI,QAAQ,GAAG,QAAQ,CAAC;YACxB,IAAI,OAAO,GAAI,CAAC,CAAC,CAAC;YAElB,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,OAAO,CAAC,MAAM,GAAG,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC;gBAC1C,MAAM,IAAI,GAAG,OAAO,CAAC,CAAC,CAAC,GAAG,GAAG,GAAG,OAAO,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC;gBAC/C,MAAM,IAAI,GAAG,IAAI,CAAC,MAAM,CAAC,GAAG,CAAC,IAAI,CAAC,CAAC;gBACnC,IAAI,IAAI,KAAK,SAAS,IAAI,IAAI,GAAG,QAAQ,EAAE,CAAC;oBACxC,QAAQ,GAAG,IAAI,CAAC;oBAChB,OAAO,GAAI,CAAC,CAAC;gBACjB,CAAC;YACL,CAAC;YAED,IAAI,OAAO,KAAK,CAAC,CAAC;gBAAE,MAAM;YAE1B,MAAM,MAAM,GAAG,OAAO,CAAC,OAAO,CAAE,GAAG,OAAO,CAAC,OAAO,GAAG,CAAC,CAAE,CAAC;YACzD,OAAO,GAAG;gBACN,GAAG,OAAO,CAAC,KAAK,CAAC,CAAC,EAAE,OAAO,CAAC;gBAC5B,MAAM;gBACN,GAAG,OAAO,CAAC,KAAK,CAAC,OAAO,GAAG,CAAC,CAAC;aAChC,CAAC;QACN,CAAC;QAED,OAAO,OAAO,CAAC;IACnB,CAAC;IAED,aAAa,CAAC,GAAa,EAAE,MAAc,EAAE,OAAgB,OAAO;QAChE,IAAI,GAAG,CAAC,MAAM,IAAI,MAAM;YAAE,OAAO,GAAG,CAAC,KAAK,CAAC,CAAC,EAAE,MAAM,CAAC,CAAC;QACtD,MAAM,KAAK,GAAG,IAAI,CAAC,KAAK,IAAI,CAAC,CAAC;QAC9B,MAAM,GAAG,GAAK,IAAI,KAAK,CAAS,MAAM,GAAG,GAAG,CAAC,MAAM,CAAC,CAAC,IAAI,CAAC,KAAK,CAAC,CAAC;QACjE,OAAO,IAAI,KAAK,OAAO,CAAC,CAAC,CAAC,CAAC,GAAG,GAAG,EAAE,GAAG,GAAG,CAAC,CAAC,CAAC,CAAC,CAAC,GAAG,GAAG,EAAE,GAAG,GAAG,CAAC,CAAC;IAClE,CAAC;IAED,IAAI,SAAS,KAAa,OAAO,IAAI,CAAC,KAAK,CAAC,IAAI,CAAC,CAAC,CAAC;CACtD"}

package/dist/training/autograd.d.ts

@@ -0,0 +1,27 @@
+/**
+ * autograd.ts – Lightweight tape-based automatic differentiation engine.
+ */
+export declare class Tensor {
+    data: GPUBuffer | null;
+    shape: number[];
+    numel: number;
+    requiresGrad: boolean;
+    grad: GPUBuffer | null;
+    _gradFn: number | null;
+    constructor(data: GPUBuffer | null, shape: number[], requiresGrad?: boolean);
+    get byteSize(): number;
+    zeroGrad(device: GPUDevice): void;
+    destroy(): void;
+}
+export declare function enableGrad(): void;
+export declare function noGrad(): void;
+export declare function clearTape(): void;
+export declare function recordOperation(backwardFn: () => void | Promise<void>): number;
+export declare function backward(): Promise<void>;
+export declare function ensureGradBuffer(device: GPUDevice, tensor: Tensor): void;
+export declare function allocateGradients(device: GPUDevice, tensors: Tensor[]): void;
+export declare function zeroGradients(device: GPUDevice, tensors: Tensor[]): void;
+export declare function onesLikeScalar(device: GPUDevice): GPUBuffer;
+export declare function crossEntropyLoss(logits: Float32Array, targetId: number): number;
+export declare function crossEntropyGrad(logits: Float32Array, targetId: number): Float32Array;
+//# sourceMappingURL=autograd.d.ts.map

package/dist/training/autograd.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"autograd.d.ts","sourceRoot":"","sources":["../../src/training/autograd.ts"],"names":[],"mappings":"AAAA;;GAEG;AAYH,qBAAa,MAAM;IACf,IAAI,EAAE,SAAS,GAAG,IAAI,CAAC;IACvB,KAAK,EAAE,MAAM,EAAE,CAAC;IAChB,KAAK,EAAE,MAAM,CAAC;IACd,YAAY,EAAE,OAAO,CAAC;IACtB,IAAI,EAAE,SAAS,GAAG,IAAI,CAAC;IACvB,OAAO,EAAE,MAAM,GAAG,IAAI,CAAC;gBAEX,IAAI,EAAE,SAAS,GAAG,IAAI,EAAE,KAAK,EAAE,MAAM,EAAE,EAAE,YAAY,UAAQ;IASzE,IAAI,QAAQ,IAAI,MAAM,CAA2B;IAEjD,QAAQ,CAAC,MAAM,EAAE,SAAS,GAAG,IAAI;IAMjC,OAAO,IAAI,IAAI;CAMlB;AAED,wBAAgB,UAAU,IAAI,IAAI,CAA2B;AAC7D,wBAAgB,MAAM,IAAI,IAAI,CAA+B;AAC7D,wBAAgB,SAAS,IAAI,IAAI,CAAkB;AAEnD,wBAAgB,eAAe,CAAC,UAAU,EAAE,MAAM,IAAI,GAAG,OAAO,CAAC,IAAI,CAAC,GAAG,MAAM,CAI9E;AAED,wBAAsB,QAAQ,IAAI,OAAO,CAAC,IAAI,CAAC,CAK9C;AAED,wBAAgB,gBAAgB,CAAC,MAAM,EAAE,SAAS,EAAE,MAAM,EAAE,MAAM,GAAG,IAAI,CAWxE;AAED,wBAAgB,iBAAiB,CAAC,MAAM,EAAE,SAAS,EAAE,OAAO,EAAE,MAAM,EAAE,GAAG,IAAI,CAI5E;AAED,wBAAgB,aAAa,CAAC,MAAM,EAAE,SAAS,EAAE,OAAO,EAAE,MAAM,EAAE,GAAG,IAAI,CAMxE;AAED,wBAAgB,cAAc,CAAC,MAAM,EAAE,SAAS,GAAG,SAAS,CAW3D;AAED,wBAAgB,gBAAgB,CAAC,MAAM,EAAE,YAAY,EAAE,QAAQ,EAAE,MAAM,GAAG,MAAM,CAW/E;AAED,wBAAgB,gBAAgB,CAAC,MAAM,EAAE,YAAY,EAAE,QAAQ,EAAE,MAAM,GAAG,YAAY,CAiBrF"}