engine/tokenizer.js

// Tokenizer loading + batch encoding, on our self-written SpmTokenizer
// (src/engine/spm_tokenizer.js) — byte-exact with the transformers.js
// tokenizer it replaced (test/tokenizer_parity.test.js).

import { SpmTokenizer } from './spm_tokenizer.js';
import { bucketFor } from './shapes.js';
import { SRC_CAP, PAD, EOS } from './constants.js';

// /model/* is served immutable (1-year HTTP cache) but the file NAMES are
// fixed, unlike the content-addressed weight shards — so when the asset
// changes (tokenizer audit v4.0.1) every browser that ever cached it keeps
// the old bytes forever. ?v=<tokRev> (from the model registry) rolls the URL
// with the content; no rev (bench/debug dev paths) keeps the bare URL.
export const tokenizerAssetUrl = (modelId, file, tokRev = null) =>
  `/model/${modelId}/${file}${tokRev ? `?v=${encodeURIComponent(tokRev)}` : ''}`;

// modelId is the directory under /model/ holding the HF tokenizer assets
// (tokenizer.json + tokenizer_config.json). Default keeps every existing
// caller (bench, debug suite) on MoxhiMT-30.
export async function loadTokenizer(modelId = 'moxhi', tokRev = null) {
  const fetchJson = async (file) => {
    const res = await fetch(tokenizerAssetUrl(modelId, file, tokRev));
    if (!res.ok) throw new Error(`fetch ${modelId}/${file}: HTTP ${res.status}`);
    return res.json();
  };
  const [tokenizerJson, config] = await Promise.all([
    fetchJson('tokenizer.json'), fetchJson('tokenizer_config.json'),
  ]);
  return new SpmTokenizer(tokenizerJson, config);
}

// Tokenize a batch of source texts into a right-padded id matrix.
//
// Each text is tokenized individually (raw ids incl. eos, no padding); rows
// longer than SRC_CAP are truncated (keeping a trailing eos). `truncated`
// lists their batch indices — the caller decides what a cut source means
// (the app blanks the row, skips its TM write and split-retries it; silent
// truncation is exactly the bug this field exists to prevent). S is the
// bucketed max row length.
//
// Returns {ids: Uint32Array [B·S] PAD-right-padded, lens: Uint32Array [B],
//          B, S, truncated: number[]}.
export async function tokenizeBatch(tok, texts) {
  const B = texts.length;
  const rows = [];
  const truncated = [];
  for (const text of texts) {
    let row = tok.encode(text);
    if (row.length > SRC_CAP) {
      console.warn(`tokenizeBatch: truncating row of ${row.length} tokens to SRC_CAP=${SRC_CAP}`);
      truncated.push(rows.length);
      row = row.slice(0, SRC_CAP);
      row[SRC_CAP - 1] = EOS;
    }
    rows.push(row);
  }
  const lens = Uint32Array.from(rows, (r) => r.length);
  const S = bucketFor(Math.max(...lens));
  const ids = new Uint32Array(B * S).fill(PAD);
  rows.forEach((row, b) => ids.set(row, b * S));
  return { ids, lens, B, S, truncated };
}