engine/spm_tokenizer.js

// Self-written tokenizer for our Marian SPM-BPE exports — replaces
// transformers.js on the app path. Loads the same HF tokenizer.json and
// reproduces @huggingface/tokenizers@0.1.3 byte-for-byte for the pipeline our
// two models use (parity enforced by test/tokenizer_parity.test.js):
//
//   normalize:  Precompiled → Strip(right) → Replace(/ {2,}/ → ▁)
//   pre-token:  Metaspace(▁, prepend always)
//   model:      BPE (byte_fallback, no fuse on encode)
//   post:       TemplateProcessing (append </s>)
//   decode:     Metaspace join, skip-special filter, no cleanup
//
// The Precompiled stage is deliberately the transformers.js APPROXIMATION of
// SentencePiece's charsmap (control-char strip + whitespace folding + NFKC
// with the ～ split quirk), not the real charsmap trie: the app has always
// tokenized through that approximation, and the golden outputs pin it.
//
// The constructor validates that tokenizer.json actually declares this exact
// pipeline and throws otherwise — a future re-export with a different shape
// must fail loudly, never mis-tokenize silently.

const need = (cond, what) => {
  if (!cond) throw new Error(`spm_tokenizer: unsupported tokenizer.json: ${what}`);
};

// Encode LRU bounds: app sentences are ≤200 chars (MAX_SENT_CHARS), so 512
// covers every real caller; 4096 entries ≈ a few MB worst case.
const CACHE_CAP = 4096;
const CACHE_MAX_CHARS = 512;

export class SpmTokenizer {
  // tokenizerJson: parsed tokenizer.json; config: parsed tokenizer_config.json
  // (only clean_up_tokenization_spaces is read from it).
  constructor(tokenizerJson, config = {}) {
    const { model, normalizer, pre_tokenizer, post_processor, decoder, added_tokens } = tokenizerJson;

    need(model?.type === 'BPE' && !Array.isArray(model.vocab), 'model must be BPE with object vocab');
    need(Array.isArray(model.merges?.[0]), 'merges must use the pair-array format');
    need(!model.ignore_merges && !model.end_of_word_suffix && !model.continuing_subword_suffix,
      'BPE suffix/ignore_merges options');
    need(normalizer?.type === 'Sequence'
      && normalizer.normalizers.map((n) => n.type).join(',') === 'Precompiled,Strip,Replace'
      && normalizer.normalizers[1].strip_left === false && normalizer.normalizers[1].strip_right === true
      && normalizer.normalizers[2].pattern?.Regex === ' {2,}',
      'normalizer must be [Precompiled, Strip(right), Replace(/ {2,}/)]');
    need(pre_tokenizer?.type === 'Metaspace' && pre_tokenizer.prepend_scheme === 'always',
      'pre_tokenizer must be Metaspace(always)');
    need(decoder?.type === 'Metaspace', 'decoder must be Metaspace');
    const single = post_processor?.single;
    need(post_processor?.type === 'TemplateProcessing' && single?.length === 2
      && single[0].Sequence?.id === 'A' && typeof single[1].SpecialToken?.id === 'string',
      'post_processor must be TemplateProcessing [A, eos]');
    need((added_tokens ?? []).every((t) => !t.normalized && !t.lstrip && !t.rstrip && !t.single_word),
      'added tokens must be unnormalized, no lstrip/rstrip/single_word');

    this.replacement = pre_tokenizer.replacement ?? '▁';
    this.spaceRun = normalizer.normalizers[2].content ?? '';
    this.eosToken = single[1].SpecialToken.id;
    this.cleanUp = config.clean_up_tokenization_spaces ?? true;

    this.vocab = new Map(Object.entries(model.vocab)); // token -> id
    this.byteFallback = model.byte_fallback ?? false;
    this.unkToken = model.unk_token ?? null;

    // rank of each merge pair, nested left -> right -> rank
    this.merges = new Map();
    model.merges.forEach(([a, b], rank) => {
      let m = this.merges.get(a);
      if (!m) this.merges.set(a, (m = new Map()));
      if (!m.has(b)) m.set(b, rank);
    });

    // Added tokens are matched as literal substrings before normalization and
    // override the vocab id for their content.
    this.addedContents = [];
    this.specialTokens = new Set();
    for (const t of added_tokens ?? []) {
      this.addedContents.push(t.content);
      this.vocab.set(t.content, t.id);
      if (t.special) this.specialTokens.add(t.content);
    }

    this.idToToken = [];
    for (const [tok, id] of this.vocab) this.idToToken[id] = tok;
    this.unkId = this.vocab.get(this.unkToken);
    this.textEncoder = new TextEncoder();
    this.cache = new Map(); // encode LRU (insertion order = recency)
  }

  // Precompiled (transformers.js approximation) → Strip(right) → Replace.
  normalize(text) {
    text = text.replace(/[\u0001-\u0008\u000B\u000E-\u001F\u007F\u008F\u009F]/g, '');
    text = text.replace(/[\u0009\u000A\u000C\u000D\u00A0\u1680\u2000-\u200F\u2028\u2029\u202F\u205F\u2581\u3000\uFEFF\uFFFD]/g, ' ');
    // NFKC would fold ～ (U+FF5E) into ~; SentencePiece keeps it, so normalize
    // around it.
    text = text.includes('～')
      ? text.split('～').map((p) => p.normalize('NFKC')).join('～')
      : text.normalize('NFKC');
    return text.trimEnd().replace(/ {2,}/gu, this.spaceRun);
  }

  // Greedy longest-match split on added-token contents; non-matching spans
  // stay as plain-text sections.
  splitOnAdded(text) {
    const sections = [];
    let start = 0;
    for (let i = 0; i < text.length; ) {
      let match = null;
      for (const c of this.addedContents) {
        if (text.startsWith(c, i) && (!match || c.length > match.length)) match = c;
      }
      if (!match) { i++; continue; }
      if (i > start) sections.push(text.slice(start, i));
      sections.push(match);
      i += match.length;
      start = i;
    }
    if (start < text.length) sections.push(text.slice(start));
    return sections;
  }

  // Merge the lowest-ranked adjacent pair (leftmost on rank ties) until no
  // pair is mergeable. `word` is an array of code points.
  bpe(word) {
    for (;;) {
      let at = -1;
      let best = Infinity;
      for (let i = 0; i + 1 < word.length; i++) {
        const rank = this.merges.get(word[i])?.get(word[i + 1]);
        if (rank !== undefined && rank < best) { best = rank; at = i; }
      }
      if (at < 0) return word;
      word.splice(at, 2, word[at] + word[at + 1]);
    }
  }

  // BPE pieces -> token strings, with <0xXX> byte fallback for OOV pieces.
  piecesToTokens(pieces) {
    const out = [];
    for (const p of pieces) {
      if (this.vocab.has(p)) {
        out.push(p);
      } else if (this.byteFallback) {
        const bytes = Array.from(this.textEncoder.encode(p),
          (x) => `<0x${x.toString(16).toUpperCase().padStart(2, '0')}>`);
        if (bytes.every((b) => this.vocab.has(b))) out.push(...bytes);
        else if (this.unkToken != null) out.push(this.unkToken);
      } else if (this.unkToken != null) {
        out.push(this.unkToken);
      }
    }
    return out;
  }

  // text -> token ids, appending eos (TemplateProcessing) by default.
  // Sentence-level LRU: novels repeat short lines (dialogue beats, sound
  // words), and the app encodes per sentence — cache whole results. Misses
  // pay one Map lookup + one array copy, ~nothing next to the BPE scan.
  encode(text, { addSpecialTokens = true } = {}) {
    const key = (addSpecialTokens ? 'S' : 'R') + text;
    const cacheable = text.length <= CACHE_MAX_CHARS;
    if (cacheable) {
      const hit = this.cache.get(key);
      if (hit !== undefined) {
        this.cache.delete(key); // refresh recency
        this.cache.set(key, hit);
        return hit.slice(); // callers own their copy — the cache stays clean
      }
    }
    const ids = this.encodeUncached(text, addSpecialTokens);
    if (cacheable) {
      this.cache.set(key, ids.slice());
      if (this.cache.size > CACHE_CAP) this.cache.delete(this.cache.keys().next().value);
    }
    return ids;
  }

  encodeUncached(text, addSpecialTokens) {
    const tokens = [];
    for (const section of this.splitOnAdded(text)) {
      if (this.addedContents.includes(section)) { tokens.push(section); continue; }
      const norm = this.normalize(section);
      if (norm.length === 0) continue;
      let pre = norm.replaceAll(' ', this.replacement);
      if (!pre.startsWith(this.replacement)) pre = this.replacement + pre;
      tokens.push(...this.piecesToTokens(this.bpe(Array.from(pre))));
    }
    if (addSpecialTokens) tokens.push(this.eosToken);
    return tokens.map((t) => this.vocab.get(t) ?? this.unkId);
  }

  // ids -> text. Option name matches the transformers.js call sites.
  decode(ids, { skip_special_tokens = false } = {}) {
    let tokens = Array.from(ids, (i) => this.idToToken[Number(i)] ?? this.unkToken);
    if (skip_special_tokens) tokens = tokens.filter((t) => !this.specialTokens.has(t));
    let text = '';
    for (let i = 0; i < tokens.length; i++) {
      let s = tokens[i].replaceAll(this.replacement, ' ');
      if (i === 0 && s.startsWith(' ')) s = s.slice(1);
      text += s;
    }
    if (this.cleanUp) {
      text = text.replace(/ \./g, '.').replace(/ \?/g, '?').replace(/ !/g, '!')
        .replace(/ ,/g, ',').replace(/ ' /g, "'").replace(/ n't/g, "n't")
        .replace(/ 'm/g, "'m").replace(/ 's/g, "'s").replace(/ 've/g, "'ve")
        .replace(/ 're/g, "'re");
    }
    return text;
  }
}