mambacode.js 1.0.0 → 1.0.1

package/src/model/{mamba_model.js → mamba_model.ts}

@@ -1,15 +1,8 @@
 /**
- * mamba_model.js – Full Mamba language model.
- *
- * Architecture (matches Qwen3.5-Coder-0.8B-style Mamba):
- *
- *   Token IDs ──► Embedding ──► [MambaBlock × numLayers] ──► RMSNorm ──► LM Head
- *
- * The LM Head is a linear projection from dModel → vocabSize.
- * All computations run on WebGPU via the kernels in src/kernels/.
+ * mamba_model.ts – Full Mamba language model.
  */
 
-import { MambaBlock } from './mamba_block.js';
+import { MambaBlock, BlockCache, BlockParam } from './mamba_block';
 import {
     createStorageBuffer,
     createEmptyStorageBuffer,
@@ -18,40 +11,60 @@ import {
     createBindGroup,
     dispatchKernel,
     readBuffer,
+    uploadBuffer,
     cdiv,
-} from '../utils/gpu_utils.js';
-import { LINEAR_FORWARD_WGSL } from '../kernels/linear_projection.js';
-import { ACTIVATIONS_WGSL }    from '../kernels/activations.js';
+} from '../utils/gpu_utils';
+import { LINEAR_FORWARD_WGSL } from '../kernels/linear_projection';
+import { ACTIVATIONS_WGSL }    from '../kernels/activations';
+
+export interface MambaModelConfig {
+  vocabSize: number;
+  dModel: number;
+  numLayers: number;
+  dState?: number;
+  dConv?: number;
+  expand?: number;
+  eosId?: number;
+}
 
-/**
- * @typedef {Object} MambaModelConfig
- * @property {number} vocabSize   – vocabulary size (Qwen3.5-Coder: 151936)
- * @property {number} dModel      – model (embedding) dimension
- * @property {number} numLayers   – number of Mamba blocks
- * @property {number} [dState]    – SSM state dimension (default 16)
- * @property {number} [dConv]     – conv kernel size (default 4)
- * @property {number} [expand]    – inner-dim expansion factor (default 2)
- */
+export interface ModelForwardResult {
+  logits: Float32Array;
+  gpuLogits: GPUBuffer;
+  caches: BlockCache[];
+}
+
+export interface SamplingOptions {
+  temperature?: number;
+  topK?: number;
+  topP?: number;
+}
 
 export class MambaModel {
-    /**
-     * @param {GPUDevice}       device
-     * @param {MambaModelConfig} config
-     */
-    constructor(device, config) {
+    device: GPUDevice;
+    config: Required<MambaModelConfig>;
+    gpuEmbedding: GPUBuffer;
+    blocks: MambaBlock[];
+    gpuFinalNorm: GPUBuffer;
+    tiedEmbedding: boolean;
+    gpuLMHeadBias: GPUBuffer;
+    private _lmHeadPipeline: GPUComputePipeline;
+    private _rmsnormPipeline: GPUComputePipeline;
+    private _embedPipeline: GPUComputePipeline;
+    private _wslaMode = false;
+
+    constructor(device: GPUDevice, config: MambaModelConfig) {
         this.device = device;
         this.config = {
             dState    : 16,
             dConv     : 4,
             expand    : 2,
+            eosId     : -1,
             ...config,
-        };
+        } as Required<MambaModelConfig>;
 
         const { vocabSize, dModel, numLayers } = this.config;
 
-        // Token embedding table: (vocabSize, dModel)
         const embedData = new Float32Array(vocabSize * dModel);
-        // Xavier-style initialisation
         const std = 1.0 / Math.sqrt(dModel);
         for (let i = 0; i < embedData.length; i++) {
             const u1 = Math.random(), u2 = Math.random();
@@ -60,7 +73,6 @@ export class MambaModel {
         }
         this.gpuEmbedding = createStorageBuffer(device, embedData, true);
 
-        // Stacked Mamba blocks
         this.blocks = Array.from({ length: numLayers }, () =>
             new MambaBlock(device, {
                 dModel,
@@ -70,36 +82,20 @@ export class MambaModel {
             })
         );
 
-        // Final RMSNorm
         const finalNormW = new Float32Array(dModel).fill(1.0);
         this.gpuFinalNorm = createStorageBuffer(device, finalNormW, true);
 
-        // LM Head: (vocabSize, dModel) – tied to embedding by default
-        // We share the embedding weight (weight tying saves memory).
         this.tiedEmbedding = true;
 
-        // Compile pipelines
         this._lmHeadPipeline  = createComputePipeline(device, LINEAR_FORWARD_WGSL, 'linear_forward');
         this._rmsnormPipeline = createComputePipeline(device, ACTIVATIONS_WGSL,    'rmsnorm_forward');
 
-        // LM Head bias (zeroed)
         this.gpuLMHeadBias = createStorageBuffer(device, new Float32Array(vocabSize), true);
 
-        // Embedding lookup pipeline (gather rows)
         this._embedPipeline = createComputePipeline(device, EMBED_LOOKUP_WGSL, 'embed_lookup');
     }
 
-    // ─── Embedding lookup ─────────────────────────────────────────────────────
-
-    /**
-     * Look up token embeddings.
-     *
-     * @param {Int32Array|Uint32Array} tokenIds  – (batch * seqLen,)
-     * @param {number} batch
-     * @param {number} seqLen
-     * @returns {GPUBuffer}  – (batch * seqLen, dModel)
-     */
-    embedTokens(tokenIds, batch, seqLen) {
+    embedTokens(tokenIds: number[] | Uint32Array, batch: number, seqLen: number): GPUBuffer {
         const { dModel } = this.config;
         const M = batch * seqLen;
 
@@ -119,27 +115,13 @@ export class MambaModel {
         return outBuf;
     }
 
-    // ─── Forward pass ─────────────────────────────────────────────────────────
-
-    /**
-     * Full model forward pass.
-     *
-     * @param {number[]|Uint32Array} tokenIds  – (batch * seqLen,) flat
-     * @param {number}  batch
-     * @param {number}  seqLen
-     * @returns {Promise<{ logits: Float32Array, gpuLogits: GPUBuffer }>}
-     *   logits   – CPU Float32Array of shape (batch * seqLen, vocabSize)
-     *   gpuLogits – GPU buffer (same data, for chained backward)
-     */
-    async forward(tokenIds, batch, seqLen) {
+    async forward(tokenIds: number[] | Uint32Array, batch: number, seqLen: number): Promise<ModelForwardResult> {
         const { dModel, vocabSize } = this.config;
         const M = batch * seqLen;
 
-        // 1. Token embedding lookup
         let hidden = this.embedTokens(tokenIds, batch, seqLen);
 
-        // 2. Mamba blocks
-        const caches = [];
+        const caches: BlockCache[] = [];
         for (const block of this.blocks) {
             const { output, cache } = block.forward(hidden, batch, seqLen);
             caches.push(cache);
@@ -147,7 +129,6 @@ export class MambaModel {
             hidden = output;
         }
 
-        // 3. Final RMSNorm
         const normOut = createEmptyStorageBuffer(this.device, M * dModel * 4, true);
         const normInv = createEmptyStorageBuffer(this.device, M * 4,          false);
         {
@@ -160,12 +141,11 @@ export class MambaModel {
             dispatchKernel(this.device, this._rmsnormPipeline, bg, [cdiv(M, 64), 1, 1]);
         }
 
-        // 4. LM Head: (M, vocabSize) = normOut @ embedding^T + bias
         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.gpuLMHeadWeight;
+            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,
@@ -175,66 +155,47 @@ export class MambaModel {
         normOut.destroy();
         normInv.destroy();
 
-        // 5. Read back logits to CPU
         const logits = await readBuffer(this.device, gpuLogits, M * vocabSize * 4);
 
         return { logits, gpuLogits, caches };
     }
 
-    /**
-     * Greedy / top-k / temperature-sampled autoregressive generation.
-     *
-     * @param {number[]} promptIds  – starting token IDs
-     * @param {number}   maxNewTokens
-     * @param {{ temperature?: number, topK?: number, topP?: number }} [samplingOpts]
-     * @returns {Promise<number[]>}  – full sequence (prompt + generated)
-     */
-    async generate(promptIds, maxNewTokens = 200, samplingOpts = {}) {
+    async generate(promptIds: number[], maxNewTokens = 200, samplingOpts: SamplingOptions = {}): Promise<number[]> {
         const { temperature = 1.0, topK = 50, topP = 0.9 } = samplingOpts;
         const { vocabSize } = this.config;
 
         let ids = [...promptIds];
 
         for (let step = 0; step < maxNewTokens; step++) {
-            // Use the full context each step (linear cost with Mamba – no kv-cache needed)
             const { logits } = await this.forward(
                 new Uint32Array(ids), 1, ids.length
             );
-            // Get logits for the last position
             const lastLogits = logits.slice((ids.length - 1) * vocabSize, ids.length * vocabSize);
 
             const nextId = sampleToken(lastLogits, { temperature, topK, topP });
             ids.push(nextId);
 
-            // Stop on EOS
             if (nextId === this.config.eosId) break;
         }
 
         return ids;
     }
 
-    /**
-     * Collect all trainable parameters across all blocks.
-     * @returns {Array<{buf: GPUBuffer, numel: number, name: string}>}
-     */
-    parameters() {
-        const params = [];
+    parameters(): BlockParam[] {
+        const params: BlockParam[] = [];
 
-        // Embedding
         params.push({
             buf  : this.gpuEmbedding,
             numel: this.config.vocabSize * this.config.dModel,
             name : 'embedding',
         });
 
-        // Blocks
         for (let i = 0; i < this.blocks.length; i++) {
-            for (const p of this.blocks[i].parameters()) {
+            for (const p of this.blocks[i]!.parameters()) {
                 params.push({ ...p, name: `block${i}.${p.name}` });
             }
         }
 
-        // Final norm
         params.push({
             buf  : this.gpuFinalNorm,
             numel: this.config.dModel,
@@ -244,19 +205,117 @@ export class MambaModel {
         return params;
     }
 
-    /**
-     * Enable WSLA (selective fine-tuning of B and C only) across all blocks.
-     * @param {boolean} enabled
-     */
-    setWSLAMode(enabled) {
+    setWSLAMode(enabled: boolean): void {
         for (const block of this.blocks) block.setWSLAMode(enabled);
         this._wslaMode = enabled;
     }
-}
 
-// ─── Embedding lookup WGSL kernel ────────────────────────────────────────────
+    /**
+     * 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(): Promise<ArrayBuffer> {
+        const params = this.parameters();
+        const nParams = params.length;
+
+        // Read all GPU buffers into CPU Float32Arrays
+        const arrays: Float32Array[] = 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);
 
-const EMBED_LOOKUP_WGSL = /* wgsl */`
+        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: ArrayBuffer): Promise<void> {
+        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: number[] = [];
+        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: string = /* wgsl */`
 struct EmbedParams {
     num_tokens : u32,
     d_model    : u32,
@@ -264,8 +323,8 @@ struct EmbedParams {
 
 @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>;  // (V, D)
-@group(0) @binding(3) var<storage, read_write> out    : array<f32>;  // (T, D)
+@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>) {
@@ -283,53 +342,38 @@ fn embed_lookup(@builtin(global_invocation_id) gid: vec3<u32>) {
 }
 `;
 
-// ─── Sampling helper ──────────────────────────────────────────────────────────
-
-/**
- * Sample a token from logits using temperature + top-k + nucleus (top-p).
- *
- * @param {Float32Array} logits
- * @param {{ temperature?: number, topK?: number, topP?: number }} opts
- * @returns {number}
- */
-function sampleToken(logits, { temperature = 1.0, topK = 50, topP = 0.9 } = {}) {
+function sampleToken(logits: Float32Array, { temperature = 1.0, topK = 50, topP = 0.9 } = {}): number {
     const n = logits.length;
 
-    // Apply temperature
     const scaled = new Float32Array(n);
-    for (let i = 0; i < n; i++) scaled[i] = logits[i] / Math.max(temperature, 1e-7);
+    for (let i = 0; i < n; i++) scaled[i] = logits[i]! / Math.max(temperature, 1e-7);
 
-    // Softmax
     let maxL = -Infinity;
-    for (let i = 0; i < n; i++) if (scaled[i] > maxL) maxL = scaled[i];
+    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]; }
+    for (let i = 0; i < n; i++) { exps[i] = Math.exp(scaled[i]! - maxL); sumE += exps[i]!; }
 
-    // Sort indices by probability (descending)
     const indices = Array.from({ length: n }, (_, i) => i)
-        .sort((a, b) => exps[b] - exps[a]);
+        .sort((a, b) => exps[b]! - exps[a]!);
 
-    // Top-K filter
     const topKIndices = indices.slice(0, topK);
 
-    // Nucleus (top-p) filter
     let cumSum = 0;
-    const nucleus = [];
+    const nucleus: number[] = [];
     for (const idx of topKIndices) {
-        cumSum += exps[idx] / sumE;
+        cumSum += exps[idx]! / sumE;
         nucleus.push(idx);
         if (cumSum >= topP) break;
     }
 
-    // Sample from nucleus
     let nucleusSum = 0;
-    for (const idx of nucleus) nucleusSum += exps[idx];
+    for (const idx of nucleus) nucleusSum += exps[idx]!;
     const threshold = Math.random() * nucleusSum;
     let acc = 0;
     for (const idx of nucleus) {
-        acc += exps[idx];
+        acc += exps[idx]!;
         if (acc >= threshold) return idx;
     }
-    return nucleus[nucleus.length - 1];
+    return nucleus[nucleus.length - 1]!;
 }

package/src/tokenizer/bpe.ts

@@ -0,0 +1,186 @@
+/**
+ * bpe.ts – Browser-side Byte Pair Encoding (BPE) tokenizer.
+ */
+
+export interface BPEEncodeOptions {
+  addBos?: boolean;
+  addEos?: boolean;
+}
+
+export type PadSide = 'right' | 'left';
+
+function buildByteEncoder(): Map<number, string> {
+    const enc = new Map<number, string>();
+    const ranges: [number, number][] = [
+        [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: 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() {
+        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: string | Record<string, number>, merges: string | string[]): Promise<void> {
+        let vocabObj: Record<string, number>;
+        if (typeof vocab === 'string') {
+            const res = await fetch(vocab);
+            vocabObj = await res.json() as Record<string, number>;
+        } 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: string[];
+        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: Record<string, number>, mergeArr: string[]): void {
+        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: string, opts: BPEEncodeOptions = {}): number[] {
+        const words = text.match(PRE_TOKENIZE_RE) ?? [];
+        const ids: number[]   = [];
+
+        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: number[]): string {
+        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: string): string[] {
+        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: number[], maxLen: number, side: PadSide = 'right'): number[] {
+        if (ids.length >= maxLen) return ids.slice(0, maxLen);
+        const padId = this.padId ?? 0;
+        const pad   = new Array<number>(maxLen - ids.length).fill(padId);
+        return side === 'right' ? [...ids, ...pad] : [...pad, ...ids];
+    }
+
+    get vocabSize(): number { return this.vocab.size; }
+}