From 1147b018ce5cdfe8d69d1a282a20dda994a7c1f9 Mon Sep 17 00:00:00 2001 From: Thomas Kosiewski Date: Tue, 3 Nov 2026 11:30:02 +0000 Subject: [PATCH 1/9] feat: add SentencePieceTokenizer for merge-list BPE vocabularies Change-Id: Id42ddd71e5349f481dbd64a6cd647a4680896be9 Signed-off-by: Thomas Kosiewski --- src/index.ts | 2 + src/sentencepiece.ts | 429 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 431 insertions(+) create mode 100644 src/sentencepiece.ts diff --git a/src/index.ts b/src/index.ts index ae3a363..6cfb5c8 100644 --- a/src/index.ts +++ b/src/index.ts @@ -5,6 +5,8 @@ import type * as encodings from "./encoding"; export default Tokenizer; export { Tokenizer }; export type { Encoding } from "./tokenizer"; +export { default as SentencePieceTokenizer } from "./sentencepiece"; +export type { SentencePieceEncoding } from "./sentencepiece"; export type ModelName = keyof typeof modelsJson; diff --git a/src/sentencepiece.ts b/src/sentencepiece.ts new file mode 100644 index 0000000..5052ea1 --- /dev/null +++ b/src/sentencepiece.ts @@ -0,0 +1,429 @@ +/** + * BPE tokenizer for SentencePiece-style vocabularies (Gemma 4, which Google + * documents as the Gemini tokenizer). + * + * It reproduces the Hugging Face `tokenizers` BPE model exactly: + * - the only normalization is " " -> "▁", + * - there is no pre-tokenization, so one text is one word, + * - a character missing from the vocabulary becomes its UTF-8 bytes as + * `<0xNN>` tokens, + * - merges run in merge-list order (the merge index is the rank). Ties go to + * the leftmost position. + * + * A word can be a whole document, so the merge loop uses a heap and a + * doubly linked symbol list. It runs in O(n log n), never O(n^2). + */ + +/** + * Encoding data format for {@link SentencePieceTokenizer}. + */ +export interface SentencePieceEncoding { + name: string; + special_tokens: Record; + /** Token text by id. A space is "▁" and a raw byte is "<0xNN>". */ + vocab: string[]; + /** + * Base64 of the merge list in merge order. For `n` merges it holds nine + * planes of `n` bytes: byte 0, 1 and 2 (little endian) of the left token + * id, then the right token id, then the merged token id. + */ + merges: string; +} + +const SPACE_MARK = "▁"; +const SPACE_MARK_CODE_POINT = 0x2581; +const REPLACEMENT_CODE_POINT = 0xfffd; +// Merge index and position share one heap key: rank * 2^32 + position. +const KEY_RADIX = 4294967296; +const BYTE_TOKEN = /^<0x([0-9A-F]{2})>$/; + +const textDecoder = new TextDecoder("utf-8"); + +function escapeRegex(str: string): string { + return str.replace(/[\\^$*+?.()|[\]{}]/g, "\\$&"); +} + +function decodeBase64(base64: string): Uint8Array { + const binary = atob(base64); + const bytes = new Uint8Array(binary.length); + for (let i = 0; i < binary.length; i++) { + bytes[i] = binary.charCodeAt(i); + } + return bytes; +} + +function hashPair(left: number, right: number): number { + let h = Math.imul(left ^ Math.imul(right, 0x85ebca6b), 0xc2b2ae35); + h ^= h >>> 16; + return h; +} + +/** + * Merge list as a hash table from a token pair to its merge index. + */ +class MergeTable { + readonly left: Int32Array; + readonly right: Int32Array; + readonly merged: Int32Array; + private readonly slots: Int32Array; + private readonly mask: number; + + constructor(base64: string) { + const bytes = decodeBase64(base64); + if (bytes.length === 0 || bytes.length % 9 !== 0) { + throw new Error("Invalid merge list: length is not a multiple of 9"); + } + const count = bytes.length / 9; + const readPlane = (plane: number): Int32Array => { + const ids = new Int32Array(count); + const base = plane * 3 * count; + for (let i = 0; i < count; i++) { + ids[i] = + bytes[base + i]! | + (bytes[base + count + i]! << 8) | + (bytes[base + 2 * count + i]! << 16); + } + return ids; + }; + this.left = readPlane(0); + this.right = readPlane(1); + this.merged = readPlane(2); + + let size = 1024; + while (size < count * 2) size *= 2; + this.mask = size - 1; + this.slots = new Int32Array(size).fill(-1); + for (let i = 0; i < count; i++) { + let slot = hashPair(this.left[i]!, this.right[i]!) & this.mask; + while (this.slots[slot] !== -1) slot = (slot + 1) & this.mask; + this.slots[slot] = i; + } + } + + /** Merge index of the pair, or -1 when the pair never merges. */ + find(left: number, right: number): number { + let slot = hashPair(left, right) & this.mask; + for (;;) { + const index = this.slots[slot]!; + if (index === -1) return -1; + if (this.left[index] === left && this.right[index] === right) { + return index; + } + slot = (slot + 1) & this.mask; + } + } +} + +/** + * Min-heap of merge candidates. Keys are exact doubles below 2^53. + */ +class MergeHeap { + private keys: Float64Array; + size = 0; + + constructor(capacity: number) { + this.keys = new Float64Array(Math.max(capacity, 16)); + } + + /** Append without restoring heap order. Call `heapify` afterwards. */ + append(key: number): void { + if (this.size === this.keys.length) this.grow(); + this.keys[this.size++] = key; + } + + heapify(): void { + for (let i = (this.size >> 1) - 1; i >= 0; i--) this.siftDown(i); + } + + push(key: number): void { + if (this.size === this.keys.length) this.grow(); + const keys = this.keys; + let i = this.size++; + while (i > 0) { + const parent = (i - 1) >> 1; + if (keys[parent]! <= key) break; + keys[i] = keys[parent]!; + i = parent; + } + keys[i] = key; + } + + pop(): number { + const keys = this.keys; + const top = keys[0]!; + const last = keys[--this.size]!; + if (this.size > 0) { + keys[0] = last; + this.siftDown(0); + } + return top; + } + + private grow(): void { + const next = new Float64Array(this.keys.length * 2); + next.set(this.keys); + this.keys = next; + } + + private siftDown(start: number): void { + const keys = this.keys; + const size = this.size; + const key = keys[start]!; + let i = start; + for (;;) { + let child = 2 * i + 1; + if (child >= size) break; + if (child + 1 < size && keys[child + 1]! < keys[child]!) child++; + if (keys[child]! >= key) break; + keys[i] = keys[child]!; + i = child; + } + keys[i] = key; + } +} + +/** + * Tokenizer for {@link SentencePieceEncoding} data. It has the same public + * counting surface as the default `Tokenizer` class. + */ +export default class SentencePieceTokenizer { + private readonly name: string; + private readonly vocab: string[]; + private readonly merges: MergeTable; + private readonly characterIds: Map; + private readonly byteIds: Int32Array; + private readonly byteValues: Int16Array; + private readonly specialTokens: Record; + private readonly specialTokenKeys: readonly string[]; + private readonly specialTokensRegex: RegExp | null; + + constructor(data: SentencePieceEncoding) { + this.name = data.name; + this.vocab = data.vocab; + this.merges = new MergeTable(data.merges); + + this.characterIds = new Map(); + this.byteIds = new Int32Array(256).fill(-1); + this.byteValues = new Int16Array(data.vocab.length).fill(-1); + for (let id = 0; id < data.vocab.length; id++) { + const piece = data.vocab[id]!; + const codePoint = piece.codePointAt(0); + if (codePoint === undefined) continue; + if (piece.length === (codePoint > 0xffff ? 2 : 1)) { + this.characterIds.set(codePoint, id); + } else if (piece.length === 6) { + const match = BYTE_TOKEN.exec(piece); + if (match !== null) { + const value = Number.parseInt(match[1]!, 16); + this.byteIds[value] = id; + this.byteValues[id] = value; + } + } + } + for (let value = 0; value < 256; value++) { + if (this.byteIds[value] === -1) { + throw new Error(`Encoding ${data.name} has no byte token for ${value}`); + } + } + if (!this.characterIds.has(REPLACEMENT_CODE_POINT)) { + throw new Error(`Encoding ${data.name} has no U+FFFD token`); + } + + this.specialTokens = { ...data.special_tokens }; + // Longest first, so an alternation prefers the longest token. + this.specialTokenKeys = Object.keys(this.specialTokens).sort( + (a, b) => b.length - a.length + ); + this.specialTokensRegex = + this.specialTokenKeys.length > 0 + ? new RegExp(this.specialTokenKeys.map(escapeRegex).join("|"), "g") + : null; + } + + get encodingName(): string { + return this.name; + } + + /** + * Encode text to token ids. The special token arguments follow + * `Tokenizer.encode`: tokens in `allowedSpecial` become their id, and + * `disallowedSpecial` throws when the text contains one of its tokens. + */ + encode( + text: string, + allowedSpecial: Array | "all" = [], + disallowedSpecial: Array | "all" = "all" + ): number[] { + const result: number[] = []; + const allowedSet = new Set( + allowedSpecial === "all" ? this.specialTokenKeys : allowedSpecial + ); + const disallowed = + disallowedSpecial === "all" + ? this.specialTokenKeys.filter((token) => !allowedSet.has(token)) + : disallowedSpecial; + for (const token of disallowed) { + if (text.includes(token)) { + throw new Error(`Text contains disallowed special token: ${token}`); + } + } + + const regex = this.specialTokensRegex; + if (regex === null || allowedSet.size === 0) { + this.encodeOrdinaryInto(text, result); + return result; + } + + let start = 0; + while (start < text.length) { + let next: RegExpExecArray | null; + let searchFrom = start; + for (;;) { + regex.lastIndex = searchFrom; + next = regex.exec(text); + if (next === null || allowedSet.has(next[0])) break; + searchFrom = next.index + 1; + } + const end = next === null ? text.length : next.index; + this.encodeOrdinaryInto(text.substring(start, end), result); + if (next === null) break; + result.push(this.specialTokens[next[0]]!); + start = next.index + next[0].length; + } + return result; + } + + /** + * Decode token ids to text. Byte tokens are joined and read as UTF-8. + */ + decode(tokens: number[]): string { + let result = ""; + let bytes: number[] = []; + for (let i = 0; i < tokens.length; i++) { + const token = tokens[i]!; + const byteValue = this.byteValues[token]; + if (byteValue !== undefined && byteValue !== -1) { + bytes.push(byteValue); + continue; + } + const piece = this.vocab[token]; + if (piece === undefined) continue; + if (bytes.length > 0) { + result += textDecoder.decode(new Uint8Array(bytes)); + bytes = []; + } + result += piece.replace(/▁/g, " "); + } + if (bytes.length > 0) { + result += textDecoder.decode(new Uint8Array(bytes)); + } + return result; + } + + /** + * Count tokens in text. Special token text counts as ordinary text. + */ + count(input: string): number { + return this.encode(input, [], []).length; + } + + private encodeOrdinaryInto(text: string, result: number[]): void { + if (text.length === 0) return; + const symbols = this.initialSymbols(text); + const count = symbols.length; + const merges = this.merges; + const mergedIds = merges.merged; + + const prev = new Int32Array(count); + const next = new Int32Array(count); + const alive = new Uint8Array(count).fill(1); + for (let i = 0; i < count; i++) { + prev[i] = i - 1; + next[i] = i + 1 < count ? i + 1 : -1; + } + + const heap = new MergeHeap(count); + for (let i = 0; i + 1 < count; i++) { + const rank = merges.find(symbols[i]!, symbols[i + 1]!); + if (rank !== -1) heap.append(rank * KEY_RADIX + i); + } + heap.heapify(); + + while (heap.size > 0) { + const key = heap.pop(); + const rank = Math.floor(key / KEY_RADIX); + const pos = key - rank * KEY_RADIX; + if (alive[pos] === 0) continue; + const right = next[pos]!; + if (right === -1) continue; + // Skip a queue entry that went stale after its neighbours merged. + const current = merges.find(symbols[pos]!, symbols[right]!); + if (current === -1 || mergedIds[current] !== mergedIds[rank]) continue; + + symbols[pos] = mergedIds[rank]!; + alive[right] = 0; + const after = next[right]!; + next[pos] = after; + if (after !== -1) prev[after] = pos; + + const before = prev[pos]!; + if (before !== -1) { + const beforeRank = merges.find(symbols[before]!, symbols[pos]!); + if (beforeRank !== -1) heap.push(beforeRank * KEY_RADIX + before); + } + if (after !== -1) { + const afterRank = merges.find(symbols[pos]!, symbols[after]!); + if (afterRank !== -1) heap.push(afterRank * KEY_RADIX + pos); + } + } + + for (let i = 0; i !== -1; i = next[i]!) result.push(symbols[i]!); + } + + /** + * One symbol per character after the space marker rule. A character the + * vocabulary lacks becomes one symbol per UTF-8 byte. + */ + private initialSymbols(text: string): Int32Array { + const symbols: number[] = []; + const length = text.length; + for (let i = 0; i < length; i++) { + let codePoint = text.charCodeAt(i); + if (codePoint >= 0xd800 && codePoint <= 0xdbff && i + 1 < length) { + const low = text.charCodeAt(i + 1); + if (low >= 0xdc00 && low <= 0xdfff) { + codePoint = 0x10000 + ((codePoint - 0xd800) << 10) + (low - 0xdc00); + i++; + } + } + // A lone surrogate becomes U+FFFD, as UTF-8 encoding would make it. + if (codePoint >= 0xd800 && codePoint <= 0xdfff) { + codePoint = REPLACEMENT_CODE_POINT; + } else if (codePoint === 0x20) { + codePoint = SPACE_MARK_CODE_POINT; + } + + const id = this.characterIds.get(codePoint); + if (id !== undefined) { + symbols.push(id); + } else if (codePoint < 0x80) { + symbols.push(this.byteIds[codePoint]!); + } else if (codePoint < 0x800) { + symbols.push(this.byteIds[0xc0 | (codePoint >> 6)]!); + symbols.push(this.byteIds[0x80 | (codePoint & 0x3f)]!); + } else if (codePoint < 0x10000) { + symbols.push(this.byteIds[0xe0 | (codePoint >> 12)]!); + symbols.push(this.byteIds[0x80 | ((codePoint >> 6) & 0x3f)]!); + symbols.push(this.byteIds[0x80 | (codePoint & 0x3f)]!); + } else { + symbols.push(this.byteIds[0xf0 | (codePoint >> 18)]!); + symbols.push(this.byteIds[0x80 | ((codePoint >> 12) & 0x3f)]!); + symbols.push(this.byteIds[0x80 | ((codePoint >> 6) & 0x3f)]!); + symbols.push(this.byteIds[0x80 | (codePoint & 0x3f)]!); + } + } + return Int32Array.from(symbols); + } +} + +export { SPACE_MARK }; From 1b15e2424e510d196d099887ebf3ef06622daa64 Mon Sep 17 00:00:00 2001 From: Thomas Kosiewski Date: Tue, 3 Nov 2026 11:30:03 +0000 Subject: [PATCH 2/9] feat(encoding): add llama4 and gemma4 encodings generated from pinned HF tokenizers Change-Id: I904ed1182063c82c2e0a35419f80d63a8850847b Signed-off-by: Thomas Kosiewski --- package.json | 24 +- scripts/generate-hf-encodings.ts | 369 + scripts/hf-sources.ts | 56 + src/encoding/gemma4.ts | 262194 ++++++++++++++++++ src/encoding/index.ts | 3 +- src/encoding/llama4.ts | 401759 ++++++++++++++++++++++++++++ 6 files changed, 664403 insertions(+), 2 deletions(-) create mode 100644 scripts/generate-hf-encodings.ts create mode 100644 scripts/hf-sources.ts create mode 100644 src/encoding/gemma4.ts create mode 100644 src/encoding/llama4.ts diff --git a/package.json b/package.json index 6af94fd..f4c1354 100644 --- a/package.json +++ b/package.json @@ -34,7 +34,9 @@ "encoding/cl100k_base": ["./dist/encoding/cl100k_base.d.ts"], "encoding/o200k_base": ["./dist/encoding/o200k_base.d.ts"], "encoding/p50k_base": ["./dist/encoding/p50k_base.d.ts"], - "encoding/claude": ["./dist/encoding/claude.d.ts"] + "encoding/claude": ["./dist/encoding/claude.d.ts"], + "encoding/gemma4": ["./dist/encoding/gemma4.d.ts"], + "encoding/llama4": ["./dist/encoding/llama4.d.ts"] } }, "exports": { @@ -107,6 +109,26 @@ "types": "./dist/encoding/claude.d.cts", "default": "./dist/encoding/claude.cjs" } + }, + "./encoding/gemma4": { + "import": { + "types": "./dist/encoding/gemma4.d.ts", + "default": "./dist/encoding/gemma4.js" + }, + "require": { + "types": "./dist/encoding/gemma4.d.cts", + "default": "./dist/encoding/gemma4.cjs" + } + }, + "./encoding/llama4": { + "import": { + "types": "./dist/encoding/llama4.d.ts", + "default": "./dist/encoding/llama4.js" + }, + "require": { + "types": "./dist/encoding/llama4.d.cts", + "default": "./dist/encoding/llama4.cjs" + } } }, "description": "A faster than tiktoken tokenizer with first-class support for Vercel's AI SDK.", diff --git a/scripts/generate-hf-encodings.ts b/scripts/generate-hf-encodings.ts new file mode 100644 index 0000000..9936e2b --- /dev/null +++ b/scripts/generate-hf-encodings.ts @@ -0,0 +1,369 @@ +/** + * Generate the llama4 and gemma4 encodings from Hugging Face tokenizer.json + * files pinned to a repository revision and a SHA-256. + * + * Usage: + * bun scripts/generate-hf-encodings.ts [llama4] [gemma4] + * + * Downloads go to .cache/hf-tokenizers/ (git-ignored). The script refuses a + * file whose hash differs from the pin, and it asserts every property of the + * tokenizer that the runtime relies on. A new revision must pass those + * assertions before its pin changes. + */ + +import * as fs from "node:fs"; +import * as path from "node:path"; +import * as prettier from "prettier"; +import { downloadTokenizerJson, SOURCES, type Source } from "./hf-sources"; + +const ENCODING_DIR = path.join(import.meta.dirname, "../src/encoding"); + +// tokenizer.json holds the regex in Rust syntax. `(?i:...)` and `\s` differ +// in JavaScript, so the JavaScript pattern is written out and the source +// pattern is asserted instead of translated. +const LLAMA4_SOURCE_PATTERN = + "[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}]*[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?|[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}]+[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n/]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+"; +// `\s` is Unicode White_Space in Rust and Oniguruma. In JavaScript it also +// matches U+FEFF and misses U+0085. +const LLAMA4_PATTERN = + "[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}]*[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}]+(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])?|[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}]+[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}]*(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])?|\\p{N}{1,3}| ?[^\\p{White_Space}\\p{L}\\p{N}]+[\\r\\n/]*|\\p{White_Space}*[\\r\\n]+|\\p{White_Space}+(?!\\P{White_Space})|\\p{White_Space}+"; + +const textEncoder = new TextEncoder(); + +function assert(condition: unknown, message: string): asserts condition { + if (!condition) throw new Error(`Assertion failed: ${message}`); +} + +function assertAddedTokens(added: any[]): Record { + const special: Record = {}; + for (const token of added) { + assert( + token.special === true && + token.normalized === false && + token.lstrip === false && + token.rstrip === false && + token.single_word === false, + `added token ${token.content} must be plain and special` + ); + assert(!(token.content in special), `duplicate ${token.content}`); + special[token.content] = token.id; + } + return special; +} + +function assertContiguous(vocab: Record): string[] { + const entries = Object.entries(vocab); + const byId = new Array(entries.length); + for (const [piece, id] of entries) { + assert(Number.isInteger(id) && id >= 0 && id < entries.length, `id ${id}`); + assert(byId[id] === undefined, `duplicate id ${id}`); + byId[id] = piece; + } + return byId; +} + +/** The GPT-2 table from byte value to the printable character in vocab keys. */ +function byteLevelDecoder(): Map { + const printable: number[] = []; + for (let b = 33; b <= 126; b++) printable.push(b); + for (let b = 161; b <= 172; b++) printable.push(b); + for (let b = 174; b <= 255; b++) printable.push(b); + const decoder = new Map(); + for (const b of printable) decoder.set(String.fromCharCode(b), b); + let extra = 0; + for (let b = 0; b < 256; b++) { + if (!printable.includes(b)) { + decoder.set(String.fromCharCode(256 + extra), b); + extra++; + } + } + assert(decoder.size === 256, "byte-level table has 256 entries"); + return decoder; +} + +/** Bytes to a string when they are valid UTF-8 and survive a round trip. */ +function bytesToString(bytes: Uint8Array): string | undefined { + let text: string; + try { + text = new TextDecoder("utf-8", { fatal: true }).decode(bytes); + } catch { + return undefined; + } + const again = textEncoder.encode(text); + if (again.length !== bytes.length) return undefined; + for (let i = 0; i < bytes.length; i++) { + if (again[i] !== bytes[i]) return undefined; + } + return text; +} + +function describeSource(source: Source): string { + return ` * Source: https://huggingface.co/${source.repo} at ${source.revision}\n * tokenizer.json sha256: ${source.sha256}`; +} + +function generateLlama4(tokenizer: any, source: Source): string { + assert(tokenizer.normalizer === null, "llama4 has no normalizer"); + const pre = tokenizer.pre_tokenizer; + assert( + pre.type === "Sequence" && pre.pretokenizers.length === 2, + "llama4 pre-tokenizer is split + byte-level" + ); + const [split, byteLevel] = pre.pretokenizers; + assert( + split.type === "Split" && + split.behavior === "Isolated" && + split.invert === false && + split.pattern.Regex === LLAMA4_SOURCE_PATTERN, + "llama4 split pattern changed" + ); + assert( + byteLevel.type === "ByteLevel" && + byteLevel.add_prefix_space === false && + byteLevel.use_regex === false, + "llama4 byte-level settings changed" + ); + const model = tokenizer.model; + assert( + model.type === "BPE" && + model.ignore_merges === true && + model.byte_fallback === false && + model.dropout === null && + model.continuing_subword_prefix === null && + model.end_of_word_suffix === null, + "llama4 model settings changed" + ); + + const special = assertAddedTokens(tokenizer.added_tokens); + const byId = assertContiguous(model.vocab); + assert(!byId.includes("__proto__"), "__proto__ would set a prototype"); + for (const id of Object.values(special)) { + assert(id >= byId.length, `special id ${id} overlaps the vocabulary`); + } + + // Rank equals id. The existing engine merges the adjacent pair with the + // lowest merged rank. Hugging Face merges the pair with the lowest merge + // index. They agree when the merge list is every splittable pair of every + // token, in token order. + const merges: Array<[string, string]> = model.merges; + const seen = new Set(); + let previousId = -1; + for (const [left, right] of merges) { + const id = model.vocab[left + right]; + assert( + id !== undefined && + model.vocab[left] !== undefined && + model.vocab[right] !== undefined, + `merge ${left} ${right} is not in the vocabulary` + ); + assert(id >= previousId, "merge list is not in token order"); + previousId = id; + seen.add(`${left} ${right}`); + } + assert(seen.size === merges.length, "duplicate merge"); + let splittable = 0; + for (const piece of byId) { + for (let i = 1; i < piece.length; i++) { + if ( + model.vocab[piece.slice(0, i)] !== undefined && + model.vocab[piece.slice(i)] !== undefined + ) { + splittable++; + } + } + } + assert(splittable === merges.length, "merge list is not the splittable set"); + + const decoderTable = byteLevelDecoder(); + const stringEntries: string[] = []; + const binaryPairs: Array<{ bytes: number[]; rank: number }> = []; + const decoderEntries: string[] = []; + for (let id = 0; id < byId.length; id++) { + const bytes = Uint8Array.from( + Array.from(byId[id]!, (ch) => { + const value = decoderTable.get(ch); + assert(value !== undefined, `byte-level char U+${ch.codePointAt(0)}`); + return value; + }) + ); + const text = bytesToString(bytes); + if (text !== undefined) { + stringEntries.push(` ${JSON.stringify(text)}: ${id}`); + decoderEntries.push(` ${id}: ${JSON.stringify(text)}`); + } else { + binaryPairs.push({ bytes: Array.from(bytes), rank: id }); + decoderEntries.push(` ${id}: new Uint8Array([${bytes.join(",")}])`); + } + } + binaryPairs.sort((a, b) => { + const minLen = Math.min(a.bytes.length, b.bytes.length); + for (let i = 0; i < minLen; i++) { + const diff = a.bytes[i]! - b.bytes[i]!; + if (diff !== 0) return diff; + } + return a.bytes.length - b.bytes.length; + }); + const binaryEntries = binaryPairs.map( + ({ bytes, rank }) => ` [new Uint8Array([${bytes.join(",")}]), ${rank}]` + ); + + return `/** + * LLAMA4 encoding + * Auto-generated by scripts/generate-hf-encodings.ts - DO NOT EDIT + * +${describeSource(source)} + * + * Optimized dual-storage format: + * - String tokens: ${stringEntries.length.toLocaleString("en-US")} (fast Map lookup) + * - Binary tokens: ${binaryPairs.length.toLocaleString("en-US")} (binary search) + * - Total tokens: ${byId.length.toLocaleString("en-US")} + * - Pattern: ${LLAMA4_PATTERN} + * - Special tokens: ${Object.keys(special).length.toLocaleString("en-US")} + */ + +export const name = "llama4"; + +export const pat_str = ${JSON.stringify(LLAMA4_PATTERN)}; + +export const special_tokens = ${JSON.stringify(special, null, 2)}; + +// String-based encoder (UTF-8 tokens) - fast object lookup. It has no +// prototype, so names such as "toString" are not found by accident. +export const stringEncoder: Record = { + __proto__: null, +${stringEntries.join(",\n")} +} as any; + +// Binary encoder (non-UTF-8 tokens) - pre-sorted for binary search +export const binaryEncoder: Array<[Uint8Array, number]> = [ +${binaryEntries.join(",\n")} +]; + +// Decoder - plain object for mobile compatibility +export const decoder: Record = { +${decoderEntries.join(",\n")} +}; +`; +} + +function generateGemma4(tokenizer: any, source: Source): string { + assert( + JSON.stringify(tokenizer.normalizer) === + JSON.stringify({ + type: "Replace", + pattern: { String: " " }, + content: "\u2581", + }), + "gemma4 normalizer changed" + ); + // After the normalizer no space is left, so this split never fires and one + // text is one word. + assert( + JSON.stringify(tokenizer.pre_tokenizer) === + JSON.stringify({ + type: "Split", + pattern: { String: " " }, + behavior: "MergedWithPrevious", + invert: false, + }), + "gemma4 pre-tokenizer changed" + ); + const model = tokenizer.model; + assert( + model.type === "BPE" && + model.byte_fallback === true && + model.ignore_merges === false && + model.dropout === null && + model.continuing_subword_prefix === null && + model.end_of_word_suffix === null, + "gemma4 model settings changed" + ); + + const special = assertAddedTokens(tokenizer.added_tokens); + const byId = assertContiguous(model.vocab); + assert(byId.length <= 1 << 18, "token ids must fit in 18 bits"); + for (const [text, id] of Object.entries(special)) { + assert(model.vocab[text] === id, `special token ${text} is in the vocab`); + } + assert( + !byId.some((piece) => /[\u2028\u2029\ud800-\udfff]/u.test(piece)), + "no line separators or lone surrogates in pieces" + ); + for (let b = 0; b < 256; b++) { + const hex = b.toString(16).toUpperCase().padStart(2, "0"); + assert(model.vocab[`<0x${hex}>`] !== undefined, `byte token ${hex}`); + } + + const merges: Array<[string, string]> = model.merges; + const count = merges.length; + const planes = new Uint8Array(count * 9); + const seen = new Set(); + merges.forEach(([left, right], i) => { + const ids: number[] = [ + model.vocab[left], + model.vocab[right], + model.vocab[left + right], + ]; + assert( + ids.every((id) => id !== undefined), + `merge ${i} has unknown parts` + ); + const key = ids[0]! * (1 << 18) + ids[1]!; + assert(!seen.has(key), `duplicate merge ${i}`); + seen.add(key); + ids.forEach((id, plane) => { + planes[(plane * 3 + 0) * count + i] = id & 0xff; + planes[(plane * 3 + 1) * count + i] = (id >> 8) & 0xff; + planes[(plane * 3 + 2) * count + i] = (id >> 16) & 0xff; + }); + }); + + return `/** + * GEMMA4 encoding + * Auto-generated by scripts/generate-hf-encodings.ts - DO NOT EDIT + * +${describeSource(source)} + * + * SentencePiece-style BPE, run by SentencePieceTokenizer: + * - Tokens: ${byId.length.toLocaleString("en-US")} + * - Merges: ${count.toLocaleString("en-US")} + * - Special tokens: ${Object.keys(special).length.toLocaleString("en-US")} + */ + +export const name = "gemma4"; + +export const special_tokens = ${JSON.stringify(special, null, 2)}; + +// Token text by id. +export const vocab: string[] = [ +${byId.map((piece) => ` ${JSON.stringify(piece)}`).join(",\n")} +]; + +// Merge list in merge order as nine byte planes, base64 encoded. See +// SentencePieceEncoding. +export const merges = ${JSON.stringify(Buffer.from(planes).toString("base64"))}; +`; +} + +const generators: Record string> = { + llama4: generateLlama4, + gemma4: generateGemma4, +}; + +const requested = process.argv.slice(2); +const names = requested.length > 0 ? requested : Object.keys(generators); +for (const name of names) { + const source = SOURCES[name]; + const generate = generators[name]; + assert(source !== undefined && generate !== undefined, `unknown ${name}`); + const tokenizer = JSON.parse( + (await downloadTokenizerJson(name)).toString("utf-8") + ); + const outputFile = path.join(ENCODING_DIR, `${name}.ts`); + // The repository formats src/ with Prettier, so write the formatted file. + const code = await prettier.format(generate(tokenizer, source), { + ...(await prettier.resolveConfig(outputFile)), + filepath: outputFile, + }); + fs.writeFileSync(outputFile, code); + console.log(`${name}: ${outputFile} (${(code.length / 1e6).toFixed(2)} MB)`); +} diff --git a/scripts/hf-sources.ts b/scripts/hf-sources.ts new file mode 100644 index 0000000..0642af7 --- /dev/null +++ b/scripts/hf-sources.ts @@ -0,0 +1,56 @@ +import { createHash } from "node:crypto"; +import * as fs from "node:fs"; +import * as path from "node:path"; + +export const CACHE_DIR = path.join( + import.meta.dirname, + "../.cache/hf-tokenizers" +); + +export interface Source { + repo: string; + revision: string; + sha256: string; +} + +export const SOURCES: Record = { + llama4: { + repo: "unsloth/Llama-4-Scout-17B-16E-Instruct", + revision: "afd8e498c87bda51c7ea8ec68ea2f7c066e6340b", + sha256: "172c9eb4beafc72601690da3ccfcede5c2e6806a8d5ec1fca33e22acea8023a4", + }, + gemma4: { + repo: "google/gemma-4-E4B-it", + revision: "ee0ef6023621cff504d758262d4e04895a5af4a2", + sha256: "cc8d3a0ce36466ccc1278bf987df5f71db1719b9ca6b4118264f45cb627bfe0f", + }, +}; + +/** + * Path of the pinned tokenizer.json. It is downloaded on first use, and its + * SHA-256 must match the pin. + */ +export async function tokenizerJsonPath(name: string): Promise { + const source = SOURCES[name]; + if (source === undefined) throw new Error(`Unknown encoding: ${name}`); + fs.mkdirSync(CACHE_DIR, { recursive: true }); + const file = path.join(CACHE_DIR, `${name}-${source.revision}.json`); + if (!fs.existsSync(file)) { + const url = `https://huggingface.co/${source.repo}/resolve/${source.revision}/tokenizer.json`; + console.log(`Downloading ${url}`); + const response = await fetch(url); + if (!response.ok) throw new Error(`${url} answered ${response.status}`); + fs.writeFileSync(file, Buffer.from(await response.arrayBuffer())); + } + const sha256 = createHash("sha256") + .update(fs.readFileSync(file)) + .digest("hex"); + if (sha256 !== source.sha256) { + throw new Error(`${file} has sha256 ${sha256}, expected ${source.sha256}`); + } + return file; +} + +export async function downloadTokenizerJson(name: string): Promise { + return fs.readFileSync(await tokenizerJsonPath(name)); +} diff --git a/src/encoding/gemma4.ts b/src/encoding/gemma4.ts new file mode 100644 index 0000000..35441d1 --- /dev/null +++ b/src/encoding/gemma4.ts @@ -0,0 +1,262194 @@ +/** + * GEMMA4 encoding + * Auto-generated by scripts/generate-hf-encodings.ts - DO NOT EDIT + * + * Source: https://huggingface.co/google/gemma-4-E4B-it at ee0ef6023621cff504d758262d4e04895a5af4a2 + * tokenizer.json sha256: cc8d3a0ce36466ccc1278bf987df5f71db1719b9ca6b4118264f45cb627bfe0f + * + * SentencePiece-style BPE, run by SentencePieceTokenizer: + * - Tokens: 262,144 + * - Merges: 514,906 + * - Special tokens: 24 + */ + +export const name = "gemma4"; + +export const special_tokens = { + "": 0, + "": 1, + "": 2, + "": 3, + "": 4, + "<|tool>": 46, + "": 47, + "<|tool_call>": 48, + "": 49, + "<|tool_response>": 50, + "": 51, + '<|"|>': 52, + "<|think|>": 98, + "<|channel>": 100, + "": 101, + "<|turn>": 105, + "": 106, + "<|image>": 255999, + "<|audio>": 256000, + "<|image|>": 258880, + "<|audio|>": 258881, + "": 258882, + "": 258883, + "<|video|>": 258884, +}; + +// Token text by id. +export const vocab: string[] = [ + "", + "", + "", + "", + "", + "[multimodal]", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "<|tool>", + "", + "<|tool_call>", + "", + "<|tool_response>", + "", + '<|"|>', + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "<|think|>", + "", + "<|channel>", + "", + "", + "", + "", + "<|turn>", + "", + "\n", + "\n\n", + "\n\n\n", + "\n\n\n\n", + "\n\n\n\n\n", + "\n\n\n\n\n\n", + "\n\n\n\n\n\n\n", + "\n\n\n\n\n\n\n\n", + "\n\n\n\n\n\n\n\n\n", + "\n\n\n\n\n\n\n\n\n\n", + "\n\n\n\n\n\n\n\n\n\n\n", + "\n\n\n\n\n\n\n\n\n\n\n\n", + "\n\n\n\n\n\n\n\n\n\n\n\n\n", + "\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "▁▁", + "▁▁▁", + "▁▁▁▁", + "▁▁▁▁▁", + "▁▁▁▁▁▁", + "▁▁▁▁▁▁▁", + "▁▁▁▁▁▁▁▁", + "▁▁▁▁▁▁▁▁▁", + "▁▁▁▁▁▁▁▁▁▁", + "▁▁▁▁▁▁▁▁▁▁▁", + "▁▁▁▁▁▁▁▁▁▁▁▁", + "▁▁▁▁▁▁▁▁▁▁▁▁▁", + "▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "", + "", + "", + "", + "", + "
", + "
", + "", + "
", + "", + "", + "", + "", + "", + "", + "", + "

", + "

", + "

", + "

", + "

", + "
", + "
", + "
", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "