public/tokenizer.js

// Byte-level BPE, read straight from a repo's tokenizer.json.
//
// The three families this build supports — GPT-2, Llama-style (SmolLM,
// TinyLlama), and Qwen2 — all use byte-level BPE, so one implementation covers
// them. A repo using anything else (SentencePiece/Unigram, WordPiece) is
// refused by name rather than approximated: a tokenizer that is merely close
// produces text that looks right and is subtly not, which is the worst
// possible failure mode to ship quietly.
(function (root) {
  "use strict";

  // GPT-2's byte<->unicode table. Byte-level BPE maps raw bytes into printable
  // codepoints so the merge table can be plain text; every byte round-trips.
  function byteTables() {
    const bs = [];
    for (let i = 33; i <= 126; i++) bs.push(i);
    for (let i = 161; i <= 172; i++) bs.push(i);
    for (let i = 174; i <= 255; i++) bs.push(i);
    const cs = bs.slice();
    let n = 0;
    for (let b = 0; b < 256; b++)
      if (!bs.includes(b)) { bs.push(b); cs.push(256 + n); n++; }
    const b2u = new Map(), u2b = new Map();
    for (let i = 0; i < bs.length; i++) {
      const ch = String.fromCodePoint(cs[i]);
      b2u.set(bs[i], ch); u2b.set(ch, bs[i]);
    }
    return { b2u, u2b };
  }
  const { b2u, u2b } = byteTables();

  // GPT-2 / Llama pretokenizer split. Keeps leading spaces attached to the
  // following word, which is what the merge table was built against.
  const SPLIT = /'s|'t|'re|'ve|'m|'ll|'d| ?[\p{L}]+| ?[\p{N}]+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+/gu;

  function build(json) {
    const model = json.model || {};
    if (String(model.type || "").toUpperCase() !== "BPE")
      throw new Error(`tokenizer type "${model.type}" is not supported — this build implements byte-level BPE ` +
                      `(GPT-2, Llama-style and Qwen2 repos). A different tokenizer would decode to subtly wrong text.`);
    const vocab = model.vocab || {};
    const ids = [];
    for (const [tokStr, id] of Object.entries(vocab)) ids[id] = tokStr;

    // merges may be "a b" strings or ["a","b"] pairs depending on the version
    const ranks = new Map();
    (model.merges || []).forEach((m, i) => {
      const pair = Array.isArray(m) ? m : String(m).split(" ");
      if (pair.length === 2) ranks.set(pair[0] + "" + pair[1], i);
    });

    // Added tokens live OUTSIDE model.vocab, so they have to be written into
    // the id table too — otherwise an id like <|endoftext|> decodes to nothing
    // and, worse, an id past the end of the vocab silently disappears instead
    // of being reported.
    const added = new Map();
    for (const t of json.added_tokens || []) { added.set(t.content, t.id); ids[t.id] = t.content; }
    const specialIds = new Set([...added.values()]);

    return { vocab, ids, ranks, added, specialIds, size: ids.length,
             name: `byte-level BPE (${ids.length} tokens)` };
  }

  function bpe(tok, word) {
    if (word.length === 1) return [word];
    let parts = [...word];
    for (;;) {
      let best = null, bestRank = Infinity, bestIdx = -1;
      for (let i = 0; i < parts.length - 1; i++) {
        const r = tok.ranks.get(parts[i] + "" + parts[i + 1]);
        if (r !== undefined && r < bestRank) { bestRank = r; best = i; bestIdx = i; }
      }
      if (best === null) break;
      parts = parts.slice(0, bestIdx).concat(parts[bestIdx] + parts[bestIdx + 1], parts.slice(bestIdx + 2));
      if (parts.length === 1) break;
    }
    return parts;
  }

  function encode(tok, text) {
    const out = [];
    // added/special tokens are matched literally before anything else, so a
    // chat template's <|im_start|> stays one token instead of being merged
    const specials = [...tok.added.keys()].sort((a, b) => b.length - a.length);
    let rest = text;
    const pieces = [];
    while (rest.length) {
      let hit = null, at = -1;
      for (const s of specials) {
        const i = rest.indexOf(s);
        if (i !== -1 && (at === -1 || i < at)) { at = i; hit = s; }
      }
      if (hit === null) { pieces.push({ text: rest }); break; }
      if (at > 0) pieces.push({ text: rest.slice(0, at) });
      pieces.push({ special: hit });
      rest = rest.slice(at + hit.length);
    }
    for (const p of pieces) {
      if (p.special) { out.push(tok.added.get(p.special)); continue; }
      for (const m of p.text.matchAll(SPLIT)) {
        let mapped = "";
        for (const b of new TextEncoder().encode(m[0])) mapped += b2u.get(b);
        for (const piece of bpe(tok, mapped)) {
          const id = tok.vocab[piece];
          if (id !== undefined) out.push(id);
          else for (const ch of piece) {         // fall back to single bytes
            const bid = tok.vocab[ch];
            if (bid !== undefined) out.push(bid);
          }
        }
      }
    }
    return Int32Array.from(out);
  }

  function decode(tok, idArr, keepSpecial) {
    const bytes = [];
    for (const id of idArr) {
      const t = tok.ids[id];
      if (t === undefined) continue;
      if (tok.specialIds.has(id) && !keepSpecial) continue;
      for (const ch of t) {
        const b = u2b.get(ch);
        // A codepoint outside the byte table means the token is literal text
        // (some added tokens are), so emit its UTF-8 rather than dropping it.
        if (b === undefined) for (const bb of new TextEncoder().encode(ch)) bytes.push(bb);
        else bytes.push(b);
      }
    }
    return new TextDecoder("utf-8", { fatal: false }).decode(Uint8Array.from(bytes));
  }

  const api = { build, encode, decode, byteTables };
  if (typeof module !== "undefined" && module.exports) module.exports = api;
  else root.Tokenizer = api;
})(typeof self !== "undefined" ? self : this);