From 7b589df4c1198727a5e7608272d79bda9c64b010 Mon Sep 17 00:00:00 2001 From: Thomas Kosiewski Date: Tue, 3 Nov 2026 10:55:02 +0000 Subject: [PATCH 1/8] test: add timing test for long single-piece inputs One run of underscores, spaces, dots or letters is a single regex piece, and bytePairMerge is quadratic in the piece length. At 100k characters the four shapes take 14 to 20 s here, so this test fails until the merge is fixed. See #6. Change-Id: I140f5a21ac255a662fa4d06a4098144c0429857f Signed-off-by: Thomas Kosiewski --- src/byte-pair-merge.test.ts | 40 +++++++++++++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) create mode 100644 src/byte-pair-merge.test.ts diff --git a/src/byte-pair-merge.test.ts b/src/byte-pair-merge.test.ts new file mode 100644 index 0000000..5962b99 --- /dev/null +++ b/src/byte-pair-merge.test.ts @@ -0,0 +1,40 @@ +/** + * Long pieces: the split pattern bounds digit runs only, so one run of + * underscores, spaces, dots or letters becomes a single piece. Merging such a + * piece must scale about linearly with its length. + * See https://github.com/coder/ai-tokenizer/issues/6. + */ + +import { describe, test, expect } from "bun:test"; +import Tokenizer from "./tokenizer"; +import * as o200k from "./encoding/o200k_base"; + +// Linear time is about 0.15 s for 100k characters. The quadratic merge took +// 10 s or more, so this bound has 10x headroom for slow CI machines. +const LONG_PIECE_CHARS = 100_000; +const LONG_PIECE_BOUND_MS = 2_000; + +const longPieces: Array<[string, string]> = [ + ["underscores", "_".repeat(LONG_PIECE_CHARS)], + ["spaces", " ".repeat(LONG_PIECE_CHARS)], + ["alternating letters", "ab".repeat(LONG_PIECE_CHARS / 2)], + ["dots", ".".repeat(LONG_PIECE_CHARS)], +]; + +describe("long pieces", () => { + test.each(longPieces)( + "%s finish within a generous time bound", + (_name, text) => { + // A fresh tokenizer, so the merge cache cannot hide the cost. + const tokenizer = new Tokenizer(o200k); + const start = performance.now(); + const tokens = tokenizer.encode(text); + const elapsedMs = performance.now() - start; + + expect(tokens.length).toBeGreaterThan(0); + expect(tokenizer.decode(tokens)).toBe(text); + expect(elapsedMs).toBeLessThan(LONG_PIECE_BOUND_MS); + }, + 60_000 + ); +}); From a704ff09c3256141e2cebc2beffd8b71f7b3065d Mon Sep 17 00:00:00 2001 From: Thomas Kosiewski Date: Tue, 3 Nov 2026 11:05:32 +0000 Subject: [PATCH 2/8] perf: merge byte pairs in O(n log n) bytePairMerge rescanned every pair rank and spliced two arrays for each merge, so one long piece (a run of underscores, spaces or dots, which the split pattern does not bound) cost O(n^2). Keep the parts in a linked list and pick merges from a min-heap ordered by (rank, position). The merges and their order are the same as before: lowest rank first, leftmost on ties. The differential test compares the new merge with a copy of the old one on random, adversarial and real pieces for all four encodings. Change-Id: I69cf8dbccdf791d5c1b7963df679631d13e4f208 Signed-off-by: Thomas Kosiewski --- src/byte-pair-merge.test.ts | 384 +++++++++++++++++++++++++++++++++++- src/tokenizer.ts | 186 ++++++++++++----- 2 files changed, 511 insertions(+), 59 deletions(-) diff --git a/src/byte-pair-merge.test.ts b/src/byte-pair-merge.test.ts index 5962b99..8d1ea70 100644 --- a/src/byte-pair-merge.test.ts +++ b/src/byte-pair-merge.test.ts @@ -1,13 +1,25 @@ /** - * Long pieces: the split pattern bounds digit runs only, so one run of - * underscores, spaces, dots or letters becomes a single piece. Merging such a - * piece must scale about linearly with its length. - * See https://github.com/coder/ai-tokenizer/issues/6. + * Byte pair merge tests. + * + * 1. Long pieces: the split pattern bounds digit runs only, so one run of + * underscores, spaces, dots or letters becomes a single piece. Merging such + * a piece must scale about linearly with its length. + * See https://github.com/coder/ai-tokenizer/issues/6. + * 2. Differential: bytePairMerge must return exactly what the original + * quadratic algorithm (copied below) returns, on random and real pieces. + * 3. End to end: Tokenizer#encode must match tiktoken on long single-piece + * inputs. */ -import { describe, test, expect } from "bun:test"; -import Tokenizer from "./tokenizer"; +import { describe, test, expect, beforeAll, afterAll } from "bun:test"; +import { readdirSync, readFileSync } from "node:fs"; +import { join } from "node:path"; +import { get_encoding, type Tiktoken, type TiktokenEncoding } from "tiktoken"; +import Tokenizer, { bytePairMerge, type Encoding } from "./tokenizer"; import * as o200k from "./encoding/o200k_base"; +import * as cl100k from "./encoding/cl100k_base"; +import * as p50k from "./encoding/p50k_base"; +import * as claude from "./encoding/claude"; // Linear time is about 0.15 s for 100k characters. The quadratic merge took // 10 s or more, so this bound has 10x headroom for slow CI machines. @@ -21,6 +33,14 @@ const longPieces: Array<[string, string]> = [ ["dots", ".".repeat(LONG_PIECE_CHARS)], ]; +// Shaped like OCR output of a scanned form: long field rules, dot leaders and +// blank runs of different lengths, about 1.1M characters. The run lengths +// differ, so the merge cache cannot answer them. +const ocrLikeDocument = Array.from( + { length: 24 }, + (_, i) => `Field ${i}: ${"_. "[i % 3]!.repeat(45_000 + i * 173)}\n` +).join(""); + describe("long pieces", () => { test.each(longPieces)( "%s finish within a generous time bound", @@ -37,4 +57,356 @@ describe("long pieces", () => { }, 60_000 ); + + test("an OCR-like document of about 1.1M characters finishes within a generous time bound", () => { + expect(ocrLikeDocument.length).toBeGreaterThan(1_000_000); + const tokenizer = new Tokenizer(o200k); + const start = performance.now(); + const count = tokenizer.count(ocrLikeDocument); + const elapsedMs = performance.now() - start; + + expect(count).toBeGreaterThan(0); + // The quadratic merge needed minutes; linear time is about 1 s. + expect(elapsedMs).toBeLessThan(20_000); + }, 60_000); +}); + +const NO_RANK = 0xffffffff; +const textEncoder = new TextEncoder(); + +/** + * The merge from 1.0.6, which this fix replaced, with its rank lookup passed + * in. Do not "improve" this copy: it is the reference the new merge must match + * exactly. + */ +function legacyBytePairMerge( + piece: Uint8Array, + getRankForSlice: (slice: Uint8Array) => number +): number[] { + const len = piece.length; + + const starts: number[] = []; + const ranks: number[] = []; + + // Inline getRank helper + const getRank = ( + startIndex: number, + pairStart?: number, + pairEnd?: number + ): number => { + if (pairStart === undefined) pairStart = starts[startIndex]; + if (pairEnd === undefined) pairEnd = starts[startIndex + 2]; + if (pairEnd === undefined) return NO_RANK; + return getRankForSlice(piece.subarray(pairStart, pairEnd)); + }; + + // Initialize starts and ranks + for (let i = 0; i <= len; i++) { + starts[i] = i; + if (i < len - 1) { + ranks[i] = getRank(i, i, i + 2); + } else { + ranks[i] = NO_RANK; + } + } + + // Main merge loop + while (starts.length > 1) { + let minRank = NO_RANK; + let minIdx = -1; + + // Find minimum rank + const ranksLen = ranks.length - 1; + for (let i = 0; i < ranksLen; i++) { + const rank = ranks[i]!; + if (rank < minRank) { + minRank = rank!; + minIdx = i; + } + } + + if (minRank === NO_RANK || minIdx === -1) break; + + // Remove elements + starts.splice(minIdx + 1, 1); + ranks.splice(minIdx, 1); + + // Update ranks + ranks[minIdx] = getRank(minIdx); + if (minIdx > 0) { + ranks[minIdx - 1] = getRank(minIdx - 1); + } + } + + // Build output + const output: number[] = []; + for (let i = 0; i < starts.length - 1; i++) { + const pairStart = starts[i]; + const pairEnd = starts[i + 1]; + const rank = getRankForSlice(piece.subarray(pairStart, pairEnd)); + if (rank !== NO_RANK) { + output.push(rank); + } + } + + return output; +} + +const bytesKey = (bytes: Uint8Array): string => + Buffer.from(bytes.buffer, bytes.byteOffset, bytes.length).toString("latin1"); + +/** + * Rank lookup built straight from the encoding data, independent of the + * tokenizer's own lookup. A slice that is valid UTF-8 is looked up as a string + * token first and as a binary token second, as in the tokenizer. + */ +function rankLookup(encoding: Encoding): (slice: Uint8Array) => number { + const stringRanks = new Map(Object.entries(encoding.stringEncoder)); + const binaryRanks = new Map(); + for (const [bytes, rank] of encoding.binaryEncoder) { + binaryRanks.set(bytesKey(bytes), rank); + } + const decoder = new TextDecoder("utf-8", { fatal: true, ignoreBOM: true }); + return (slice) => { + let text: string | undefined; + try { + text = decoder.decode(slice); + } catch { + text = undefined; + } + return ( + (text === undefined ? undefined : stringRanks.get(text)) ?? + binaryRanks.get(bytesKey(slice)) ?? + NO_RANK + ); + }; +} + +// U+FEFF. Before the tokenizer decoded with ignoreBOM, a slice starting with +// these bytes was looked up without them, which the lookup above does not do. +const hasBom = (bytes: Uint8Array): boolean => + bytes.some( + (byte, i) => byte === 0xef && bytes[i + 1] === 0xbb && bytes[i + 2] === 0xbf + ); + +/** Small seeded generator (mulberry32): the corpus is the same on every run. */ +function seededRandom(seed: number): () => number { + let state = seed >>> 0; + return () => { + state = (state + 0x6d2b79f5) >>> 0; + let t = state; + t = Math.imul(t ^ (t >>> 15), t | 1); + t ^= t + Math.imul(t ^ (t >>> 7), t | 61); + return ((t ^ (t >>> 14)) >>> 0) / 4294967296; + }; +} + +const repoRoot = join(import.meta.dir, ".."); +const corpusFiles = [ + join(repoRoot, "README.md"), + join(repoRoot, "LICENSE"), + ...["src", "bench", "scripts", "site"].flatMap((dir) => + readdirSync(join(repoRoot, dir), { withFileTypes: true }) + .filter( + (entry) => entry.isFile() && /\.(ts|js|html|css)$/.test(entry.name) + ) + .map((entry) => join(repoRoot, dir, entry.name)) + ), +]; +const multilingual = [ + "你好世界,这是一个没有空格的很长的中文句子,用来测试分词器在长片段上的表现。".repeat( + 30 + ), + "こんにちは世界、これは長い日本語の文章です。スペースがありません。".repeat( + 30 + ), + "مرحبا بالعالم هذه جملة طويلة باللغة العربية لاختبار المحلل".repeat(30), + "Привет мир! Это длинное предложение на русском языке для проверки. ".repeat( + 30 + ), + "Hello 👋 🌍 🚀 naïve café Ünïcödé ".repeat(60), + "😀".repeat(500) + "世".repeat(500) + "é".repeat(500), +]; +const corpusTexts = [ + ...corpusFiles.map((file) => readFileSync(file, "utf8")), + ...multilingual, +]; + +const encodings = { + o200k_base: o200k, + cl100k_base: cl100k, + p50k_base: p50k, + claude, +}; + +/** Pieces that stress the merge, from several generators. */ +function syntheticPieces(encoding: Encoding, seed: number): Uint8Array[] { + const random = seededRandom(seed); + const int = (max: number) => Math.floor(random() * max); + const pieces: Uint8Array[] = []; + + // Random bytes, mostly invalid UTF-8. + for (let i = 0; i < 400; i++) { + pieces.push(Uint8Array.from({ length: 1 + int(200) }, () => int(256))); + } + + // One repeated byte at every length: ties everywhere. + for (const byte of [0x5f, 0x20, 0x2e, 0x61, 0x0a, 0x30, 0xff, 0xe4]) { + for (let length = 1; length <= 300; length++) { + pieces.push(new Uint8Array(length).fill(byte)); + } + } + + // Two or three letters: many equal ranks. + for (let i = 0; i < 300; i++) { + const alphabet = 2 + int(2); + pieces.push( + Uint8Array.from({ length: 1 + int(400) }, () => 97 + int(alphabet)) + ); + } + + // Concatenated vocabulary tokens, cut at a random byte: merge-rich input. + const stringTokens = Object.keys(encoding.stringEncoder); + const binaryTokens = encoding.binaryEncoder.map(([bytes]) => bytes); + for (let i = 0; i < 300; i++) { + const target = 1 + int(i % 20 === 0 ? 2000 : 400); + const bytes: number[] = []; + while (bytes.length < target) { + const token = + int(8) === 0 + ? binaryTokens[int(binaryTokens.length)]! + : textEncoder.encode(stringTokens[int(stringTokens.length)]!); + bytes.push(...token); + } + pieces.push(Uint8Array.from(bytes.slice(0, target))); + } + + // Random code points from scripts with 2, 3 and 4 byte UTF-8. + const ranges: Array<[number, number]> = [ + [0xa0, 0xff], + [0x600, 0x6ff], + [0x4e00, 0x9fff], + [0x1f300, 0x1f64f], + ]; + for (let i = 0; i < 200; i++) { + const [low, high] = ranges[int(ranges.length)]!; + const chars = Array.from({ length: 1 + int(300) }, () => + String.fromCodePoint(low + int(high - low + 1)) + ); + pieces.push(textEncoder.encode(chars.join(""))); + } + + return pieces; +} + +/** Every distinct pattern match in the corpus, as UTF-8 bytes. */ +function corpusPieces(encoding: Encoding): Uint8Array[] { + const pattern = new RegExp(encoding.pat_str, "ug"); + const seen = new Set(); + for (const text of corpusTexts) { + pattern.lastIndex = 0; + let match: RegExpExecArray | null; + while ((match = pattern.exec(text)) !== null) seen.add(match[0]); + } + return [...seen].map((piece) => textEncoder.encode(piece)); +} + +describe("differential: bytePairMerge matches the original algorithm", () => { + let total = 0; + let identical = 0; + let skipped = 0; + + afterAll(() => { + console.log( + `differential: encodings=${Object.keys(encodings).length} pieces=${total} identical=${identical} skipped_bom=${skipped}` + ); + }); + + describe.each(Object.entries(encodings))("%s", (name, encoding) => { + const tokenizer = new Tokenizer(encoding); + const lookup = rankLookup(encoding); + const stringRanks = tokenizer["stringRankEncoder"]; + const binaryIndex = tokenizer["binaryFirstByteIndex"]; + + const check = (kind: string, pieces: Uint8Array[]) => { + const mismatches: string[] = []; + for (const piece of pieces) { + if (hasBom(piece)) { + skipped++; + continue; + } + total++; + const expected = legacyBytePairMerge(piece, lookup); + const actual = bytePairMerge(piece, stringRanks, binaryIndex); + if ( + actual.length === expected.length && + actual.every((rank, i) => rank === expected[i]) + ) { + identical++; + } else if (mismatches.length < 3) { + mismatches.push( + `${Buffer.from(piece).toString("hex").slice(0, 80)} expected ${expected.slice(0, 8)} got ${actual.slice(0, 8)}` + ); + } + } + console.log(`${name} ${kind}: pieces=${pieces.length}`); + expect(mismatches).toEqual([]); + }; + + test("random and adversarial pieces", () => { + check("synthetic", syntheticPieces(encoding, 1234)); + }, 60_000); + + test("pieces of real text", () => { + const pieces = corpusPieces(encoding); + expect(pieces.length).toBeGreaterThan(1_000); + check("corpus", pieces); + }, 60_000); + + test("empty and single-byte pieces", () => { + check("tiny", [ + new Uint8Array(0), + ...Array.from({ length: 256 }, (_, i) => Uint8Array.of(i)), + ]); + }); + }); +}); + +describe("end to end: long single-piece inputs match tiktoken", () => { + const references: Array<[string, TiktokenEncoding, Encoding]> = [ + ["o200k_base", "o200k_base", o200k], + ["cl100k_base", "cl100k_base", cl100k], + ["p50k_base", "p50k_base", p50k], + ]; + const runs = ["_", " ", ".", "a", "ab", "\n", "世", "😀", "é", "1"]; + const lengths = [ + 1, 2, 3, 5, 31, 100, 127, 128, 129, 255, 256, 257, 1_000, 4_097, + ]; + let compared = 0; + + afterAll(() => { + console.log(`end to end: inputs=${compared} identical=${compared}`); + }); + + describe.each(references)("%s", (_name, tiktokenName, encoding) => { + let reference: Tiktoken; + const tokenizer = new Tokenizer(encoding); + + beforeAll(() => { + reference = get_encoding(tiktokenName); + }); + afterAll(() => reference.free()); + + test("runs of one repeated character", () => { + for (const run of runs) { + for (const length of lengths) { + const text = run.repeat(length); + expect(tokenizer.encode(text)).toEqual( + Array.from(reference.encode(text)) + ); + compared++; + } + } + }, 60_000); + }); }); diff --git a/src/tokenizer.ts b/src/tokenizer.ts index 4c89ce7..62a5c13 100644 --- a/src/tokenizer.ts +++ b/src/tokenizer.ts @@ -478,15 +478,18 @@ function escapeRegex(str: string): string { } /** - * Hyper-optimized BPE merge - * Matches gpt-tokenizer's approach exactly but with our optimizations + * Byte pair merge in O(n log n) for a piece of n bytes. + * + * The parts form a linked list, and a min-heap of adjacent pairs ordered by + * (rank, position) picks the next merge: the lowest rank first, the leftmost + * pair among equal ranks. Exported for tests; not part of the public API. */ -function bytePairMerge( +export function bytePairMerge( piece: Uint8Array, stringRanks: Record, binaryFirstByteIndex: Array | null> ): number[] { - const len = piece.length; + const n = piece.length; // Ultra-fast inlined rank lookup const getRankForSlice = (slice: Uint8Array): number => { @@ -506,65 +509,142 @@ function bytePairMerge( return NO_RANK; }; - const starts: number[] = []; - const ranks: number[] = []; - - // Inline getRank helper - const getRank = ( - startIndex: number, - pairStart?: number, - pairEnd?: number - ): number => { - if (pairStart === undefined) pairStart = starts[startIndex]; - if (pairEnd === undefined) pairEnd = starts[startIndex + 2]; - if (pairEnd === undefined) return NO_RANK; - return getRankForSlice(piece.subarray(pairStart, pairEnd)); + const output: number[] = []; + if (n === 0) return output; + + // A part is named by the index of its first byte. next[i] is where the part + // after i starts (n for the last part), or -1 once i was merged into the part + // on its left. partRank[i] is the rank of i's part if it was merged. + const next = new Int32Array(n + 1); + const prev = new Int32Array(n); + const partRank = new Uint32Array(n).fill(NO_RANK); + for (let i = 0; i < n; i++) { + next[i] = i + 1; + prev[i] = i - 1; + } + next[n] = -1; + + // Heap entry: merge the part at `pos` with the part after it; together they + // span [pos, end). An entry is live while pos still has a right neighbour + // whose own right edge is `end`. Every merge adds at most two entries and + // removes one, so 2n slots are enough. + const capacity = 2 * n; + const heapRank = new Uint32Array(capacity); + const heapPos = new Int32Array(capacity); + const heapEnd = new Int32Array(capacity); + let size = 0; + + const siftUp = (index: number): void => { + const rank = heapRank[index]!; + const pos = heapPos[index]!; + const end = heapEnd[index]!; + while (index > 0) { + const parent = (index - 1) >> 1; + const parentRank = heapRank[parent]!; + if ( + parentRank < rank || + (parentRank === rank && heapPos[parent]! < pos) + ) { + break; + } + heapRank[index] = parentRank; + heapPos[index] = heapPos[parent]!; + heapEnd[index] = heapEnd[parent]!; + index = parent; + } + heapRank[index] = rank; + heapPos[index] = pos; + heapEnd[index] = end; }; - // Initialize starts and ranks - for (let i = 0; i <= len; i++) { - starts[i] = i; - if (i < len - 1) { - ranks[i] = getRank(i, i, i + 2); - } else { - ranks[i] = NO_RANK; + const siftDown = (index: number): void => { + const rank = heapRank[index]!; + const pos = heapPos[index]!; + const end = heapEnd[index]!; + const half = size >> 1; + while (index < half) { + let child = 2 * index + 1; + const right = child + 1; + if ( + right < size && + (heapRank[right]! < heapRank[child]! || + (heapRank[right] === heapRank[child] && + heapPos[right]! < heapPos[child]!)) + ) { + child = right; + } + const childRank = heapRank[child]!; + if (rank < childRank || (rank === childRank && pos < heapPos[child]!)) { + break; + } + heapRank[index] = childRank; + heapPos[index] = heapPos[child]!; + heapEnd[index] = heapEnd[child]!; + index = child; + } + heapRank[index] = rank; + heapPos[index] = pos; + heapEnd[index] = end; + }; + + // Only a pair whose rank is below NO_RANK can merge. + for (let i = 0; i + 2 <= n; i++) { + const rank = getRankForSlice(piece.subarray(i, i + 2)); + if (rank < NO_RANK) { + heapRank[size] = rank; + heapPos[size] = i; + heapEnd[size] = i + 2; + size++; } } + for (let i = (size >> 1) - 1; i >= 0; i--) siftDown(i); + + while (size > 0) { + const rank = heapRank[0]!; + const pos = heapPos[0]!; + const end = heapEnd[0]!; + size--; + if (size > 0) { + heapRank[0] = heapRank[size]!; + heapPos[0] = heapPos[size]!; + heapEnd[0] = heapEnd[size]!; + siftDown(0); + } - // Main merge loop - while (starts.length > 1) { - let minRank = NO_RANK; - let minIdx = -1; - - // Find minimum rank - const ranksLen = ranks.length - 1; - for (let i = 0; i < ranksLen; i++) { - const rank = ranks[i]!; - if (rank < minRank) { - minRank = rank!; - minIdx = i; + const right = next[pos]!; + if (right < 0 || next[right] !== end) continue; + + next[pos] = end; + next[right] = -1; + partRank[pos] = rank; + if (end < n) { + prev[end] = pos; + const afterEnd = next[end]!; + const rankWithNext = getRankForSlice(piece.subarray(pos, afterEnd)); + if (rankWithNext < NO_RANK) { + heapRank[size] = rankWithNext; + heapPos[size] = pos; + heapEnd[size] = afterEnd; + siftUp(size++); } } - - if (minRank === NO_RANK || minIdx === -1) break; - - // Remove elements - starts.splice(minIdx + 1, 1); - ranks.splice(minIdx, 1); - - // Update ranks - ranks[minIdx] = getRank(minIdx); - if (minIdx > 0) { - ranks[minIdx - 1] = getRank(minIdx - 1); + const left = prev[pos]!; + if (left >= 0) { + const rankWithLeft = getRankForSlice(piece.subarray(left, end)); + if (rankWithLeft < NO_RANK) { + heapRank[size] = rankWithLeft; + heapPos[size] = left; + heapEnd[size] = end; + siftUp(size++); + } } } - // Build output - const output: number[] = []; - for (let i = 0; i < starts.length - 1; i++) { - const pairStart = starts[i]; - const pairEnd = starts[i + 1]; - const rank = getRankForSlice(piece.subarray(pairStart, pairEnd)); + for (let i = 0; i < n; i = next[i]!) { + let rank = partRank[i]!; + if (rank === NO_RANK) { + rank = getRankForSlice(piece.subarray(i, i + 1)); + } if (rank !== NO_RANK) { output.push(rank); } From b3241f007f2b0f31d722932070ef111b3297f263 Mon Sep 17 00:00:00 2001 From: Thomas Kosiewski Date: Tue, 3 Nov 2026 11:07:06 +0000 Subject: [PATCH 3/8] ci: typecheck the README example against the packed tarball Packs the package, installs it into a fresh project and typechecks the README example under bundler, node16 (ESM and CJS) and node10 resolution. It fails today: index.d.ts and sdk.d.ts each declare their own Tokenizer class, so a Tokenizer from 'ai-tokenizer' is not assignable to the one 'ai-tokenizer/sdk' expects. See #1. Change-Id: I60e43f6f077fdee94e19394b9bb943fb4910edc4 Signed-off-by: Thomas Kosiewski --- .github/workflows/ci.yml | 5 ++ scripts/check-package-types.ts | 154 +++++++++++++++++++++++++++++++++ 2 files changed, 159 insertions(+) create mode 100644 scripts/check-package-types.ts diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 9322aaa..3981a28 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -33,3 +33,8 @@ jobs: - name: Build run: bun run build --dts=false + - name: Build with type declarations + run: bun run build + + - name: Check packed types + run: bun scripts/check-package-types.ts diff --git a/scripts/check-package-types.ts b/scripts/check-package-types.ts new file mode 100644 index 0000000..ada3946 --- /dev/null +++ b/scripts/check-package-types.ts @@ -0,0 +1,154 @@ +/** + * Typechecks the README example against the packed tarball, the way a + * consumer installs it. Run `bun run build` first: it needs the declarations. + * + * Regression check for https://github.com/coder/ai-tokenizer/issues/1: a + * Tokenizer from "ai-tokenizer" must be assignable to the Tokenizer that + * `count` from "ai-tokenizer/sdk" expects. + * + * bun scripts/check-package-types.ts + */ + +import { execFileSync } from "node:child_process"; +import { mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; + +const repoRoot = join(import.meta.dir, ".."); +const installedVersion = (name: string): string => + JSON.parse( + readFileSync(join(repoRoot, "node_modules", name, "package.json"), "utf8") + ).version; + +// The README example, with `messages` and `tools` declared. +const example = `import Tokenizer, { models } from "ai-tokenizer"; +import { count } from "ai-tokenizer/sdk"; +import * as encoding from "ai-tokenizer/encoding"; +import type { ModelMessage, ToolSet } from "ai"; + +declare const messages: ModelMessage[]; +declare const tools: ToolSet; + +const model = models["openai/gpt-5"]; +const tokenizer = new Tokenizer(encoding[model.encoding]); +export const result = count({ tokenizer, model, messages, tools }); +export const total: number = result.total; +`; + +const resolutions: Array<{ + name: string; + file: string; + module: string; + moduleResolution: string; +}> = [ + { + name: "bundler", + file: "index.ts", + module: "esnext", + moduleResolution: "bundler", + }, + { + name: "node16 (ESM)", + file: "index.ts", + module: "node16", + moduleResolution: "node16", + }, + { + name: "node16 (CJS)", + file: "index.cts", + module: "node16", + moduleResolution: "node16", + }, + { + name: "node10", + file: "index.ts", + module: "commonjs", + moduleResolution: "node10", + }, +]; + +const run = (cwd: string, command: string, args: string[]): string => + execFileSync(command, args, { cwd, encoding: "utf8", stdio: "pipe" }); + +const workDir = mkdtempSync(join(tmpdir(), "ai-tokenizer-types-")); +let failed = 0; +try { + const packDir = join(workDir, "pack"); + const consumerDir = join(workDir, "consumer"); + run(workDir, "mkdir", [packDir, consumerDir]); + + const tarball = join( + packDir, + JSON.parse( + run(repoRoot, "npm", ["pack", "--json", "--pack-destination", packDir]) + )[0].filename + ); + + writeFileSync( + join(consumerDir, "package.json"), + JSON.stringify({ name: "consumer", private: true, type: "module" }) + ); + run(consumerDir, "npm", [ + "install", + "--no-audit", + "--no-fund", + `typescript@${installedVersion("typescript")}`, + `ai@${installedVersion("ai")}`, + "@types/node", + tarball, + ]); + + for (const resolution of resolutions) { + writeFileSync(join(consumerDir, resolution.file), example); + const tsconfig = join( + consumerDir, + `tsconfig.${resolution.moduleResolution}.json` + ); + writeFileSync( + tsconfig, + JSON.stringify({ + compilerOptions: { + strict: true, + noEmit: true, + target: "es2022", + module: resolution.module, + moduleResolution: resolution.moduleResolution, + types: ["node"], + }, + files: [resolution.file], + }) + ); + + let output: string; + try { + output = run(consumerDir, "node", [ + join("node_modules", "typescript", "bin", "tsc"), + "-p", + tsconfig, + ]); + } catch (error) { + output = String((error as { stdout?: string }).stdout ?? error); + } + // Declarations of other packages are not ours to check. + const errors = output + .split("\n") + .filter((line) => + /^(index\.c?ts|node_modules\/ai-tokenizer\/)/.test(line) + ); + + if (errors.length === 0) { + console.log(`PASS ${resolution.name}`); + } else { + failed++; + console.log(`FAIL ${resolution.name}`); + for (const line of errors) console.log(` ${line}`); + } + } +} finally { + rmSync(workDir, { recursive: true, force: true }); +} + +console.log( + `package types: ${resolutions.length - failed} of ${resolutions.length} resolutions typecheck` +); +process.exit(failed === 0 ? 0 : 1); From 2eafae8f329084cba948162fa7e5c5a2e24dbb18 Mon Sep 17 00:00:00 2001 From: Thomas Kosiewski Date: Tue, 3 Nov 2026 11:09:10 +0000 Subject: [PATCH 4/8] build: share one Tokenizer declaration between entry points index and sdk were two separate tsdown builds, so each bundled its own copy of the Tokenizer class and its .d.ts. TypeScript treats classes with private members as nominal, so a Tokenizer from 'ai-tokenizer' was not assignable to the one 'ai-tokenizer/sdk' expects. Build both entries together so the bundler emits the class and its declaration once, in a shared chunk. Fixes #1. PR #2 found the cause and proposed externalising a shared tokenizer module, which also needs a tsdown upgrade and a new ./tokenizer export. A single multi-entry build needs neither. Change-Id: I849322c12091fa284281a23a8cfa0ded1833712e Signed-off-by: Thomas Kosiewski --- tsdown.config.ts | 47 ++++++++++++++++++++++++----------------------- 1 file changed, 24 insertions(+), 23 deletions(-) diff --git a/tsdown.config.ts b/tsdown.config.ts index f01d5ac..0387848 100644 --- a/tsdown.config.ts +++ b/tsdown.config.ts @@ -1,25 +1,26 @@ import { defineConfig } from "tsdown"; -export default defineConfig([{ - entry: "./src/index.ts", - outDir: "./dist", - dts: true, - format: ["esm", "cjs"], - target: "es2020", - sourcemap: true, -}, { - entry: "./src/sdk.ts", - outDir: "./dist", - dts: true, - format: ["esm", "cjs"], - target: "es2020", - sourcemap: true, -}, { - entry: "./src/encoding/*", - outDir: "./dist/encoding/", - format: ["esm", "cjs"], - target: "es2020", - dts: true, - sourcemap: false, - minify: false, -}]); +export default defineConfig([ + { + // One build for both entries, so code they share (the Tokenizer class and + // its declaration) lands in one chunk instead of a copy in each entry. + entry: { + index: "./src/index.ts", + sdk: "./src/sdk.ts", + }, + outDir: "./dist", + dts: true, + format: ["esm", "cjs"], + target: "es2020", + sourcemap: true, + }, + { + entry: "./src/encoding/*", + outDir: "./dist/encoding/", + format: ["esm", "cjs"], + target: "es2020", + dts: true, + sourcemap: false, + minify: false, + }, +]); From a93d79321318ed4561b5d501fe986a117d9c42eb Mon Sep 17 00:00:00 2001 From: Thomas Kosiewski Date: Tue, 3 Nov 2026 12:36:11 +0000 Subject: [PATCH 5/8] ci: fail the package types check on any TypeScript diagnostic The check only counted error lines from the example file and ai-tokenizer. An option error in the temporary tsconfig makes tsc skip the type check of the example, and that error line was dropped, so the check printed PASS. Count every 'error TS' line except those in other packages' node_modules, and treat a failed tsc run with no diagnostic as an error. Change-Id: I66a0c4e16331914e044d902178a91d56074c8132 Signed-off-by: Thomas Kosiewski --- scripts/check-package-types.ts | 25 +++++++++++++++++-------- 1 file changed, 17 insertions(+), 8 deletions(-) diff --git a/scripts/check-package-types.ts b/scripts/check-package-types.ts index ada3946..3f93a3e 100644 --- a/scripts/check-package-types.ts +++ b/scripts/check-package-types.ts @@ -119,7 +119,8 @@ try { }) ); - let output: string; + let output = ""; + let threw = false; try { output = run(consumerDir, "node", [ join("node_modules", "typescript", "bin", "tsc"), @@ -127,14 +128,22 @@ try { tsconfig, ]); } catch (error) { - output = String((error as { stdout?: string }).stdout ?? error); + threw = true; + const { stdout, stderr } = error as { stdout?: string; stderr?: string }; + output = `${stdout ?? ""}${stderr ?? ""}` || String(error); } - // Declarations of other packages are not ours to check. - const errors = output - .split("\n") - .filter((line) => - /^(index\.c?ts|node_modules\/ai-tokenizer\/)/.test(line) - ); + // Every diagnostic counts, including option errors in the tsconfig, which + // make tsc skip the check of the example. Only declarations of other + // packages are not ours to check. A failed run with no diagnostic at all + // is an error too. + const lines = output.split("\n").filter((line) => line.trim() !== ""); + const diagnostics = lines.filter((line) => /error TS\d+/.test(line)); + const errors = + threw && diagnostics.length === 0 + ? lines + : diagnostics.filter( + (line) => !/^node_modules\/(?!ai-tokenizer\/)/.test(line) + ); if (errors.length === 0) { console.log(`PASS ${resolution.name}`); From 9bca535c911497eacef7b3fb362b1373b5fce22a Mon Sep 17 00:00:00 2001 From: Thomas Kosiewski Date: Tue, 3 Nov 2026 12:36:12 +0000 Subject: [PATCH 6/8] perf: assert the merge heap never overflows Typed arrays drop writes past their end, so a wrong capacity bound would give wrong tokens instead of an error. Route the two pushes made during merging through one function that throws a RangeError when the heap is full. Change-Id: Ib0753ac41e67c820320340b6a46f2b73dc88edf8 Signed-off-by: Thomas Kosiewski --- src/tokenizer.ts | 23 +++++++++++------------ 1 file changed, 11 insertions(+), 12 deletions(-) diff --git a/src/tokenizer.ts b/src/tokenizer.ts index 62a5c13..615e107 100644 --- a/src/tokenizer.ts +++ b/src/tokenizer.ts @@ -557,6 +557,15 @@ export function bytePairMerge( heapEnd[index] = end; }; + // Typed arrays drop writes past their end, so check the bound here. + const push = (rank: number, pos: number, end: number): void => { + if (size >= capacity) throw new RangeError("bytePairMerge: heap is full"); + heapRank[size] = rank; + heapPos[size] = pos; + heapEnd[size] = end; + siftUp(size++); + }; + const siftDown = (index: number): void => { const rank = heapRank[index]!; const pos = heapPos[index]!; @@ -621,22 +630,12 @@ export function bytePairMerge( prev[end] = pos; const afterEnd = next[end]!; const rankWithNext = getRankForSlice(piece.subarray(pos, afterEnd)); - if (rankWithNext < NO_RANK) { - heapRank[size] = rankWithNext; - heapPos[size] = pos; - heapEnd[size] = afterEnd; - siftUp(size++); - } + if (rankWithNext < NO_RANK) push(rankWithNext, pos, afterEnd); } const left = prev[pos]!; if (left >= 0) { const rankWithLeft = getRankForSlice(piece.subarray(left, end)); - if (rankWithLeft < NO_RANK) { - heapRank[size] = rankWithLeft; - heapPos[size] = left; - heapEnd[size] = end; - siftUp(size++); - } + if (rankWithLeft < NO_RANK) push(rankWithLeft, left, end); } } From 0d28d1b7465b7ece2f7188824fa27b178f504923 Mon Sep 17 00:00:00 2001 From: Thomas Kosiewski Date: Tue, 3 Nov 2026 13:38:21 +0000 Subject: [PATCH 7/8] test: compare pieces that contain U+FEFF in the differential test The tokenizer now decodes with ignoreBOM, so its rank lookup and the test's agree on pieces that contain U+FEFF. Drop the skip for them. Change-Id: I75bfdefdf3eb11a5257521a810e4a8d4f8f0bb64 Signed-off-by: Thomas Kosiewski --- src/byte-pair-merge.test.ts | 14 +------------- 1 file changed, 1 insertion(+), 13 deletions(-) diff --git a/src/byte-pair-merge.test.ts b/src/byte-pair-merge.test.ts index 8d1ea70..b5d4561 100644 --- a/src/byte-pair-merge.test.ts +++ b/src/byte-pair-merge.test.ts @@ -182,13 +182,6 @@ function rankLookup(encoding: Encoding): (slice: Uint8Array) => number { }; } -// U+FEFF. Before the tokenizer decoded with ignoreBOM, a slice starting with -// these bytes was looked up without them, which the lookup above does not do. -const hasBom = (bytes: Uint8Array): boolean => - bytes.some( - (byte, i) => byte === 0xef && bytes[i + 1] === 0xbb && bytes[i + 2] === 0xbf - ); - /** Small seeded generator (mulberry32): the corpus is the same on every run. */ function seededRandom(seed: number): () => number { let state = seed >>> 0; @@ -314,11 +307,10 @@ function corpusPieces(encoding: Encoding): Uint8Array[] { describe("differential: bytePairMerge matches the original algorithm", () => { let total = 0; let identical = 0; - let skipped = 0; afterAll(() => { console.log( - `differential: encodings=${Object.keys(encodings).length} pieces=${total} identical=${identical} skipped_bom=${skipped}` + `differential: encodings=${Object.keys(encodings).length} pieces=${total} identical=${identical}` ); }); @@ -331,10 +323,6 @@ describe("differential: bytePairMerge matches the original algorithm", () => { const check = (kind: string, pieces: Uint8Array[]) => { const mismatches: string[] = []; for (const piece of pieces) { - if (hasBom(piece)) { - skipped++; - continue; - } total++; const expected = legacyBytePairMerge(piece, lookup); const actual = bytePairMerge(piece, stringRanks, binaryIndex); From 1125a602644f99e8822f42b915f799c1822dcc8e Mon Sep 17 00:00:00 2001 From: Thomas Kosiewski Date: Wed, 4 Nov 2026 01:40:02 +0000 Subject: [PATCH 8/8] ci: typecheck a claude_v5 tokenizer through the packed types Encodings with a kind are typed as CodecEncoding, a second member of the Tokenizer constructor's parameter type. Pass one through the shared Tokenizer declaration in the packed-tarball check too. Change-Id: Ie0977de7912f32be1156c837d222c49863477acc Signed-off-by: Thomas Kosiewski --- scripts/check-package-types.ts | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/scripts/check-package-types.ts b/scripts/check-package-types.ts index 3f93a3e..06b7186 100644 --- a/scripts/check-package-types.ts +++ b/scripts/check-package-types.ts @@ -33,6 +33,11 @@ const model = models["openai/gpt-5"]; const tokenizer = new Tokenizer(encoding[model.encoding]); export const result = count({ tokenizer, model, messages, tools }); export const total: number = result.total; +export const claudeResult = count({ + tokenizer: new Tokenizer(encoding.claude_v5), + model, + messages, +}); `; const resolutions: Array<{