diff --git a/.github/workflows/build-maturin.yml b/.github/workflows/build-maturin.yml index db63613..23c370a 100644 --- a/.github/workflows/build-maturin.yml +++ b/.github/workflows/build-maturin.yml @@ -75,7 +75,7 @@ jobs: if: ${{ matrix.platform.ft }} with: target: ${{ matrix.platform.target }} - args: --release --out dist -i python3.13t + args: --release --out dist -i python3.14t sccache: "false" manylinux: auto before-script-linux: | @@ -122,7 +122,7 @@ jobs: uses: PyO3/maturin-action@v1 with: target: ${{ matrix.platform.target }} - args: --release --out dist -i python3.13t + args: --release --out dist -i python3.14t sccache: "false" manylinux: musllinux_1_2 - name: Upload wheels @@ -175,7 +175,7 @@ jobs: uses: PyO3/maturin-action@v1 with: target: ${{ matrix.platform.target }} - args: --release --out dist -i python3.13t + args: --release --out dist -i python3.14t sccache: "false" - name: Upload wheels uses: actions/upload-artifact@v4 diff --git a/.github/workflows/build-release-fastokens.yml b/.github/workflows/build-release-fastokens.yml index ede70cc..98f4e56 100644 --- a/.github/workflows/build-release-fastokens.yml +++ b/.github/workflows/build-release-fastokens.yml @@ -48,9 +48,9 @@ jobs: - uses: actions/checkout@v6 - uses: dtolnay/rust-toolchain@stable - name: Verify package - run: cargo publish -p fastokens --dry-run --locked + run: cargo publish -p fastokens-b10 --dry-run --locked - name: Publish if: ${{ github.event_name == 'workflow_dispatch' && fromJSON(github.event.inputs.publish) }} env: CARGO_REGISTRY_TOKEN: ${{ secrets.CARGO_REGISTRY_TOKEN }} - run: cargo publish -p fastokens --locked + run: cargo publish -p fastokens-b10 --locked diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 558fe4e..3f0ab59 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -5,9 +5,11 @@ on: branches: [main] pull_request: branches: [main] + workflow_dispatch: env: CARGO_TERM_COLOR: always + HF_TOKEN: ${{ secrets.HF_TOKEN }} jobs: test: @@ -49,6 +51,9 @@ jobs: - name: Build and install wheel run: maturin build --release --out dist && pip install dist/*.whl + - name: Install Python test dependencies + run: pip install pytest tiktoken + - name: Smoke test run: | python - <<'EOF' @@ -61,6 +66,9 @@ jobs: print("smoke test passed") EOF + - name: Run Python tests + run: pytest python/tests -v + patch-transformers: runs-on: ubuntu-latest strategy: @@ -85,7 +93,7 @@ jobs: run: maturin build --release --out dist && pip install dist/*.whl - name: Install transformers ${{ matrix.transformers }} and pytest - run: pip install transformers==${{ matrix.transformers }} pytest + run: pip install transformers==${{ matrix.transformers }} pytest jinja2 - name: Test patch_transformers run: pytest python/tests/test_patch_transformers.py -v diff --git a/.github/workflows/python-ci.yml b/.github/workflows/python-ci.yml new file mode 100644 index 0000000..3fb20db --- /dev/null +++ b/.github/workflows/python-ci.yml @@ -0,0 +1,62 @@ +name: Python CI + +on: + push: + branches: [main] + pull_request: + branches: [main] + workflow_dispatch: + +env: + CARGO_TERM_COLOR: always + HF_TOKEN: ${{ secrets.HF_TOKEN }} + +permissions: + contents: read + +jobs: + test: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v6 + + - uses: dtolnay/rust-toolchain@stable + + - uses: Swatinem/rust-cache@v2 + + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + cache: pip + + - name: Cache Hugging Face tokenizers + uses: actions/cache@v4 + with: + path: | + ~/.cache/huggingface/hub + ~/.cache/huggingface/xet + ~/.cache/tokenizers + key: ${{ runner.os }}-py3.12-hf-tokenizers-${{ hashFiles('pyproject.toml', 'python/Cargo.toml', 'Cargo.lock', '.github/workflows/python-ci.yml') }} + restore-keys: | + ${{ runner.os }}-py3.12-hf-tokenizers- + ${{ runner.os }}-hf-tokenizers- + + - name: Configure Hugging Face token + run: | + if [ -n "${HF_TOKEN:-}" ]; then + mkdir -p ~/.cache/huggingface + printf '%s' "$HF_TOKEN" > ~/.cache/huggingface/token + fi + + - name: Install Python test dependencies + run: python -m pip install --upgrade pip maturin pytest transformers==4.57.1 jinja2 + + - name: Build and install fastokens + run: maturin build --release --out dist && pip install dist/*.whl + + - name: Run Python tests + run: pytest python/tests -v + + - name: Remove Hugging Face token before cache save + if: always() + run: rm -f ~/.cache/huggingface/token diff --git a/.gitignore b/.gitignore index ee53101..cfa605a 100644 --- a/.gitignore +++ b/.gitignore @@ -4,4 +4,5 @@ __pycache__/ *.pyd dist/ -target/ \ No newline at end of file +target/ +vendored_tokenizers/private_* \ No newline at end of file diff --git a/Cargo.lock b/Cargo.lock index f41b036..57fd4e7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -31,6 +31,15 @@ dependencies = [ "memchr", ] +[[package]] +name = "android_system_properties" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311" +dependencies = [ + "libc", +] + [[package]] name = "anstream" version = "0.6.21" @@ -183,6 +192,17 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" +[[package]] +name = "chrono" +version = "0.4.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1aa79e62e7697b8e29b513a68abacf485adcd1fe8284a4316c5ae868e6633327" +dependencies = [ + "iana-time-zone", + "num-traits", + "windows-link", +] + [[package]] name = "clap" version = "4.5.58" @@ -270,6 +290,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "core-foundation-sys" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" + [[package]] name = "crc32fast" version = "1.5.0" @@ -450,6 +476,12 @@ version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "34aa73646ffb006b8f5147f3dc182bd4bcb190227ce861fc4a4844bf8e3cb2c0" +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + [[package]] name = "esaxx-rs" version = "0.1.10" @@ -471,18 +503,22 @@ dependencies = [ ] [[package]] -name = "fastokens" -version = "0.2.1" +name = "fastokens-b10" +version = "0.2.6" dependencies = [ "anyhow", + "chrono", "clap", "csv", "daachorse", "fancy-regex", "hf-hub", "icu_normalizer", + "icu_properties", "indicatif 0.18.4", "memchr", + "minijinja", + "minijinja-contrib", "pcre2", "rayon", "serde", @@ -494,10 +530,13 @@ dependencies = [ [[package]] name = "fastokens-python" -version = "0.2.1" +version = "0.2.6" dependencies = [ - "fastokens", + "fastokens-b10", + "numpy", "pyo3", + "pyo3-async-runtimes", + "pythonize", "rayon", "serde_json", ] @@ -533,6 +572,21 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "futures" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b147ee9d1f6d097cef9ce628cd2ee62288d963e16fb287bd9286455b241382d" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + [[package]] name = "futures-channel" version = "0.3.32" @@ -540,6 +594,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "07bbe89c50d7a535e539b8c17bc0b49bdb77747034daa8087407d655f3f7cc1d" dependencies = [ "futures-core", + "futures-sink", ] [[package]] @@ -548,6 +603,17 @@ version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" +[[package]] +name = "futures-executor" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baf29c38818342a3b26b5b923639e7b1f4a61fc5e76102d4b1981c6dc7a7579d" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + [[package]] name = "futures-io" version = "0.3.32" @@ -583,6 +649,7 @@ version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" dependencies = [ + "futures-channel", "futures-core", "futures-io", "futures-macro", @@ -620,6 +687,12 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + [[package]] name = "heck" version = "0.5.0" @@ -746,6 +819,30 @@ dependencies = [ "tracing", ] +[[package]] +name = "iana-time-zone" +version = "0.1.65" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e31bc9ad994ba00e440a8aa5c9ef0ec67d5cb5e5cb0cc7f8b744a35b389cc470" +dependencies = [ + "android_system_properties", + "core-foundation-sys", + "iana-time-zone-haiku", + "js-sys", + "log", + "wasm-bindgen", + "windows-core", +] + +[[package]] +name = "iana-time-zone-haiku" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f" +dependencies = [ + "cc", +] + [[package]] name = "icu_collections" version = "2.1.1" @@ -857,6 +954,16 @@ dependencies = [ "icu_properties", ] +[[package]] +name = "indexmap" +version = "2.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +dependencies = [ + "equivalent", + "hashbrown", +] + [[package]] name = "indicatif" version = "0.17.11" @@ -999,12 +1106,28 @@ version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "670fdfda89751bc4a84ac13eaa63e205cf0fd22b4c9a5fbfa085b63c1f1d3a30" +[[package]] +name = "matrixmultiply" +version = "0.3.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a06de3016e9fae57a36fd14dba131fccf49f74b40b7fbdb472f96e361ec71a08" +dependencies = [ + "autocfg", + "rawpointer", +] + [[package]] name = "memchr" version = "2.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" +[[package]] +name = "memo-map" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38d1115007560874e373613744c6fba374c17688327a71c1476d1a5954cc857b" + [[package]] name = "memoffset" version = "0.9.1" @@ -1014,6 +1137,27 @@ dependencies = [ "autocfg", ] +[[package]] +name = "minijinja" +version = "2.21.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb3d648e68cea56d9858d535ee28f9538404e2dd8cb08ed0bd05dca379477f39" +dependencies = [ + "indexmap", + "memo-map", + "serde", +] + +[[package]] +name = "minijinja-contrib" +version = "2.21.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85342f6fac0be8ccd5bd00d9066be538f34f393f577b75d81b17c8398a6b43bb" +dependencies = [ + "minijinja", + "serde", +] + [[package]] name = "minimal-lexical" version = "0.2.1" @@ -1063,6 +1207,21 @@ dependencies = [ "syn", ] +[[package]] +name = "ndarray" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "882ed72dce9365842bf196bdeedf5055305f11fc8c03dee7bb0194a6cad34841" +dependencies = [ + "matrixmultiply", + "num-complex", + "num-integer", + "num-traits", + "portable-atomic", + "portable-atomic-util", + "rawpointer", +] + [[package]] name = "nom" version = "7.1.3" @@ -1073,12 +1232,55 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-integer" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + [[package]] name = "number_prefix" version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3" +[[package]] +name = "numpy" +version = "0.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29f1dee9aa8d3f6f8e8b9af3803006101bb3653866ef056d530d53ae68587191" +dependencies = [ + "libc", + "ndarray", + "num-complex", + "num-integer", + "num-traits", + "pyo3", + "pyo3-build-config", + "rustc-hash", +] + [[package]] name = "once_cell" version = "1.21.3" @@ -1177,6 +1379,15 @@ version = "1.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" +[[package]] +name = "portable-atomic-util" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a106d1259c23fac8e543272398ae0e3c0b8d33c88ed73d0cc71b0f1d902618" +dependencies = [ + "portable-atomic", +] + [[package]] name = "potential_utf" version = "0.1.4" @@ -1221,6 +1432,19 @@ dependencies = [ "unindent", ] +[[package]] +name = "pyo3-async-runtimes" +version = "0.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d73cc6b1b7d8b3cef02101d37390dbdfe7e450dfea14921cae80a9534ba59ef2" +dependencies = [ + "futures", + "once_cell", + "pin-project-lite", + "pyo3", + "tokio", +] + [[package]] name = "pyo3-build-config" version = "0.25.1" @@ -1266,6 +1490,16 @@ dependencies = [ "syn", ] +[[package]] +name = "pythonize" +version = "0.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "597907139a488b22573158793aa7539df36ae863eba300c75f3a0d65fc475e27" +dependencies = [ + "pyo3", + "serde", +] + [[package]] name = "quinn" version = "0.11.9" @@ -1365,6 +1599,12 @@ dependencies = [ "getrandom 0.3.4", ] +[[package]] +name = "rawpointer" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" + [[package]] name = "rayon" version = "1.11.0" @@ -1581,6 +1821,7 @@ version = "1.0.149" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "83fc039473c5595ace860d8c4fafa220ff474b3fc6bfdb4293327f1a37e94d86" dependencies = [ + "indexmap", "itoa", "memchr", "serde", @@ -2188,12 +2429,65 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" +[[package]] +name = "windows-core" +version = "0.62.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" +dependencies = [ + "windows-implement", + "windows-interface", + "windows-link", + "windows-result", + "windows-strings", +] + +[[package]] +name = "windows-implement" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "windows-interface" +version = "0.59.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "windows-link" version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" +[[package]] +name = "windows-result" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-strings" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091" +dependencies = [ + "windows-link", +] + [[package]] name = "windows-sys" version = "0.52.0" diff --git a/Cargo.toml b/Cargo.toml index 304aa30..018c41e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -2,7 +2,7 @@ members = [".", "python"] [workspace.package] -version = "0.2.1" +version = "0.2.6" edition = "2024" license = "Apache-2.0" repository = "https://github.com/Atero-ai/fastokens" @@ -10,7 +10,7 @@ homepage = "https://github.com/Atero-ai/fastokens" documentation = "https://github.com/Atero-ai/fastokens/blob/main/README.md" [package] -name = "fastokens" +name = "fastokens-b10" version.workspace = true edition.workspace = true description = "Fast Tokenizer" @@ -18,6 +18,10 @@ license.workspace = true repository.workspace = true homepage.workspace = true documentation.workspace = true +exclude = ["vendored_tokenizers/**"] + +[lib] +name = "fastokens" [features] default = ["hf-hub"] @@ -27,14 +31,27 @@ hf-hub = ["dep:hf-hub"] daachorse = "1.0.0" fancy-regex = "0.17.0" memchr = "2" +chrono = { version = "0.4", default-features = false, features = ["clock"] } +minijinja = { version = "2.21.0", default-features = false, features = [ + "builtins", + "deserialization", + "loop_controls", + "macros", + "multi_template", + "preserve_order", + "serde", + "std_collections", +] } +minijinja-contrib = { version = "2.21.0", features = ["pycompat"] } pcre2 = "0.2" hf-hub = { version = "0.4.3", optional = true, default-features = false, features = ["rustls-tls", "ureq"] } rayon = "1.11.0" serde = { version = "1.0.228", features = ["derive"] } -serde_json = "1.0.149" +serde_json = { version = "1.0.149", features = ["preserve_order"] } strum = { version = "0.27.2", features = ["derive"] } thiserror = "2.0.18" icu_normalizer = "2" +icu_properties = "2" [dev-dependencies] anyhow = "1" @@ -62,3 +79,7 @@ required-features = ["hf-hub"] [[example]] name = "simple_bench" required-features = ["hf-hub"] + +[[example]] +name = "long_context_bench" +required-features = ["hf-hub"] diff --git a/docs/perf/specialized-split-scanners.md b/docs/perf/specialized-split-scanners.md new file mode 100644 index 0000000..bdaa56d --- /dev/null +++ b/docs/perf/specialized-split-scanners.md @@ -0,0 +1,122 @@ +# Specialized split scanners + +Hand-written pre-tokenization scanners for split regexes known ahead of time, +dispatched by byte-for-byte pattern match at `Split` construction +(`src/pre_tokenizers/fast_split.rs`). Technique adapted from +[gigatoken](https://github.com/marcelroed/gigatoken) (MIT): first-byte +dispatch + SWAR ASCII-letter runs + a packed 2-bit-per-codepoint Unicode class +table, no regex engine. + +Tier order in `Split::pre_tokenize` (fastest first): +specialized scanner → PCRE2 JIT (parallel + incremental cache) → fancy-regex. +The scanner activates only on an exact recognized pattern with `Isolated` +behavior and `invert == false`; everything else is untouched. + +## Measured (2026-07-25, EPYC 208-core box, `taskset`-pinned, Qwen3) + +The benchmark that matters for serving is **long-context, per-request, cold +cache**: each iteration encodes a *distinct* context so the PCRE2 tier's +incremental split cache (keyed on shared prefix with the previous call) never +hits. That cache helps streaming/growing sequences, not independent requests — +re-encoding one fixed string, as an earlier draft of this benchmark did, +overstates PCRE2 by making its split nearly free and hides the real per-request +cost. Reproduce with `examples/serving_split_bench.rs`, A/B by checking out +`main` (PCRE2) vs this branch (scanner). + +Full `encode()`, median, distinct inputs: + +| context | scanner (8 core) | PCRE2 (8 core) | scanner (1 core) | PCRE2 (1 core) | +|---------|-----------------:|---------------:|-----------------:|---------------:| +| 10k ch | 0.11 ms | 0.20 ms (1.8x) | 0.079 ms | 0.170 ms (2.15x) | +| 50k ch | 0.29 ms | 0.59 ms (2.0x) | 0.39 ms | *panic\** | +| 200k ch | 0.9 ms | 1.8 ms (1.9x) | 1.54 ms | *panic\** | + +**~1.8–2.15x faster full encode across 10k–200k char prompts**, single- or +multi-threaded. The gain is per-core split throughput; at full encode the BPE +stage is shared, so this is a split-bound-workload / serving-latency win, not a +bulk-throughput one. On a single re-encoded string (warm incremental cache) the +two tiers tie — that regime is not representative of independent requests. + +\* PCRE2's parallel authority-zone path (`split.rs`) panics on >16 KiB inputs +under `RAYON_NUM_THREADS=1` — a pre-existing bug on `main`, unrelated to this +change. The scanner replaces that path for recognized patterns, so Qwen3 +sidesteps it; other patterns still hit it. + +### Kimi (o200k-family) — `fast_split_o200k.rs` + +The Kimi (moonshotai K2) pattern uses `[\p{Han}]+` runs and `&&[^\p{Han}]` +class intersections that PCRE2 can't parse, so the PCRE2 tier runs a +hand-rewritten pattern with a per-codepoint `(?!\p{Han})` negative lookahead on +every letter-bracket char — slow on both ASCII and CJK. The scanner classifies +via one packed table load. A/B, distinct inputs, full `encode()` on the +vendored Kimi tokenizer: + +| workload | scanner | PCRE2 | speedup | +|-------------|---------:|---------:|---------| +| ASCII 10k | 0.095 ms | 0.215 ms | ~2.2x | +| ASCII 50k | 0.22 ms | 0.61 ms | ~2.7x | +| CJK 10k | 0.14 ms | 0.33 ms | ~2.3x | +| CJK 50k | 0.26 ms | 0.74 ms | ~2.8x | + +Enabling Kimi also lets the hand-maintained PCRE2 intersection-rewrite special +case in `split.rs` eventually retire (once o200k/Nemotron are ported too — they +share `fast_split_o200k`'s parameterized `advance_pos`). + +## Correctness + +Differential vs `fancy-regex` on the exact pattern (fancy-regex handles the +`&&` intersections, so it is a faithful oracle): Qwen2 has 60+ hand-picked edge +cases + a 4000-round fuzz; Kimi has 30+ cases (camelCase phase automaton, +`HTTPResponse`, contractions incl. U+017F, Han runs, Han numerals U+3007, Han +symbols/marks U+16FF0, mixed CJK/Latin, combining marks) + a 6000-round fuzz +over a CJK-heavy pool. Both have an integration test asserting the scanner path +— including the parallel chunked path — equals the regex path through +`Split::pre_tokenize`. Parallel chunk boundaries are placed only after a `\n` +followed by a non-whitespace ASCII byte, a point no alternative can cross. + +### Qwen2N3 (GLM / dolma2 / OLMo) + +Same as Qwen2 but with `\p{N}{1,3}` number runs. Enabled. Full `encode()` on +the vendored GLM-5 tokenizer, distinct inputs, A/B vs PCRE2: + +| context | scanner | PCRE2 | speedup | +|---------|---------:|---------:|---------| +| 10k ch | 0.085 ms | 0.170 ms | ~2.0x | +| 50k ch | 0.315 ms | 0.651 ms | ~2.1x | +| 200k ch | 0.926 ms | 2.098 ms | ~2.3x | + +An earlier draft gated this off after a same-text-4 MiB microbenchmark showed +the GLM scanner ~5x slower than the Qwen2 scanner. That was a **first-touch / +page-warming artifact of re-encoding one fixed buffer**: swapping the A/B order +reversed which scheme was "slow", and on distinct inputs (real serving) the two +scanners are within noise in both the sequential and parallel paths. No +scheme-specific slowdown exists. + +## Status + +Every vendored tokenizer is now scanner-backed, one scheme each: + +- **`Qwen2`** (Qwen2/Qwen3 and derivatives): **enabled**. +- **`Qwen2N3`** (GLM / dolma2 / OLMo, `\p{N}{1,3}`): **enabled**. +- **`Kimi`** (moonshotai K2 family, o200k-family with Han runs): **enabled**. + Also catches **knext / Kimi K3**, whose tiktoken `pat_str` is byte-identical + and reaches `from_pattern` verbatim through the tiktoken→tokenizer.json + conversion (`python/fastokens/tiktoken.py`) — a `recognizes_only_known_patterns` + assertion pins that exact string so an upstream regex tweak fails loudly + instead of silently falling back to PCRE2. +- **`O200k`** (o200k_base: GPT-4o, gpt-oss): **enabled**. Same walker as Kimi + with `HAN=false` and a `[\r\n/]*` punct tail (`SLASH=true`). + +Recognition is byte-exact: a model whose split regex differs by one character +silently uses PCRE2 (safe — correctness is never at risk — but unaccelerated). + +## Next steps (ranked) + +1. Nemotron-3 scheme (`fast_split_o200k::advance_pos::`) + — a pattern const + one dispatch line + fuzz. With o200k done this also lets + the hand-maintained PCRE2 intersection-rewrite special case in `split.rs` + retire once nothing else needs it. +2. Fix the PCRE2 `RAYON_NUM_THREADS=1` panic independently (affects all + non-specialized patterns). +3. Later scalar/SIMD tiers from gigatoken once split is the bottleneck: + dual-cursor ILP (+25% measured by them), NEON/AVX-512 mask scanner. diff --git a/examples/long_context_bench.rs b/examples/long_context_bench.rs new file mode 100644 index 0000000..65d1e1d --- /dev/null +++ b/examples/long_context_bench.rs @@ -0,0 +1,90 @@ +//! Cold-cache benchmark for single 200k- and 1M-token contexts. +//! +//! The input is the repository's multilingual corpus repeated to the requested +//! token count. A separate tokenizer calibrates input size; timed iterations +//! construct a fresh tokenizer before encoding so BPE caches start cold. +//! +//! RAYON_NUM_THREADS=8 cargo run --release --example long_context_bench +//! RAYON_NUM_THREADS=8 cargo run --release --example long_context_bench -- kimi-k2.5 5 + +use std::time::Instant; + +fn tokenizer_json(name: &str) -> serde_json::Value { + let repo = env!("CARGO_MANIFEST_DIR"); + let path = format!("{repo}/vendored_tokenizers/{name}/tokenizer.json"); + serde_json::from_str(&std::fs::read_to_string(path).expect("read tokenizer.json")) + .expect("parse tokenizer.json") +} + +fn build_context(tokenizer: &fastokens::Tokenizer, corpus: &str, target_tokens: usize) -> String { + let corpus_tokens = tokenizer.encode(corpus).expect("calibrate corpus").len(); + let estimated_bytes = corpus + .len() + .saturating_mul(target_tokens) + .div_ceil(corpus_tokens); + let repeats = estimated_bytes.div_ceil(corpus.len()).max(1); + let mut text = corpus.repeat(repeats); + let mut cutoff = estimated_bytes.min(text.len()); + while !text.is_char_boundary(cutoff) { + cutoff -= 1; + } + text.truncate(cutoff); + text +} + +fn percentile(sorted: &[f64], percentile: usize) -> f64 { + sorted[(sorted.len() - 1) * percentile / 100] +} + +fn main() { + let repo = env!("CARGO_MANIFEST_DIR"); + let name = std::env::args() + .nth(1) + .unwrap_or_else(|| "kimi-k2.5".to_string()); + let rounds = std::env::args() + .nth(2) + .and_then(|arg| arg.parse().ok()) + .unwrap_or(5usize); + let raw_corpus = + std::fs::read_to_string(format!("{repo}/tests/fixtures/corpus_multilingual.txt")) + .expect("read multilingual corpus"); + let corpus = raw_corpus.replace("<|endoftext|>", " "); + let json = tokenizer_json(&name); + let calibration = + fastokens::Tokenizer::from_json(json.clone()).expect("construct calibration tokenizer"); + + println!( + "{name}, RAYON_NUM_THREADS={}, {rounds} cold rounds", + std::env::var("RAYON_NUM_THREADS").unwrap_or_else(|_| "default".to_string()) + ); + for target in [200_000usize, 1_000_000] { + let text = build_context(&calibration, &corpus, target); + let actual_tokens = calibration.encode(&text).expect("count tokens").len(); + let mib = text.len() as f64 / (1 << 20) as f64; + let mut times = Vec::with_capacity(rounds); + let mut construction_times = Vec::with_capacity(rounds); + + for _ in 0..rounds { + let construction_start = Instant::now(); + let tokenizer = + fastokens::Tokenizer::from_json(json.clone()).expect("construct tokenizer"); + construction_times.push(construction_start.elapsed().as_secs_f64() * 1e3); + let start = Instant::now(); + let ids = tokenizer.encode(&text).expect("encode context"); + times.push(start.elapsed().as_secs_f64() * 1e3); + assert_eq!(ids.len(), actual_tokens); + } + + times.sort_by(f64::total_cmp); + construction_times.sort_by(f64::total_cmp); + let median = percentile(×, 50); + let p90 = percentile(×, 90); + let construction_median = percentile(&construction_times, 50); + println!( + "{target:>9} target | {actual_tokens:>9} actual | {mib:>6.2} MiB | \ + p50 {median:>8.2} ms | p90 {p90:>8.2} ms | {throughput:>6.1} MiB/s | \ + load {construction_median:>7.1} ms", + throughput = mib / (median / 1e3), + ); + } +} diff --git a/examples/serving_split_bench.rs b/examples/serving_split_bench.rs new file mode 100644 index 0000000..7ffcf8b --- /dev/null +++ b/examples/serving_split_bench.rs @@ -0,0 +1,95 @@ +//! Long-context, single-request encode benchmark — the serving latency that +//! matters (10k+ char prompts, ms scale). +//! +//! Each iteration encodes a DISTINCT context so the PCRE2 tier's incremental +//! split cache (keyed on shared prefix with the previous call) never hits — +//! that cache helps streaming/growing sequences, not independent requests, so +//! measuring repeated identical input overstates it. Runs single-threaded +//! (`RAYON_NUM_THREADS=1`) to model a loaded server where each concurrent +//! request effectively gets one core. +//! +//! A/B by checking out `main` (PCRE2 tier) vs this branch (scanner): +//! RAYON_NUM_THREADS=1 cargo run --release --example serving_split_bench + +use std::time::Instant; + +fn build_context(chars: usize, seed: u64) -> String { + let mut rng = seed | 1; + let words = [ + "the", + "tokenizer", + "performance", + "streaming", + "request", + "context", + "attention", + "under", + "load", + "server", + "batch", + "and", + "for", + "with", + "It's", + "don't", + "2024", + "1.5e-7", + "GPU", + "throughput,", + "(latency)", + "https://example.test/path", + "fn", + "encode()", + "->", + "Vec", + "0x7f", + "многоязычный", + "文脈", + "コンテキスト", + ]; + let mut s = String::with_capacity(chars + 32); + while s.chars().count() < chars { + rng = rng + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + s.push_str(words[(rng >> 33) as usize % words.len()]); + s.push(if rng & 0xF == 0 { '\n' } else { ' ' }); + } + s +} + +fn main() { + let path = std::env::args() + .nth(1) + .unwrap_or_else(|| "vendored_tokenizers/qwen3-8b/tokenizer.json".to_string()); + let raw = std::fs::read_to_string(&path).expect("read tokenizer.json"); + let value: serde_json::Value = serde_json::from_str(&raw).unwrap(); + let tok = fastokens::Tokenizer::from_json(value).unwrap(); + + for chars in [10_000usize, 50_000, 200_000] { + // Pool of distinct contexts; rotate so no two consecutive encodes + // share a prefix (defeats the incremental split cache). + let pool: Vec = (0..16).map(|i| build_context(chars, 0x1234 + i)).collect(); + for c in &pool { + let _ = tok.encode(c); + } + let iters = (40_000_000 / chars).max(64); + let mut times = Vec::with_capacity(iters); + for i in 0..iters { + let text = &pool[i % pool.len()]; + let t = Instant::now(); + std::hint::black_box(tok.encode(text).unwrap()); + times.push(t.elapsed().as_secs_f64() * 1e3); + } + times.sort_by(|a, b| a.partial_cmp(b).unwrap()); + let median = times[times.len() / 2]; + let bytes = pool[0].len(); + println!( + "{:>7} chars ({:>8} bytes): {:7.3} ms ({:6.1} MiB/s)", + chars, + bytes, + median, + bytes as f64 / (1 << 20) as f64 / (median / 1e3) + ); + } +} diff --git a/pyproject.toml b/pyproject.toml index 7373e4a..b1ca7dd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,9 +3,13 @@ requires = ["maturin>=1.0,<2.0"] build-backend = "maturin" [project] -name = "fastokens" +name = "fastokens-b10" dynamic = ["version"] -requires-python = ">=3.9" +requires-python = ">=3.10" +dependencies = ["numpy>=1.23.0"] + +[project.optional-dependencies] +tiktoken = ["tiktoken>=0.13.0"] [tool.maturin] manifest-path = "python/Cargo.toml" diff --git a/python/Cargo.toml b/python/Cargo.toml index 7b91f04..53d7beb 100644 --- a/python/Cargo.toml +++ b/python/Cargo.toml @@ -13,7 +13,10 @@ name = "_native" crate-type = ["cdylib"] [dependencies] -fastokens = { path = ".." } -pyo3 = {version = "0.25", features = ["extension-module", "abi3-py39"] } +fastokens = { package = "fastokens-b10", path = ".." } +numpy = "0.25" +pyo3 = {version = "0.25", features = ["extension-module", "abi3-py310"] } +pyo3-async-runtimes = { version = "0.25", features = ["tokio-runtime"] } +pythonize = "0.25" rayon = "1" -serde_json = "1" +serde_json = { version = "1", features = ["preserve_order"] } diff --git a/python/fastokens/__init__.py b/python/fastokens/__init__.py index f03d25e..7b29b4f 100644 --- a/python/fastokens/__init__.py +++ b/python/fastokens/__init__.py @@ -1,10 +1,28 @@ -from fastokens._native import Tokenizer +from __future__ import annotations -__all__ = ["Tokenizer", "patch_transformers", "unpatch_transformers"] +import os +import warnings +from fastokens._native import StructuralTokenConfig, Tokenizer +from fastokens.tiktoken import ( + tiktoken_model_to_tokenizer_json, + tiktoken_to_tokenizer_json, +) + +__all__ = [ + "StructuralTokenConfig", + "Tokenizer", + "patch_transformers", + "tiktoken_model_to_tokenizer_json", + "tiktoken_to_tokenizer_json", + "unpatch_transformers", +] + +APPLY_CHAT_TEMPLATE_ENV = "FASTOKENS_APPLY_CHAT_TEMPLATE" _patched = False _originals: dict = {} +_native_chat_fallback_warned = False def _swap_backend(tokenizer, shim_cls): @@ -15,7 +33,100 @@ def _swap_backend(tokenizer, shim_cls): return tokenizer -def patch_transformers() -> None: +def _env_flag(name: str) -> bool | None: + """Parse a boolean-like environment variable; ``None`` when unset.""" + value = os.environ.get(name) + if value is None: + return None + normalized = value.strip().lower() + if normalized in ("1", "true", "yes", "on", "native"): + return True + if normalized in ("", "0", "false", "no", "off"): + return False + raise ValueError(f"{name} must be boolean-like (0/1/true/false), got {value!r}") + + +def _warn_native_chat_fallback(exc: Exception) -> None: + global _native_chat_fallback_warned + if not _native_chat_fallback_warned: + _native_chat_fallback_warned = True + warnings.warn( + "[fastokens] native chat template rendering failed; falling back to " + f"transformers' jinja2 renderer for this and future failures: {exc}", + RuntimeWarning, + stacklevel=3, + ) + + +def _patched_apply_chat_template(self, conversation, *args, **kwargs): + """Route render-only ``apply_chat_template`` calls to the fastokens shim. + + Anything the native path does not support byte-for-byte — positional + arguments, ``tokenize=True``, dict outputs, batched conversations, a + non-shim backend — falls through to transformers' own implementation, + as does any error raised by the native renderer (warned once). + """ + from fastokens._compat import _TokenizerShim + + original = _originals["PreTrainedTokenizerBase.apply_chat_template"] + backend = getattr(self, "_tokenizer", None) + if ( + args + or kwargs.get("tokenize", True) + or kwargs.get("return_dict") + or kwargs.get("return_assistant_tokens_mask") + or not isinstance(backend, _TokenizerShim) + or ( + isinstance(conversation, (list, tuple)) + and conversation + and isinstance(conversation[0], (list, tuple)) + ) + ): + return original(self, conversation, *args, **kwargs) + + try: + get_chat_template = getattr(self, "get_chat_template", None) + if get_chat_template is not None: + chat_template = get_chat_template(kwargs.get("chat_template"), kwargs.get("tools")) + else: + chat_template = kwargs.get("chat_template") or self.chat_template + if not isinstance(chat_template, str): + return original(self, conversation, **kwargs) + + native_kwargs = dict(kwargs) + for consumed in ( + "tokenize", + "chat_template", + "tools", + "documents", + "add_generation_prompt", + "continue_final_message", + ): + native_kwargs.pop(consumed, None) + # transformers silently ignores kwargs the template does not read; + # keep that contract on the patched path unless the caller opts in. + native_kwargs.setdefault("fastokens_strict_template", False) + special_tokens = { + key: [str(item) for item in value] if isinstance(value, (list, tuple)) else str(value) + for key, value in self.special_tokens_map.items() + } + return backend.apply_chat_template( + conversation, + chat_template=chat_template, + tokenize=False, + add_generation_prompt=kwargs.get("add_generation_prompt", False), + continue_final_message=kwargs.get("continue_final_message") or False, + tools=kwargs.get("tools"), + documents=kwargs.get("documents"), + special_tokens=special_tokens, + **native_kwargs, + ) + except Exception as exc: # noqa: BLE001 - fallback must equal the status quo + _warn_native_chat_fallback(exc) + return original(self, conversation, **kwargs) + + +def patch_transformers(apply_chat_template: bool = False) -> None: """ Monkey-patch ``tokenizers.Tokenizer`` so that the ``transformers`` library uses fastokens for encoding. @@ -33,12 +144,22 @@ def patch_transformers() -> None: Supports both transformers v4 (``tokenization_utils_fast``) and v5+ (``tokenization_utils_tokenizers``). + + ``apply_chat_template=True`` additionally routes render-only + ``tokenizer.apply_chat_template(..., tokenize=False)`` calls to the + fastokens native renderer (falling back to transformers' jinja2 on any + unsupported input or error). The ``FASTOKENS_APPLY_CHAT_TEMPLATE`` + environment variable, when set, overrides this argument in either + direction — it doubles as a fleet-wide kill switch. """ global _patched if _patched: print("[fastokens] patch_transformers: already patched.") return + env_override = _env_flag(APPLY_CHAT_TEMPLATE_ENV) + route_chat_template = env_override if env_override is not None else apply_chat_template + from fastokens._compat import _TokenizerShim from fastokens._native import DecodeStream @@ -97,13 +218,36 @@ def _patched_from_pretrained(cls, *args, **kwargs): _originals["DecodeStream"] = _td.DecodeStream _td.DecodeStream = DecodeStream + if route_chat_template: + from transformers import PreTrainedTokenizerBase + + _originals["PreTrainedTokenizerBase.apply_chat_template"] = ( + PreTrainedTokenizerBase.apply_chat_template + ) + PreTrainedTokenizerBase.apply_chat_template = _patched_apply_chat_template + _patched = True from importlib.metadata import version + # Assuming transformers is installed. # If not, this will raise an error, which is fine since patching won't work without it. transformers_version = version("transformers") - print(f"[fastokens] patch_transformers: successfully patched transformers v{transformers_version}") + print( + f"[fastokens] patch_transformers: successfully patched transformers v{transformers_version}" + ) + if route_chat_template: + source = "env" if env_override is not None else "arg" + print( + "[fastokens] chat template engine: fastokens-native for " + f"tokenize=False (via {source})" + ) + else: + print( + "[fastokens] chat template engine: transformers jinja2 " + "(enable native with patch_transformers(apply_chat_template=True) " + f"or {APPLY_CHAT_TEMPLATE_ENV}=1)" + ) def unpatch_transformers() -> None: @@ -136,6 +280,13 @@ def unpatch_transformers() -> None: mod, original_cls = _originals["tokenization_utils_fast"] mod.TokenizerFast = original_cls + if "PreTrainedTokenizerBase.apply_chat_template" in _originals: + from transformers import PreTrainedTokenizerBase + + PreTrainedTokenizerBase.apply_chat_template = _originals[ + "PreTrainedTokenizerBase.apply_chat_template" + ] + _td.DecodeStream = _originals["DecodeStream"] _originals.clear() diff --git a/python/fastokens/_compat.py b/python/fastokens/_compat.py index 6d4dbd6..9b43a1a 100644 --- a/python/fastokens/_compat.py +++ b/python/fastokens/_compat.py @@ -20,10 +20,17 @@ _Encoding = Encoding +def _reject_unsupported_kwargs(kwargs: dict) -> None: + if kwargs: + names = ", ".join(sorted(kwargs)) + raise TypeError(f"unsupported fastokens encode option(s): {names}") + + # --------------------------------------------------------------------------- # _TokenizerShim # --------------------------------------------------------------------------- + class _TokenizerShim: """ Complete replacement for ``tokenizers.Tokenizer``. @@ -40,16 +47,30 @@ def __init__(self, src) -> None: elif isinstance(src, _TokenizerShim): self._json = src._json self._fast = Tokenizer.from_json_str(src._json) + self._encode_special_tokens = src._encode_special_tokens + self._chat_template = src._chat_template + self._special_tokens = getattr(src, "_special_tokens", None) elif hasattr(src, "to_str"): # Accept a real tokenizers.Tokenizer (e.g. from convert_slow_tokenizer). self._json = src.to_str() self._fast = Tokenizer.from_json_str(self._json) + self._encode_special_tokens = getattr(src, "encode_special_tokens", False) else: raise TypeError( f"expected JSON string, _TokenizerShim, or tokenizers.Tokenizer; " f"got {type(src).__name__}" ) - self._encode_special_tokens: bool = False + if not hasattr(self, "_encode_special_tokens"): + self._encode_special_tokens = False + if not hasattr(self, "_chat_template"): + self._chat_template = None + if not hasattr(self, "_native_chat_template"): + self._native_chat_template = None + if not hasattr(self, "_special_tokens"): + self._special_tokens = None + self._set_native_chat_template(self._chat_template) + if self._special_tokens is not None: + self._fast.set_special_tokens(self._special_tokens) # -- Pickle / copy -------------------------------------------------- @@ -59,6 +80,8 @@ def __getstate__(self): self._fast.truncation, self._fast.padding, self._encode_special_tokens, + self._chat_template, + self._special_tokens, ) def __setstate__(self, state) -> None: @@ -66,13 +89,27 @@ def __setstate__(self, state) -> None: # Backwards compat: old pickles stored just the JSON string. self.__init__(state) # type: ignore[misc] else: - json_str, trunc, pad, enc_special = state + special_tokens = None + if len(state) == 4: + json_str, trunc, pad, enc_special = state + chat_template = None + elif len(state) == 5: + json_str, trunc, pad, enc_special, chat_template = state + else: + json_str, trunc, pad, enc_special, chat_template, special_tokens = state self.__init__(json_str) # type: ignore[misc] if trunc is not None: self._fast.enable_truncation(**trunc) if pad is not None: - self._fast.enable_padding(**{k: v for k, v in pad.items() if v is not None}) - self._encode_special_tokens = enc_special + self._fast.enable_padding( + **{k: v for k, v in pad.items() if v is not None} + ) + self.encode_special_tokens = enc_special + self._chat_template = chat_template + self._native_chat_template = None + self._set_native_chat_template(chat_template) + if special_tokens is not None: + self.set_special_tokens(special_tokens) def __deepcopy__(self, memo): new = object.__new__(_TokenizerShim) @@ -86,6 +123,12 @@ def __deepcopy__(self, memo): if pad is not None: new._fast.enable_padding(**{k: v for k, v in pad.items() if v is not None}) new._encode_special_tokens = self._encode_special_tokens + new._chat_template = self._chat_template + new._native_chat_template = None + new._set_native_chat_template(new._chat_template) + new._special_tokens = dict(self._special_tokens) if self._special_tokens else None + if new._special_tokens is not None: + new._fast.set_special_tokens(new._special_tokens) if hasattr(self, "_special_prefix"): new._special_prefix = list(self._special_prefix) new._special_suffix = list(self._special_suffix) @@ -111,7 +154,10 @@ def from_pretrained( from huggingface_hub import hf_hub_download path = hf_hub_download( - identifier, "tokenizer.json", revision=revision, token=token, + identifier, + "tokenizer.json", + revision=revision, + token=token, ) return cls.from_file(path) @@ -164,7 +210,9 @@ def padding(self, value: dict | None) -> None: if value is None: self._fast.no_padding() else: - self._fast.enable_padding(**{k: v for k, v in value.items() if v is not None}) + self._fast.enable_padding( + **{k: v for k, v in value.items() if v is not None} + ) def enable_truncation( self, @@ -195,7 +243,11 @@ def enable_padding( pad_type_id=pad_type_id, pad_token=pad_token, **({"length": length} if length is not None else {}), - **({"pad_to_multiple_of": pad_to_multiple_of} if pad_to_multiple_of is not None else {}), + **( + {"pad_to_multiple_of": pad_to_multiple_of} + if pad_to_multiple_of is not None + else {} + ), ) def no_padding(self) -> None: @@ -209,38 +261,186 @@ def encode( pair: str | None = None, is_pretokenized: bool = False, add_special_tokens: bool = True, + split_special_tokens: bool | None = None, **kwargs, ) -> Encoding: + _reject_unsupported_kwargs(kwargs) if pair is not None: raise NotImplementedError("pair encoding is not supported by fastokens") if is_pretokenized: - raise NotImplementedError("pre-tokenized input is not supported by fastokens") - return self._fast.encode(sequence, add_special_tokens=add_special_tokens) + raise NotImplementedError( + "pre-tokenized input is not supported by fastokens" + ) + if split_special_tokens is None: + split_special_tokens = self._encode_special_tokens + return self._fast.encode( + sequence, + add_special_tokens=add_special_tokens, + split_special_tokens=split_special_tokens, + ) def encode_batch( self, inputs: list, is_pretokenized: bool = False, add_special_tokens: bool = True, + split_special_tokens: bool | None = None, **kwargs, ) -> list[Encoding]: + _reject_unsupported_kwargs(kwargs) if is_pretokenized or any(isinstance(inp, (list, tuple)) for inp in inputs): raise NotImplementedError( "pair/pre-tokenized batch encoding is not supported by fastokens" ) - return self._fast.encode_batch(inputs, add_special_tokens=add_special_tokens) + if split_special_tokens is None: + split_special_tokens = self._encode_special_tokens + return self._fast.encode_batch( + inputs, + add_special_tokens=add_special_tokens, + split_special_tokens=split_special_tokens, + ) def encode_batch_fast( self, inputs: list[str], is_pretokenized: bool = False, add_special_tokens: bool = True, + split_special_tokens: bool | None = None, ) -> list[Encoding]: if is_pretokenized or any(isinstance(inp, (list, tuple)) for inp in inputs): raise NotImplementedError( "pair/pre-tokenized batch encoding is not supported by fastokens" ) - return self._fast.encode_batch(inputs, add_special_tokens=add_special_tokens) + if split_special_tokens is None: + split_special_tokens = self._encode_special_tokens + return self._fast.encode_batch( + inputs, + add_special_tokens=add_special_tokens, + split_special_tokens=split_special_tokens, + ) + + def encode_with_structural_tokens( + self, + sequence: str, + structural_config, + placeholder_map: dict[str, str] | None = None, + add_special_tokens: bool = False, + ) -> Encoding: + """Encode rendered template text with structural-token boundaries.""" + return self._fast.encode_with_structural_tokens( + sequence, + structural_config, + placeholder_map or {}, + add_special_tokens=add_special_tokens, + ) + + def encode_segments( + self, + segments: list[tuple[str, bool]], + add_special_tokens: bool = False, + tiktoken_safe: bool = True, + ) -> Encoding: + """Encode ``(text, allow_special)`` segments in one backend call. + + ``tiktoken_safe=True`` is the default and reproduces legacy tiktoken + tokenizer chunking for token-ID parity. Pass ``False`` only when you + explicitly want whole-segment BPE encoding. + """ + return self._fast.encode_segments( + segments, + add_special_tokens=add_special_tokens, + tiktoken_safe=tiktoken_safe, + ) + + def set_chat_template(self, chat_template: str | None) -> None: + self._chat_template = chat_template + self._set_native_chat_template(chat_template) + + def set_special_tokens(self, special_tokens: dict | None) -> None: + self._special_tokens = dict(special_tokens) if special_tokens is not None else None + self._fast.set_special_tokens(special_tokens) + + def _set_native_chat_template(self, chat_template: str | None) -> None: + if self._native_chat_template != chat_template: + self._fast.set_chat_template(chat_template) + self._native_chat_template = chat_template + + def apply_chat_template( + self, + messages, + chat_template: str | None = None, + tokenize: bool = False, + add_generation_prompt: bool = False, + continue_final_message: bool | str = False, + add_special_tokens: bool = False, + tools=None, + documents=None, + special_tokens=None, + **kwargs, + ): + """Render a HuggingFace-style chat template using the Rust backend.""" + if chat_template is None: + self._set_native_chat_template(self._chat_template) + return self._fast.apply_chat_template( + messages, + chat_template=chat_template, + tokenize=tokenize, + add_generation_prompt=add_generation_prompt, + continue_final_message=continue_final_message, + add_special_tokens=add_special_tokens, + tools=tools, + documents=documents, + special_tokens=special_tokens, + **kwargs, + ) + + async def async_apply_chat_template( + self, + messages, + chat_template: str | None = None, + tokenize: bool = False, + add_generation_prompt: bool = False, + continue_final_message: bool | str = False, + add_special_tokens: bool = False, + tools=None, + documents=None, + special_tokens=None, + **kwargs, + ): + """Render a chat template on a background thread; awaitable.""" + if chat_template is None: + self._set_native_chat_template(self._chat_template) + return await self._fast.async_apply_chat_template( + messages, + chat_template=chat_template, + tokenize=tokenize, + add_generation_prompt=add_generation_prompt, + continue_final_message=continue_final_message, + add_special_tokens=add_special_tokens, + tools=tools, + documents=documents, + special_tokens=special_tokens, + **kwargs, + ) + + async def async_encode_batch( + self, + inputs: list[str], + is_pretokenized: bool = False, + add_special_tokens: bool = True, + split_special_tokens: bool | None = None, + ) -> list[Encoding]: + if is_pretokenized or any(isinstance(inp, (list, tuple)) for inp in inputs): + raise NotImplementedError( + "pair/pre-tokenized batch encoding is not supported by fastokens" + ) + if split_special_tokens is None: + split_special_tokens = self._encode_special_tokens + return await self._fast.async_encode_batch( + inputs, + add_special_tokens=add_special_tokens, + split_special_tokens=split_special_tokens, + ) # -- Post-processing ------------------------------------------------ @@ -265,7 +465,18 @@ def decode_batch( sequences: list[list[int]], skip_special_tokens: bool = True, ) -> list[str]: - return self._fast.decode_batch(sequences, skip_special_tokens=skip_special_tokens) + return self._fast.decode_batch( + sequences, skip_special_tokens=skip_special_tokens + ) + + async def async_decode_batch( + self, + sequences: list[list[int]], + skip_special_tokens: bool = True, + ) -> list[str]: + return await self._fast.async_decode_batch( + sequences, skip_special_tokens=skip_special_tokens + ) # -- Vocabulary ----------------------------------------------------- @@ -356,12 +567,11 @@ def post_processor(self, value) -> None: self._fast.post_processor = value - - # --------------------------------------------------------------------------- # Helper classes # --------------------------------------------------------------------------- + class _DecoderShim: """ Stand-in for ``tokenizers.decoders.Decoder``. diff --git a/python/fastokens/_native.pyi b/python/fastokens/_native.pyi index 0292a16..1336eb7 100644 --- a/python/fastokens/_native.pyi +++ b/python/fastokens/_native.pyi @@ -1,4 +1,4 @@ -from typing import Optional +from typing import Any, Optional class Encoding: """Rust-backed encoding returned by ``Tokenizer.encode`` / ``encode_batch``.""" @@ -15,47 +15,55 @@ class Encoding: ids: list[int], attention_mask: Optional[list[int]] = None, ) -> "Encoding": ... - def __len__(self) -> int: ... def __repr__(self) -> str: ... - + def into_numpy( + self, + ids: bool = True, + attention_mask: bool = False, + type_ids: bool = False, + special_tokens_mask: bool = False, + ) -> dict[str, Any]: + """Move selected fields into numpy.uint32 arrays and drain the encoding.""" + ... # Properties that raise NotImplementedError @property def tokens(self) -> list[str]: ... @tokens.setter def tokens(self, value: list[str]) -> None: ... - @property def offsets(self) -> list[tuple[int, int]]: ... @offsets.setter def offsets(self, value: list[tuple[int, int]]) -> None: ... - @property def sequence_ids(self) -> list[Optional[int]]: ... @sequence_ids.setter def sequence_ids(self, value: list[Optional[int]]) -> None: ... - @property def word_ids(self) -> list[Optional[int]]: ... @word_ids.setter def word_ids(self, value: list[Optional[int]]) -> None: ... - @property def words(self) -> list[Optional[int]]: ... @words.setter def words(self, value: list[Optional[int]]) -> None: ... - def set_sequence_id(self, sequence_id: int) -> None: ... - - def char_to_token(self, char_pos: int, sequence_index: int = 0) -> Optional[int]: ... + def char_to_token( + self, char_pos: int, sequence_index: int = 0 + ) -> Optional[int]: ... def char_to_word(self, char_pos: int, sequence_index: int = 0) -> Optional[int]: ... def token_to_chars(self, token_index: int) -> Optional[tuple[int, int]]: ... def token_to_sequence(self, token_index: int) -> Optional[int]: ... def token_to_word(self, token_index: int) -> Optional[int]: ... - def word_to_chars(self, word_index: int, sequence_index: int = 0) -> Optional[tuple[int, int]]: ... - def word_to_tokens(self, word_index: int, sequence_index: int = 0) -> Optional[tuple[int, int]]: ... - - def truncate(self, max_length: int, stride: int = 0, direction: str = "right") -> None: ... + def word_to_chars( + self, word_index: int, sequence_index: int = 0 + ) -> Optional[tuple[int, int]]: ... + def word_to_tokens( + self, word_index: int, sequence_index: int = 0 + ) -> Optional[tuple[int, int]]: ... + def truncate( + self, max_length: int, stride: int = 0, direction: str = "right" + ) -> None: ... def pad( self, length: int, @@ -64,34 +72,42 @@ class Encoding: pad_type_id: int = 0, pad_token: str = "[PAD]", ) -> None: ... - @staticmethod - def merge(encodings: list["Encoding"], growing_offsets: bool = True) -> "Encoding": ... + def merge( + encodings: list["Encoding"], growing_offsets: bool = True + ) -> "Encoding": ... + +class StructuralTokenConfig: + """Constant structural-token state for rendered-prompt encoding. + + Include every token string the rendered template may use as a structural + boundary, including tag-like non-special added tokens. + """ + def __new__( + cls, + structural_tokens: list[str] | set[str] | tuple[str, ...], + non_special_added_tokens: Optional[ + set[str] | list[str] | tuple[str, ...] + ] = None, + ) -> "StructuralTokenConfig": ... class Tokenizer: """An LLM tokenizer backed by ``tokenizer.json``.""" def __new__(cls, model: str) -> "Tokenizer": ... - @staticmethod def from_file(path: str) -> "Tokenizer": ... - @staticmethod def from_json_str(json: str) -> "Tokenizer": ... - @staticmethod def from_model(model: str) -> "Tokenizer": ... - @property def vocab_size(self) -> int: ... - @property def truncation(self) -> Optional[dict]: ... - @property def padding(self) -> Optional[dict]: ... - def enable_truncation( self, max_length: int, @@ -99,9 +115,7 @@ class Tokenizer: strategy: str = "longest_first", direction: str = "right", ) -> None: ... - def no_truncation(self) -> None: ... - def enable_padding( self, direction: str = "right", @@ -111,30 +125,116 @@ class Tokenizer: length: Optional[int] = None, pad_to_multiple_of: Optional[int] = None, ) -> None: ... - def no_padding(self) -> None: ... - def post_process( self, encoding: Encoding, pair: Optional[Encoding] = None, add_special_tokens: bool = True, ) -> Encoding: ... - def num_special_tokens_to_add(self, is_pair: bool) -> int: ... - - def encode(self, input: str, add_special_tokens: bool = False) -> Encoding: ... + def encode( + self, + input: str, + add_special_tokens: bool = False, + split_special_tokens: bool = False, + ) -> Encoding: ... + def encode_segments( + self, + segments: list[tuple[str, bool]], + add_special_tokens: bool = False, + tiktoken_safe: bool = True, + ) -> Encoding: + """Encode ``(text, allow_special)`` segments in one backend call. + + ``tiktoken_safe=True`` is the default and reproduces legacy tiktoken + tokenizer chunking for token-ID parity. Pass ``False`` only when you + explicitly want whole-segment BPE encoding. + """ + ... def encode_batch( - self, inputs: list[str], add_special_tokens: bool = False + self, + inputs: list[str], + add_special_tokens: bool = False, + split_special_tokens: bool = False, ) -> list[Encoding]: ... + def encode_with_structural_tokens( + self, + input: str, + structural_config: StructuralTokenConfig, + placeholder_map: Optional[dict[str, str]] = None, + add_special_tokens: bool = False, + ) -> Encoding: + """Encode rendered template text with structural-token boundaries. + + ``placeholder_map`` is per request and maps placeholder text to the + original user text. Keep ``add_special_tokens=False`` when replacing an + existing rendered chat-template encode path. Backend truncation and + padding settings are intentionally not applied. + """ + ... + + def set_chat_template(self, chat_template: Optional[str]) -> None: ... + def set_special_tokens(self, special_tokens: Optional[dict[str, object]]) -> None: ... + @property + def chat_template_variables(self) -> Optional[list[str]]: + """Variables the configured chat template reads from the render context. + + Derived statically from the template AST; ``None`` when no chat + template is set. Keyword arguments to ``apply_chat_template`` outside + this set are rejected unless ``fastokens_strict_template=False``. + """ + ... + + def apply_chat_template( + self, + messages: list[dict], + chat_template: Optional[str] = None, + tokenize: bool = False, + add_generation_prompt: bool = False, + continue_final_message: bool | str = False, + add_special_tokens: bool = False, + tools: Optional[list[dict]] = None, + documents: Optional[list[dict]] = None, + special_tokens: Optional[dict[str, object]] = None, + **kwargs, + ) -> str: ... + async def async_apply_chat_template( + self, + messages: list[dict], + chat_template: Optional[str] = None, + tokenize: bool = False, + add_generation_prompt: bool = False, + continue_final_message: bool | str = False, + add_special_tokens: bool = False, + tools: Optional[list[dict]] = None, + documents: Optional[list[dict]] = None, + special_tokens: Optional[dict[str, object]] = None, + **kwargs, + ) -> str: + """Like ``apply_chat_template`` but renders on a background thread. + + Input conversion still runs synchronously at call time; only the + template render is offloaded, so awaiting this keeps large renders + off the event loop. + """ + ... + + async def async_encode_batch( + self, + inputs: list[str], + add_special_tokens: bool = False, + split_special_tokens: bool = False, + ) -> list[Encoding]: ... def decode_tokens(self, tokens: list[str]) -> str: ... def decode(self, ids: list[int], skip_special_tokens: bool = False) -> str: ... - def decode_batch( self, sentences: list[list[int]], skip_special_tokens: bool = False ) -> list[str]: ... - + async def async_decode_batch( + self, sentences: list[list[int]], skip_special_tokens: bool = False + ) -> list[str]: ... def token_to_id(self, token: str) -> Optional[int]: ... def id_to_token(self, id: int) -> Optional[str]: ... diff --git a/python/fastokens/tiktoken.py b/python/fastokens/tiktoken.py new file mode 100644 index 0000000..c237e6e --- /dev/null +++ b/python/fastokens/tiktoken.py @@ -0,0 +1,415 @@ +from __future__ import annotations + +import json +import logging +from collections.abc import Mapping +from pathlib import Path +from typing import Any, Iterable, NamedTuple +from urllib.parse import urlparse + +DEFAULT_TIKTOKEN_PATTERN = ( + r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|" + r"\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|" + r"\s+(?!\S)|\s+" +) +KIMI_TIKTOKEN_PATTERN = "|".join( + [ + r"[\p{Han}]+", + r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?", + r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?", + r"\p{N}{1,3}", + r" ?[^\s\p{L}\p{N}]+[\r\n]*", + r"\s*[\r\n]+", + r"\s+(?!\S)", + r"\s+", + ] +) +PRINTABLE_ASCII_START = 33 +PRINTABLE_ASCII_END = 126 +LATIN1_DIRECT_START = 0xA1 +LATIN1_DIRECT_END_EXCLUSIVE = 0xAD +LATIN1_DIRECT_RESUME = 0xAE + +logger = logging.getLogger(__name__) + + +class _MergeCandidate(NamedTuple): + rank: int + left_rank: int + right_rank: int + left: bytes + right: bytes + + +class _AddedToken(NamedTuple): + id: int + content: str + single_word: bool = False + lstrip: bool = False + rstrip: bool = False + normalized: bool = False + special: bool = True + + +def _byte_to_unicode() -> dict[int, str]: + table: dict[int, str] = {} + next_codepoint = 256 + for byte in range(256): + # Match the GPT-2 byte-level alphabet: printable ASCII plus most + # Latin-1 bytes are kept as-is, while whitespace/control bytes map + # to code points above the byte range. + should_use_direct_mapping = ( + PRINTABLE_ASCII_START <= byte <= PRINTABLE_ASCII_END + or LATIN1_DIRECT_START <= byte < LATIN1_DIRECT_END_EXCLUSIVE + or byte >= LATIN1_DIRECT_RESUME + ) + if should_use_direct_mapping: + table[byte] = chr(byte) + else: + table[byte] = chr(next_codepoint) + next_codepoint += 1 + return table + + +_BYTE_ENCODER = _byte_to_unicode() + + +def _token_bytes_to_string(token: bytes) -> str: + return "".join(_BYTE_ENCODER[byte] for byte in token) + + +def _extract_encoding(encoding: Any) -> Any: + if isinstance(encoding, str): + try: + import tiktoken + except ImportError: + return None + return tiktoken.get_encoding(encoding) + return encoding + + +def _get_rank_pair_for_sorting(split: tuple[int, int, bytes, bytes]) -> tuple[int, int]: + left_rank, right_rank, _, _ = split + return left_rank, right_rank + + +def _extract_vocab_and_merges( + mergeable_ranks: dict[bytes, int], +) -> tuple[dict[str, int], list[list[str]]]: + vocab = { + _token_bytes_to_string(token): rank + for token, rank in sorted(mergeable_ranks.items(), key=lambda item: item[1]) + } + + merge_candidates: list[_MergeCandidate] = [] + for token, rank in mergeable_ranks.items(): + if len(token) == 1: + continue + split_candidates: list[tuple[int, int, bytes, bytes]] = [] + for index in range(1, len(token)): + left = token[:index] + right = token[index:] + if left in mergeable_ranks and right in mergeable_ranks: + split_candidates.append( + (mergeable_ranks[left], mergeable_ranks[right], left, right) + ) + split_candidates.sort(key=_get_rank_pair_for_sorting) + merge_candidates.extend( + _MergeCandidate( + rank=rank, + left_rank=left_rank, + right_rank=right_rank, + left=left, + right=right, + ) + for left_rank, right_rank, left, right in split_candidates + ) + + merge_candidates.sort(key=lambda item: (item.rank, item.left_rank, item.right_rank)) + merges = [ + [ + _token_bytes_to_string(candidate.left), + _token_bytes_to_string(candidate.right), + ] + for candidate in merge_candidates + ] + return vocab, merges + + +def _normalise_special_tokens( + special_tokens: Mapping[str, int] | None, +) -> list[_AddedToken]: + if special_tokens is None: + return [] + return [ + _AddedToken(id=token_id, content=token, special=True) + for token, token_id in sorted( + dict(special_tokens).items(), key=lambda item: item[1] + ) + ] + + +def _bool_config(value: Any, default: bool) -> bool: + return value if isinstance(value, bool) else default + + +def _auto_tokenizer_entries(auto_tokenizer: Any) -> list[str]: + if isinstance(auto_tokenizer, str): + return [auto_tokenizer] + if isinstance(auto_tokenizer, list): + return [item for item in auto_tokenizer if isinstance(item, str)] + return [] + + +def _is_kimi_tiktoken_config(config: dict[str, Any]) -> bool: + auto_map = config.get("auto_map", {}) + if isinstance(auto_map, dict): + if any( + "tokenization_kimi.TikTokenTokenizer" in item + for item in _auto_tokenizer_entries(auto_map.get("AutoTokenizer")) + ): + return True + if config.get("tokenizer_class") != "TikTokenTokenizer": + return False + added_tokens_decoder = config.get("added_tokens_decoder", {}) + if not isinstance(added_tokens_decoder, dict): + return False + kimi_markers = { + "<|open|>", + "<|close|>", + "<|sep|>", + "<|end_of_msg|>", + "<|im_user|>", + "<|im_assistant|>", + "<|im_middle|>", + "<|tool_calls_section_begin|>", + } + for token_config in added_tokens_decoder.values(): + if ( + isinstance(token_config, dict) + and token_config.get("content") in kimi_markers + ): + return True + return False + + +def _has_tokenizer_identity(config: dict[str, Any]) -> bool: + return "tokenizer_class" in config or "auto_map" in config + + +def _config_added_tokens( + config: dict[str, Any], base_vocab_size: int +) -> list[_AddedToken]: + added_tokens: dict[int, _AddedToken] = {} + added_tokens_decoder = config.get("added_tokens_decoder", {}) + if isinstance(added_tokens_decoder, dict): + for token_id, token_config in added_tokens_decoder.items(): + try: + token_id_int = int(token_id) + except (TypeError, ValueError): + continue + if isinstance(token_config, dict): + content = token_config.get("content") + special = _bool_config(token_config.get("special"), False) + single_word = _bool_config(token_config.get("single_word"), False) + lstrip = _bool_config(token_config.get("lstrip"), False) + rstrip = _bool_config(token_config.get("rstrip"), False) + normalized = _bool_config(token_config.get("normalized"), False) + else: + content = token_config + special = False + single_word = False + lstrip = False + rstrip = False + normalized = False + if isinstance(content, str): + added_tokens[token_id_int] = _AddedToken( + id=token_id_int, + content=content, + single_word=single_word, + lstrip=lstrip, + rstrip=rstrip, + normalized=normalized, + special=special, + ) + + if _is_kimi_tiktoken_config(config) and added_tokens: + max_token_id = max(added_tokens) + for token_id in range(base_vocab_size, max_token_id + 1): + added_tokens.setdefault( + token_id, + _AddedToken( + id=token_id, content=f"<|reserved_token_{token_id}|>", special=True + ), + ) + return [added_tokens[token_id] for token_id in sorted(added_tokens)] + + +def _load_model_config(model: str | Path) -> tuple[str, dict[str, Any]]: + if isinstance(model, str) and urlparse(model).scheme in { + "http", + "https", + "s3", + "gs", + "az", + "file", + }: + return model, {} + model_path = Path(model) + if model_path.is_dir(): + config_path = model_path / "tokenizer_config.json" + config = json.loads(config_path.read_text()) if config_path.exists() else {} + return str(model_path / "tiktoken.model"), config + return str(model), {} + + +def _tokenizer_json_from_parts( + mergeable_ranks: dict[bytes, int], + pattern: str, + added_tokens: Iterable[_AddedToken], + *, + pretty: bool = False, +) -> str: + vocab, merges = _extract_vocab_and_merges(mergeable_ranks) + added_tokens = [ + { + "id": token.id, + "content": token.content, + "single_word": token.single_word, + "lstrip": token.lstrip, + "rstrip": token.rstrip, + "normalized": token.normalized, + "special": token.special, + } + for token in added_tokens + ] + + tokenizer_json = { + "version": "1.0", + "truncation": None, + "padding": None, + "added_tokens": added_tokens, + "normalizer": None, + "pre_tokenizer": { + "type": "Sequence", + "pretokenizers": [ + { + "type": "Split", + "pattern": {"Regex": pattern}, + "behavior": "Isolated", + "invert": False, + }, + { + "type": "ByteLevel", + "add_prefix_space": False, + "trim_offsets": True, + "use_regex": False, + }, + ], + }, + "post_processor": { + "type": "ByteLevel", + "add_prefix_space": False, + "trim_offsets": False, + "use_regex": False, + }, + "decoder": {"type": "ByteLevel"}, + "model": { + "type": "BPE", + "dropout": None, + "unk_token": None, + "continuing_subword_prefix": "", + "end_of_word_suffix": "", + "fuse_unk": False, + "byte_fallback": False, + "ignore_merges": True, + "vocab": vocab, + "merges": merges, + }, + } + if pretty: + return json.dumps(tokenizer_json, indent=2, ensure_ascii=False) + return json.dumps(tokenizer_json, ensure_ascii=False) + + +def tiktoken_model_to_tokenizer_json( + model: str | Path, + *, + pattern: str | None = None, + special_tokens: Mapping[str, int] | None = None, + pretty: bool = False, +) -> str | None: + """ + Convert a ``tiktoken.model`` BPE file to a Hugging Face ``tokenizer.json`` string. + + ``model`` may point to a local ``tiktoken.model`` file, a URL accepted by + ``tiktoken.load.load_tiktoken_bpe``, or a local directory containing + ``tiktoken.model`` and optionally ``tokenizer_config.json``. Passing a model + path returns ``None`` if the optional ``tiktoken`` package is not installed. + """ + try: + from tiktoken.load import load_tiktoken_bpe + except ImportError: + return None + + model_path, config = _load_model_config(model) + if pattern is None: + config_pattern = config.get("pat_str") + if not isinstance(config_pattern, str): + config_pattern = config.get("pattern") + if isinstance(config_pattern, str): + pattern = config_pattern + elif _is_kimi_tiktoken_config(config): + pattern = KIMI_TIKTOKEN_PATTERN + elif _has_tokenizer_identity(config): + logger.warning( + "fastokens does not support converting tiktoken tokenizer config " + "%r: tokenizer_config.json does not define pat_str or pattern, " + "and the tokenizer type is unknown", + model_path, + ) + return None + else: + pattern = DEFAULT_TIKTOKEN_PATTERN + try: + mergeable_ranks = load_tiktoken_bpe(model_path) + except ValueError as exc: + raise ValueError(f"invalid tiktoken model {model_path!r}: {exc}") from exc + added_tokens = ( + _config_added_tokens(config, len(mergeable_ranks)) + if special_tokens is None + else _normalise_special_tokens(special_tokens) + ) + return _tokenizer_json_from_parts( + mergeable_ranks, pattern, added_tokens, pretty=pretty + ) + + +def tiktoken_to_tokenizer_json(encoding: Any, *, pretty: bool = False) -> str | None: + """ + Convert a ``tiktoken`` encoding to a Hugging Face ``tokenizer.json`` string. + + ``encoding`` may be either a ``tiktoken.Encoding`` instance or an encoding + name accepted by ``tiktoken.get_encoding``. Passing an encoding name returns + ``None`` if the optional ``tiktoken`` package is not installed. The returned + JSON can be passed directly to ``fastokens.Tokenizer.from_json_str``. + """ + encoding = _extract_encoding(encoding) + if encoding is None: + return None + try: + mergeable_ranks = encoding._mergeable_ranks + pattern = encoding._pat_str + special_tokens = encoding._special_tokens + except AttributeError as exc: + raise TypeError( + "expected a tiktoken.Encoding or encoding name with " + "_mergeable_ranks, _pat_str, and _special_tokens" + ) from exc + + return _tokenizer_json_from_parts( + mergeable_ranks, + pattern, + _normalise_special_tokens(special_tokens), + pretty=pretty, + ) diff --git a/python/src/lib.rs b/python/src/lib.rs index 45ff6ab..b94b07e 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -1,10 +1,366 @@ -use std::sync::RwLock; - -use pyo3::exceptions::{PyNotImplementedError, PyValueError}; +use std::{ + collections::{HashMap, HashSet}, + sync::{Arc, RwLock}, +}; + +use fastokens::chat_template::minijinja::{Value as TemplateValue, value::ValueKind}; +use numpy::IntoPyArray; +use pyo3::exceptions::{PyNotImplementedError, PyTypeError, PyValueError}; use pyo3::prelude::*; -use pyo3::types::{PyDict, PyList}; +use pyo3::types::{PyAny, PyDict, PyList, PyString}; +use pyo3_async_runtimes::tokio::future_into_py; +use rayon::prelude::*; use serde_json::Value; +fn extract_string_vec(value: &Bound<'_, PyAny>, name: &str) -> PyResult> { + if value.downcast::().is_ok() { + return Err(PyTypeError::new_err(format!( + "{name} must be a sequence or set of strings, not a single string" + ))); + } + + if let Ok(values) = value.extract::>() { + return Ok(values); + } + if let Ok(values) = value.extract::>() { + return Ok(values.into_iter().collect()); + } + Err(PyTypeError::new_err(format!( + "{name} must be a sequence or set of strings" + ))) +} + +/// Convert a Python object into a `serde_json::Value` without a Python JSON +/// serialization round-trip. +fn py_to_json(value: &Bound<'_, PyAny>) -> PyResult { + pythonize::depythonize(value) + .map_err(|e| PyTypeError::new_err(format!("value is not JSON-compatible: {e}"))) +} + +/// Convert optional Python keyword arguments into a JSON object map. +/// +/// `None` becomes an empty map. A provided value must be a dictionary after +/// JSON conversion; otherwise a `TypeError` is raised. +fn py_dict_to_json_map( + dict: Option<&Bound<'_, PyDict>>, +) -> PyResult> { + let Some(dict) = dict else { + return Ok(serde_json::Map::new()); + }; + match py_to_json(dict.as_any())? { + Value::Object(map) => Ok(map), + _ => Err(PyTypeError::new_err("expected a dictionary")), + } +} + +fn pop_bool_context_arg( + map: &mut serde_json::Map, + key: &str, +) -> PyResult> { + match map.remove(key) { + Some(Value::Bool(value)) => Ok(Some(value)), + Some(Value::Null) | None => Ok(None), + Some(value) => Err(PyTypeError::new_err(format!( + "{key} must be a bool, got {value}" + ))), + } +} + +fn consume_render_only_tokenizer_kwargs(map: &mut serde_json::Map) -> PyResult<()> { + for key in [ + "messages", + "conversation", + "add_generation_prompt", + "continue_final_message", + "tools", + "documents", + "special_tokens", + ] { + if map.contains_key(key) { + return Err(PyTypeError::new_err(format!( + "apply_chat_template got multiple values for keyword argument `{key}`" + ))); + } + } + + if pop_bool_context_arg(map, "return_assistant_tokens_mask")?.unwrap_or(false) { + return Err(PyValueError::new_err( + "`return_assistant_tokens_mask=True` requires `return_dict=True` and `tokenize=True`", + )); + } + + if pop_bool_context_arg(map, "return_dict")?.unwrap_or(false) { + return Err(PyValueError::new_err( + "`return_dict=True` requires `tokenize=True`", + )); + } + for key in [ + "padding", + "truncation", + "max_length", + "return_tensors", + "tokenizer_kwargs", + ] { + map.remove(key); + } + Ok(()) +} + +fn is_batched_chat(messages: &Value) -> bool { + messages + .as_array() + .and_then(|messages| messages.first()) + .is_some_and(Value::is_array) +} + +/// Convert a Python object straight into a MiniJinja value, skipping the +/// intermediate `serde_json::Value` tree that `py_to_json` builds. On large +/// conversations that second tree dominates `apply_chat_template` cost. +fn py_to_template_value(value: &Bound<'_, PyAny>) -> PyResult { + pythonize::depythonize(value) + .map_err(|e| PyTypeError::new_err(format!("value is not JSON-compatible: {e}"))) +} + +/// Messages ready for rendering. `continue_final_message` needs to mutate the +/// final message before rendering, which only the JSON representation +/// supports; every other call takes the single-hop MiniJinja conversion. +enum PreparedMessages { + Value(TemplateValue), + Json(Value), +} + +fn is_batched_chat_value(messages: &TemplateValue) -> bool { + messages.kind() == ValueKind::Seq + && messages + .try_iter() + .ok() + .and_then(|mut items| items.next()) + .is_some_and(|first| first.kind() == ValueKind::Seq) +} + +/// A failed `apply_chat_template` call, split by the Python exception type it +/// should surface as. +enum ChatTemplateCallError { + /// Unknown keyword argument rejected by strict validation -> `TypeError`. + UnknownKwarg(String), + /// Template compilation or rendering failure -> `ValueError`. + Render(String), +} + +impl From for PyErr { + fn from(err: ChatTemplateCallError) -> PyErr { + match err { + ChatTemplateCallError::UnknownKwarg(message) => PyTypeError::new_err(message), + ChatTemplateCallError::Render(message) => PyValueError::new_err(message), + } + } +} + +/// An `apply_chat_template` call after all GIL-held preparation. +/// +/// Owns no Python references, so the CPU-bound [`render`](Self::render) half +/// can run wherever the caller likes: under `allow_threads` for the sync +/// binding, on the Tokio blocking pool for the async one. Keeping the split +/// here is what lets both bindings share every line except the launch. +struct PreparedChatTemplateCall { + template: String, + renderer: Option>, + messages: PreparedMessages, + options: fastokens::ChatTemplateOptions, + continuation: Option, + strict_template: bool, +} + +impl PreparedChatTemplateCall { + fn render(self) -> Result { + let Self { + template, + renderer, + messages, + options, + continuation, + strict_template, + } = self; + let run = |renderer: &fastokens::ChatTemplateRenderer| { + if strict_template { + validate_template_kwargs(renderer, &options.extra_context)?; + } + match messages { + PreparedMessages::Value(messages) => renderer.render_value(messages, options), + PreparedMessages::Json(messages) => renderer.render(messages, options), + } + .map_err(|e| ChatTemplateCallError::Render(e.to_string())) + }; + let rendered = match renderer { + Some(renderer) => run(&renderer)?, + None => match fastokens::ChatTemplateRenderer::new(&template) { + Ok(renderer) => run(&renderer)?, + Err(e) => return Err(ChatTemplateCallError::Render(e.to_string())), + }, + }; + match continuation { + Some(final_message) => trim_continue_final_message(rendered, &final_message) + .map_err(ChatTemplateCallError::Render), + None => Ok(rendered), + } + } +} + +/// Reject keyword arguments the template never reads. +/// +/// Typos like `enable_thinkng` otherwise render silently against the +/// template's default branch, producing a wrong prompt with no signal. +fn validate_template_kwargs( + renderer: &fastokens::ChatTemplateRenderer, + extra_context: &serde_json::Map, +) -> Result<(), ChatTemplateCallError> { + let known = renderer.undeclared_variables(); + for key in extra_context.keys() { + if !known.contains(key) { + let suggestion = closest_variable(key, known) + .map(|candidate| format!(". Did you mean `{candidate}`?")) + .unwrap_or_default(); + return Err(ChatTemplateCallError::UnknownKwarg(format!( + "apply_chat_template got an unexpected keyword argument `{key}`: the chat \ + template does not reference it{suggestion} Pass \ + fastokens_strict_template=False to allow unused keyword arguments." + ))); + } + } + Ok(()) +} + +fn closest_variable<'a>( + key: &str, + candidates: &'a std::collections::HashSet, +) -> Option<&'a str> { + let max_distance = 1 + key.len() / 4; + candidates + .iter() + .map(|candidate| (levenshtein(key, candidate), candidate)) + .filter(|(distance, _)| *distance <= max_distance) + .min() + .map(|(_, candidate)| candidate.as_str()) +} + +fn levenshtein(left: &str, right: &str) -> usize { + let left: Vec = left.chars().collect(); + let right: Vec = right.chars().collect(); + let mut row: Vec = (0..=right.len()).collect(); + for (i, l) in left.iter().enumerate() { + let mut previous_diagonal = row[0]; + row[0] = i + 1; + for (j, r) in right.iter().enumerate() { + let substitution = previous_diagonal + usize::from(l != r); + previous_diagonal = row[j + 1]; + row[j + 1] = substitution.min(row[j] + 1).min(previous_diagonal + 1); + } + } + row[right.len()] +} + +const CONTINUE_FINAL_MESSAGE_TAG: &str = "CONTINUE_FINAL_MESSAGE_TAG "; + +fn parse_continue_final_message(value: Option<&Bound<'_, PyAny>>) -> PyResult> { + let Some(value) = value else { + return Ok(None); + }; + if value.is_none() { + return Ok(None); + } + if let Ok(continue_final_message) = value.extract::() { + return Ok(continue_final_message.then(|| "content".to_string())); + } + if let Ok(field) = value.extract::() { + return Ok((!field.is_empty()).then_some(field)); + } + Err(PyTypeError::new_err( + "continue_final_message must be a bool or string", + )) +} + +fn prepare_continue_final_message( + messages: &mut Value, + field: &str, + chat_template: &str, +) -> PyResult { + if !chat_template.contains(field) { + return Err(PyValueError::new_err(format!( + "continue_final_message is set to \"{field}\" but this is not an accepted field in the chat_template" + ))); + } + + let final_message = messages + .as_array_mut() + .and_then(|messages| messages.last_mut()) + .and_then(Value::as_object_mut) + .ok_or_else(|| { + PyValueError::new_err("continue_final_message requires a non-empty message list") + })?; + let final_value = final_message.get_mut(field).ok_or_else(|| { + PyValueError::new_err(format!( + "continue_final_message is set but the final message has no \"{field}\" to continue" + )) + })?; + + match final_value { + Value::String(content) => { + let original = content.clone(); + content.push_str(CONTINUE_FINAL_MESSAGE_TAG); + Ok(original) + } + Value::Array(blocks) => { + for block in blocks.iter_mut().rev() { + if let Some(text_value) = block + .as_object_mut() + .and_then(|block| block.get_mut("text")) + { + if let Some(text) = text_value.as_str().map(str::to_string) { + *text_value = Value::String(format!("{text}{CONTINUE_FINAL_MESSAGE_TAG}")); + return Ok(text); + } + } + } + Err(PyValueError::new_err( + "continue_final_message is set but we could not find any text to continue in the final message", + )) + } + _ => Err(PyValueError::new_err(format!( + "continue_final_message is set but final message field \"{field}\" is not a string or text block list" + ))), + } +} + +fn trim_continue_final_message(rendered: String, final_message: &str) -> Result { + let final_message = final_message.trim(); + let tag = CONTINUE_FINAL_MESSAGE_TAG.trim(); + if !rendered.contains(final_message) || !rendered.contains(tag) { + let final_message = truncate_for_error(final_message, 512); + let truncated = truncate_for_error(&rendered, 4096); + return Err(format!( + "continue_final_message is set but the final message does not appear in the chat after applying the chat template. Final message to continue: {final_message}\nRendered chat:\n{truncated}" + )); + } + let tag_loc = rendered + .rfind(tag) + .expect("tag presence was checked before rfind"); + if rendered[tag_loc..].starts_with(CONTINUE_FINAL_MESSAGE_TAG) { + Ok(rendered[..tag_loc].to_string()) + } else { + Ok(rendered[..tag_loc].trim_end().to_string()) + } +} + +fn truncate_for_error(value: &str, max_chars: usize) -> String { + let mut chars = value.chars(); + let truncated: String = chars.by_ref().take(max_chars).collect(); + if chars.next().is_some() { + format!("{truncated}...") + } else { + truncated + } +} + // --------------------------------------------------------------------------- // PyEncoding // --------------------------------------------------------------------------- @@ -17,72 +373,132 @@ use serde_json::Value; /// raise `NotImplementedError` to match the HuggingFace API surface. #[pyclass(name = "Encoding")] pub struct PyEncoding { - #[pyo3(get, set)] pub ids: Vec, - #[pyo3(get, set)] - pub attention_mask: Vec, - #[pyo3(get, set)] - pub type_ids: Vec, - #[pyo3(get, set)] - pub special_tokens_mask: Vec, + attention_mask: Option>, + type_ids: Option>, + special_tokens_mask: Option>, #[pyo3(get, set)] pub n_sequences: usize, // Backing storage for set-only properties. - _sequence_ids: Vec>, - _word_ids: Vec>, + _sequence_ids: Option>>, + _word_ids: Option>>, } impl PyEncoding { - pub fn make(ids: Vec, attention_mask: Vec) -> Self { - let n = ids.len(); + pub fn make(ids: Vec) -> Self { Self { - type_ids: vec![0u32; n], - special_tokens_mask: vec![0u32; n], + attention_mask: None, + type_ids: None, + special_tokens_mask: None, n_sequences: 1, - _sequence_ids: vec![Some(0); n], - _word_ids: vec![None; n], + _sequence_ids: None, + _word_ids: None, ids, - attention_mask, } } + fn materialized_attention_mask(&self) -> Vec { + self.attention_mask + .clone() + .unwrap_or_else(|| vec![1u32; self.ids.len()]) + } + + fn materialized_type_ids(&self) -> Vec { + self.type_ids + .clone() + .unwrap_or_else(|| vec![0u32; self.ids.len()]) + } + + fn materialized_special_tokens_mask(&self) -> Vec { + self.special_tokens_mask + .clone() + .unwrap_or_else(|| vec![0u32; self.ids.len()]) + } + + fn materialized_sequence_ids(&self) -> Vec> { + self._sequence_ids + .clone() + .unwrap_or_else(|| vec![Some(0); self.ids.len()]) + } + + fn materialized_word_ids(&self) -> Vec> { + self._word_ids + .clone() + .unwrap_or_else(|| vec![None; self.ids.len()]) + } + fn apply_slice(&mut self, start: usize, end: usize) { self.ids = self.ids[start..end].to_vec(); - self.attention_mask = self.attention_mask[start..end].to_vec(); - self.type_ids = self.type_ids[start..end].to_vec(); - self.special_tokens_mask = self.special_tokens_mask[start..end].to_vec(); - self._sequence_ids = self._sequence_ids[start..end].to_vec(); - self._word_ids = self._word_ids[start..end].to_vec(); + if let Some(values) = &self.attention_mask { + self.attention_mask = Some(values[start..end].to_vec()); + } + if let Some(values) = &self.type_ids { + self.type_ids = Some(values[start..end].to_vec()); + } + if let Some(values) = &self.special_tokens_mask { + self.special_tokens_mask = Some(values[start..end].to_vec()); + } + if let Some(values) = &self._sequence_ids { + self._sequence_ids = Some(values[start..end].to_vec()); + } + if let Some(values) = &self._word_ids { + self._word_ids = Some(values[start..end].to_vec()); + } } fn extend_right(&mut self, pad_id: u32, pad_type_id: u32, count: usize) { + let n = self.ids.len(); + let target = n + count; self.ids.extend(vec![pad_id; count]); - self.attention_mask.extend(vec![0u32; count]); - self.type_ids.extend(vec![pad_type_id; count]); - self.special_tokens_mask.extend(vec![0u32; count]); - self._sequence_ids.extend(vec![None; count]); - self._word_ids.extend(vec![None; count]); + let mut attention_mask = self.attention_mask.take().unwrap_or_else(|| vec![1u32; n]); + attention_mask.resize(target, 0u32); + self.attention_mask = Some(attention_mask); + + let mut type_ids = self.type_ids.take().unwrap_or_else(|| vec![0u32; n]); + type_ids.resize(target, pad_type_id); + self.type_ids = Some(type_ids); + + if let Some(mut special_tokens_mask) = self.special_tokens_mask.take() { + special_tokens_mask.resize(target, 0u32); + self.special_tokens_mask = Some(special_tokens_mask); + } + if let Some(mut sequence_ids) = self._sequence_ids.take() { + sequence_ids.resize(target, None); + self._sequence_ids = Some(sequence_ids); + } + if let Some(mut word_ids) = self._word_ids.take() { + word_ids.resize(target, None); + self._word_ids = Some(word_ids); + } } fn extend_left(&mut self, pad_id: u32, pad_type_id: u32, count: usize) { + let n = self.ids.len(); let mut ids = vec![pad_id; count]; ids.extend_from_slice(&self.ids); let mut mask = vec![0u32; count]; - mask.extend_from_slice(&self.attention_mask); + mask.extend_from_slice(&self.attention_mask.take().unwrap_or_else(|| vec![1u32; n])); let mut type_ids = vec![pad_type_id; count]; - type_ids.extend_from_slice(&self.type_ids); - let mut special = vec![0u32; count]; - special.extend_from_slice(&self.special_tokens_mask); - let mut seq_ids = vec![None; count]; - seq_ids.extend_from_slice(&self._sequence_ids); - let mut word_ids = vec![None; count]; - word_ids.extend_from_slice(&self._word_ids); + type_ids.extend_from_slice(&self.type_ids.take().unwrap_or_else(|| vec![0u32; n])); self.ids = ids; - self.attention_mask = mask; - self.type_ids = type_ids; - self.special_tokens_mask = special; - self._sequence_ids = seq_ids; - self._word_ids = word_ids; + self.attention_mask = Some(mask); + self.type_ids = Some(type_ids); + + if let Some(special_tokens_mask) = self.special_tokens_mask.take() { + let mut special = vec![0u32; count]; + special.extend_from_slice(&special_tokens_mask); + self.special_tokens_mask = Some(special); + } + if let Some(sequence_ids) = self._sequence_ids.take() { + let mut seq_ids = vec![None; count]; + seq_ids.extend_from_slice(&sequence_ids); + self._sequence_ids = Some(seq_ids); + } + if let Some(word_ids) = self._word_ids.take() { + let mut word_ids_padded = vec![None; count]; + word_ids_padded.extend_from_slice(&word_ids); + self._word_ids = Some(word_ids_padded); + } } } @@ -91,9 +507,45 @@ impl PyEncoding { #[new] #[pyo3(signature = (ids, attention_mask = None))] fn new(ids: Vec, attention_mask: Option>) -> Self { - let n = ids.len(); - let mask = attention_mask.unwrap_or_else(|| vec![1u32; n]); - Self::make(ids, mask) + let mut encoding = Self::make(ids); + encoding.attention_mask = attention_mask; + encoding + } + + #[getter] + fn ids(&self) -> Vec { + self.ids.clone() + } + #[setter] + fn set_ids(&mut self, value: Vec) { + self.ids = value; + } + + #[getter] + fn attention_mask(&self) -> Vec { + self.materialized_attention_mask() + } + #[setter] + fn set_attention_mask(&mut self, value: Vec) { + self.attention_mask = Some(value); + } + + #[getter] + fn type_ids(&self) -> Vec { + self.materialized_type_ids() + } + #[setter] + fn set_type_ids(&mut self, value: Vec) { + self.type_ids = Some(value); + } + + #[getter] + fn special_tokens_mask(&self) -> Vec { + self.materialized_special_tokens_mask() + } + #[setter] + fn set_special_tokens_mask(&mut self, value: Vec) { + self.special_tokens_mask = Some(value); } fn __len__(&self) -> usize { @@ -104,6 +556,62 @@ impl PyEncoding { format!("Encoding(num_tokens={})", self.ids.len()) } + /// Move selected fields into NumPy uint32 arrays. + /// + /// This drains the encoding's per-token fields. The returned dict + /// contains only requested arrays; unrequested fields are cleared. + #[pyo3(signature = ( + ids = true, + attention_mask = false, + type_ids = false, + special_tokens_mask = false + ))] + fn into_numpy<'py>( + &mut self, + py: Python<'py>, + ids: bool, + attention_mask: bool, + type_ids: bool, + special_tokens_mask: bool, + ) -> PyResult> { + if !(ids || attention_mask || type_ids || special_tokens_mask) { + return Err(PyValueError::new_err( + "at least one field must be selected for into_numpy()", + )); + } + + let out = PyDict::new(py); + let n = self.ids.len(); + let ids_vec = std::mem::take(&mut self.ids); + let attention_mask_vec = self.attention_mask.take(); + let type_ids_vec = self.type_ids.take(); + let special_tokens_mask_vec = self.special_tokens_mask.take(); + self._sequence_ids = None; + self._word_ids = None; + self.n_sequences = 0; + + if ids { + out.set_item("ids", ids_vec.into_pyarray(py))?; + } + if attention_mask { + let attention_mask_vec = attention_mask_vec.unwrap_or_else(|| vec![1u32; n]); + out.set_item("attention_mask", attention_mask_vec.into_pyarray(py))?; + } + if type_ids { + let type_ids_vec = type_ids_vec.unwrap_or_else(|| vec![0u32; n]); + out.set_item("type_ids", type_ids_vec.into_pyarray(py))?; + } + if special_tokens_mask { + let special_tokens_mask_vec = special_tokens_mask_vec.unwrap_or_else(|| vec![0u32; n]); + out.set_item( + "special_tokens_mask", + special_tokens_mask_vec.into_pyarray(py), + )?; + } + + Ok(out) + } + // -- Properties that raise NotImplementedError ---------------------- #[getter] @@ -133,7 +641,7 @@ impl PyEncoding { } #[setter] fn set_sequence_ids(&mut self, value: Vec>) { - self._sequence_ids = value; + self._sequence_ids = Some(value); } #[getter] @@ -144,7 +652,7 @@ impl PyEncoding { } #[setter] fn set_word_ids(&mut self, value: Vec>) { - self._word_ids = value; + self._word_ids = Some(value); } #[getter] @@ -155,7 +663,7 @@ impl PyEncoding { } #[setter] fn set_words(&mut self, value: Vec>) { - self._word_ids = value; + self._word_ids = Some(value); } /// Always empty — fastokens does not produce overflowing sequences. @@ -170,7 +678,7 @@ impl PyEncoding { fn set_sequence_id(&mut self, sequence_id: i64) { let n = self.ids.len(); - self._sequence_ids = vec![Some(sequence_id); n]; + self._sequence_ids = Some(vec![Some(sequence_id); n]); } // -- Positional mapping (all raise NotImplementedError) ------------- @@ -291,22 +799,22 @@ impl PyEncoding { for enc_py in &encodings { let enc = enc_py.borrow(py); ids.extend_from_slice(&enc.ids); - attention_mask.extend_from_slice(&enc.attention_mask); - type_ids.extend_from_slice(&enc.type_ids); - special_tokens_mask.extend_from_slice(&enc.special_tokens_mask); + attention_mask.extend_from_slice(&enc.materialized_attention_mask()); + type_ids.extend_from_slice(&enc.materialized_type_ids()); + special_tokens_mask.extend_from_slice(&enc.materialized_special_tokens_mask()); n_sequences += enc.n_sequences; - seq_ids.extend_from_slice(&enc._sequence_ids); - word_ids.extend_from_slice(&enc._word_ids); + seq_ids.extend_from_slice(&enc.materialized_sequence_ids()); + word_ids.extend_from_slice(&enc.materialized_word_ids()); } PyEncoding { ids, - attention_mask, - type_ids, - special_tokens_mask, + attention_mask: Some(attention_mask), + type_ids: Some(type_ids), + special_tokens_mask: Some(special_tokens_mask), n_sequences, - _sequence_ids: seq_ids, - _word_ids: word_ids, + _sequence_ids: Some(seq_ids), + _word_ids: Some(word_ids), } } } @@ -315,6 +823,7 @@ impl PyEncoding { // TruncationParams / PaddingParams // --------------------------------------------------------------------------- +#[derive(Clone)] struct TruncationParams { max_length: usize, stride: usize, @@ -322,6 +831,7 @@ struct TruncationParams { direction: String, } +#[derive(Clone)] struct PaddingParams { direction: String, pad_id: u32, @@ -332,8 +842,7 @@ struct PaddingParams { } fn build_encoding(ids: Vec, pad: Option<&PaddingParams>, target: usize) -> PyEncoding { - let n = ids.len(); - let mut enc = PyEncoding::make(ids, vec![1u32; n]); + let mut enc = PyEncoding::make(ids); if let Some(p) = pad { enc.pad(target, &p.direction, p.pad_id, p.pad_type_id, &p.pad_token); } @@ -380,6 +889,9 @@ struct TokenizerState { pad: Option, /// Cached JSON of the current post-processor (for the getter). post_processor_json: Option, + chat_template: Option, + chat_template_renderer: Option>, + special_tokens: serde_json::Map, } impl TokenizerState { @@ -404,6 +916,43 @@ impl TokenizerState { } } + fn encode_batch_encodings( + &self, + inputs: &[String], + add_special_tokens: bool, + split_special_tokens: bool, + ) -> Result, String> { + let mut batch: Vec> = inputs + .par_iter() + .map(|s| { + self.inner + .encode_with_options(s.as_str(), add_special_tokens, split_special_tokens) + .map_err(|e| e.to_string()) + }) + .collect::, _>>()?; + + for ids in &mut batch { + self.do_truncate(ids); + } + + let pad_target: Option = self.pad.as_ref().map(|p| { + let max_len = batch.iter().map(|ids| ids.len()).max().unwrap_or(0); + let base = p.length.unwrap_or(max_len).max(max_len); + match p.pad_to_multiple_of { + Some(m) if m > 0 => (base + m - 1) / m * m, + _ => base, + } + }); + + Ok(batch + .into_iter() + .map(|ids| { + let target = pad_target.unwrap_or(ids.len()); + build_encoding(ids, self.pad.as_ref(), target) + }) + .collect()) + } + /// Parse `json`, update the Rust post-processor in place, and cache the JSON. fn update_post_processor_json(&mut self, json: &str) -> PyResult<()> { use fastokens::json_structs::PostProcessorConfig; @@ -421,10 +970,43 @@ impl TokenizerState { } } +#[pyclass(name = "StructuralTokenConfig")] +struct PyStructuralTokenConfig { + inner: Arc, +} + +#[pymethods] +impl PyStructuralTokenConfig { + /// Build constant structural-token state for rendered-prompt encoding. + /// + /// Include every token string that the rendered template may use as a + /// structural boundary, including tag-like non-special added tokens. + #[new] + #[pyo3(signature = (structural_tokens, non_special_added_tokens = None))] + fn new( + structural_tokens: &Bound<'_, PyAny>, + non_special_added_tokens: Option<&Bound<'_, PyAny>>, + ) -> PyResult { + let structural_tokens = extract_string_vec(structural_tokens, "structural_tokens")?; + let non_special_added_tokens = match non_special_added_tokens { + Some(value) => extract_string_vec(value, "non_special_added_tokens")? + .into_iter() + .collect(), + None => HashSet::new(), + }; + let inner = + fastokens::StructuralTokenConfig::new(&structural_tokens, &non_special_added_tokens) + .map_err(|e| PyValueError::new_err(e.to_string()))?; + Ok(Self { + inner: Arc::new(inner), + }) + } +} + /// An LLM tokenizer backed by `tokenizer.json`. #[pyclass(name = "Tokenizer")] struct PyTokenizer { - state: RwLock, + state: Arc>, } impl PyTokenizer { @@ -436,6 +1018,103 @@ impl PyTokenizer { self.state.write().expect("PyTokenizer state lock poisoned") } + /// GIL-held half of `apply_chat_template`, shared by the sync and async + /// bindings: validate arguments, convert Python inputs to owned Rust + /// values, and resolve the template. The returned call owns everything + /// the CPU-bound render needs, so the bindings differ only in where they + /// run it. + #[allow(clippy::too_many_arguments)] + fn prepare_chat_template_call( + &self, + py: Python<'_>, + messages: &Bound<'_, PyAny>, + chat_template: Option, + tokenize: bool, + add_generation_prompt: bool, + continue_final_message: Option>, + tools: Option<&Bound<'_, PyAny>>, + documents: Option<&Bound<'_, PyAny>>, + special_tokens: Option<&Bound<'_, PyDict>>, + kwargs: Option<&Bound<'_, PyDict>>, + ) -> PyResult { + if tokenize { + return Err(PyNotImplementedError::new_err( + "apply_chat_template(tokenize=True) is not supported by fastokens; render with tokenize=False and call encode separately", + )); + } + + let continue_field = parse_continue_final_message( + continue_final_message.as_ref().map(|value| value.bind(py)), + )?; + if add_generation_prompt && continue_field.is_some() { + return Err(PyValueError::new_err( + "continue_final_message and add_generation_prompt are not compatible", + )); + } + + let mut messages = if continue_field.is_some() { + PreparedMessages::Json(py_to_json(messages)?) + } else { + PreparedMessages::Value(py_to_template_value(messages)?) + }; + let batched = match &messages { + PreparedMessages::Json(messages) => is_batched_chat(messages), + PreparedMessages::Value(messages) => is_batched_chat_value(messages), + }; + if batched { + return Err(PyNotImplementedError::new_err( + "batched conversations are not supported by fastokens apply_chat_template", + )); + } + let tools = tools.map(py_to_json).transpose()?; + let documents = documents.map(py_to_json).transpose()?; + let per_call_special_tokens = py_dict_to_json_map(special_tokens)?; + let mut extra_context = py_dict_to_json_map(kwargs)?; + consume_render_only_tokenizer_kwargs(&mut extra_context)?; + let strict_template = + pop_bool_context_arg(&mut extra_context, "fastokens_strict_template")?.unwrap_or(true); + + let state = self.read(); + let (template, renderer) = match chat_template { + Some(template) => (template, None), + None => { + let template = state + .chat_template + .clone() + .ok_or_else(|| PyValueError::new_err("chat_template must be provided"))?; + let renderer = state + .chat_template_renderer + .clone() + .ok_or_else(|| PyValueError::new_err("chat_template must be provided"))?; + (template, Some(renderer)) + } + }; + let continuation = match (&continue_field, &mut messages) { + (Some(field), PreparedMessages::Json(messages)) => { + Some(prepare_continue_final_message(messages, field, &template)?) + } + _ => None, + }; + let mut options = fastokens::ChatTemplateOptions { + add_generation_prompt, + continue_final_message: continue_field.is_some(), + tools, + documents, + special_tokens: state.special_tokens.clone(), + extra_context, + }; + options.special_tokens.extend(per_call_special_tokens); + + Ok(PreparedChatTemplateCall { + template, + renderer, + messages, + options, + continuation, + strict_template, + }) + } + /// Build from a raw JSON string, extracting the post-processor field so /// the getter can return it without needing to re-serialize. fn build_from_str(json: &str, py: Python<'_>) -> PyResult { @@ -449,12 +1128,15 @@ impl PyTokenizer { .allow_threads(|| fastokens::Tokenizer::from_json(value).map_err(|e| e.to_string())) .map_err(PyValueError::new_err)?; Ok(Self { - state: RwLock::new(TokenizerState { + state: Arc::new(RwLock::new(TokenizerState { inner, trunc: None, pad: None, post_processor_json, - }), + chat_template: None, + chat_template_renderer: None, + special_tokens: fastokens::chat_template::default_special_tokens(), + })), }) } } @@ -496,6 +1178,145 @@ impl PyTokenizer { Self::build_from_str(&json, py) } + /// Set the default chat template used by `apply_chat_template`. + fn set_chat_template(&self, chat_template: Option) -> PyResult<()> { + let renderer = chat_template + .as_deref() + .map(fastokens::ChatTemplateRenderer::new) + .transpose() + .map(|renderer| renderer.map(Arc::new)) + .map_err(|e| PyValueError::new_err(e.to_string()))?; + + let mut state = self.write(); + state.chat_template = chat_template; + state.chat_template_renderer = renderer; + Ok(()) + } + + /// Variables the configured chat template reads from the render context, + /// or ``None`` if no chat template is set. + /// + /// Derived statically from the template, so dynamic lookups are invisible. + /// Useful for validating or documenting per-model template kwargs. + #[getter] + fn chat_template_variables(&self) -> Option> { + let state = self.read(); + let renderer = state.chat_template_renderer.as_ref()?; + let mut variables: Vec = renderer.undeclared_variables().iter().cloned().collect(); + variables.sort(); + Some(variables) + } + + /// Set persistent special-token variables used by `apply_chat_template`. + fn set_special_tokens(&self, special_tokens: Option<&Bound<'_, PyDict>>) -> PyResult<()> { + let special_tokens = match special_tokens { + Some(special_tokens) => py_dict_to_json_map(Some(special_tokens))?, + None => fastokens::chat_template::default_special_tokens(), + }; + + let mut state = self.write(); + state.special_tokens = special_tokens; + Ok(()) + } + + /// Render a HuggingFace-style chat template. + #[pyo3(signature = ( + messages, + chat_template = None, + tokenize = false, + add_generation_prompt = false, + continue_final_message = None, + add_special_tokens = false, + tools = None, + documents = None, + special_tokens = None, + **kwargs + ))] + fn apply_chat_template( + &self, + messages: &Bound<'_, PyAny>, + chat_template: Option, + tokenize: bool, + add_generation_prompt: bool, + continue_final_message: Option>, + add_special_tokens: bool, + tools: Option<&Bound<'_, PyAny>>, + documents: Option<&Bound<'_, PyAny>>, + special_tokens: Option<&Bound<'_, PyDict>>, + kwargs: Option<&Bound<'_, PyDict>>, + py: Python<'_>, + ) -> PyResult { + let _ = add_special_tokens; + let call = self.prepare_chat_template_call( + py, + messages, + chat_template, + tokenize, + add_generation_prompt, + continue_final_message, + tools, + documents, + special_tokens, + kwargs, + )?; + let rendered = py.allow_threads(|| call.render()).map_err(PyErr::from)?; + Ok(rendered.into_pyobject(py)?.unbind().into_any()) + } + + /// Render a HuggingFace-style chat template without blocking the event + /// loop: returns an awaitable that renders on a background thread. + /// + /// Input conversion still happens synchronously at call time (it reads + /// live Python objects); only the render is offloaded. + #[pyo3(signature = ( + messages, + chat_template = None, + tokenize = false, + add_generation_prompt = false, + continue_final_message = None, + add_special_tokens = false, + tools = None, + documents = None, + special_tokens = None, + **kwargs + ))] + fn async_apply_chat_template<'py>( + &self, + messages: &Bound<'py, PyAny>, + chat_template: Option, + tokenize: bool, + add_generation_prompt: bool, + continue_final_message: Option>, + add_special_tokens: bool, + tools: Option<&Bound<'py, PyAny>>, + documents: Option<&Bound<'py, PyAny>>, + special_tokens: Option<&Bound<'py, PyDict>>, + kwargs: Option<&Bound<'py, PyDict>>, + py: Python<'py>, + ) -> PyResult> { + let _ = add_special_tokens; + let call = self.prepare_chat_template_call( + py, + messages, + chat_template, + tokenize, + add_generation_prompt, + continue_final_message, + tools, + documents, + special_tokens, + kwargs, + )?; + future_into_py(py, async move { + let rendered = pyo3_async_runtimes::tokio::get_runtime() + .spawn_blocking(move || call.render()) + .await + .map_err(|e| PyValueError::new_err(e.to_string()))? + .map_err(PyErr::from)?; + Python::with_gil(|py| Ok(rendered.into_pyobject(py)?.unbind().into_any())) + }) + } + // ── Post-processor ──────────────────────────────────────────────── /// The current post-processor, or ``None`` if none is configured. @@ -626,70 +1447,167 @@ impl PyTokenizer { /// /// Truncation and padding configured via `enable_truncation` / /// `enable_padding` are applied before returning. - #[pyo3(signature = (input, add_special_tokens = false))] + #[pyo3(signature = (input, add_special_tokens = false, split_special_tokens = false))] fn encode( &self, input: &str, add_special_tokens: bool, + split_special_tokens: bool, py: Python<'_>, ) -> PyResult> { - let state = self.read(); - let mut ids = state - .inner - .encode_with_special_tokens(input, add_special_tokens) - .map_err(|e| PyValueError::new_err(e.to_string()))?; - state.do_truncate(&mut ids); - let target = state.single_pad_target(ids.len()); + let encoding = py + .allow_threads(|| { + let state = self.read(); + let mut ids = state + .inner + .encode_with_options(input, add_special_tokens, split_special_tokens) + .map_err(|e| e.to_string())?; + state.do_truncate(&mut ids); + let target = state.single_pad_target(ids.len()); + Ok::(build_encoding(ids, state.pad.as_ref(), target)) + }) + .map_err(PyValueError::new_err)?; - Py::new(py, build_encoding(ids, state.pad.as_ref(), target)) + Py::new(py, encoding) + } + + /// Encode text segments with per-segment control over added-token matching. + /// + /// Each segment is ``(text, allow_special)``. When ``allow_special`` is + /// false, all added-token matching is bypassed for that segment, matching + /// ``tiktoken.encode(..., disallowed_special=())`` for user/tool text. + /// + /// ``tiktoken_safe=True`` is the default and reproduces legacy tiktoken + /// tokenizer chunking for token-ID parity. Pass ``False`` only when you + /// explicitly want whole-segment BPE encoding. + #[pyo3(signature = (segments, add_special_tokens = false, tiktoken_safe = true))] + fn encode_segments( + &self, + segments: Vec<(String, bool)>, + add_special_tokens: bool, + tiktoken_safe: bool, + py: Python<'_>, + ) -> PyResult> { + let encoding = py + .allow_threads(|| { + let state = self.read(); + let segment_iter = segments.iter().map(|(text, allow)| (text, *allow)); + let mut ids = if tiktoken_safe { + state + .inner + .encode_segments_tiktoken_safe(segment_iter, add_special_tokens) + } else { + state + .inner + .encode_segments(segment_iter, add_special_tokens) + } + .map_err(|e| e.to_string())?; + state.do_truncate(&mut ids); + let target = state.single_pad_target(ids.len()); + Ok::(build_encoding(ids, state.pad.as_ref(), target)) + }) + .map_err(PyValueError::new_err)?; + + Py::new(py, encoding) + } + + /// Encode a rendered prompt containing structural token strings. + /// + /// Structural tokens are emitted as token IDs. Text spans are encoded with + /// split-special-token behavior, after restoring placeholders to their + /// original user text. Keep `add_special_tokens=false` when replacing an + /// existing rendered chat-template encode path. + #[pyo3(signature = ( + input, + structural_config, + placeholder_map = None, + add_special_tokens = false + ))] + fn encode_with_structural_tokens( + &self, + input: &str, + structural_config: PyRef<'_, PyStructuralTokenConfig>, + placeholder_map: Option>, + add_special_tokens: bool, + py: Python<'_>, + ) -> PyResult> { + let placeholder_map = placeholder_map.unwrap_or_default(); + let structural_config = Arc::clone(&structural_config.inner); + let ids = py + .allow_threads(|| { + let state = self.read(); + state + .inner + .encode_with_structural_tokens( + input, + &structural_config, + &placeholder_map, + add_special_tokens, + ) + .map_err(|e| e.to_string()) + }) + .map_err(PyValueError::new_err)?; + Py::new(py, PyEncoding::make(ids)) } /// Encode a batch of inputs in parallel. /// /// Truncation is applied per-sequence; padding (if enabled) pads the /// batch to a uniform length. - #[pyo3(signature = (inputs, add_special_tokens = false))] + #[pyo3(signature = (inputs, add_special_tokens = false, split_special_tokens = false))] fn encode_batch( &self, inputs: Vec, add_special_tokens: bool, + split_special_tokens: bool, py: Python<'_>, ) -> PyResult>> { - use rayon::prelude::*; - - let state = self.read(); - let mut batch: Vec> = inputs - .par_iter() - .map(|s| { - state - .inner - .encode_with_special_tokens(s.as_str(), add_special_tokens) - .map_err(|e| PyValueError::new_err(e.to_string())) + let encodings = py + .allow_threads(|| { + let state = self.read(); + state.encode_batch_encodings(&inputs, add_special_tokens, split_special_tokens) }) - .collect::>>()?; - - for ids in &mut batch { - state.do_truncate(ids); - } - - let pad_target: Option = state.pad.as_ref().map(|p| { - let max_len = batch.iter().map(|ids| ids.len()).max().unwrap_or(0); - let base = p.length.unwrap_or(max_len).max(max_len); - match p.pad_to_multiple_of { - Some(m) if m > 0 => (base + m - 1) / m * m, - _ => base, - } - }); - - batch + .map_err(PyValueError::new_err)?; + encodings .into_iter() - .map(|ids| { - let target = pad_target.unwrap_or(ids.len()); - Py::new(py, build_encoding(ids, state.pad.as_ref(), target)) - }) + .map(|encoding| Py::new(py, encoding)) .collect() } + /// Encode a batch of inputs in parallel and return a Python awaitable. + /// + /// Truncation is applied per-sequence; padding (if enabled) pads the + /// batch to a uniform length. + #[pyo3(signature = (inputs, add_special_tokens = false, split_special_tokens = false))] + fn async_encode_batch<'py>( + &self, + py: Python<'py>, + inputs: Vec, + add_special_tokens: bool, + split_special_tokens: bool, + ) -> PyResult> { + let state = Arc::clone(&self.state); + + future_into_py(py, async move { + let encodings = pyo3_async_runtimes::tokio::get_runtime() + .spawn_blocking(move || { + let state = state.read().expect("PyTokenizer state lock poisoned"); + state.encode_batch_encodings(&inputs, add_special_tokens, split_special_tokens) + }) + .await + .map_err(|e| PyValueError::new_err(e.to_string()))? + .map_err(PyValueError::new_err)?; + + Python::with_gil(|py| { + let encodings: PyResult>> = encodings + .into_iter() + .map(|encoding| Py::new(py, encoding)) + .collect(); + Ok(encodings?.into_pyobject(py)?.unbind().into_any()) + }) + }) + } + // ── Post-processing ─────────────────────────────────────────────── /// Apply the post-processor to an existing encoding. @@ -714,8 +1632,7 @@ impl PyTokenizer { } let ids = encoding.borrow(py).ids.clone(); let new_ids = self.read().inner.post_process(ids, true); - let n = new_ids.len(); - Py::new(py, PyEncoding::make(new_ids, vec![1u32; n])) + Py::new(py, PyEncoding::make(new_ids)) } /// Return the number of special tokens added for a single or pair sequence. @@ -766,6 +1683,34 @@ impl PyTokenizer { .map_err(|e| PyValueError::new_err(e.to_string())) } + /// Decode a batch of token ID sequences and return a Python awaitable. + #[pyo3(signature = (sentences, skip_special_tokens = false))] + fn async_decode_batch<'py>( + &self, + py: Python<'py>, + sentences: Vec>, + skip_special_tokens: bool, + ) -> PyResult> { + let state = Arc::clone(&self.state); + + future_into_py(py, async move { + let decoded = pyo3_async_runtimes::tokio::get_runtime() + .spawn_blocking(move || { + let state = state.read().expect("PyTokenizer state lock poisoned"); + let refs: Vec<&[u32]> = sentences.iter().map(Vec::as_slice).collect(); + state + .inner + .decode_batch(&refs, skip_special_tokens) + .map_err(|e| e.to_string()) + }) + .await + .map_err(|e| PyValueError::new_err(e.to_string()))? + .map_err(PyValueError::new_err)?; + + Python::with_gil(|py| Ok(decoded.into_pyobject(py)?.unbind().into_any())) + }) + } + // ── Vocabulary ──────────────────────────────────────────────────── /// Look up the token ID for a string. @@ -802,9 +1747,9 @@ mod tests { enc.pad(5, "right", 0u32, 1u32, "[PAD]"); assert_eq!(enc.ids, vec![10u32, 20, 30, 0, 0]); - assert_eq!(enc.attention_mask, vec![1u32, 1, 1, 0, 0]); + assert_eq!(enc.materialized_attention_mask(), vec![1u32, 1, 1, 0, 0]); assert_eq!( - enc.type_ids, + enc.materialized_type_ids(), vec![0u32, 0, 0, 1, 1], "padded positions should carry pad_type_id=1 in type_ids" ); @@ -825,8 +1770,8 @@ mod tests { let enc = build_encoding(vec![10u32, 20, 30], Some(&pad), 5); assert_eq!(enc.ids, vec![10u32, 20, 30, 0, 0]); - assert_eq!(enc.attention_mask, vec![1u32, 1, 1, 0, 0]); - assert_eq!(enc.type_ids, vec![0u32, 0, 0, 1, 1]); + assert_eq!(enc.materialized_attention_mask(), vec![1u32, 1, 1, 0, 0]); + assert_eq!(enc.materialized_type_ids(), vec![0u32, 0, 0, 1, 1]); } #[test] @@ -842,8 +1787,8 @@ mod tests { let enc = build_encoding(vec![10u32, 20, 30], Some(&pad), 5); assert_eq!(enc.ids, vec![0u32, 0, 10, 20, 30]); - assert_eq!(enc.attention_mask, vec![0u32, 0, 1, 1, 1]); - assert_eq!(enc.type_ids, vec![7u32, 7, 0, 0, 0]); + assert_eq!(enc.materialized_attention_mask(), vec![0u32, 0, 1, 1, 1]); + assert_eq!(enc.materialized_type_ids(), vec![7u32, 7, 0, 0, 0]); } } @@ -905,6 +1850,7 @@ impl PyDecodeStream { fn _native(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; Ok(()) diff --git a/python/tests/test_async_stub.py b/python/tests/test_async_stub.py new file mode 100644 index 0000000..0026ec4 --- /dev/null +++ b/python/tests/test_async_stub.py @@ -0,0 +1,106 @@ +import inspect +import unittest + +from fastokens._compat import _TokenizerShim +from fastokens._native import Tokenizer + + +TOKENIZER_JSON = """ +{ + "version": "1.0", + "added_tokens": [], + "normalizer": null, + "pre_tokenizer": null, + "post_processor": null, + "decoder": null, + "model": { + "type": "BPE", + "dropout": null, + "unk_token": null, + "continuing_subword_prefix": "", + "end_of_word_suffix": "", + "fuse_unk": false, + "byte_fallback": false, + "vocab": { + "h": 0, + "e": 1, + "l": 2, + "o": 3, + " ": 4, + "w": 5, + "r": 6, + "d": 7, + "he": 8, + "hel": 9, + "hell": 10, + "hello": 11, + "wo": 12, + "wor": 13, + "worl": 14, + "world": 15 + }, + "merges": [ + "h e", + "he l", + "hel l", + "hell o", + "w o", + "wo r", + "wor l", + "worl d" + ] + } +} +""" + +INPUTS = ["hello", "hello world", "world"] + + +def _ids(encodings): + return [encoding.ids for encoding in encodings] + + +class NativeAsyncStubTests(unittest.IsolatedAsyncioTestCase): + async def test_async_encode_batch_matches_sync(self) -> None: + tokenizer = Tokenizer.from_json_str(TOKENIZER_JSON) + + sync_encodings = tokenizer.encode_batch(INPUTS) + awaitable = tokenizer.async_encode_batch(INPUTS) + async_encodings = await awaitable + + self.assertTrue(inspect.isawaitable(awaitable)) + self.assertEqual(_ids(sync_encodings), _ids(async_encodings)) + self.assertEqual(sync_encodings[1].ids, [11, 4, 15]) + + async def test_async_decode_batch_matches_sync(self) -> None: + tokenizer = Tokenizer.from_json_str(TOKENIZER_JSON) + encoded = tokenizer.encode_batch(INPUTS) + batch_ids = [encoding.ids for encoding in encoded] + + sync_decoded = tokenizer.decode_batch(batch_ids) + awaitable = tokenizer.async_decode_batch(batch_ids) + async_decoded = await awaitable + + self.assertTrue(inspect.isawaitable(awaitable)) + self.assertEqual(sync_decoded, async_decoded) + self.assertEqual(async_decoded, INPUTS) + + +class CompatAsyncStubTests(unittest.IsolatedAsyncioTestCase): + async def test_async_shim_methods_are_usable(self) -> None: + tokenizer = _TokenizerShim.from_str(TOKENIZER_JSON) + + encode_awaitable = tokenizer.async_encode_batch(INPUTS) + async_encodings = await encode_awaitable + batch_ids = [encoding.ids for encoding in async_encodings] + decode_awaitable = tokenizer.async_decode_batch(batch_ids) + async_decoded = await decode_awaitable + + self.assertTrue(inspect.isawaitable(encode_awaitable)) + self.assertTrue(inspect.isawaitable(decode_awaitable)) + self.assertEqual(batch_ids, [[11], [11, 4, 15], [15]]) + self.assertEqual(async_decoded, INPUTS) + + +if __name__ == "__main__": + unittest.main() diff --git a/python/tests/test_chat_template.py b/python/tests/test_chat_template.py new file mode 100644 index 0000000..3835047 --- /dev/null +++ b/python/tests/test_chat_template.py @@ -0,0 +1,553 @@ +from pathlib import Path + +import pytest + +from fastokens._compat import _TokenizerShim +from fastokens._native import Tokenizer +from test_async_stub import TOKENIZER_JSON + +REPO_ROOT = Path(__file__).resolve().parents[2] + +CHAT_TEMPLATE = ( + "{%- for message in messages -%}" + "{%- if message['role'] == 'user' -%}" + "user: {{ message['content'] }}" + "{%- elif message['role'] == 'assistant' -%}" + "assistant: {{ message.get('content') }}" + "{%- endif -%}" + "{%- endfor -%}" + "{%- if add_generation_prompt -%}assistant: {%- endif -%}" +) + + +def test_native_apply_chat_template_renders_string() -> None: + tokenizer = Tokenizer.from_json_str(TOKENIZER_JSON) + + rendered = tokenizer.apply_chat_template( + [{"role": "user", "content": "hello"}], + chat_template=CHAT_TEMPLATE, + tokenize=False, + add_generation_prompt=True, + ) + + assert rendered == "user: helloassistant:" + + +def test_native_apply_chat_template_rejects_tokenize_true() -> None: + tokenizer = Tokenizer.from_json_str(TOKENIZER_JSON) + + with pytest.raises(NotImplementedError, match="tokenize=True"): + tokenizer.apply_chat_template( + [{"role": "user", "content": "hello"}], + chat_template="{{ messages[0]['content'] }}", + tokenize=True, + ) + + +def test_native_apply_chat_template_consumes_render_only_tokenizer_kwargs() -> None: + tokenizer = Tokenizer.from_json_str(TOKENIZER_JSON) + template = ( + "{% if return_dict is defined %}return_dict leaked{% endif %}" + "{% if tokenizer_kwargs is defined %}tokenizer_kwargs leaked{% endif %}" + "{% if padding is defined %}padding leaked{% endif %}" + "{{ messages[0]['content'] }}" + ) + + rendered = tokenizer.apply_chat_template( + [{"role": "user", "content": "hello"}], + chat_template=template, + tokenize=False, + return_dict=False, + tokenizer_kwargs={"add_prefix_space": True}, + padding=False, + truncation=False, + max_length=128, + return_tensors=None, + ) + + assert rendered == "hello" + + +def test_native_apply_chat_template_omits_null_special_tokens() -> None: + tokenizers = pytest.importorskip("tokenizers") + transformers = pytest.importorskip("transformers") + template = ( + "{% if pad_token is defined %}defined{% else %}undefined{% endif %}" + "{{ pad_token }}" + ) + messages = [{"role": "user", "content": "hello"}] + + hf_tokenizer = transformers.PreTrainedTokenizerFast( + tokenizer_object=tokenizers.Tokenizer.from_str(TOKENIZER_JSON), + bos_token="", + eos_token="", + ) + fast_tokenizer = Tokenizer.from_json_str(TOKENIZER_JSON) + + expected = hf_tokenizer.apply_chat_template( + messages, + chat_template=template, + tokenize=False, + ) + rendered = fast_tokenizer.apply_chat_template( + messages, + chat_template=template, + tokenize=False, + ) + + assert rendered == expected == "undefined" + + +def test_native_apply_chat_template_uses_persistent_special_tokens() -> None: + tokenizer = Tokenizer.from_json_str(TOKENIZER_JSON) + tokenizer.set_special_tokens({"bos_token": "", "eos_token": ""}) + + rendered = tokenizer.apply_chat_template( + [{"role": "user", "content": "hello"}], + chat_template="{{ bos_token }}{{ messages[0]['content'] }}{{ eos_token }}", + tokenize=False, + ) + + assert rendered == "hello" + + +def test_native_apply_chat_template_rejects_reserved_kwargs() -> None: + tokenizer = Tokenizer.from_json_str(TOKENIZER_JSON) + + with pytest.raises(TypeError, match="multiple values"): + tokenizer.apply_chat_template( + [{"role": "user", "content": "hello"}], + chat_template="{{ messages }}", + tokenize=False, + messages="shadow", + ) + + +def test_native_apply_chat_template_rejects_assistant_mask_without_tokenize() -> None: + tokenizer = Tokenizer.from_json_str(TOKENIZER_JSON) + + with pytest.raises(ValueError, match="return_assistant_tokens_mask=True"): + tokenizer.apply_chat_template( + [{"role": "user", "content": "hello"}], + chat_template="{{ messages[0]['content'] }}", + tokenize=False, + return_assistant_tokens_mask=True, + ) + + +def test_native_apply_chat_template_rejects_batched_conversations() -> None: + tokenizer = Tokenizer.from_json_str(TOKENIZER_JSON) + + with pytest.raises(NotImplementedError, match="batched conversations"): + tokenizer.apply_chat_template( + [[{"role": "user", "content": "hello"}]], + chat_template="{{ messages[0]['content'] }}", + tokenize=False, + ) + + +def test_native_apply_chat_template_supports_kwargs_and_special_tokens() -> None: + tokenizer = Tokenizer.from_json_str(TOKENIZER_JSON) + template = ( + "{{ bos_token }}" + "{{ messages[0].content.upper() }}" + "{{ payload | tojson }}" + "{{ eos_token }}" + ) + + rendered = tokenizer.apply_chat_template( + [{"role": "user", "content": "hello"}], + chat_template=template, + tokenize=False, + payload={"text": "&'"}, + special_tokens={"bos_token": "", "eos_token": ""}, + ) + + assert rendered == r"""HELLO{"text": "&'"}""" + + +def test_native_apply_chat_template_rejects_continue_and_generation_prompt() -> None: + tokenizers = pytest.importorskip("tokenizers") + transformers = pytest.importorskip("transformers") + template = "{{ messages[0]['content'] }}" + messages = [{"role": "assistant", "content": "partial"}] + hf_tokenizer = transformers.PreTrainedTokenizerFast( + tokenizer_object=tokenizers.Tokenizer.from_str(TOKENIZER_JSON), + ) + fast_tokenizer = Tokenizer.from_json_str(TOKENIZER_JSON) + + with pytest.raises(ValueError, match="continue_final_message"): + hf_tokenizer.apply_chat_template( + messages, + chat_template=template, + tokenize=False, + add_generation_prompt=True, + continue_final_message=True, + ) + with pytest.raises(ValueError, match="continue_final_message"): + fast_tokenizer.apply_chat_template( + messages, + chat_template=template, + tokenize=False, + add_generation_prompt=True, + continue_final_message=True, + ) + + +def test_native_apply_chat_template_strips_generation_markers() -> None: + tokenizer = Tokenizer.from_json_str(TOKENIZER_JSON) + + rendered = tokenizer.apply_chat_template( + [], + chat_template="a{% generation %}b{% endgeneration %}c", + tokenize=False, + ) + + assert rendered == "abc" + + +def test_native_apply_chat_template_supports_generation_whitespace_control() -> None: + tokenizer = Tokenizer.from_json_str(TOKENIZER_JSON) + + rendered = tokenizer.apply_chat_template( + [], + chat_template="a\n{%- generation -%}\nb\n{%- endgeneration -%}\nc", + tokenize=False, + ) + + assert rendered == "abc" + + +def test_native_apply_chat_template_rejects_invalid_strftime_without_panic() -> None: + tokenizer = Tokenizer.from_json_str(TOKENIZER_JSON) + + with pytest.raises(ValueError, match="invalid strftime_now format"): + tokenizer.apply_chat_template( + [], + chat_template="{{ strftime_now('%Q %#z') }}", + tokenize=False, + ) + + +def test_continue_final_message_error_truncates_rendered_prompt() -> None: + tokenizer = Tokenizer.from_json_str(TOKENIZER_JSON) + message = "x" * 10_000 + + with pytest.raises(ValueError) as exc_info: + tokenizer.apply_chat_template( + [{"role": "assistant", "content": message}], + chat_template="{{ messages[0]['content'] | replace('x', 'y') }}", + tokenize=False, + continue_final_message=True, + ) + + error = str(exc_info.value) + assert len(error) < 5_000 + assert "..." in error + + +@pytest.mark.parametrize( + ("fixture", "kwargs"), + [ + ("glm-5.2", {"enable_thinking": False}), + ("gpt-oss-120b", {"builtin_tools": []}), + ], +) +def test_vendored_chat_templates_match_transformers(fixture: str, kwargs: dict) -> None: + tokenizers = pytest.importorskip("tokenizers") + transformers = pytest.importorskip("transformers") + template = (REPO_ROOT / "vendored_tokenizers" / fixture / "chat_template.jinja").read_text() + messages = [{"role": "user", "content": "hello"}] + + hf_tokenizer = transformers.PreTrainedTokenizerFast( + tokenizer_object=tokenizers.Tokenizer.from_str(TOKENIZER_JSON), + bos_token="", + eos_token="", + ) + fast_tokenizer = Tokenizer.from_json_str(TOKENIZER_JSON) + + expected = hf_tokenizer.apply_chat_template( + messages, + chat_template=template, + tokenize=False, + **kwargs, + ) + rendered = fast_tokenizer.apply_chat_template( + messages, + chat_template=template, + tokenize=False, + special_tokens={"bos_token": "", "eos_token": ""}, + **kwargs, + ) + + assert rendered == expected + + +def test_kimi_25_template_matches_transformers_with_tools_and_thinking() -> None: + tokenizers = pytest.importorskip("tokenizers") + transformers = pytest.importorskip("transformers") + template = ( + REPO_ROOT / "vendored_tokenizers" / "kimi-k2.5" / "chat_template.jinja" + ).read_text() + messages = [ + {"role": "system", "content": "Be terse."}, + { + "role": "user", + "name": "alice", + "content": [ + {"type": "text", "text": "look this up"}, + {"type": "image_url", "image_url": {"url": "https://example.test/i.png"}}, + ], + }, + { + "role": "assistant", + "content": "I will check.", + "tool_calls": [ + { + "id": "call_1", + "function": { + "name": "search", + "arguments": {"query": "rust tokenizer", "top_k": 2}, + }, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "result text"}, + { + "role": "assistant", + "content": "final answer", + "reasoning_content": "private reasoning", + }, + ] + tools = [ + { + "type": "function", + "function": { + "name": "search", + "description": "Search docs", + "parameters": { + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], + }, + }, + } + ] + + hf_tokenizer = transformers.PreTrainedTokenizerFast( + tokenizer_object=tokenizers.Tokenizer.from_str(TOKENIZER_JSON), + bos_token="", + eos_token="", + ) + fast_tokenizer = Tokenizer.from_json_str(TOKENIZER_JSON) + + expected = hf_tokenizer.apply_chat_template( + messages, + tools=tools, + chat_template=template, + tokenize=False, + add_generation_prompt=True, + thinking=False, + tools_ts_str="", + ) + rendered = fast_tokenizer.apply_chat_template( + messages, + tools=tools, + chat_template=template, + tokenize=False, + add_generation_prompt=True, + thinking=False, + tools_ts_str="", + special_tokens={"bos_token": "", "eos_token": ""}, + ) + + assert rendered == expected + + +@pytest.mark.parametrize( + ("messages", "continue_final_message"), + [ + ([{"role": "assistant", "content": "partial answer"}], True), + ( + [ + { + "role": "assistant", + "content": [ + {"type": "image_url", "image_url": {"url": "https://example.test/i.png"}}, + {"type": "text", "text": "partial answer"}, + ], + } + ], + "content", + ), + ], +) +def test_continue_final_message_matches_transformers( + messages: list[dict], continue_final_message: bool | str +) -> None: + tokenizers = pytest.importorskip("tokenizers") + transformers = pytest.importorskip("transformers") + template = ( + "{%- for message in messages -%}" + "{%- if message.content is string -%}" + "{{ message.content }}{{ eos_token }}" + "{%- else -%}" + "{%- for content in message.content -%}" + "{%- if content.text is defined -%}{{ content.text }}{{ eos_token }}{%- endif -%}" + "{%- endfor -%}" + "{%- endif -%}" + "{%- endfor -%}" + ) + + hf_tokenizer = transformers.PreTrainedTokenizerFast( + tokenizer_object=tokenizers.Tokenizer.from_str(TOKENIZER_JSON), + bos_token="", + eos_token="", + ) + fast_tokenizer = Tokenizer.from_json_str(TOKENIZER_JSON) + + expected = hf_tokenizer.apply_chat_template( + messages, + chat_template=template, + tokenize=False, + continue_final_message=continue_final_message, + ) + rendered = fast_tokenizer.apply_chat_template( + messages, + chat_template=template, + tokenize=False, + continue_final_message=continue_final_message, + special_tokens={"bos_token": "", "eos_token": ""}, + ) + + assert rendered == expected == "partial answer" + + +def test_strict_template_rejects_unknown_kwargs_with_suggestion() -> None: + tokenizer = Tokenizer.from_json_str(TOKENIZER_JSON) + template = "{% if enable_thinking %}think: {% endif %}{{ messages[0]['content'] }}" + + with pytest.raises(TypeError, match=r"enable_thinkng.*Did you mean `enable_thinking`"): + tokenizer.apply_chat_template( + [{"role": "user", "content": "hello"}], + chat_template=template, + tokenize=False, + enable_thinkng=True, + ) + + +def test_strict_template_rejects_unknown_kwargs_on_cached_renderer() -> None: + tokenizer = Tokenizer.from_json_str(TOKENIZER_JSON) + tokenizer.set_chat_template("{{ messages[0]['content'] }}") + + with pytest.raises(TypeError, match="unexpected keyword argument `enable_thinking`"): + tokenizer.apply_chat_template( + [{"role": "user", "content": "hello"}], + tokenize=False, + enable_thinking=False, + ) + + +def test_strict_template_opt_out_allows_unused_kwargs() -> None: + tokenizer = Tokenizer.from_json_str(TOKENIZER_JSON) + + rendered = tokenizer.apply_chat_template( + [{"role": "user", "content": "hello"}], + chat_template="{{ messages[0]['content'] }}", + tokenize=False, + enable_thinkng=True, + fastokens_strict_template=False, + ) + + assert rendered == "hello" + + +def test_chat_template_variables_property() -> None: + tokenizer = Tokenizer.from_json_str(TOKENIZER_JSON) + assert tokenizer.chat_template_variables is None + + tokenizer.set_chat_template("{% if enable_thinking %}{{ messages }}{% endif %}") + + assert tokenizer.chat_template_variables == ["enable_thinking", "messages"] + + +def test_shim_apply_chat_template_forwards_to_native() -> None: + tokenizer = _TokenizerShim.from_str(TOKENIZER_JSON) + + rendered = tokenizer.apply_chat_template( + [{"role": "assistant", "content": "world"}], + chat_template=CHAT_TEMPLATE, + ) + + assert rendered == "assistant: world" + + +def test_async_apply_chat_template_matches_sync() -> None: + import asyncio + + tokenizer = Tokenizer.from_json_str(TOKENIZER_JSON) + messages = [{"role": "user", "content": "hello"}] + sync_rendered = tokenizer.apply_chat_template( + messages, chat_template=CHAT_TEMPLATE, tokenize=False, add_generation_prompt=True + ) + + async def render() -> str: + return await tokenizer.async_apply_chat_template( + messages, chat_template=CHAT_TEMPLATE, tokenize=False, add_generation_prompt=True + ) + + assert asyncio.run(render()) == sync_rendered == "user: helloassistant:" + + +def test_async_apply_chat_template_propagates_errors() -> None: + import asyncio + + tokenizer = Tokenizer.from_json_str(TOKENIZER_JSON) + + async def render_tokenize_true() -> str: + return await tokenizer.async_apply_chat_template( + [], chat_template="x", tokenize=True + ) + + with pytest.raises(NotImplementedError, match="tokenize=True"): + asyncio.run(render_tokenize_true()) + + async def render_unknown_kwarg() -> str: + return await tokenizer.async_apply_chat_template( + [], chat_template="{{ messages }}", tokenize=False, bogus_kwarg=1 + ) + + with pytest.raises(TypeError, match="unexpected keyword argument"): + asyncio.run(render_unknown_kwarg()) + + +def test_shim_async_apply_chat_template_forwards_to_native() -> None: + import asyncio + + tokenizer = _TokenizerShim.from_str(TOKENIZER_JSON) + tokenizer.set_chat_template(CHAT_TEMPLATE) + + async def render() -> str: + return await tokenizer.async_apply_chat_template( + [{"role": "assistant", "content": "world"}] + ) + + assert asyncio.run(render()) == "assistant: world" + + +def test_shim_preserves_special_tokens_across_copy_and_pickle() -> None: + import copy + import pickle + + tokenizer = _TokenizerShim.from_str(TOKENIZER_JSON) + tokenizer.set_chat_template("{{ bos_token }}x") + tokenizer.set_special_tokens({"bos_token": ""}) + assert tokenizer.apply_chat_template([], tokenize=False) == "x" + + clones = { + "deepcopy": copy.deepcopy(tokenizer), + "pickle": pickle.loads(pickle.dumps(tokenizer)), + "from_shim": _TokenizerShim(tokenizer), + } + for name, clone in clones.items(): + assert clone.apply_chat_template([], tokenize=False) == "x", name diff --git a/python/tests/test_patch_transformers.py b/python/tests/test_patch_transformers.py index dbb6ebe..5f40ba4 100644 --- a/python/tests/test_patch_transformers.py +++ b/python/tests/test_patch_transformers.py @@ -55,3 +55,86 @@ def test_unpatch_restores_backend(): assert not isinstance(tok._tokenizer, _TokenizerShim), ( "backend should be original tokenizers.Tokenizer after unpatch" ) + + +def test_apply_chat_template_not_routed_by_default(): + """Default patch must leave transformers' apply_chat_template untouched.""" + import fastokens + + original = transformers.PreTrainedTokenizerBase.apply_chat_template + fastokens.patch_transformers() + + assert transformers.PreTrainedTokenizerBase.apply_chat_template is original + + +def test_apply_chat_template_opt_in_routes_native(monkeypatch): + """With apply_chat_template=True, render-only calls go through the shim.""" + import fastokens + + calls = [] + original_shim_apply = _TokenizerShim.apply_chat_template + + def spy(self, *args, **kwargs): + calls.append(kwargs.get("chat_template") is not None) + return original_shim_apply(self, *args, **kwargs) + + monkeypatch.setattr(_TokenizerShim, "apply_chat_template", spy) + fastokens.patch_transformers(apply_chat_template=True) + + tok = transformers.AutoTokenizer.from_pretrained(MODEL) + messages = [{"role": "user", "content": "hi"}] + rendered = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + + assert calls, "shim apply_chat_template was not called" + + fastokens.unpatch_transformers() + reference = transformers.AutoTokenizer.from_pretrained(MODEL) + expected = reference.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + assert rendered == expected + + +def test_apply_chat_template_tokenize_true_falls_through(monkeypatch): + """tokenize=True is unsupported natively and must use transformers' path.""" + import fastokens + + fastokens.patch_transformers(apply_chat_template=True) + + tok = transformers.AutoTokenizer.from_pretrained(MODEL) + encoded = tok.apply_chat_template([{"role": "user", "content": "hi"}], tokenize=True) + # v4 returns a list of ids, v5 returns a BatchEncoding mapping. + ids = encoded["input_ids"] if not isinstance(encoded, list) else encoded + assert len(ids) > 0 + + +def test_apply_chat_template_env_var_overrides_arg(monkeypatch): + """FASTOKENS_APPLY_CHAT_TEMPLATE=0 wins over apply_chat_template=True.""" + import fastokens + + original = transformers.PreTrainedTokenizerBase.apply_chat_template + monkeypatch.setenv(fastokens.APPLY_CHAT_TEMPLATE_ENV, "0") + fastokens.patch_transformers(apply_chat_template=True) + + assert transformers.PreTrainedTokenizerBase.apply_chat_template is original + + +def test_apply_chat_template_env_var_enables(monkeypatch): + """FASTOKENS_APPLY_CHAT_TEMPLATE=1 enables routing without the arg.""" + import fastokens + + original = transformers.PreTrainedTokenizerBase.apply_chat_template + monkeypatch.setenv(fastokens.APPLY_CHAT_TEMPLATE_ENV, "1") + fastokens.patch_transformers() + + assert transformers.PreTrainedTokenizerBase.apply_chat_template is not original + + +def test_unpatch_restores_apply_chat_template(): + """Unpatch must restore the original transformers method.""" + import fastokens + + original = transformers.PreTrainedTokenizerBase.apply_chat_template + fastokens.patch_transformers(apply_chat_template=True) + assert transformers.PreTrainedTokenizerBase.apply_chat_template is not original + + fastokens.unpatch_transformers() + assert transformers.PreTrainedTokenizerBase.apply_chat_template is original diff --git a/python/tests/test_split_special_tokens.py b/python/tests/test_split_special_tokens.py new file mode 100644 index 0000000..b3d00dd --- /dev/null +++ b/python/tests/test_split_special_tokens.py @@ -0,0 +1,431 @@ +import json +import unittest +from copy import deepcopy + +import numpy as np + +from fastokens._compat import _TokenizerShim +from fastokens._native import Encoding, StructuralTokenConfig, Tokenizer + + +BOS_TOKEN = "[BOS]" +BOS_ID = 100 +SPECIAL_ID = 101 +NON_SPECIAL_ID = 102 + + +def _tokenizer_json(special_token: str, non_special_token: str) -> str: + chars = sorted(set(BOS_TOKEN + special_token + non_special_token + " hello")) + vocab = {ch: i for i, ch in enumerate(chars)} + return json.dumps( + { + "version": "1.0", + "added_tokens": [ + { + "id": BOS_ID, + "content": BOS_TOKEN, + "single_word": False, + "lstrip": False, + "rstrip": False, + "normalized": False, + "special": True, + }, + { + "id": SPECIAL_ID, + "content": special_token, + "single_word": False, + "lstrip": False, + "rstrip": False, + "normalized": False, + "special": True, + }, + { + "id": NON_SPECIAL_ID, + "content": non_special_token, + "single_word": False, + "lstrip": False, + "rstrip": False, + "normalized": False, + "special": False, + }, + ], + "normalizer": None, + "pre_tokenizer": None, + "post_processor": { + "type": "TemplateProcessing", + "single": [ + {"SpecialToken": {"id": BOS_TOKEN, "type_id": 0}}, + {"Sequence": {"id": "A", "type_id": 0}}, + ], + "pair": [], + "special_tokens": { + BOS_TOKEN: { + "id": BOS_TOKEN, + "ids": [BOS_ID], + "tokens": [BOS_TOKEN], + } + }, + }, + "decoder": None, + "model": { + "type": "BPE", + "dropout": None, + "unk_token": None, + "continuing_subword_prefix": "", + "end_of_word_suffix": "", + "fuse_unk": False, + "byte_fallback": False, + "vocab": vocab, + "merges": [], + }, + } + ) + + +def _char_ids(tokenizer_json: str, text: str) -> list[int]: + vocab = json.loads(tokenizer_json)["model"]["vocab"] + return [vocab[ch] for ch in text] + + +class SplitSpecialTokensTests(unittest.TestCase): + CASES = [ + ( + "kimi2.5", + "", + "<|tool_calls_section_begin|>", + ), + ( + "ascii-endoftext", + "<|endoftext|>", + "", + ), + ( + "ascii-short", + "", + "", + ), + ] + + def test_native_split_special_tokens_matches_kimi_style_expectations(self) -> None: + for model_name, special_token, non_special_token in self.CASES: + with self.subTest(model_name=model_name): + tokenizer_json = _tokenizer_json(special_token, non_special_token) + tokenizer = Tokenizer.from_json_str(tokenizer_json) + expected_text_ids = _char_ids(tokenizer_json, special_token) + + default_ids = tokenizer.encode( + special_token, + add_special_tokens=False, + ).ids + split_ids = tokenizer.encode( + special_token, + add_special_tokens=False, + split_special_tokens=True, + ).ids + non_special_ids = tokenizer.encode( + non_special_token, + add_special_tokens=False, + split_special_tokens=True, + ).ids + with_bos_ids = tokenizer.encode( + special_token, + add_special_tokens=True, + split_special_tokens=True, + ).ids + + self.assertEqual(default_ids, [SPECIAL_ID]) + self.assertEqual(split_ids, expected_text_ids) + self.assertEqual(non_special_ids, [NON_SPECIAL_ID]) + self.assertEqual(with_bos_ids, [BOS_ID, *expected_text_ids]) + + def test_native_split_special_tokens_handles_specials_inside_text(self) -> None: + for model_name, special_token, _non_special_token in self.CASES: + with self.subTest(model_name=model_name): + tokenizer_json = _tokenizer_json(special_token, "") + tokenizer = Tokenizer.from_json_str(tokenizer_json) + text = f"hello{special_token} hello" + + default_ids = tokenizer.encode( + text, + add_special_tokens=False, + ).ids + split_ids = tokenizer.encode( + text, + add_special_tokens=False, + split_special_tokens=True, + ).ids + + self.assertIn(SPECIAL_ID, default_ids) + self.assertEqual(split_ids, _char_ids(tokenizer_json, text)) + + def test_shim_forwards_split_special_tokens_keyword(self) -> None: + for model_name, special_token, _non_special_token in self.CASES: + with self.subTest(model_name=model_name): + tokenizer_json = _tokenizer_json(special_token, "") + tokenizer = _TokenizerShim.from_str(tokenizer_json) + + self.assertEqual( + tokenizer.encode( + special_token, + add_special_tokens=False, + ).ids, + [SPECIAL_ID], + ) + self.assertEqual( + tokenizer.encode( + special_token, + add_special_tokens=False, + split_special_tokens=True, + ).ids, + _char_ids(tokenizer_json, special_token), + ) + + def test_shim_respects_encode_special_tokens_property(self) -> None: + for model_name, special_token, _non_special_token in self.CASES: + with self.subTest(model_name=model_name): + tokenizer_json = _tokenizer_json(special_token, "") + tokenizer = _TokenizerShim.from_str(tokenizer_json) + tokenizer.encode_special_tokens = True + + self.assertEqual( + tokenizer.encode(special_token, add_special_tokens=False).ids, + _char_ids(tokenizer_json, special_token), + ) + self.assertEqual( + tokenizer.encode_batch( + [special_token], + add_special_tokens=False, + )[0].ids, + _char_ids(tokenizer_json, special_token), + ) + + def test_shim_explicit_split_special_tokens_overrides_property(self) -> None: + for model_name, special_token, _non_special_token in self.CASES: + with self.subTest(model_name=model_name): + tokenizer_json = _tokenizer_json(special_token, "") + tokenizer = _TokenizerShim.from_str(tokenizer_json) + tokenizer.encode_special_tokens = True + + self.assertEqual( + tokenizer.encode( + special_token, + add_special_tokens=False, + split_special_tokens=False, + ).ids, + [SPECIAL_ID], + ) + + def test_shim_copy_preserves_encode_special_tokens_property(self) -> None: + for model_name, special_token, _non_special_token in self.CASES: + with self.subTest(model_name=model_name): + tokenizer_json = _tokenizer_json(special_token, "") + tokenizer = _TokenizerShim.from_str(tokenizer_json) + tokenizer.encode_special_tokens = True + + wrapped = _TokenizerShim(tokenizer) + copied = deepcopy(tokenizer) + + self.assertTrue(wrapped.encode_special_tokens) + self.assertEqual( + wrapped.encode(special_token, add_special_tokens=False).ids, + _char_ids(tokenizer_json, special_token), + ) + self.assertTrue(copied.encode_special_tokens) + self.assertEqual( + copied.encode(special_token, add_special_tokens=False).ids, + _char_ids(tokenizer_json, special_token), + ) + + def test_native_encode_with_structural_tokens(self) -> None: + tokenizer_json = _tokenizer_json("", "") + tokenizer = Tokenizer.from_json_str(tokenizer_json) + think_placeholder = "\ue000STRUCTTOK_0\ue000" + tool_placeholder = "\ue000STRUCTTOK_1\ue000" + structural_config = StructuralTokenConfig( + {"", ""}, + {""}, + ) + + ids = tokenizer.encode_with_structural_tokens( + f"hello {think_placeholder} {tool_placeholder}", + structural_config, + { + think_placeholder: "", + tool_placeholder: "", + }, + ).ids + + self.assertEqual( + ids, + [ + *_char_ids(tokenizer_json, "hello "), + SPECIAL_ID, + *_char_ids(tokenizer_json, " "), + *_char_ids(tokenizer_json, ""), + *_char_ids(tokenizer_json, " "), + NON_SPECIAL_ID, + *_char_ids(tokenizer_json, " "), + *_char_ids(tokenizer_json, ""), + ], + ) + + def test_encoding_into_numpy_moves_selected_fields(self) -> None: + encoding = Encoding([1, 2, 3]) + + arrays = encoding.into_numpy( + attention_mask=True, + type_ids=True, + special_tokens_mask=True, + ) + + self.assertEqual( + set(arrays), + {"ids", "attention_mask", "type_ids", "special_tokens_mask"}, + ) + ids = arrays["ids"] + self.assertEqual(ids.dtype, np.dtype("uint32")) + np.testing.assert_array_equal(ids, np.array([1, 2, 3], dtype=np.uint32)) + np.testing.assert_array_equal( + arrays["attention_mask"], + np.array([1, 1, 1], dtype=np.uint32), + ) + np.testing.assert_array_equal( + arrays["type_ids"], + np.array([0, 0, 0], dtype=np.uint32), + ) + np.testing.assert_array_equal( + arrays["special_tokens_mask"], + np.array([0, 0, 0], dtype=np.uint32), + ) + self.assertEqual(encoding.ids, []) + self.assertEqual(encoding.attention_mask, []) + self.assertEqual(encoding.type_ids, []) + self.assertEqual(encoding.special_tokens_mask, []) + self.assertEqual(encoding.n_sequences, 0) + + def test_encoding_into_numpy_can_skip_ids(self) -> None: + encoding = Encoding([1, 2, 3]) + + arrays = encoding.into_numpy(ids=False, attention_mask=True) + + self.assertEqual(set(arrays), {"attention_mask"}) + np.testing.assert_array_equal( + arrays["attention_mask"], + np.array([1, 1, 1], dtype=np.uint32), + ) + self.assertEqual(encoding.ids, []) + self.assertEqual(encoding.attention_mask, []) + + def test_encoding_into_numpy_rejects_empty_selection(self) -> None: + encoding = Encoding([1, 2, 3]) + + with self.assertRaisesRegex(ValueError, "at least one field"): + encoding.into_numpy( + ids=False, + attention_mask=False, + type_ids=False, + special_tokens_mask=False, + ) + + def test_native_encode_with_structural_tokens_can_add_special_tokens(self) -> None: + tokenizer_json = _tokenizer_json("", "") + tokenizer = Tokenizer.from_json_str(tokenizer_json) + structural_config = StructuralTokenConfig({""}) + + ids = tokenizer.encode_with_structural_tokens( + "hello ", + structural_config, + add_special_tokens=True, + ).ids + + self.assertEqual( + ids, + [BOS_ID, *_char_ids(tokenizer_json, "hello "), SPECIAL_ID], + ) + + def test_structural_token_config_rejects_single_string(self) -> None: + with self.assertRaisesRegex(TypeError, "not a single string"): + StructuralTokenConfig("") + + with self.assertRaisesRegex(TypeError, "not a single string"): + StructuralTokenConfig({""}, "") + + def test_native_encode_with_structural_tokens_keeps_bare_non_special_added_tokens( + self, + ) -> None: + tokenizer_json = _tokenizer_json("", "magic") + tokenizer = Tokenizer.from_json_str(tokenizer_json) + structural_config = StructuralTokenConfig({""}) + + ids = tokenizer.encode_with_structural_tokens( + "hello magic ", + structural_config, + ).ids + + self.assertEqual( + ids, + [ + *_char_ids(tokenizer_json, "hello "), + NON_SPECIAL_ID, + *_char_ids(tokenizer_json, " "), + SPECIAL_ID, + ], + ) + + def test_native_encode_with_structural_tokens_ignores_padding_and_truncation( + self, + ) -> None: + tokenizer_json = _tokenizer_json("", "") + tokenizer = Tokenizer.from_json_str(tokenizer_json) + tokenizer.enable_truncation(max_length=2) + tokenizer.enable_padding(length=8, pad_id=99) + structural_config = StructuralTokenConfig({""}) + + encoding = tokenizer.encode_with_structural_tokens( + "hello ", + structural_config, + ) + + self.assertEqual( + encoding.ids, + [*_char_ids(tokenizer_json, "hello "), SPECIAL_ID], + ) + self.assertEqual(encoding.attention_mask, [1] * len(encoding.ids)) + + def test_shim_forwards_encode_with_structural_tokens(self) -> None: + tokenizer_json = _tokenizer_json("", "") + tokenizer = _TokenizerShim.from_str(tokenizer_json) + think_placeholder = "\ue000STRUCTTOK_0\ue000" + structural_config = StructuralTokenConfig({""}) + + ids = tokenizer.encode_with_structural_tokens( + f"{think_placeholder}", + structural_config, + {think_placeholder: ""}, + ).ids + + self.assertEqual( + ids, + [SPECIAL_ID, *_char_ids(tokenizer_json, "")], + ) + + def test_shim_rejects_unknown_encode_kwargs(self) -> None: + tokenizer = _TokenizerShim.from_str(_tokenizer_json("", "")) + + with self.assertRaisesRegex(TypeError, "return_offsets_mapping"): + tokenizer.encode( + "hello", + add_special_tokens=False, + return_offsets_mapping=True, + ) + + with self.assertRaisesRegex(TypeError, "return_attention_mask"): + tokenizer.encode_batch( + ["hello"], + add_special_tokens=False, + return_attention_mask=True, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/python/tests/test_tiktoken_conversion.py b/python/tests/test_tiktoken_conversion.py new file mode 100644 index 0000000..d927f42 --- /dev/null +++ b/python/tests/test_tiktoken_conversion.py @@ -0,0 +1,474 @@ +import base64 +import builtins +from collections import Counter +import gzip +import json +import logging +from pathlib import Path +import sys +import types + +import pytest + +from fastokens import ( + Tokenizer, + tiktoken_model_to_tokenizer_json, + tiktoken_to_tokenizer_json, +) + +VENDORED_DIR = Path(__file__).parents[2] / "vendored_tokenizers" + + +@pytest.fixture +def hide_tiktoken(monkeypatch: pytest.MonkeyPatch) -> None: + original_import = builtins.__import__ + + def import_without_tiktoken(name, *args, **kwargs): + if name == "tiktoken" or name.startswith("tiktoken."): + raise ImportError("No module named 'tiktoken'") + return original_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", import_without_tiktoken) + + +def _create_test_encoding(): + tiktoken = pytest.importorskip("tiktoken") + return tiktoken.Encoding( + name="fastokens-test", + pat_str=r"(?s).+", + mergeable_ranks={ + b"a": 0, + b"b": 1, + b"c": 2, + b"d": 3, + b"ab": 4, + b"cd": 5, + b"abcd": 6, + b" ": 7, + b"!": 8, + b"c!": 9, + }, + special_tokens={"<|end|>": 100}, + ) + + +def _write_tiktoken_model(path: Path, encoding) -> None: + path.write_text( + "\n".join( + f"{base64.b64encode(token).decode()} {rank}" + for token, rank in encoding._mergeable_ranks.items() + ) + ) + + +def _merge_counter(merges: list[list[str]]) -> Counter[tuple[str, str]]: + return Counter(tuple(merge) for merge in merges) + + +def _special_token_map(added_tokens: list[dict[str, object]]) -> dict[str, int]: + return {t["content"]: t["id"] for t in added_tokens if t.get("special")} + + +def test_tiktoken_to_tokenizer_json_matches_encoding() -> None: + encoding = _create_test_encoding() + tokenizer = Tokenizer.from_json_str(tiktoken_to_tokenizer_json(encoding)) + + texts = [ + "abcd", + "ab cd", + "abc!", + ] + for text in texts: + assert tokenizer.encode(text, add_special_tokens=False).ids == encoding.encode( + text + ) + + +def test_tiktoken_to_tokenizer_json_with_encoding_name( + monkeypatch: pytest.MonkeyPatch, +) -> None: + tiktoken = pytest.importorskip("tiktoken") + encoding = _create_test_encoding() + monkeypatch.setattr(tiktoken, "get_encoding", lambda name: encoding) + + tokenizer_json = tiktoken_to_tokenizer_json("fastokens-test") + tokenizer = Tokenizer.from_json_str(tokenizer_json) + + assert tokenizer.encode("abcd", add_special_tokens=False).ids == [6] + + +def test_tiktoken_to_tokenizer_json_preserves_special_tokens() -> None: + encoding = _create_test_encoding() + tokenizer_json = tiktoken_to_tokenizer_json(encoding) + config = json.loads(tokenizer_json) + special_token = "<|end|>" + special_id = 100 + + assert { + "id": special_id, + "content": special_token, + "single_word": False, + "lstrip": False, + "rstrip": False, + "normalized": False, + "special": True, + } in config["added_tokens"] + + tokenizer = Tokenizer.from_json_str(tokenizer_json) + assert tokenizer.encode(special_token, add_special_tokens=False).ids == [special_id] + + +def test_tiktoken_to_tokenizer_json_returns_none_without_optional_tiktoken( + hide_tiktoken, +) -> None: + assert tiktoken_to_tokenizer_json("cl100k_base") is None + + +def test_tiktoken_model_to_tokenizer_json_matches_model_file(tmp_path) -> None: + pytest.importorskip("tiktoken") + encoding = _create_test_encoding() + model_path = tmp_path / "tiktoken.model" + _write_tiktoken_model(model_path, encoding) + + tokenizer_json = tiktoken_model_to_tokenizer_json( + model_path, + pattern=encoding._pat_str, + special_tokens={"<|end|>": 100}, + ) + tokenizer = Tokenizer.from_json_str(tokenizer_json) + + for text in ["abcd", "ab cd", "abc!"]: + assert tokenizer.encode(text, add_special_tokens=False).ids == encoding.encode( + text + ) + assert tokenizer.encode("<|end|>", add_special_tokens=False).ids == [100] + + +def test_tiktoken_model_to_tokenizer_json_reads_model_directory(tmp_path) -> None: + tiktoken = pytest.importorskip("tiktoken") + encoding = _create_test_encoding() + _write_tiktoken_model(tmp_path / "tiktoken.model", encoding) + (tmp_path / "tokenizer_config.json").write_text( + json.dumps( + { + "pat_str": encoding._pat_str, + "added_tokens_decoder": { + "100": { + "content": "<|end|>", + "special": True, + } + }, + } + ) + ) + + tokenizer_json = tiktoken_model_to_tokenizer_json(tmp_path) + tokenizer = Tokenizer.from_json_str(tokenizer_json) + + assert tokenizer.encode("abcd", add_special_tokens=False).ids == [6] + assert tokenizer.encode("<|end|>", add_special_tokens=False).ids == [100] + + +def test_tiktoken_model_to_tokenizer_json_expands_sparse_kimi_added_tokens( + tmp_path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + mergeable_ranks = {b"a": 0, b"b": 1, b"c": 2, b"d": 3} + tiktoken_module = types.ModuleType("tiktoken") + tiktoken_load_module = types.ModuleType("tiktoken.load") + tiktoken_load_module.load_tiktoken_bpe = lambda path: mergeable_ranks + tiktoken_module.load = tiktoken_load_module + monkeypatch.setitem(sys.modules, "tiktoken", tiktoken_module) + monkeypatch.setitem(sys.modules, "tiktoken.load", tiktoken_load_module) + + (tmp_path / "tiktoken.model").write_text("") + (tmp_path / "tokenizer_config.json").write_text( + json.dumps( + { + "tokenizer_class": "TikTokenTokenizer", + "auto_map": { + "AutoTokenizer": "tokenization_kimi.TikTokenTokenizer", + }, + "added_tokens_decoder": { + "6": { + "content": "<|tool_call_begin|>", + "special": False, + "single_word": False, + "lstrip": False, + "rstrip": False, + "normalized": False, + }, + "8": { + "content": "[PAD]", + "special": True, + "single_word": False, + "lstrip": False, + "rstrip": False, + "normalized": False, + }, + }, + } + ) + ) + + tokenizer_json = tiktoken_model_to_tokenizer_json(tmp_path) + config = json.loads(tokenizer_json) + added = config["added_tokens"] + + assert [token["id"] for token in added] == [4, 5, 6, 7, 8] + assert [token["content"] for token in added] == [ + "<|reserved_token_4|>", + "<|reserved_token_5|>", + "<|tool_call_begin|>", + "<|reserved_token_7|>", + "[PAD]", + ] + assert [token["special"] for token in added] == [True, True, False, True, True] + + +def test_kimi_k2_5_config_conversion_preserves_reserved_added_token_range( + tmp_path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + config_path = VENDORED_DIR / "kimi-k2.5" / "tokenizer_config.json" + if not config_path.exists(): + pytest.skip("vendored kimi-k2.5 tokenizer_config.json not found") + + class FakeMergeableRanks(dict): + def __len__(self) -> int: + return 163584 + + mergeable_ranks = FakeMergeableRanks({b"a": 0, b"b": 1, b"c": 2, b"d": 3}) + tiktoken_module = types.ModuleType("tiktoken") + tiktoken_load_module = types.ModuleType("tiktoken.load") + tiktoken_load_module.load_tiktoken_bpe = lambda path: mergeable_ranks + tiktoken_module.load = tiktoken_load_module + monkeypatch.setitem(sys.modules, "tiktoken", tiktoken_module) + monkeypatch.setitem(sys.modules, "tiktoken.load", tiktoken_load_module) + + (tmp_path / "tiktoken.model").write_text("") + (tmp_path / "tokenizer_config.json").write_text(config_path.read_text()) + + tokenizer_json = tiktoken_model_to_tokenizer_json(tmp_path) + config = json.loads(tokenizer_json) + added_by_id = {token["id"]: token for token in config["added_tokens"]} + vendored = json.loads((VENDORED_DIR / "kimi-k2.5" / "tokenizer.json").read_text()) + vendored_pattern = vendored["pre_tokenizer"]["pretokenizers"][0]["pattern"]["Regex"] + + assert list(added_by_id) == list(range(163584, 163840)) + assert added_by_id[163589] == { + "id": 163589, + "content": "<|reserved_token_163589|>", + "single_word": False, + "lstrip": False, + "rstrip": False, + "normalized": False, + "special": True, + } + assert added_by_id[163595]["content"] == "<|tool_calls_section_begin|>" + assert added_by_id[163595]["special"] is False + assert added_by_id[163606]["content"] == "" + assert added_by_id[163606]["special"] is False + assert added_by_id[163838]["content"] == "[UNK]" + assert added_by_id[163839]["content"] == "[PAD]" + assert ( + config["pre_tokenizer"]["pretokenizers"][0]["pattern"]["Regex"] + == vendored_pattern + ) + + +def test_kimi_k2_5_conversion_encode_matches_hf_tokenizers(tmp_path) -> None: + pytest.importorskip("tiktoken") + hf_tokenizers = pytest.importorskip("tokenizers") + + kimi_dir = VENDORED_DIR / "kimi-k2.5" + gz_path = kimi_dir / "tiktoken.model.gz" + config_path = kimi_dir / "tokenizer_config.json" + if not gz_path.exists() or not config_path.exists(): + pytest.skip("vendored kimi-k2.5 assets not found") + + model_path = tmp_path / "tiktoken.model" + model_path.write_bytes(gzip.decompress(gz_path.read_bytes())) + (tmp_path / "tokenizer_config.json").write_text(config_path.read_text()) + + tokenizer_json = tiktoken_model_to_tokenizer_json(tmp_path) + assert tokenizer_json is not None, "tiktoken_model_to_tokenizer_json returned None" + + tokenizer_config = json.loads(tokenizer_json) + tokenizer_config.pop("decoder", None) + tokenizer_json = json.dumps(tokenizer_config, ensure_ascii=False) + fastokens_tok = Tokenizer.from_json_str(tokenizer_json) + hf_tok = hf_tokenizers.Tokenizer.from_str(tokenizer_json) + + corpus = [ + "[BOS]", + "<|im_user|>", + "<|tool_calls_section_begin|>", + "", + "<|reserved_token_163589|>", + "<|im_user|>user<|im_middle|>Hello<|im_end|>", + ] + for text in corpus: + fastokens_ids = fastokens_tok.encode(text, add_special_tokens=False).ids + hf_ids = hf_tok.encode(text, add_special_tokens=False).ids + assert fastokens_ids == hf_ids, ( + f"encoding mismatch for {text!r}:\n" + f" fastokens : {fastokens_ids}\n" + f" tokenizers: {hf_ids}" + ) + + +def test_tiktoken_model_to_tokenizer_json_rejects_unknown_patternless_config( + tmp_path, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + tiktoken_module = types.ModuleType("tiktoken") + tiktoken_load_module = types.ModuleType("tiktoken.load") + tiktoken_load_module.load_tiktoken_bpe = lambda path: {b"a": 0} + tiktoken_module.load = tiktoken_load_module + monkeypatch.setitem(sys.modules, "tiktoken", tiktoken_module) + monkeypatch.setitem(sys.modules, "tiktoken.load", tiktoken_load_module) + + (tmp_path / "tiktoken.model").write_text("") + (tmp_path / "tokenizer_config.json").write_text( + json.dumps( + { + "tokenizer_class": "TikTokenTokenizer", + "auto_map": { + "AutoTokenizer": [ + "tokenization_glm.TikTokenTokenizer", + None, + ], + }, + "added_tokens_decoder": {}, + } + ) + ) + + with caplog.at_level(logging.WARNING, logger="fastokens.tiktoken"): + assert tiktoken_model_to_tokenizer_json(tmp_path) is None + assert "does not support converting tiktoken tokenizer config" in caplog.text + assert "does not define pat_str or pattern" in caplog.text + + +def test_tiktoken_model_to_tokenizer_json_returns_none_without_optional_tiktoken( + hide_tiktoken, +) -> None: + assert tiktoken_model_to_tokenizer_json("tiktoken.model") is None + + +def test_tiktoken_model_to_tokenizer_json_raises_for_missing_file(tmp_path) -> None: + pytest.importorskip("tiktoken") + + with pytest.raises(FileNotFoundError): + tiktoken_model_to_tokenizer_json(tmp_path / "missing.model") + + +def test_tiktoken_model_to_tokenizer_json_raises_for_invalid_file(tmp_path) -> None: + pytest.importorskip("tiktoken") + model_path = tmp_path / "tiktoken.model" + model_path.write_text("not-a-valid-model-line") + + with pytest.raises(ValueError, match="invalid tiktoken model"): + tiktoken_model_to_tokenizer_json(model_path) + + +def test_kimi_k2_5_tiktoken_gz_conversion_matches_vendored_tokenizer_json( + tmp_path, +) -> None: + """Convert the vendored Kimi K2.5 tiktoken.model.gz on the fly and verify + that encoding results match the vendored tokenizer.json for a corpus of + test strings. Requires tiktoken; skipped otherwise.""" + pytest.importorskip("tiktoken") + + kimi_dir = VENDORED_DIR / "kimi-k2.5" + gz_path = kimi_dir / "tiktoken.model.gz" + vendored_json_path = kimi_dir / "tokenizer.json" + + if not gz_path.exists() or not vendored_json_path.exists(): + pytest.skip("vendored kimi-k2.5 assets not found") + + # Decompress the tiktoken model to a temporary plain file. + model_path = tmp_path / "tiktoken.model" + model_path.write_bytes(gzip.decompress(gz_path.read_bytes())) + + # Extract the pre-tokenizer pattern and special tokens from the vendored + # tokenizer.json so the conversion uses exactly the same configuration. + vendored = json.loads(vendored_json_path.read_text()) + pattern: str = vendored["pre_tokenizer"]["pretokenizers"][0]["pattern"]["Regex"] + special_tokens: dict[str, int] = { + t["content"]: t["id"] for t in vendored["added_tokens"] if t.get("special") + } + + # Convert the tiktoken model to a tokenizer JSON and build a Tokenizer. + converted_json = tiktoken_model_to_tokenizer_json( + model_path, + pattern=pattern, + special_tokens=special_tokens, + ) + assert converted_json is not None, "tiktoken_model_to_tokenizer_json returned None" + converted_tok = Tokenizer.from_json_str(converted_json) + converted = json.loads(converted_json) + + # Also build a Tokenizer directly from the vendored tokenizer.json. + vendored_tok = Tokenizer.from_json_str(vendored_json_path.read_text()) + + # Rough structural parity checks (order-insensitive). + converted_model = converted["model"] + vendored_model = vendored["model"] + assert converted_model["type"] == vendored_model["type"] == "BPE" + assert len(converted_model["vocab"]) == len(vendored_model["vocab"]) + assert converted_model["vocab"] == vendored_model["vocab"] + assert len(converted_model["merges"]) == len(vendored_model["merges"]) + assert _merge_counter(converted_model["merges"]) == _merge_counter( + vendored_model["merges"] + ) + + converted_special_tokens = _special_token_map(converted["added_tokens"]) + vendored_special_tokens = _special_token_map(vendored["added_tokens"]) + assert converted_special_tokens == vendored_special_tokens + + corpus = [ + "Hello, world!", + "1 + 2 = 3", + "The quick brown fox jumps over the lazy dog.", + "你好，世界", + "fastokens converts tiktoken models on the fly", + ] + for text in corpus: + converted_ids = converted_tok.encode(text, add_special_tokens=False).ids + vendored_ids = vendored_tok.encode(text, add_special_tokens=False).ids + assert converted_ids == vendored_ids, ( + f"encoding mismatch for {text!r}:\n" + f" converted : {converted_ids}\n" + f" vendored : {vendored_ids}" + ) + + +def test_converted_tiktoken_tokenizer_encodes_segments_with_special_control() -> None: + encoding = _create_test_encoding() + tokenizer = Tokenizer.from_json_str(tiktoken_to_tokenizer_json(encoding)) + + structural_ids = tokenizer.encode_segments( + [("ab", False), ("<|end|>", True), ("cd", False)] + ).ids + text_ids = tokenizer.encode_segments([("ab cd", False)]).ids + + assert structural_ids == [4, 100, 5] + assert text_ids == encoding.encode("ab cd", disallowed_special=()) + + +def test_converted_tiktoken_tokenizer_tiktoken_safe_segments_match_manual_chunks() -> ( + None +): + encoding = _create_test_encoding() + tokenizer = Tokenizer.from_json_str(tiktoken_to_tokenizer_json(encoding)) + text = "aaaa" + + safe_ids = tokenizer.encode_segments( + [(text, False)], + tiktoken_safe=True, + ).ids + + assert safe_ids == encoding.encode(text, disallowed_special=()) diff --git a/python/tests/test_unsupported_encode_options.py b/python/tests/test_unsupported_encode_options.py new file mode 100644 index 0000000..a853fa8 --- /dev/null +++ b/python/tests/test_unsupported_encode_options.py @@ -0,0 +1,77 @@ +import pickle + +import pytest + +from fastokens._compat import _TokenizerShim +from test_async_stub import TOKENIZER_JSON + + +HELLO_TOKEN_ID = 11 + + +def test_encode_accepts_split_special_tokens_true_as_noop_without_added_tokens() -> None: + tokenizer = _TokenizerShim.from_str(TOKENIZER_JSON) + + plain = tokenizer.encode("hello") + split = tokenizer.encode("hello", split_special_tokens=True) + + assert split.ids == plain.ids + + +def test_encode_batch_accepts_split_special_tokens_true_as_noop_without_added_tokens() -> None: + tokenizer = _TokenizerShim.from_str(TOKENIZER_JSON) + + plain = tokenizer.encode_batch(["hello"]) + split = tokenizer.encode_batch(["hello"], split_special_tokens=True) + + assert split[0].ids == plain[0].ids + + +def test_encode_accepts_split_special_tokens_false_as_noop() -> None: + tokenizer = _TokenizerShim.from_str(TOKENIZER_JSON) + + plain = tokenizer.encode("hello") + explicit_false = tokenizer.encode("hello", split_special_tokens=False) + + assert explicit_false.ids == plain.ids + + +def test_encode_rejects_unknown_kwargs() -> None: + tokenizer = _TokenizerShim.from_str(TOKENIZER_JSON) + + with pytest.raises(TypeError, match="return_offsets_mapping"): + tokenizer.encode("hello", return_offsets_mapping=True) + + +def test_encode_special_tokens_true_is_supported() -> None: + tokenizer = _TokenizerShim.from_str(TOKENIZER_JSON) + + tokenizer.encode_special_tokens = True + + assert tokenizer.encode_special_tokens is True + assert tokenizer.encode("hello").ids == [HELLO_TOKEN_ID] + + +def test_encode_special_tokens_false_is_noop() -> None: + tokenizer = _TokenizerShim.from_str(TOKENIZER_JSON) + + tokenizer.encode_special_tokens = False + + assert tokenizer.encode_special_tokens is False + + +def test_pickle_restores_stored_encode_special_tokens_true() -> None: + tokenizer = _TokenizerShim.from_str(TOKENIZER_JSON) + state = ( + tokenizer.to_str(), + tokenizer.truncation, + tokenizer.padding, + True, + ) + + payload = pickle.dumps(state) + restored = _TokenizerShim.from_str(TOKENIZER_JSON) + + restored.__setstate__(pickle.loads(payload)) + + assert restored.encode_special_tokens is True diff --git a/src/added_tokens.rs b/src/added_tokens.rs index 8ddbd1f..c7eb764 100644 --- a/src/added_tokens.rs +++ b/src/added_tokens.rs @@ -12,15 +12,9 @@ use crate::json_structs::AddedTokenConfig; /// Matched spans are assigned their token IDs directly; unmatched spans pass /// through normalization, pre-tokenization and the model as usual. pub struct AddedTokens { - daac: DoubleArrayAhoCorasick, - /// Token lengths (in bytes) indexed by token ID, for matched tokens only. - /// Non-added token IDs map to 0. - token_lens: Vec, - /// Distinct first bytes of all added token strings. Used to quickly skip - /// positions that cannot start any token via SIMD memchr. - start_bytes: Vec, - /// Longest added token in bytes. Limits the DAAC scan window. - max_token_len: usize, + all: AddedTokenMatcher, + has_special_tokens: bool, + has_non_special_tokens: bool, /// Mapping from token ID to token content string. id_to_content: HashMap, /// Reverse mapping: token content string → token ID. @@ -48,49 +42,29 @@ pub struct AddedTokenInfo<'a> { pub special: bool, } -impl AddedTokens { - /// Build from the `added_tokens` array in `tokenizer.json`. - /// - /// Returns `None` if there are no added tokens. - pub fn from_configs(configs: &[AddedTokenConfig]) -> Result, String> { - if configs.is_empty() { +struct AddedTokenMatcher { + daac: DoubleArrayAhoCorasick, + /// Distinct first bytes of all added token strings. Used to quickly skip + /// positions that cannot start any token via SIMD memchr. + start_bytes: Vec, + /// Longest added token in bytes. Limits the DAAC scan window. + max_token_len: usize, +} + +impl AddedTokenMatcher { + fn new(patterns: Vec<(&str, u32)>) -> Result, String> { + if patterns.is_empty() { return Ok(None); } - let max_id = configs.iter().map(|c| c.id).max().unwrap_or(0); - let mut token_lens = vec![0usize; (max_id + 1) as usize]; - - let mut id_to_content = HashMap::with_capacity(configs.len()); - let mut special_ids = HashSet::new(); - - let mut content_to_id = HashMap::with_capacity(configs.len()); - - let patterns: Vec<(&str, u32)> = configs - .iter() - .map(|c| { - token_lens[c.id as usize] = c.content.len(); - id_to_content.insert(c.id, c.content.clone()); - content_to_id.insert(c.content.clone(), c.id); - if c.special { - special_ids.insert(c.id); - } - (c.content.as_str(), c.id) - }) - .collect(); - - let daac = DoubleArrayAhoCorasickBuilder::new() - .match_kind(daachorse::MatchKind::LeftmostLongest) - .build_with_values(patterns) - .map_err(|e| format!("error building added-tokens DAAC: {e}"))?; - // Collect distinct first bytes for memchr prefilter. let mut start_set = [false; 256]; let mut max_token_len = 0; - for c in configs { - if let Some(&b) = c.content.as_bytes().first() { + for (content, _) in &patterns { + if let Some(&b) = content.as_bytes().first() { start_set[b as usize] = true; } - max_token_len = max_token_len.max(c.content.len()); + max_token_len = max_token_len.max(content.len()); } let start_bytes: Vec = start_set .iter() @@ -99,63 +73,25 @@ impl AddedTokens { .map(|(i, _)| i as u8) .collect(); + let daac = DoubleArrayAhoCorasickBuilder::new() + .match_kind(daachorse::MatchKind::LeftmostLongest) + .build_with_values(patterns) + .map_err(|e| format!("error building added-tokens DAAC: {e}"))?; + Ok(Some(Self { daac, - token_lens, start_bytes, max_token_len, - id_to_content, - content_to_id, - special_ids, })) } - /// Look up the string content of an added token by ID. - pub fn id_to_token(&self, id: u32) -> Option<&str> { - self.id_to_content.get(&id).map(String::as_str) - } - - /// Look up the token ID for a content string. - pub fn token_to_id(&self, token: &str) -> Option { - self.content_to_id.get(token).copied() - } - - /// Check if a token ID is a special added token. - pub fn is_special(&self, id: u32) -> bool { - self.special_ids.contains(&id) - } - - /// Return the number of added tokens. - pub fn len(&self) -> usize { - self.id_to_content.len() - } - - /// Return whether there are no added tokens. - pub fn is_empty(&self) -> bool { - self.id_to_content.is_empty() - } - - /// Iterate over added-token entries. - /// - /// The iteration order is unspecified. Callers that need a stable order - /// should sort by `id` themselves. - pub fn iter(&self) -> impl Iterator> { - self.id_to_content - .iter() - .map(|(&id, content)| AddedTokenInfo { - id, - content: content.as_str(), - special: self.special_ids.contains(&id), - }) - } - /// Split `input` into segments: spans matching added tokens and spans of /// regular text. /// /// Added tokens are matched leftmost-longest. Non-overlapping matches are /// emitted as [`Segment::Token`]; the gaps between them as /// [`Segment::Text`]. - pub fn split<'a>(&self, input: &'a str) -> Vec> { + fn split<'a>(&self, input: &'a str) -> Vec> { // When there are few distinct start bytes, use SIMD memchr to skip // positions that cannot start any added token. This avoids scanning // the full input through the Aho-Corasick automaton. @@ -243,11 +179,167 @@ impl AddedTokens { } } +impl AddedTokens { + /// Build from the `added_tokens` array in `tokenizer.json`. + /// + /// Returns `None` if there are no added tokens. + pub fn from_configs(configs: &[AddedTokenConfig]) -> Result, String> { + if configs.is_empty() { + return Ok(None); + } + + let mut id_to_content = HashMap::with_capacity(configs.len()); + let mut special_ids = HashSet::new(); + + let mut content_to_id = HashMap::with_capacity(configs.len()); + + let patterns: Vec<(&str, u32)> = configs + .iter() + .map(|c| { + id_to_content.insert(c.id, c.content.clone()); + content_to_id.insert(c.content.clone(), c.id); + if c.special { + special_ids.insert(c.id); + } + (c.content.as_str(), c.id) + }) + .collect(); + + let Some(all) = AddedTokenMatcher::new(patterns)? else { + debug_assert!(false, "non-empty configs should produce token patterns"); + return Ok(None); + }; + let has_special_tokens = !special_ids.is_empty(); + let has_non_special_tokens = special_ids.len() != configs.len(); + + Ok(Some(Self { + all, + has_special_tokens, + has_non_special_tokens, + id_to_content, + content_to_id, + special_ids, + })) + } + + /// Look up the string content of an added token by ID. + pub fn id_to_token(&self, id: u32) -> Option<&str> { + self.id_to_content.get(&id).map(String::as_str) + } + + /// Look up the token ID for a content string. + pub fn token_to_id(&self, token: &str) -> Option { + self.content_to_id.get(token).copied() + } + + /// Check if a token ID is a special added token. + pub fn is_special(&self, id: u32) -> bool { + self.special_ids.contains(&id) + } + + /// Return the number of added tokens. + pub fn len(&self) -> usize { + self.id_to_content.len() + } + + /// Return whether there are no added tokens. + pub fn is_empty(&self) -> bool { + self.id_to_content.is_empty() + } + + /// Iterate over added-token entries. + /// + /// The iteration order is unspecified. Callers that need a stable order + /// should sort by `id` themselves. + pub fn iter(&self) -> impl Iterator> { + self.id_to_content + .iter() + .map(|(&id, content)| AddedTokenInfo { + id, + content: content.as_str(), + special: self.special_ids.contains(&id), + }) + } + + /// Split `input` into segments: spans matching added tokens and spans of + /// regular text. + /// + /// Added tokens are matched leftmost-longest. Non-overlapping matches are + /// emitted as [`Segment::Token`]; the gaps between them as + /// [`Segment::Text`]. + pub fn split<'a>(&self, input: &'a str) -> Vec> { + self.all.split(input) + } + + /// Split `input` into segments while leaving special added tokens as text. + /// + /// This mirrors HuggingFace's `split_special_tokens=True` encode option: + /// special-token strings are tokenized through the normal model pipeline, + /// while non-special added tokens can still match as added tokens. + pub fn split_special_as_text<'a>(&self, input: &'a str) -> Vec> { + if input.is_empty() { + return Vec::new(); + } + + if !self.has_special_tokens { + return self.all.split(input); + } + if !self.has_non_special_tokens { + return vec![Segment::Text(input)]; + } + + let mut segments = Vec::new(); + let mut cursor = 0; + let mut pending_text_start: Option = None; + + for segment in self.all.split(input) { + match segment { + Segment::Text(text) => { + if !text.is_empty() && pending_text_start.is_none() { + pending_text_start = Some(cursor); + } + cursor += text.len(); + } + Segment::Token(id) => { + let content = self + .id_to_content + .get(&id) + .expect("matcher id must exist in added-token content map"); + let start = cursor; + let end = start + content.len(); + + if self.special_ids.contains(&id) { + if pending_text_start.is_none() { + pending_text_start = Some(start); + } + } else { + if let Some(text_start) = pending_text_start.take() + && text_start < start + { + segments.push(Segment::Text(&input[text_start..start])); + } + segments.push(Segment::Token(id)); + } + + cursor = end; + } + } + } + + if let Some(text_start) = pending_text_start + && text_start < cursor + { + segments.push(Segment::Text(&input[text_start..cursor])); + } + + segments + } +} + impl fmt::Debug for AddedTokens { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - let count = self.token_lens.iter().filter(|&&len| len > 0).count(); f.debug_struct("AddedTokens") - .field("count", &count) + .field("count", &self.len()) .finish() } } @@ -452,6 +544,68 @@ mod tests { assert!(!at.is_special(99)); // unknown id } + #[test] + fn split_special_as_text_keeps_non_special_matches() { + let mut special = make_config(1, ""); + special.special = true; + let non_special = make_config(2, "<|tool_calls_section_begin|>"); + let at = AddedTokens::from_configs(&[special, non_special]) + .unwrap() + .unwrap(); + + assert_eq!( + at.split_special_as_text("ab<|tool_calls_section_begin|>c"), + vec![ + Segment::Text("ab"), + Segment::Token(2), + Segment::Text("c"), + ] + ); + } + + #[test] + fn split_special_as_text_without_non_special_tokens_returns_text() { + let mut special = make_config(1, ""); + special.special = true; + let mut end_special = make_config(2, ""); + end_special.special = true; + let at = AddedTokens::from_configs(&[special, end_special]) + .unwrap() + .unwrap(); + + assert_eq!( + at.split_special_as_text("abc"), + vec![Segment::Text("abc")] + ); + assert!(at.split_special_as_text("").is_empty()); + } + + #[test] + fn split_special_as_text_with_no_special_tokens_uses_all_matches() { + let plain = make_config(1, ""); + let at = AddedTokens::from_configs(&[plain]).unwrap().unwrap(); + + assert_eq!( + at.split_special_as_text("ab"), + vec![Segment::Text("a"), Segment::Token(1), Segment::Text("b")] + ); + } + + #[test] + fn split_special_as_text_blocks_non_special_match_inside_special_text() { + let mut special = make_config(1, ""); + special.special = true; + let non_special = make_config(2, "ask>"); + let at = AddedTokens::from_configs(&[special, non_special]) + .unwrap() + .unwrap(); + + assert_eq!( + at.split_special_as_text(""), + vec![Segment::Text("")] + ); + } + #[test] fn iter_exposes_id_content_and_special_flag() { let mut special = make_config(1, ""); diff --git a/src/chat_template/mod.rs b/src/chat_template/mod.rs new file mode 100644 index 0000000..3fd75ec --- /dev/null +++ b/src/chat_template/mod.rs @@ -0,0 +1,295 @@ +mod runtime; +mod special_tokens; + +use minijinja::Error as MiniJinjaError; +use serde_json::{Map, Value}; + +pub use runtime::ChatTemplateRenderer; +pub use special_tokens::default_special_tokens; + +// Re-exported so bindings can deserialize directly into the same +// `minijinja::Value` the renderer consumes (see `ChatTemplateRenderer::render_value`) +// without depending on a possibly-mismatched minijinja version themselves. +pub use minijinja; + +/// Options for rendering a HuggingFace-style chat template. +#[derive(Clone, Debug, Default)] +pub struct ChatTemplateOptions { + pub add_generation_prompt: bool, + pub continue_final_message: bool, + pub tools: Option, + pub documents: Option, + pub special_tokens: Map, + pub extra_context: Map, +} + +/// Render a HuggingFace-style Jinja chat template. +/// +/// This covers the hot serving path for `apply_chat_template(..., tokenize=False)`. +/// Python bindings can tokenize the rendered string directly afterward, avoiding +/// the Python-side template loop while keeping the tokenizer API surface small. +pub fn apply_chat_template( + chat_template: &str, + messages: Value, + options: ChatTemplateOptions, +) -> Result { + ChatTemplateRenderer::new(chat_template)?.render(messages, options) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::{Value, json}; + + #[test] + fn renders_common_hf_style_template() { + let template = r#" +{%- for message in messages -%} +{%- if message['role'] == 'user' -%} +{{ bos_token }}user: {{ message['content'] }}{{ eos_token }} +{%- elif message['role'] == 'assistant' -%} +assistant: {{ message.get('content') }} +{%- endif -%} +{%- endfor -%} +{%- if add_generation_prompt -%}assistant:{%- endif -%} +"#; + let messages = json!([ + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "hi"} + ]); + let special_tokens = Map::from_iter([ + ("bos_token".to_string(), json!("[BOS]")), + ("eos_token".to_string(), json!("[EOS]")), + ]); + let options = ChatTemplateOptions { + add_generation_prompt: true, + special_tokens, + ..Default::default() + }; + + let rendered = apply_chat_template(template, messages, options).unwrap(); + + assert_eq!(rendered, "[BOS]user: hello[EOS]assistant: hiassistant:"); + } + + #[test] + fn omits_null_special_tokens_from_context() { + let mut special_tokens = default_special_tokens(); + special_tokens.insert("custom_token".to_string(), Value::Null); + + let rendered = apply_chat_template( + "{% if pad_token is defined %}defined{% else %}undefined{% endif %}{{ pad_token }}", + json!([]), + ChatTemplateOptions { + special_tokens, + ..Default::default() + }, + ) + .unwrap(); + + assert_eq!(rendered, "undefined"); + } + + #[test] + fn supports_tools_documents_tojson_and_kwargs() { + let template = r#"{{ tools[0].name }} {{ documents[0].title }} {{ custom | tojson }}"#; + let mut extra = Map::new(); + extra.insert("custom".to_string(), json!({"x": 1})); + let options = ChatTemplateOptions { + tools: Some(json!([{"name": "search"}])), + documents: Some(json!([{"title": "doc"}])), + extra_context: extra, + ..Default::default() + }; + + let rendered = apply_chat_template(template, json!([]), options).unwrap(); + + assert_eq!(rendered, r#"search doc {"x": 1}"#); + } + + #[test] + fn supports_python_compat_methods() { + let template = r#"{{ message.get('missing', 'fallback') }} {{ 'abc'.upper() }} {{ [1, 1, 2].count(1) }}"#; + let mut extra = Map::new(); + extra.insert("message".to_string(), json!({"content": "hello"})); + + let rendered = apply_chat_template( + template, + json!([]), + ChatTemplateOptions { + extra_context: extra, + ..Default::default() + }, + ) + .unwrap(); + + assert_eq!(rendered, "fallback ABC 2"); + } + + #[test] + fn tojson_matches_transformers_defaults() { + let template = r#"{{ value | tojson }}"#; + let mut extra = Map::new(); + extra.insert("value".to_string(), json!({"text": "&'"})); + + let rendered = apply_chat_template( + template, + json!([]), + ChatTemplateOptions { + extra_context: extra, + ..Default::default() + }, + ) + .unwrap(); + + assert_eq!(rendered, r#"{"text": "&'"}"#); + } + + #[test] + fn tojson_supports_compact_separators_and_ascii_escape() { + let template = r#"{{ value | tojson(separators=(',', ':'), ensure_ascii=True) }}"#; + let mut extra = Map::new(); + extra.insert("value".to_string(), json!({"text": "é"})); + + let rendered = apply_chat_template( + template, + json!([]), + ChatTemplateOptions { + extra_context: extra, + ..Default::default() + }, + ) + .unwrap(); + + assert_eq!(rendered, r#"{"text":"\u00e9"}"#); + } + + #[test] + fn strips_generation_markers() { + let rendered = apply_chat_template( + "a{% generation %}b{% endgeneration %}c", + json!([]), + ChatTemplateOptions::default(), + ) + .unwrap(); + + assert_eq!(rendered, "abc"); + } + + #[test] + fn rewrites_generation_markers_with_whitespace_control() { + let rendered = apply_chat_template( + "a\n{%- generation -%}\nb\n{%- endgeneration -%}\nc", + json!([]), + ChatTemplateOptions::default(), + ) + .unwrap(); + + assert_eq!(rendered, "abc"); + } + + #[test] + fn tojson_indent_matches_json_dumps() { + let mut extra = Map::new(); + extra.insert("value".to_string(), json!({"b": [1, 2], "a": {}, "c": "x"})); + + let rendered = apply_chat_template( + "{{ value | tojson(indent=2) }}", + json!([]), + ChatTemplateOptions { + extra_context: extra, + ..Default::default() + }, + ) + .unwrap(); + + // json.dumps(value, indent=2) + assert_eq!( + rendered, + "{\n \"b\": [\n 1,\n 2\n ],\n \"a\": {},\n \"c\": \"x\"\n}" + ); + } + + #[test] + fn tojson_indent_honors_separators() { + let mut extra = Map::new(); + extra.insert("value".to_string(), json!({"b": [1, 2], "a": {}, "c": "x"})); + + let rendered = apply_chat_template( + "{{ value | tojson(indent=2, separators=(',', ':')) }}", + json!([]), + ChatTemplateOptions { + extra_context: extra, + ..Default::default() + }, + ) + .unwrap(); + + // json.dumps(value, indent=2, separators=(",", ":")) + assert_eq!( + rendered, + "{\n \"b\":[\n 1,\n 2\n ],\n \"a\":{},\n \"c\":\"x\"\n}" + ); + } + + #[test] + fn tojson_indent_keeps_separator_whitespace_like_json_dumps() { + let mut extra = Map::new(); + extra.insert("value".to_string(), json!({"b": [1, 2], "a": {}, "c": "x"})); + + let rendered = apply_chat_template( + "{{ value | tojson(indent=2, separators=(', ', ': '), sort_keys=True) }}", + json!([]), + ChatTemplateOptions { + extra_context: extra, + ..Default::default() + }, + ) + .unwrap(); + + // json.dumps keeps the item separator verbatim before the newline, + // trailing space included. + assert_eq!( + rendered, + "{\n \"a\": {}, \n \"b\": [\n 1, \n 2\n ], \n \"c\": \"x\"\n}" + ); + } + + #[test] + fn exposes_undeclared_template_variables() { + let renderer = ChatTemplateRenderer::new( + "{% for m in messages %}{{ m.content }}{% endfor %}{% if enable_thinking %}t{% endif %}", + ) + .unwrap(); + + let variables = renderer.undeclared_variables(); + + assert!(variables.contains("messages")); + assert!(variables.contains("enable_thinking")); + assert!(!variables.contains("m")); + } + + #[test] + fn raise_exception_returns_error() { + let err = apply_chat_template( + "{{ raise_exception('bad role') }}", + json!([]), + ChatTemplateOptions::default(), + ) + .unwrap_err(); + + assert!(err.to_string().contains("bad role")); + } + + #[test] + fn strftime_now_invalid_format_returns_error() { + let err = apply_chat_template( + "{{ strftime_now('%Q %#z') }}", + json!([]), + ChatTemplateOptions::default(), + ) + .unwrap_err(); + + assert!(err.to_string().contains("invalid strftime_now format")); + } +} diff --git a/src/chat_template/runtime.rs b/src/chat_template/runtime.rs new file mode 100644 index 0000000..7161762 --- /dev/null +++ b/src/chat_template/runtime.rs @@ -0,0 +1,424 @@ +use chrono::{DateTime, Local}; +use minijinja::{Environment, Error, ErrorKind, Value as MiniJinjaValue, value::Kwargs}; +use serde::Serialize; +use serde_json::Value; +use std::collections::HashSet; +use std::{fmt::Write as _, io}; + +use super::ChatTemplateOptions; + +/// Compiled chat-template renderer. +/// +/// Reuse this for repeated `apply_chat_template` calls with the same template; +/// constructing a MiniJinja environment and parsing the template is avoidable +/// work on the hot serving path. +pub struct ChatTemplateRenderer { + env: Environment<'static>, + undeclared_variables: HashSet, +} + +impl ChatTemplateRenderer { + pub fn new(chat_template: &str) -> Result { + let mut env = Environment::new(); + env.set_lstrip_blocks(true); + env.set_trim_blocks(true); + env.set_unknown_method_callback(minijinja_contrib::pycompat::unknown_method_callback); + env.add_filter("tojson", tojson); + env.add_function("raise_exception", raise_exception); + env.add_function("strftime_now", strftime_now); + env.add_template_owned("chat", strip_known_non_jinja_tags(chat_template))?; + let undeclared_variables = env.get_template("chat")?.undeclared_variables(false); + + Ok(Self { + env, + undeclared_variables, + }) + } + + /// Top-level variables the template reads from the render context. + /// + /// Computed statically from the template AST at compile time, so dynamic + /// lookups (e.g. `context[name]`) are invisible. Use this to reject + /// keyword arguments the template can never see, not to require ones it + /// might. + pub fn undeclared_variables(&self) -> &HashSet { + &self.undeclared_variables + } + + pub fn render(&self, messages: Value, options: ChatTemplateOptions) -> Result { + self.render_value(json_to_minijinja_value(messages), options) + } + + /// Render with messages already converted to a MiniJinja value. + /// + /// Bindings that can deserialize straight into [`minijinja::Value`] should + /// prefer this over [`render`](Self::render): it skips the intermediate + /// `serde_json::Value` tree, which dominates per-call cost on large + /// conversations. + pub fn render_value( + &self, + messages: MiniJinjaValue, + options: ChatTemplateOptions, + ) -> Result { + self.render_context(build_context(messages, options)) + } + + pub fn render_context(&self, context: MiniJinjaValue) -> Result { + self.env.get_template("chat")?.render(context) + } +} + +fn build_context(messages: MiniJinjaValue, options: ChatTemplateOptions) -> MiniJinjaValue { + let mut context_entries = + Vec::with_capacity(6 + options.special_tokens.len() + options.extra_context.len()); + context_entries.push(("messages".to_string(), messages)); + context_entries.push(( + "add_generation_prompt".to_string(), + MiniJinjaValue::from(options.add_generation_prompt), + )); + context_entries.push(( + "continue_final_message".to_string(), + MiniJinjaValue::from(options.continue_final_message), + )); + context_entries.push(( + "tools".to_string(), + options + .tools + .map(json_to_minijinja_value) + .unwrap_or_else(|| MiniJinjaValue::from(())), + )); + context_entries.push(( + "documents".to_string(), + options + .documents + .map(json_to_minijinja_value) + .unwrap_or_else(|| MiniJinjaValue::from(())), + )); + context_entries.extend( + options + .special_tokens + .into_iter() + .filter(|(_, value)| !value.is_null()) + .map(|(key, value)| (key, json_to_minijinja_value(value))), + ); + context_entries.extend( + options + .extra_context + .into_iter() + .map(|(key, value)| (key, json_to_minijinja_value(value))), + ); + context_entries.into_iter().collect() +} + +fn json_to_minijinja_value(value: Value) -> MiniJinjaValue { + match value { + Value::Null => MiniJinjaValue::from(()), + Value::Bool(value) => MiniJinjaValue::from(value), + Value::Number(value) => { + if let Some(value) = value.as_i64() { + MiniJinjaValue::from(value) + } else if let Some(value) = value.as_u64() { + MiniJinjaValue::from(value) + } else if let Some(value) = value.as_f64() { + MiniJinjaValue::from(value) + } else { + MiniJinjaValue::from(()) + } + } + Value::String(value) => MiniJinjaValue::from(value), + Value::Array(values) => MiniJinjaValue::from( + values + .into_iter() + .map(json_to_minijinja_value) + .collect::>(), + ), + Value::Object(values) => values + .into_iter() + .map(|(key, value)| (key, json_to_minijinja_value(value))) + .collect(), + } +} + +fn strip_known_non_jinja_tags(template: &str) -> String { + let mut stripped = String::with_capacity(template.len()); + let mut rest = template; + while let Some(start) = rest.find("{%") { + stripped.push_str(&rest[..start]); + rest = &rest[start..]; + let Some(end) = rest.find("%}") else { + stripped.push_str(rest); + return stripped; + }; + let tag = &rest[..end + 2]; + let inner = tag + .strip_prefix("{%") + .and_then(|tag| tag.strip_suffix("%}")) + .unwrap_or_default(); + let inner = inner.strip_prefix('-').unwrap_or(inner); + let inner = inner.strip_suffix('-').unwrap_or(inner).trim(); + + match inner { + "generation" => stripped.push_str(&rewrite_block_tag(tag, "if true")), + "endgeneration" => stripped.push_str(&rewrite_block_tag(tag, "endif")), + _ => stripped.push_str(tag), + } + rest = &rest[end + 2..]; + } + stripped.push_str(rest); + stripped +} + +fn rewrite_block_tag(tag: &str, replacement: &str) -> String { + let left_trim = tag.starts_with("{%-"); + let right_trim = tag.ends_with("-%}"); + format!( + "{{%{} {replacement} {}%}}", + if left_trim { "-" } else { "" }, + if right_trim { "-" } else { "" } + ) +} + +fn raise_exception(message: String) -> Result { + Err(Error::new(ErrorKind::InvalidOperation, message)) +} + +fn strftime_now(format: &str) -> Result { + let now: DateTime = Local::now(); + let mut rendered = String::new(); + write!(&mut rendered, "{}", now.format(format)).map_err(|err| { + Error::new(ErrorKind::InvalidOperation, "invalid strftime_now format").with_source(err) + })?; + Ok(MiniJinjaValue::from_safe_string(rendered)) +} + +fn tojson(value: MiniJinjaValue, kwargs: Kwargs) -> Result { + let ensure_ascii: Option = kwargs.get("ensure_ascii")?; + let indent: Option = kwargs.get("indent")?; + let separators: Option> = kwargs.get("separators")?; + let sort_keys: Option = kwargs.get("sort_keys")?; + kwargs.assert_all_used()?; + + let serialized = if sort_keys.unwrap_or(false) { + let mut value = serde_json::to_value(&value).map_err(json_error)?; + sort_json_value(&mut value); + serialize_json(&value, indent, separators)? + } else { + serialize_json(&value, indent, separators)? + }; + + let serialized = if ensure_ascii.unwrap_or(false) { + escape_non_ascii(&serialized) + } else { + serialized + }; + + Ok(MiniJinjaValue::from_safe_string(serialized)) +} + +fn serialize_json( + value: &T, + indent: Option, + separators: Option>, +) -> Result { + let (item_separator, key_separator) = json_separators(separators, indent.is_some())?; + let formatter = HfJsonFormatter { + item_separator, + key_separator, + indent: indent.map(|width| b" ".repeat(width)), + current_indent: 0, + has_value: false, + }; + let mut buf = Vec::new(); + let mut serializer = serde_json::Serializer::with_formatter(&mut buf, formatter); + value.serialize(&mut serializer).map_err(json_error)?; + String::from_utf8(buf).map_err(|err| { + Error::new(ErrorKind::BadSerialization, "cannot serialize to JSON").with_source(err) + }) +} + +/// Serializer matching Python's `json.dumps` output, including its indented +/// form: separators apply verbatim even when `indent` is set (so a custom +/// item separator keeps its trailing whitespace before the newline, exactly +/// like Python), and empty containers stay `[]`/`{}`. +struct HfJsonFormatter { + item_separator: Vec, + key_separator: Vec, + indent: Option>, + current_indent: usize, + has_value: bool, +} + +impl HfJsonFormatter { + fn write_newline_indent(&self, writer: &mut W) -> io::Result<()> + where + W: ?Sized + io::Write, + { + if let Some(indent) = &self.indent { + writer.write_all(b"\n")?; + for _ in 0..self.current_indent { + writer.write_all(indent)?; + } + } + Ok(()) + } + + fn begin_container(&mut self, writer: &mut W, open: &[u8]) -> io::Result<()> + where + W: ?Sized + io::Write, + { + if self.indent.is_some() { + self.current_indent += 1; + self.has_value = false; + } + writer.write_all(open) + } + + fn end_container(&mut self, writer: &mut W, close: &[u8]) -> io::Result<()> + where + W: ?Sized + io::Write, + { + if self.indent.is_some() { + self.current_indent -= 1; + if self.has_value { + self.write_newline_indent(writer)?; + } + } + writer.write_all(close) + } +} + +impl serde_json::ser::Formatter for HfJsonFormatter { + fn begin_array(&mut self, writer: &mut W) -> io::Result<()> + where + W: ?Sized + io::Write, + { + self.begin_container(writer, b"[") + } + + fn end_array(&mut self, writer: &mut W) -> io::Result<()> + where + W: ?Sized + io::Write, + { + self.end_container(writer, b"]") + } + + fn begin_object(&mut self, writer: &mut W) -> io::Result<()> + where + W: ?Sized + io::Write, + { + self.begin_container(writer, b"{") + } + + fn end_object(&mut self, writer: &mut W) -> io::Result<()> + where + W: ?Sized + io::Write, + { + self.end_container(writer, b"}") + } + + fn begin_array_value(&mut self, writer: &mut W, first: bool) -> io::Result<()> + where + W: ?Sized + io::Write, + { + if !first { + writer.write_all(&self.item_separator)?; + } + self.write_newline_indent(writer) + } + + fn end_array_value(&mut self, _writer: &mut W) -> io::Result<()> + where + W: ?Sized + io::Write, + { + self.has_value = true; + Ok(()) + } + + fn begin_object_key(&mut self, writer: &mut W, first: bool) -> io::Result<()> + where + W: ?Sized + io::Write, + { + if !first { + writer.write_all(&self.item_separator)?; + } + self.write_newline_indent(writer) + } + + fn begin_object_value(&mut self, writer: &mut W) -> io::Result<()> + where + W: ?Sized + io::Write, + { + writer.write_all(&self.key_separator) + } + + fn end_object_value(&mut self, _writer: &mut W) -> io::Result<()> + where + W: ?Sized + io::Write, + { + self.has_value = true; + Ok(()) + } +} + +fn json_separators( + separators: Option>, + indented: bool, +) -> Result<(Vec, Vec), Error> { + let Some(separators) = separators else { + // json.dumps defaults: (", ", ": ") normally, (",", ": ") when + // indent is given (the newline replaces the space after commas). + let item_separator = if indented { + b",".to_vec() + } else { + b", ".to_vec() + }; + return Ok((item_separator, b": ".to_vec())); + }; + if separators.len() != 2 { + return Err(Error::new( + ErrorKind::InvalidOperation, + "tojson separators must contain exactly two strings", + )); + } + Ok(( + separators[0].as_bytes().to_vec(), + separators[1].as_bytes().to_vec(), + )) +} + +fn sort_json_value(value: &mut Value) { + match value { + Value::Array(values) => { + for value in values { + sort_json_value(value); + } + } + Value::Object(map) => { + let mut entries: Vec<_> = std::mem::take(map).into_iter().collect(); + entries.sort_by(|(left, _), (right, _)| left.cmp(right)); + for (key, mut value) in entries { + sort_json_value(&mut value); + map.insert(key, value); + } + } + _ => {} + } +} + +fn json_error(err: serde_json::Error) -> Error { + Error::new(ErrorKind::BadSerialization, "cannot serialize to JSON").with_source(err) +} + +fn escape_non_ascii(serialized: &str) -> String { + let mut escaped = String::with_capacity(serialized.len()); + for c in serialized.chars() { + if c.is_ascii() { + escaped.push(c); + } else { + for unit in c.encode_utf16(&mut [0; 2]) { + write!(&mut escaped, "\\u{unit:04x}").expect("writing to string cannot fail"); + } + } + } + escaped +} diff --git a/src/chat_template/special_tokens.rs b/src/chat_template/special_tokens.rs new file mode 100644 index 0000000..cd68ee2 --- /dev/null +++ b/src/chat_template/special_tokens.rs @@ -0,0 +1,15 @@ +use serde_json::{Map, Value}; + +/// Return baseline special-token variables for standalone template rendering. +/// +/// Python callers can override these with tokenizer-specific values. Tokens +/// default to null and are omitted from the template context so absent tokenizer +/// metadata behaves like Transformers' undefined variables. +pub fn default_special_tokens() -> Map { + Map::from_iter([ + ("bos_token".to_string(), Value::Null), + ("eos_token".to_string(), Value::Null), + ("unk_token".to_string(), Value::Null), + ("pad_token".to_string(), Value::Null), + ]) +} diff --git a/src/lib.rs b/src/lib.rs index b931528..a6f1611 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,4 +1,5 @@ pub mod added_tokens; +pub mod chat_template; pub mod decoders; pub mod json_structs; pub mod models; @@ -7,13 +8,19 @@ pub mod post_processors; pub mod pre_tokenized; pub mod pre_tokenizers; -use std::{fs, path::Path}; +use std::{ + collections::{HashMap, HashSet}, + fs, + path::Path, +}; +use daachorse::{DoubleArrayAhoCorasick, DoubleArrayAhoCorasickBuilder}; use rayon::prelude::*; use serde_json::Value; pub use self::{ added_tokens::{AddedTokenInfo, AddedTokens}, + chat_template::{ChatTemplateOptions, ChatTemplateRenderer, apply_chat_template}, json_structs::{ AddedTokenConfig, DecoderConfig, DecoderKind, ModelConfig, ModelKind, NormalizerConfig, NormalizerKind, PostProcessorConfig, PostProcessorKind, PreTokenizerConfig, @@ -31,6 +38,14 @@ use self::{ pre_tokenized::{PreTokenizedString, Split as PtSplit}, }; +// Upper bound a single tiktoken `encode` call can handle without the regex +// pre-tokenizer panicking across the pyo3 boundary (`pyo3_runtime. +// PanicException`). Legacy tiktoken tokenizers chunk inputs at this size. +const TIKTOKEN_MAX_ENCODE_CHARS: usize = 400_000; +// Guard against catastrophic backtracking in tiktoken's regex on very long +// whitespace or non-whitespace runs (tiktoken issue #195). +const TIKTOKEN_MAX_CONSECUTIVE_WHITESPACE_OR_NONWHITESPACE_CHARS: usize = 25_000; + #[cfg(feature = "hf-hub")] mod hf_hub_support { pub use hf_hub::api::sync::ApiError; @@ -107,6 +122,9 @@ pub enum Error { #[error("decoder error: {0}")] Decoder(#[from] decoders::Error), + #[error("chat template error: {0}")] + ChatTemplate(#[from] minijinja::Error), + #[error("model error: {0}")] Model(String), @@ -260,6 +278,21 @@ impl Tokenizer { &self, input: &str, add_special_tokens: bool, + ) -> Result, Error> { + self.encode_with_options(input, add_special_tokens, false) + } + + /// Run the full encoding pipeline with control over special token insertion + /// and whether special-token strings should be split as ordinary text. + /// + /// When `split_special_tokens` is true, special added-token strings are not + /// emitted directly as their special token IDs. Non-special added tokens can + /// still match. + pub fn encode_with_options( + &self, + input: &str, + add_special_tokens: bool, + split_special_tokens: bool, ) -> Result, Error> { if input.is_empty() { return if add_special_tokens { @@ -269,18 +302,37 @@ impl Tokenizer { }; } + let ids = self.encode_without_post_process(input, split_special_tokens, false)?; + + // 4. Post-process. + Ok(self.post_process(ids, add_special_tokens)) + } + + fn encode_without_post_process( + &self, + input: &str, + split_special_tokens: bool, + bypass_added_tokens: bool, + ) -> Result, Error> { + if input.is_empty() { + return Ok(Vec::new()); + } + // 1. Split on added tokens + normalize into a single buffer. - let mut pts = self.build_pre_tokenized(input); + let mut pts = self.build_pre_tokenized_with_added_token_options( + input, + split_special_tokens, + bypass_added_tokens, + ); // Fused path: run only Split, then batch-tokenize with inline ByteLevel. if let Some(ref split) = self.split_only { split.pre_tokenize(&mut pts)?; - let ids = pts + return pts .tokenize_batched(|buf, splits, out| { self.model.tokenize_batch_fused(buf, splits, out) }) - .map_err(Error::Model)?; - return Ok(self.post_process(ids, add_special_tokens)); + .map_err(Error::Model); } // 2. Pre-tokenize (refine splits in place). @@ -293,7 +345,70 @@ impl Tokenizer { .tokenize(|text, out| self.model.tokenize_into(text, out)) .map_err(Error::Model)?; - // 4. Post-process. + Ok(ids) + } + + /// Encode already-rendered text segments with per-segment control over + /// added-token matching. + /// + /// When `allow_special` is true, the segment uses normal added-token + /// matching. When false, all added-token matching is bypassed, so strings + /// like a control-token string inside user text are encoded as ordinary BPE + /// text. + /// Post-processing is applied once after all segment IDs are concatenated. + pub fn encode_segments( + &self, + segments: I, + add_special_tokens: bool, + ) -> Result, Error> + where + I: IntoIterator, + S: AsRef, + { + let mut ids = Vec::new(); + for (text, allow_special) in segments { + ids.extend(self.encode_without_post_process(text.as_ref(), false, !allow_special)?); + } + Ok(self.post_process(ids, add_special_tokens)) + } + + /// Encode segments while preserving custom tiktoken tokenizer chunking. + /// + /// This splits every segment into at most 400k-character chunks and further + /// splits very long whitespace/non-whitespace runs at 25k characters before + /// BPE encoding. Legacy tiktoken tokenizers use those limits to avoid + /// tiktoken regex pre-tokenizer panics/pathologies. fastokens does not need + /// them for robustness, but the artificial boundaries affect BPE merges, so + /// callers migrating from those tokenizers need this path for exact parity on + /// very long inputs. + pub fn encode_segments_tiktoken_safe( + &self, + segments: I, + add_special_tokens: bool, + ) -> Result, Error> + where + I: IntoIterator, + S: AsRef, + { + let mut ids = Vec::new(); + for (text, allow_special) in segments { + let text = text.as_ref(); + let chunks = tiktoken_safe_chunks(text); + if chunks.len() <= 1 { + for chunk in chunks { + ids.extend(self.encode_without_post_process(chunk, false, !allow_special)?); + } + continue; + } + + let chunk_ids = chunks + .par_iter() + .map(|chunk| self.encode_without_post_process(chunk, false, !allow_special)) + .collect::, _>>()?; + for mut chunk_ids in chunk_ids { + ids.append(&mut chunk_ids); + } + } Ok(self.post_process(ids, add_special_tokens)) } @@ -302,13 +417,141 @@ impl Tokenizer { &self, inputs: &[S], add_special_tokens: bool, + ) -> Result>, Error> { + self.encode_batch_with_options(inputs, add_special_tokens, false) + } + + /// Encode a batch of inputs with full encode options. + pub fn encode_batch_with_options + Sync>( + &self, + inputs: &[S], + add_special_tokens: bool, + split_special_tokens: bool, ) -> Result>, Error> { inputs .par_iter() - .map(|input| self.encode_with_special_tokens(input.as_ref(), add_special_tokens)) + .map(|input| { + self.encode_with_options(input.as_ref(), add_special_tokens, split_special_tokens) + }) .collect() } + /// Encode a rendered prompt that contains structural token strings. + /// + /// Structural token strings are emitted as their token IDs when possible. + /// Non-structural text is encoded with `split_special_tokens=true`, after + /// restoring placeholders to their original text. This is intended for + /// chat-template output where template control tokens must remain + /// structural, but escaped user content that spells a control token must be + /// tokenized as ordinary text. + /// + /// `placeholder_map` is request-specific and maps placeholder text to the + /// original user text. `add_special_tokens` defaults to false in the Python + /// binding; callers replacing a rendered chat-template encode path should + /// keep it false unless they explicitly want post-processor tokens. + pub fn encode_with_structural_tokens( + &self, + input: &str, + structural_tokens: &StructuralTokenConfig, + placeholder_map: &HashMap, + add_special_tokens: bool, + ) -> Result, Error> { + if structural_tokens.is_empty() { + return self.encode_with_options(input, add_special_tokens, false); + } + + let mut ids = Vec::new(); + + if input.is_empty() { + return Ok(self.post_process(ids, add_special_tokens)); + } + + let placeholder_matcher = if placeholder_map.is_empty() { + None + } else { + let placeholder_tokens: Vec = placeholder_map.keys().cloned().collect(); + Some(PatternMatcher::new( + &placeholder_tokens, + "placeholder matcher", + )?) + }; + + for part in structural_tokens.structural_matcher.split(input) { + if part.is_match { + match self.token_to_id(part.text) { + Some(id) => ids.push(id), + None => ids.extend(self.encode_with_options(part.text, false, false)?), + } + } else { + self.encode_structural_text_part( + part.text, + placeholder_map, + placeholder_matcher.as_ref(), + &structural_tokens.non_special_added_tokens, + &mut ids, + )?; + } + } + + Ok(self.post_process(ids, add_special_tokens)) + } + + fn encode_structural_text_part( + &self, + text: &str, + placeholder_map: &HashMap, + placeholder_matcher: Option<&PatternMatcher>, + non_special_added_tokens: &HashSet, + ids: &mut Vec, + ) -> Result<(), Error> { + if text.is_empty() { + return Ok(()); + } + + let Some(placeholder_matcher) = placeholder_matcher else { + ids.extend(self.encode_with_options(text, false, true)?); + return Ok(()); + }; + + for part in placeholder_matcher.split(text) { + if !part.is_match { + ids.extend(self.encode_with_options(part.text, false, true)?); + continue; + } + + let Some(original) = placeholder_map.get(part.text) else { + ids.extend(self.encode_with_options(part.text, false, true)?); + continue; + }; + + if non_special_added_tokens.contains(original) { + self.encode_literal_structural_token(original, ids)?; + } else { + ids.extend(self.encode_with_options(original, false, true)?); + } + } + + Ok(()) + } + + fn encode_literal_structural_token( + &self, + token: &str, + ids: &mut Vec, + ) -> Result<(), Error> { + let Some(inner) = token.strip_prefix('<').and_then(|s| s.strip_suffix('>')) else { + ids.extend(self.encode_with_options(token, false, true)?); + return Ok(()); + }; + + ids.extend(self.encode_with_options("<", false, true)?); + if !inner.is_empty() { + ids.extend(self.encode_with_options(inner, false, true)?); + } + ids.extend(self.encode_with_options(">", false, true)?); + Ok(()) + } + /// Replace the post-processor. Called when transformers dynamically /// updates the post-processor (e.g. for `add_bos_token=True`). pub fn set_post_processor(&mut self, pp: Option) { @@ -376,6 +619,17 @@ impl Tokenizer { .collect() } + /// Render a HuggingFace-style chat template. + pub fn apply_chat_template( + &self, + chat_template: &str, + messages: Value, + options: ChatTemplateOptions, + ) -> Result { + chat_template::apply_chat_template(chat_template, messages, options) + .map_err(Error::ChatTemplate) + } + // ── Vocabulary access ──────────────────────────────────────────── /// Look up the string for a token ID, checking added tokens first, @@ -421,11 +675,39 @@ impl Tokenizer { /// Build a [`PreTokenizedString`] by splitting on added tokens and /// normalizing text segments into a single contiguous buffer. pub fn build_pre_tokenized(&self, input: &str) -> PreTokenizedString { + self.build_pre_tokenized_with_options(input, false) + } + + /// Build a [`PreTokenizedString`] with encode-time added-token options. + pub fn build_pre_tokenized_with_options( + &self, + input: &str, + split_special_tokens: bool, + ) -> PreTokenizedString { + self.build_pre_tokenized_with_added_token_options(input, split_special_tokens, false) + } + + fn build_pre_tokenized_with_added_token_options( + &self, + input: &str, + split_special_tokens: bool, + bypass_added_tokens: bool, + ) -> PreTokenizedString { let segments = match &self.added_tokens { + Some(_) if bypass_added_tokens => vec![Segment::Text(input)], + Some(at) if split_special_tokens => at.split_special_as_text(input), Some(at) => at.split(input), None => vec![Segment::Text(input)], }; + self.build_pre_tokenized_from_segments(input, segments) + } + + fn build_pre_tokenized_from_segments<'a>( + &self, + input: &'a str, + segments: Vec>, + ) -> PreTokenizedString { // Fast path: if there's exactly one Text segment (no added token matches) // and normalization returns Cow::Borrowed, we just need a string copy. if segments.len() == 1 @@ -485,6 +767,165 @@ impl Tokenizer { } } +/// Constant structural-token state for rendered-prompt encoding. +/// +/// Build this once per tokenizer/chat-template setup. `structural_tokens` +/// should include every token string that the rendered template may use as a +/// structural boundary, including tag-like non-special added tokens such as +/// `` in addition to tokens marked special by the tokenizer. +pub struct StructuralTokenConfig { + structural_matcher: PatternMatcher, + non_special_added_tokens: HashSet, +} + +impl StructuralTokenConfig { + pub fn new( + structural_tokens: &[String], + non_special_added_tokens: &HashSet, + ) -> Result { + Ok(Self { + structural_matcher: PatternMatcher::new(structural_tokens, "structural-token matcher")?, + non_special_added_tokens: non_special_added_tokens.clone(), + }) + } + + fn is_empty(&self) -> bool { + self.structural_matcher.is_empty() + } +} + +fn tiktoken_safe_chunks(text: &str) -> Vec<&str> { + if text.is_empty() { + return Vec::new(); + } + + let mut chunks = Vec::new(); + let mut outer_start = 0; + let mut outer_chars = 0; + for (idx, _) in text.char_indices() { + if outer_chars == TIKTOKEN_MAX_ENCODE_CHARS { + push_whitespace_limited_chunks(&text[outer_start..idx], &mut chunks); + outer_start = idx; + outer_chars = 0; + } + outer_chars += 1; + } + push_whitespace_limited_chunks(&text[outer_start..], &mut chunks); + chunks +} + +fn push_whitespace_limited_chunks<'a>(text: &'a str, chunks: &mut Vec<&'a str>) { + if text.is_empty() { + return; + } + + let mut slice_start = 0; + let mut current_slice_len = 0; + let mut current_slice_is_space = text.chars().next().is_some_and(char::is_whitespace); + + for (idx, char) in text.char_indices() { + let is_now_space = char.is_whitespace(); + if current_slice_is_space ^ is_now_space { + current_slice_len = 1; + current_slice_is_space = is_now_space; + } else { + current_slice_len += 1; + if current_slice_len > TIKTOKEN_MAX_CONSECUTIVE_WHITESPACE_OR_NONWHITESPACE_CHARS { + if slice_start < idx { + chunks.push(&text[slice_start..idx]); + } + slice_start = idx; + current_slice_len = 1; + } + } + } + + if slice_start < text.len() { + chunks.push(&text[slice_start..]); + } +} + +struct MatchedPart<'a> { + text: &'a str, + is_match: bool, +} + +struct PatternMatcher { + matcher: Option>, +} + +impl PatternMatcher { + fn new(patterns: &[String], name: &str) -> Result { + let mut unique_patterns: Vec<&str> = patterns + .iter() + .map(String::as_str) + .filter(|pattern| !pattern.is_empty()) + .collect(); + unique_patterns.sort_unstable(); + unique_patterns.dedup(); + + if unique_patterns.is_empty() { + return Ok(Self { matcher: None }); + } + + let pattern_values: Vec<(&str, usize)> = unique_patterns + .iter() + .enumerate() + .map(|(idx, pattern)| (*pattern, idx)) + .collect(); + let matcher = DoubleArrayAhoCorasickBuilder::new() + .match_kind(daachorse::MatchKind::LeftmostLongest) + .build_with_values(pattern_values) + .map_err(|e| Error::Model(format!("error building {name}: {e}")))?; + + Ok(Self { + matcher: Some(matcher), + }) + } + + fn is_empty(&self) -> bool { + self.matcher.is_none() + } + + fn split<'a>(&self, input: &'a str) -> Vec> { + let Some(matcher) = &self.matcher else { + if input.is_empty() { + return Vec::new(); + } + return vec![MatchedPart { + text: input, + is_match: false, + }]; + }; + + let mut parts = Vec::new(); + let mut prev_end = 0; + + for m in matcher.leftmost_find_iter(input) { + if m.start() > prev_end { + parts.push(MatchedPart { + text: &input[prev_end..m.start()], + is_match: false, + }); + } + parts.push(MatchedPart { + text: &input[m.start()..m.end()], + is_match: true, + }); + prev_end = m.end(); + } + + if prev_end < input.len() { + parts.push(MatchedPart { + text: &input[prev_end..], + is_match: false, + }); + } + + parts + } +} + // --------------------------------------------------------------------------- // Streaming decode // --------------------------------------------------------------------------- @@ -597,6 +1038,8 @@ pub fn decode_stream_step( mod tests { use crate::hf_hub_support::make_api; + use std::str::FromStr; + use super::*; const HF_MODELS: &[&str] = &[ @@ -832,26 +1275,46 @@ mod tests { "| col1 | col2 |\n|------|------|\n| a | b |", ]; - /// Helper: compare both encoding and decoding of every input in `corpus` - /// between our tokenizer and the HuggingFace tokenizer for a given model. + /// Helper: compare encoding of every input in `corpus` in both default + /// and split-special-token modes, and compare decoding of the default IDs. + /// Special-token samples from each model are included so split mode covers + /// model-specific added-token strings, not just generic text. /// Returns a list of failure descriptions (empty = all passed). fn compare_encode_decode(model_name: &str, corpus: &[&str]) -> Vec { - let hf = tokenizers::Tokenizer::from_pretrained(model_name, None) + let mut hf = tokenizers::Tokenizer::from_pretrained(model_name, None) .unwrap_or_else(|e| panic!("{model_name}: HF load failed: {e}")); let ours = Tokenizer::from_model(model_name) .unwrap_or_else(|e| panic!("{model_name}: fastokens load failed: {e}")); + let special_samples: Vec = ours + .added_tokens() + .map(|added| { + added + .iter() + .filter(|entry| entry.special) + .take(4) + .flat_map(|entry| { + [ + entry.content.to_string(), + format!("hello {} world", entry.content), + ] + }) + .collect() + }) + .unwrap_or_default(); + let mut failures = Vec::new(); - for &input in corpus { + let mut compare_input = |input: &str, compare_decode: bool| { + hf.set_encode_special_tokens(false); let hf_enc = hf .encode(input, false) .unwrap_or_else(|e| panic!("{model_name}: HF encode({input:?}): {e}")); let hf_ids = hf_enc.get_ids().to_vec(); - let our_ids = match ours.encode(input) { + let our_ids = match ours.encode_with_options(input, false, false) { Ok(ids) => ids, Err(e) => { failures.push(format!(" encode error on {input:?}: {e}")); - continue; + return; } }; if our_ids != hf_ids { @@ -866,19 +1329,45 @@ mod tests { )); } + hf.set_encode_special_tokens(true); + let hf_split_ids = hf + .encode(input, false) + .unwrap_or_else(|e| panic!("{model_name}: HF split encode({input:?}): {e}")) + .get_ids() + .to_vec(); + let our_split_ids = match ours.encode_with_options(input, false, true) { + Ok(ids) => ids, + Err(e) => { + failures.push(format!(" split encode error on {input:?}: {e}")); + return; + } + }; + if our_split_ids != hf_split_ids { + failures.push(format!( + " split encode mismatch on {input:?}: got {} tokens, expected {}\n\ + \x20 ours: {:?}\n\ + \x20 hf: {:?}", + our_split_ids.len(), + hf_split_ids.len(), + &our_split_ids[..our_split_ids.len().min(20)], + &hf_split_ids[..hf_split_ids.len().min(20)], + )); + } + // Decode comparison (skip empty inputs / empty token sequences). - if input.is_empty() || hf_ids.is_empty() { - continue; + if !compare_decode || input.is_empty() || hf_ids.is_empty() { + return; } + hf.set_encode_special_tokens(false); let hf_decoded = match hf.decode(&hf_ids, false) { Ok(d) => d, - Err(_) => continue, + Err(_) => return, }; let our_decoded = match ours.decode(&hf_ids, false) { Ok(d) => d, Err(e) => { failures.push(format!(" decode error on {input:?}: {e}")); - continue; + return; } }; if our_decoded != hf_decoded { @@ -890,6 +1379,13 @@ mod tests { &hf_decoded[..hf_decoded.len().min(100)], )); } + }; + + for &input in corpus { + compare_input(input, true); + } + for input in &special_samples { + compare_input(input, false); } failures } @@ -992,6 +1488,347 @@ mod tests { assert!(f.is_empty(), "hoangquan456/Kimi-K2.5:\n{}", f.join("\n")); } + #[test] + fn split_special_tokens_kimi_k2_5_matches_hf_tokenizers() { + let model = "hoangquan456/Kimi-K2.5"; + let mut hf = tokenizers::Tokenizer::from_pretrained(model, None) + .unwrap_or_else(|e| panic!("{model}: HF load failed: {e}")); + let ours = + Tokenizer::from_model(model).unwrap_or_else(|e| panic!("{model}: load failed: {e}")); + + let inputs = &[ + "", + "hello world", + "🤔final answer", + "<|tool_calls_section_begin|>{\"name\":\"search\"}", + ]; + + hf.set_encode_special_tokens(true); + for input in inputs { + let hf_ids = hf + .encode(*input, false) + .unwrap_or_else(|e| panic!("{model}: HF encode({input:?}): {e}")) + .get_ids() + .to_vec(); + let our_ids = ours + .encode_with_options(input, false, true) + .unwrap_or_else(|e| panic!("{model}: fastokens encode({input:?}): {e}")); + assert_eq!( + our_ids, hf_ids, + "split_special_tokens=true mismatch on {input:?}" + ); + } + + hf.set_encode_special_tokens(false); + let input = ""; + let hf_ids = hf + .encode(input, false) + .unwrap_or_else(|e| panic!("{model}: HF encode({input:?}): {e}")) + .get_ids() + .to_vec(); + let our_ids = ours + .encode_with_options(input, false, false) + .unwrap_or_else(|e| panic!("{model}: fastokens encode({input:?}): {e}")); + assert_eq!(our_ids, hf_ids, "default special-token matching mismatch"); + } + + /// Verify that `split_special_tokens=true` matches HuggingFace + /// `tokenizers` when a non-special added token's content is contained + /// within a special added token's content. + /// + /// Config: special `` (id 6) + non-special `mask` (id 7), input + /// ``. HF runs the full leftmost-longest matcher, so `` + /// (special) wins over `mask` (non-special) and — with + /// `encode_special_tokens=true` — the matched span is tokenized as + /// ordinary text, yielding the char-level ids `[0,1,2,3,4,5]`. fastokens + /// should do the same rather than letting `mask` match inside the special + /// span as added-token id 7. + /// + /// The added-token ids are chosen as `vocab_size + index` (6, 7) because HF + /// `tokenizers` ignores the `id` field in `added_tokens` and assigns ids + /// sequentially from the model vocab size; fastokens respects the `id` + /// field, so the two only agree when the JSON ids match HF's assignment. + /// + /// The default-path baseline (`split_special_tokens=false`) confirms the + /// tokenizer JSON is valid for both backends. + #[test] + fn split_special_tokens_overlapping_non_special_matches_hf() { + let json = r#"{ + "version": "1.0", + "added_tokens": [ + {"id": 6, "content": "", "single_word": false, "lstrip": false, "rstrip": false, "normalized": false, "special": true}, + {"id": 7, "content": "mask", "single_word": false, "lstrip": false, "rstrip": false, "normalized": false, "special": false} + ], + "normalizer": null, + "pre_tokenizer": null, + "post_processor": null, + "decoder": null, + "model": { + "type": "BPE", + "dropout": null, + "unk_token": null, + "continuing_subword_prefix": "", + "end_of_word_suffix": "", + "fuse_unk": false, + "byte_fallback": false, + "vocab": {"<": 0, "m": 1, "a": 2, "s": 3, "k": 4, ">": 5}, + "merges": [] + } + }"#; + + let mut hf = tokenizers::Tokenizer::from_str(json).expect("HF load failed"); + let ours = Tokenizer::from_json(serde_json::from_str(json).unwrap()) + .expect("fastokens load failed"); + + // Baseline: default path (split=false) — both emit the special id. + hf.set_encode_special_tokens(false); + let hf_default: Vec = hf.encode("", false).unwrap().get_ids().to_vec(); + let our_default = ours.encode_with_options("", false, false).unwrap(); + assert_eq!( + our_default, hf_default, + "default path baseline: fastokens must match HF" + ); + assert_eq!( + our_default, + vec![6], + "default path should emit the special id" + ); + + // Split path (split=true): HF tokenizes the whole as text. + hf.set_encode_special_tokens(true); + let hf_split: Vec = hf.encode("", false).unwrap().get_ids().to_vec(); + let our_split = ours.encode_with_options("", false, true).unwrap(); + + assert_eq!( + hf_split, + vec![0, 1, 2, 3, 4, 5], + "HF should tokenize the special span as ordinary text" + ); + assert_eq!( + our_split, hf_split, + "split_special_tokens=true should match HF by tokenizing the whole \ + special span as text" + ); + } + + fn structural_test_tokenizer() -> Tokenizer { + let json = r#"{ + "version": "1.0", + "added_tokens": [ + {"id": 100, "content": "[BOS]", "single_word": false, "lstrip": false, "rstrip": false, "normalized": false, "special": true}, + {"id": 101, "content": "", "single_word": false, "lstrip": false, "rstrip": false, "normalized": false, "special": true}, + {"id": 102, "content": "", "single_word": false, "lstrip": false, "rstrip": false, "normalized": false, "special": false}, + {"id": 103, "content": "magic", "single_word": false, "lstrip": false, "rstrip": false, "normalized": false, "special": false} + ], + "normalizer": null, + "pre_tokenizer": null, + "post_processor": { + "type": "TemplateProcessing", + "single": [ + {"SpecialToken": {"id": "[BOS]", "type_id": 0}}, + {"Sequence": {"id": "A", "type_id": 0}} + ], + "pair": [], + "special_tokens": { + "[BOS]": {"id": "[BOS]", "ids": [100], "tokens": ["[BOS]"]} + } + }, + "decoder": null, + "model": { + "type": "BPE", + "dropout": null, + "unk_token": null, + "continuing_subword_prefix": "", + "end_of_word_suffix": "", + "fuse_unk": false, + "byte_fallback": false, + "vocab": { + " ": 0, "<": 1, ">": 2, "e": 3, "h": 4, "i": 5, + "k": 6, "l": 7, "n": 8, "o": 9, "t": 10, + "a": 11, "c": 12, "g": 13, "m": 14 + }, + "merges": [] + } + }"#; + Tokenizer::from_json(serde_json::from_str(json).unwrap()).unwrap() + } + + #[test] + fn encode_with_structural_tokens_mixes_template_ids_and_literal_placeholders() { + let tok = structural_test_tokenizer(); + let structural_config = StructuralTokenConfig::new( + &["".to_string(), "".to_string()], + &HashSet::from(["".to_string()]), + ) + .unwrap(); + let think_placeholder = "\u{e000}STRUCTTOK_0\u{e000}".to_string(); + let tool_placeholder = "\u{e000}STRUCTTOK_1\u{e000}".to_string(); + let placeholder_map = HashMap::from([ + (think_placeholder.clone(), "".to_string()), + (tool_placeholder.clone(), "".to_string()), + ]); + + let ids = tok + .encode_with_structural_tokens( + &format!("hello {think_placeholder} {tool_placeholder}"), + &structural_config, + &placeholder_map, + false, + ) + .unwrap(); + + assert_eq!( + ids, + vec![ + 4, 3, 7, 7, 9, 0, // "hello " + 101, // structural "" + 0, 1, 10, 4, 5, 8, 6, 2, // literal " " + 0, 102, // structural " " + 0, 1, 10, 9, 9, 7, 2, // literal " " + ] + ); + } + + #[test] + fn encode_segments_bypasses_added_tokens_for_text_segments() { + let tok = structural_test_tokenizer(); + + let ids = tok + .encode_segments( + [ + ("", true), + (" hello ", false), + (" ", true), + (" ", false), + ], + false, + ) + .unwrap(); + + assert_eq!( + ids, + vec![ + 102, // structural non-special "" + 0, 4, 3, 7, 7, 9, 0, 1, 10, 9, 9, 7, 2, // literal " hello " + 0, 101, // structural " " + 0, 1, 10, 4, 5, 8, 6, 2, // literal " " + ] + ); + } + + #[test] + fn encode_segments_tiktoken_safe_splits_long_runs_like_legacy_tiktoken() { + let tok = structural_test_tokenizer(); + let long = format!( + "{}{}", + "h".repeat(TIKTOKEN_MAX_CONSECUTIVE_WHITESPACE_OR_NONWHITESPACE_CHARS + 3), + "e".repeat(4), + ); + + let ids = tok + .encode_segments_tiktoken_safe([(long.as_str(), false)], false) + .unwrap(); + + assert_eq!( + ids.len(), + TIKTOKEN_MAX_CONSECUTIVE_WHITESPACE_OR_NONWHITESPACE_CHARS + 3 + "".len() + 4 + ); + assert!( + !ids.contains(&102), + "literal must not become added token" + ); + } + + #[test] + fn encode_with_structural_tokens_can_add_post_processor_tokens() { + let tok = structural_test_tokenizer(); + let structural_config = + StructuralTokenConfig::new(&["".to_string()], &HashSet::new()).unwrap(); + + let ids = tok + .encode_with_structural_tokens( + "hello ", + &structural_config, + &HashMap::new(), + true, + ) + .unwrap(); + + assert_eq!(ids, vec![100, 4, 3, 7, 7, 9, 0, 101]); + } + + #[test] + fn encode_with_structural_tokens_empty_config_uses_normal_added_token_matching() { + let tok = structural_test_tokenizer(); + let structural_config = StructuralTokenConfig::new(&[], &HashSet::new()).unwrap(); + + let ids = tok + .encode_with_structural_tokens( + " magic", + &structural_config, + &HashMap::new(), + false, + ) + .unwrap(); + + assert_eq!(ids, vec![101, 0, 103]); + } + + #[test] + fn encode_with_structural_tokens_keeps_bare_non_special_added_tokens() { + let tok = structural_test_tokenizer(); + let structural_config = + StructuralTokenConfig::new(&["".to_string()], &HashSet::new()).unwrap(); + + let ids = tok + .encode_with_structural_tokens( + "hello magic ", + &structural_config, + &HashMap::new(), + false, + ) + .unwrap(); + + assert_eq!( + ids, + vec![ + 4, 3, 7, 7, 9, 0, // "hello " + 103, // bare non-special added token "magic" + 0, 101, // structural " " + ] + ); + } + + #[test] + fn encode_with_structural_tokens_keeps_restored_bare_non_special_added_tokens() { + let tok = structural_test_tokenizer(); + let structural_config = StructuralTokenConfig::new( + &["".to_string()], + &HashSet::from(["magic".to_string()]), + ) + .unwrap(); + let placeholder = "\u{e000}STRUCTTOK_0\u{e000}".to_string(); + let placeholder_map = HashMap::from([(placeholder.clone(), "magic".to_string())]); + + let ids = tok + .encode_with_structural_tokens( + &format!("hello {placeholder} "), + &structural_config, + &placeholder_map, + false, + ) + .unwrap(); + + assert_eq!( + ids, + vec![ + 4, 3, 7, 7, 9, 0, // "hello " + 103, // restored bare non-special added token "magic" + 0, 101, // structural " " + ] + ); + } + // ── Cache consistency ──────────────────────────────────────────── /// Verify that encoding the same input twice produces identical results, diff --git a/src/models/bpe.rs b/src/models/bpe.rs index ad81eee..d81160d 100644 --- a/src/models/bpe.rs +++ b/src/models/bpe.rs @@ -22,6 +22,11 @@ type Vocab = HashMap; const INVALID_TOKEN: u32 = u32::MAX; +/// Pretokens with at most this many initial (per-byte) symbols use the +/// stack-resident linear-scan merge instead of the heap. Sized so the stack +/// arrays fit in registers/L1 and `u8` linked-list indices stay in range. +const SMALL_MERGE_MAX: usize = 32; + /// Open-addressing hash table for merge lookups. #[derive(Clone, PartialEq)] struct MergeMap { @@ -103,6 +108,7 @@ fn fx_hash(key: u64) -> u64 { } /// FxHash-based [`BuildHasher`] for the token cache. +#[derive(Clone, Default)] struct FxBuildHasher; impl std::hash::BuildHasher for FxBuildHasher { @@ -142,8 +148,6 @@ impl std::hash::Hasher for FxStrHasher { type FxHashMap = HashMap; const FLAT_CACHE_BITS: usize = 16; -const FLAT_CACHE_SIZE: usize = 1 << FLAT_CACHE_BITS; -const FLAT_CACHE_MASK: usize = FLAT_CACHE_SIZE - 1; const EMPTY_SLOT: u64 = 0; #[derive(Clone, Copy)] @@ -156,11 +160,10 @@ struct CacheSlot { key_offset: u32, } -/// Maximum load factor before the cache is cleared. -const FLAT_CACHE_MAX_LOAD: usize = FLAT_CACHE_SIZE * 3 / 4; - struct FlatCache { bpe_id: usize, + mask: usize, + max_load: usize, slots: Vec, pool: Vec, key_pool: Vec, @@ -169,8 +172,17 @@ struct FlatCache { impl FlatCache { fn new() -> Self { + Self::with_bits(FLAT_CACHE_BITS) + } + + /// A flat cache with `1 << bits` slots. The thread-local L1 uses + /// [`FLAT_CACHE_BITS`]; the shared-cache shards use fewer bits each. + fn with_bits(bits: usize) -> Self { + let size = 1usize << bits; Self { bpe_id: 0, + mask: size - 1, + max_load: size * 3 / 4, slots: vec![ CacheSlot { hash: EMPTY_SLOT, @@ -179,10 +191,10 @@ impl FlatCache { key_len: 0, key_offset: 0, }; - FLAT_CACHE_SIZE + size ], - pool: Vec::with_capacity(256 * 1024), - key_pool: Vec::with_capacity(512 * 1024), + pool: Vec::new(), + key_pool: Vec::new(), count: 0, } } @@ -222,7 +234,7 @@ impl FlatCache { fn get(&self, key: &str, out: &mut Vec) -> bool { let hash = Self::hash_str(key); let key_bytes = key.as_bytes(); - let mut idx = hash as usize & FLAT_CACHE_MASK; + let mut idx = hash as usize & self.mask; loop { let slot = unsafe { self.slots.get_unchecked(idx) }; if slot.hash == hash { @@ -238,18 +250,18 @@ impl FlatCache { if slot.hash == EMPTY_SLOT { return false; } - idx = (idx + 1) & FLAT_CACHE_MASK; + idx = (idx + 1) & self.mask; } } #[inline(always)] fn insert(&mut self, key: &str, ids: &[u32]) { - if self.count >= FLAT_CACHE_MAX_LOAD { + if self.count >= self.max_load { self.clear(); } let hash = Self::hash_str(key); let key_bytes = key.as_bytes(); - let mut idx = hash as usize & FLAT_CACHE_MASK; + let mut idx = hash as usize & self.mask; loop { let slot = unsafe { self.slots.get_unchecked(idx) }; let h = slot.hash; @@ -288,7 +300,7 @@ impl FlatCache { return; } } - idx = (idx + 1) & FLAT_CACHE_MASK; + idx = (idx + 1) & self.mask; } } } @@ -300,45 +312,50 @@ thread_local! { const CACHE_SHARDS: usize = 64; +/// Slot bits per shared-cache shard: 64 shards x 4096 slots = 256k entries. +const SHARED_SHARD_BITS: usize = 12; + +/// Cross-thread token cache: [`CACHE_SHARDS`] mutex-guarded [`FlatCache`] +/// shards. Compared to the previous `HashMap>` shards this +/// is allocation-free per insert (keys and ids are copied into per-shard +/// pools that retain capacity across clears) and bounded (a shard clears at +/// 3/4 load instead of growing forever), which both removes the two heap +/// allocations per cold pretoken and fixes unbounded memory growth on +/// diverse long-running traffic. struct SharedCache { - shards: Vec>>>, + shards: Vec>, } impl SharedCache { fn new() -> Self { Self { shards: (0..CACHE_SHARDS) - .map(|_| Mutex::new(HashMap::with_hasher(FxBuildHasher))) + .map(|_| Mutex::new(FlatCache::with_bits(SHARED_SHARD_BITS))) .collect(), } } + /// Shard selector. Uses the TOP bits of the same hash `FlatCache` + /// indexes slots with (low bits), so the two are independent. #[inline] fn shard_index(key: &str) -> usize { - let bytes = key.as_bytes(); - let mut h: u64 = bytes.len() as u64; - for &b in &bytes[..bytes.len().min(8)] { - h = h.wrapping_add(b as u64).wrapping_mul(0x9E3779B97F4A7C15); - } - h as usize & (CACHE_SHARDS - 1) + (FlatCache::hash_str(key) >> (64 - 6)) as usize & (CACHE_SHARDS - 1) } #[inline] fn get_into(&self, key: &str, out: &mut Vec) -> bool { - let shard = self.shards[Self::shard_index(key)].lock().unwrap(); - if let Some(ids) = shard.get(key) { - out.extend_from_slice(ids); - true - } else { - false - } + self.shards[Self::shard_index(key)] + .lock() + .unwrap() + .get(key, out) } - fn insert(&self, key: String, value: Vec) { - self.shards[Self::shard_index(&key)] + #[inline] + fn insert(&self, key: &str, ids: &[u32]) { + self.shards[Self::shard_index(key)] .lock() .unwrap() - .insert(key, value); + .insert(key, ids); } } @@ -588,7 +605,7 @@ pub struct Bpe { shared_cache: SharedCache, fused_shared_cache: SharedCache, id_to_token: Vec, - token_to_id: HashMap, + token_to_id: FxHashMap, byte_to_initial_token: [u32; 256], ranked_merge_map: RankedMergeMap, byte_pair_initial: Vec<(u32, u32)>, @@ -802,7 +819,11 @@ impl Bpe { shared_cache: SharedCache::new(), fused_shared_cache: SharedCache::new(), id_to_token, - token_to_id: vocab.clone(), + token_to_id: { + let mut m = HashMap::with_capacity_and_hasher(vocab.len(), FxBuildHasher); + m.extend(vocab.iter().map(|(k, v)| (k.clone(), *v))); + m + }, byte_to_initial_token, ranked_merge_map, byte_pair_initial, @@ -908,7 +929,7 @@ impl Bpe { } c.insert(input, ids); }); - self.shared_cache.insert(input.to_string(), ids.to_vec()); + self.shared_cache.insert(input, ids); Ok(()) } @@ -955,12 +976,154 @@ impl Bpe { }) } + /// Linear-scan BPE merge for short pretokens (`n <= SMALL_MERGE_MAX` + /// initial symbols). Avoids the `BinaryHeap` entirely: a stack-resident + /// doubly-linked list plus a per-position rank array, find-min by a short + /// scan over stack `u32`s, merge (O(1) pointer update), then refresh only + /// the two neighbor pairs. At these sizes this beats the heap's + /// sift/stale-entry traffic and does zero heap allocation. + /// + /// Produces the identical token sequence as [`run_merge_loop`]: both + /// process the globally lowest-`(rank, pos)` active pair each step (the + /// heap's `MergeEntry` key is `(rank << 32) | pos`; the scan's strict + /// `<` keeps the leftmost/lowest-`pos` position on ties). Enforced by + /// `merge_small_matches_heap` differential test. `ids[..n]` are the + /// per-byte initial token ids. + fn merge_small_raw( + &self, + bytes: &[u8], + ids: &mut [u32; SMALL_MERGE_MAX], + n: usize, + out: &mut Vec, + ) { + let mut next = [0u8; SMALL_MERGE_MAX]; + let mut prev = [0u8; SMALL_MERGE_MAX]; + let mut ranks = [u32::MAX; SMALL_MERGE_MAX]; + let mut new_ids = [0u32; SMALL_MERGE_MAX]; + for i in 0..n { + next[i] = (i + 1) as u8; + prev[i] = (i as u8).wrapping_sub(1); // prev[0] = 255 (>= n): sentinel + } + // Round-1 ranks via the dense byte-pair table: one direct-indexed + // load per pair instead of a CSR neighbor scan. + for i in 0..n - 1 { + let (rank, new_id) = + self.byte_pair_initial[bytes[i] as usize * 256 + bytes[i + 1] as usize]; + if rank != u32::MAX { + ranks[i] = rank; + new_ids[i] = new_id; + } + } + loop { + let mut best = u32::MAX; + let mut best_i = 0usize; + for (i, &rank) in ranks[..n - 1].iter().enumerate() { + if rank < best { + best = rank; + best_i = i; + } + } + if best == u32::MAX { + break; + } + let i = best_i; + ids[i] = new_ids[i]; + let dead = next[i] as usize; + let new_right = next[dead] as usize; + next[i] = new_right as u8; + ranks[dead] = u32::MAX; + if new_right < n { + prev[new_right] = i as u8; + match self.merge_adj.get(ids[i], ids[new_right]) { + Some((rank, new_id)) => { + ranks[i] = rank; + new_ids[i] = new_id; + } + None => ranks[i] = u32::MAX, + } + } else { + ranks[i] = u32::MAX; + } + let left = prev[i] as usize; + if left < n { + match self.merge_adj.get(ids[left], ids[i]) { + Some((rank, new_id)) => { + ranks[left] = rank; + new_ids[left] = new_id; + } + None => ranks[left] = u32::MAX, + } + } + } + let mut i = 0usize; + while i < n { + out.push(ids[i]); + i = next[i] as usize; + } + } + + /// `ignore_merges` whole-pretoken lookup: is the ByteLevel-encoded form + /// of the entire pretoken a single vocab token? Encodes into a stack + /// buffer for short pretokens (`BYTE_TO_CHAR` codepoints are <= U+0143, + /// so <= 2 UTF-8 bytes per input byte) instead of allocating a `String` + /// per cold pretoken; falls back to a heap `String` only for long ones. + #[inline] + fn whole_pretoken_id(&self, raw: &str) -> Option { + const STACK: usize = 128; + let bytes = raw.as_bytes(); + if bytes.len() * 2 <= STACK { + let mut buf = [0u8; STACK]; + let mut n = 0; + for &b in bytes { + n += BYTE_TO_CHAR[b as usize].encode_utf8(&mut buf[n..]).len(); + } + // SAFETY: buf[..n] is a concatenation of `char::encode_utf8` + // outputs, hence valid UTF-8. + let encoded = unsafe { std::str::from_utf8_unchecked(&buf[..n]) }; + self.token_to_id.get(encoded).copied() + } else { + let mut encoded = String::with_capacity(bytes.len() * 2); + for &b in bytes { + encoded.push(BYTE_TO_CHAR[b as usize]); + } + self.token_to_id.get(encoded.as_str()).copied() + } + } + /// Priority-queue BPE merge on raw (pre-ByteLevel) bytes. fn merge_all_raw_into(&self, raw_input: &str, out: &mut Vec) -> Result<()> { if raw_input.is_empty() { return Ok(()); } + // Short pretokens (the common case) skip the heap + thread-local + // scratch and use the stack-resident linear-scan merge. `n` initial + // symbols == byte length, so the tier is known up front. + let bytes = raw_input.as_bytes(); + if bytes.len() <= SMALL_MERGE_MAX { + let n = bytes.len(); + let mut ids = [0u32; SMALL_MERGE_MAX]; + for (i, &byte) in bytes.iter().enumerate() { + let id = self.byte_to_initial_token[byte as usize]; + if id == INVALID_TOKEN { + return Err(format!("byte 0x{byte:02x} has no token in vocabulary")); + } + ids[i] = id; + } + if n == 1 { + out.push(ids[0]); + } else { + self.merge_small_raw(bytes, &mut ids, n, out); + } + return Ok(()); + } + + self.merge_all_raw_heap_into(raw_input, out) + } + + /// Reference heap implementation used for long pretokens and as the + /// correctness oracle for the short linear-scan tier. + fn merge_all_raw_heap_into(&self, raw_input: &str, out: &mut Vec) -> Result<()> { TL_MERGE_SCRATCH.with(|s| { let mut scratch = s.borrow_mut(); scratch.symbols.clear(); @@ -1124,15 +1287,11 @@ impl Bpe { return Ok(()); } - if self.ignore_merges { - let mut encoded = String::with_capacity(raw_input.len()); - for &byte in raw_input.as_bytes() { - encoded.push(BYTE_TO_CHAR[byte as usize]); - } - if let Some(&id) = self.token_to_id.get(encoded.as_str()) { - out.push(id); - return Ok(()); - } + if self.ignore_merges + && let Some(id) = self.whole_pretoken_id(raw_input) + { + out.push(id); + return Ok(()); } self.merge_all_raw_into(raw_input, out)?; @@ -1146,8 +1305,7 @@ impl Bpe { } c.insert(raw_input, ids); }); - self.fused_shared_cache - .insert(raw_input.to_string(), ids.to_vec()); + self.fused_shared_cache.insert(raw_input, ids); Ok(()) } @@ -1185,24 +1343,18 @@ impl Bpe { continue; } - if self.ignore_merges { - let mut encoded = String::with_capacity(text.len()); - for &byte in text.as_bytes() { - encoded.push(BYTE_TO_CHAR[byte as usize]); - } - if let Some(&id) = self.token_to_id.get(encoded.as_str()) { - out.push(id); - cache.insert(text, &out[start..]); - continue; - } + if self.ignore_merges + && let Some(id) = self.whole_pretoken_id(text) + { + out.push(id); + cache.insert(text, &out[start..]); + continue; } self.merge_all_raw_into(text, out)?; cache.insert(text, &out[start..]); - let key = text.to_string(); - let val = out[start..].to_vec(); - self.fused_shared_cache.insert(key, val); + self.fused_shared_cache.insert(text, &out[start..]); } } Ok(()) @@ -1331,6 +1483,32 @@ mod tests { assert_eq!(bpe.tokenize("abab").unwrap(), vec![4, 4]); } + #[test] + fn merge_small_matches_heap() { + let bpe = test_bpe(); + let alphabet = [b'a', b'b', b'c', b'd']; + let mut state = 0x9e37_79b9_7f4a_7c15u64; + + for len in 1..=SMALL_MERGE_MAX { + for _ in 0..512 { + let mut bytes = Vec::with_capacity(len); + for _ in 0..len { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + bytes.push(alphabet[state as usize & 3]); + } + let input = std::str::from_utf8(&bytes).unwrap(); + let mut small = Vec::new(); + let mut heap = Vec::new(); + + bpe.merge_all_raw_into(input, &mut small).unwrap(); + bpe.merge_all_raw_heap_into(input, &mut heap).unwrap(); + assert_eq!(small, heap, "short merge mismatch for {input:?}"); + } + } + } + #[test] fn deserialize_from_json() { let json = serde_json::json!({ diff --git a/src/pre_tokenizers.rs b/src/pre_tokenizers.rs index 30cc65e..f43c578 100644 --- a/src/pre_tokenizers.rs +++ b/src/pre_tokenizers.rs @@ -1,4 +1,6 @@ pub(crate) mod byte_level; +pub(crate) mod fast_split; +pub(crate) mod fast_split_o200k; mod split; use crate::{ diff --git a/src/pre_tokenizers/fast_split.rs b/src/pre_tokenizers/fast_split.rs new file mode 100644 index 0000000..042eded --- /dev/null +++ b/src/pre_tokenizers/fast_split.rs @@ -0,0 +1,897 @@ +//! Hand-specialized scanners for known pre-tokenization split regexes. +//! +//! The split regexes of common tokenizers are known ahead of time, so instead +//! of running a general regex engine (fancy-regex / PCRE2) we can recognize +//! the exact pattern string at [`Split`](super::Split) construction and +//! dispatch to a scanner hand-written for that pattern: a first-byte dispatch +//! plus SWAR run scans, roughly an order of magnitude faster than PCRE2 JIT +//! on the same pattern. +//! +//! Scalar scanning technique and the Qwen2 walker are adapted from +//! gigatoken (), +//! MIT License, Copyright (c) 2026 Marcel Rød. +//! +//! Correctness contract: for a recognized pattern, iterating +//! [`FastSplitScheme::advance`] from 0 must produce exactly the spans of +//! `regex.find_iter(text)` for the original pattern (whose alternations +//! cover every char, so consecutive matches tile the text with no gaps). +//! The differential tests at the bottom of this file enforce that against +//! fancy-regex on hand-picked edge cases and a deterministic fuzz corpus; +//! any new scheme added here needs the same treatment. + +use std::sync::LazyLock; + +/// The Qwen2/Qwen3 split pattern, as it appears in `tokenizer.json`. +pub(crate) const QWEN2_PATTERN: &str = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"; + +/// The Qwen2 scheme with number runs of up to three chars (`\p{N}{1,3}` +/// instead of `\p{N}`): GLM-4/GLM-5, and dolma2 / OLMo 2/3. +pub(crate) const QWEN2_N3_PATTERN: &str = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"; + +/// A split regex with a hand-specialized scanner. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum FastSplitScheme { + /// [`QWEN2_PATTERN`]: Qwen2/Qwen3 and derivatives. + Qwen2, + /// [`QWEN2_N3_PATTERN`]: Qwen2 with `\p{N}{1,3}` — GLM, dolma2/OLMo. + Qwen2N3, + /// The Kimi (moonshotai K2 family) o200k-family pattern. + Kimi, + /// The o200k_base pattern (GPT-4o, gpt-oss). + O200k, +} + +impl FastSplitScheme { + /// Recognize a split regex source string. Only byte-for-byte known + /// patterns qualify; anything else takes the generic engines. + pub fn from_pattern(pattern: &str) -> Option { + match pattern { + QWEN2_PATTERN => Some(Self::Qwen2), + QWEN2_N3_PATTERN => Some(Self::Qwen2N3), + super::fast_split_o200k::KIMI_PATTERN => Some(Self::Kimi), + super::fast_split_o200k::O200K_PATTERN => Some(Self::O200k), + _ => None, + } + } + + /// Advance past one pretoken starting at `pos`. Requires + /// `pos < bytes.len()`; always returns `end > pos` and `end <= len`, + /// on a UTF-8 char boundary for valid UTF-8 input. + #[inline(always)] + pub fn advance(self, bytes: &[u8], pos: usize) -> usize { + match self { + Self::Qwen2 => qwen2_advance::<1>(bytes, pos), + Self::Qwen2N3 => qwen2_advance::<3>(bytes, pos), + Self::Kimi => super::fast_split_o200k::kimi_advance(bytes, pos), + Self::O200k => super::fast_split_o200k::o200k_advance(bytes, pos), + } + } + + /// A chunk boundary near `target` that no pretoken can cross, for + /// parallel scanning of large inputs: right after a `\n` whose next byte + /// is non-whitespace ASCII. Safe because no alternative reaches backward + /// across such a point — letter-run prefixes exclude `\r\n`, a whitespace + /// run containing a newline always ends at its last newline, and a + /// punctuation run's absorbed tail is only `[\r\n]` (all whitespace, so + /// excluded by the non-whitespace check). + /// + /// The one exception is the **O200k** scheme, whose punct tail is + /// `[\r\n/]*`: a punct token can absorb the `\n` *and* a following `/` + /// (e.g. `.\n/` is one token), so a boundary placed between the `\n` and + /// the `/` would be crossed — the preceding chunk would scan past it and + /// the next chunk would re-scan the `/`, producing overlapping ranges. + /// So `/` is not a safe boundary-follower for O200k. Returns `None` when + /// no safe boundary exists in `[target, limit)`. + pub fn find_safe_boundary(self, bytes: &[u8], target: usize, limit: usize) -> Option { + // O200k's `[\r\n/]*` tail lets a punct token cross a newline into a `/`. + let slash_unsafe = matches!(self, Self::O200k); + let mut p = target.max(1); + while p < limit { + if bytes[p - 1] == b'\n' { + let b = bytes[p]; + if b < 0x80 && !is_ascii_ws(b) && !(slash_unsafe && b == b'/') { + return Some(p); + } + } + p += 1; + } + None + } +} + +// --------------------------------------------------------------------------- +// Unicode classification: packed 2-bit class per codepoint +// --------------------------------------------------------------------------- + +/// Character class as used by the split regexes: `\p{L}`, `\p{N}`, `\s` +/// (the Unicode White_Space property), and everything else. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[repr(u8)] +enum CharClass { + Letter = 0, + Number = 1, + Whitespace = 2, + Other = 3, +} + +/// 2 bits per codepoint, 4 codepoints per byte (~272 KiB). One L1/L2 load +/// replaces an ICU trie walk per non-ASCII char; only the cache lines for +/// scripts actually present in the input stay resident. +static CLASS_TABLE: LazyLock> = LazyLock::new(build_class_table); + +fn build_class_table() -> Box<[u8]> { + use icu_properties::props::{GeneralCategory, GeneralCategoryGroup, WhiteSpace}; + use icu_properties::{CodePointMapData, CodePointSetData}; + + const N: usize = 0x11_0000; + let mut classes = vec![CharClass::Other as u8; N]; + let gc = CodePointMapData::::new(); + for (group, class) in [ + (GeneralCategoryGroup::Letter, CharClass::Letter), + (GeneralCategoryGroup::Number, CharClass::Number), + ] { + for range in gc.iter_ranges_for_group(group) { + classes[*range.start() as usize..=*range.end() as usize].fill(class as u8); + } + } + // White_Space is disjoint from GC Letter/Number, so fill order is moot. + for range in CodePointSetData::new::().iter_ranges() { + classes[*range.start() as usize..=*range.end() as usize].fill(CharClass::Whitespace as u8); + } + + let mut packed = vec![0u8; N / 4].into_boxed_slice(); + for (cp, class) in classes.iter().enumerate() { + packed[cp >> 2] |= class << ((cp & 3) * 2); + } + packed +} + +/// Class of a codepoint. `cp` must be `<= 0x10FFFF`, which [`decode_cp`] +/// guarantees via its clamp. +#[inline(always)] +fn class_of(cp: u32) -> CharClass { + debug_assert!(cp <= 0x10_FFFF); + let table = &**CLASS_TABLE; + // SAFETY: cp <= 0x10FFFF, so cp >> 2 < table.len() == 0x110000 / 4. + let byte = unsafe { *table.get_unchecked((cp >> 2) as usize) }; + match (byte >> ((cp & 3) * 2)) & 3 { + 0 => CharClass::Letter, + 1 => CharClass::Number, + 2 => CharClass::Whitespace, + _ => CharClass::Other, + } +} + +// --------------------------------------------------------------------------- +// Branchless byte predicates +// --------------------------------------------------------------------------- + +#[inline(always)] +pub(super) fn is_letter(b: u8) -> bool { + (b | 0x20).wrapping_sub(b'a') < 26 +} + +#[inline(always)] +pub(super) fn is_digit(b: u8) -> bool { + b.wrapping_sub(b'0') < 10 +} + +#[inline(always)] +pub(super) fn is_ascii_ws(b: u8) -> bool { + b == b' ' || b.wrapping_sub(9) < 5 +} + +// --------------------------------------------------------------------------- +// UTF-8 decode tolerant of invalid input +// --------------------------------------------------------------------------- + +/// The codepoint reported for byte sequences that cannot be decoded within +/// bounds (truncated tails) or that assemble past the Unicode range. +/// U+10FFFF is an unassigned noncharacter: class `Other`. +const CP_INVALID: u32 = 0x10_FFFF; + +/// Decode one non-ASCII scalar. Requires only `pos < bytes.len()` and +/// `bytes[pos] >= 0x80`; invalid bytes decode deterministically, never read +/// past `bytes.len()`, and the returned length never overruns it. +#[inline(always)] +pub(super) unsafe fn decode_cp(bytes: &[u8], pos: usize) -> (u32, usize) { + if pos + 4 > bytes.len() { + return decode_cp_near_end(bytes, pos); + } + unsafe { + let b0 = *bytes.get_unchecked(pos) as u32; + let b1 = (*bytes.get_unchecked(pos + 1) & 0x3F) as u32; + if b0 < 0xE0 { + return (((b0 & 0x1F) << 6) | b1, 2); + } + let b2 = (*bytes.get_unchecked(pos + 2) & 0x3F) as u32; + if b0 < 0xF0 { + return (((b0 & 0x0F) << 12) | (b1 << 6) | b2, 3); + } + let b3 = (*bytes.get_unchecked(pos + 3) & 0x3F) as u32; + ( + (((b0 & 0x07) << 18) | (b1 << 12) | (b2 << 6) | b3).min(CP_INVALID), + 4, + ) + } +} + +/// [`decode_cp`]'s slow path for `pos + 4 > bytes.len()`: identical results +/// for complete sequences; a sequence truncated by the buffer end consumes +/// exactly the remaining bytes and yields [`CP_INVALID`]. +#[cold] +#[inline(never)] +fn decode_cp_near_end(bytes: &[u8], pos: usize) -> (u32, usize) { + let len = bytes.len(); + let b0 = bytes[pos] as u32; + let need = if b0 < 0xE0 { + 2 + } else if b0 < 0xF0 { + 3 + } else { + 4 + }; + if pos + need > len { + return (CP_INVALID, len - pos); + } + let b1 = (bytes[pos + 1] & 0x3F) as u32; + if need == 2 { + return (((b0 & 0x1F) << 6) | b1, 2); + } + let b2 = (bytes[pos + 2] & 0x3F) as u32; + if need == 3 { + return (((b0 & 0x0F) << 12) | (b1 << 6) | b2, 3); + } + let b3 = (bytes[pos + 3] & 0x3F) as u32; + ( + (((b0 & 0x07) << 18) | (b1 << 12) | (b2 << 6) | b3).min(CP_INVALID), + 4, + ) +} + +// --------------------------------------------------------------------------- +// SWAR letter scan +// --------------------------------------------------------------------------- + +const HI: u64 = 0x8080_8080_8080_8080; + +/// High bit set in each lane that is NOT an ASCII letter. +#[inline(always)] +fn swar64_letter_nonmask(word: u64) -> u64 { + let lowered = word | 0x2020_2020_2020_2020; + let ge_a = (lowered | HI).wrapping_sub(0x6161_6161_6161_6161); + let le_z = 0xFAFA_FAFA_FAFA_FAFA_u64.wrapping_sub(lowered); + !(ge_a & le_z) & HI +} + +/// Advance `pos` past ASCII letters, 8 bytes per iteration. +#[inline(always)] +fn swar_scan_letters(bytes: &[u8], mut pos: usize) -> usize { + let len = bytes.len(); + while pos + 8 <= len { + // SAFETY: pos + 8 <= len. + let word = unsafe { (bytes.as_ptr().add(pos) as *const u64).read_unaligned() }; + if word & HI != 0 { + break; + } + let nonletter = swar64_letter_nonmask(word); + if nonletter != 0 { + return pos + nonletter.to_le().trailing_zeros() as usize / 8; + } + pos += 8; + } + while pos < len { + // SAFETY: pos < len. + let b = unsafe { *bytes.get_unchecked(pos) }; + if is_letter(b) { + pos += 1; + } else { + break; + } + } + pos +} + +// --------------------------------------------------------------------------- +// Shared run scans +// --------------------------------------------------------------------------- + +/// `\p{L}+` continuation: past ASCII letters via SWAR, then non-ASCII +/// letters via the class table. +#[inline(always)] +fn scan_letters_from(bytes: &[u8], pos: usize) -> usize { + let len = bytes.len(); + let mut p = pos; + loop { + p = swar_scan_letters(bytes, p); + if p < len && unsafe { *bytes.get_unchecked(p) } >= 0x80 { + let (cp, l) = unsafe { decode_cp(bytes, p) }; + if class_of(cp) == CharClass::Letter { + p += l; + continue; + } + } + return p; + } +} + +/// `[^\s\p{L}\p{N}]+` continuation. +#[inline(always)] +fn scan_other_from(bytes: &[u8], pos: usize) -> usize { + let len = bytes.len(); + let mut p = pos; + loop { + while p < len { + // SAFETY: p < len. + let b = unsafe { *bytes.get_unchecked(p) }; + if b >= 0x80 { + break; + } + if is_letter(b) || is_digit(b) || is_ascii_ws(b) { + return p; + } + p += 1; + } + if p < len { + let (cp, l) = unsafe { decode_cp(bytes, p) }; + if class_of(cp) == CharClass::Other { + p += l; + continue; + } + } + return p; + } +} + +/// `[\r\n]*`: past a run of CR/LF bytes. +#[inline(always)] +fn scan_newlines(bytes: &[u8], mut pos: usize) -> usize { + while pos < bytes.len() { + // SAFETY: pos < len. + let b = unsafe { *bytes.get_unchecked(pos) }; + if b == b'\r' || b == b'\n' { + pos += 1; + } else { + break; + } + } + pos +} + +/// `\p{N}{1,MAX}` continuation: extend a number run that already matched +/// `consumed` chars to at most `MAX` chars total. ASCII digits and Unicode +/// `\p{N}` count alike. +#[inline(always)] +pub(super) fn scan_numbers_max( + bytes: &[u8], + mut pos: usize, + mut consumed: u32, +) -> usize { + let len = bytes.len(); + while consumed < MAX && pos < len { + // SAFETY: pos < len. + let b = unsafe { *bytes.get_unchecked(pos) }; + if is_digit(b) { + pos += 1; + consumed += 1; + continue; + } + if b >= 0x80 { + let (cp, l) = unsafe { decode_cp(bytes, pos) }; + if class_of(cp) == CharClass::Number { + pos += l; + consumed += 1; + continue; + } + } + break; + } + pos +} + +/// If the char at `pos` is `\p{L}`, return the offset just past it. +#[inline(always)] +fn letter_end_at(bytes: &[u8], pos: usize) -> Option { + let &b = bytes.get(pos)?; + if is_letter(b) { + return Some(pos + 1); + } + if b >= 0x80 { + let (cp, l) = unsafe { decode_cp(bytes, pos) }; + if class_of(cp) == CharClass::Letter { + return Some(pos + l); + } + } + None +} + +// --------------------------------------------------------------------------- +// Qwen2 scheme +// --------------------------------------------------------------------------- + +/// Whitespace-led token starting at `start`: the alternatives +/// `\s*[\r\n]+` | `\s+(?!\S)` | `\s+`, in that priority. +/// Precondition: the letter-prefix (`[^\r\n\p{L}\p{N}]?\p{L}+`) and +/// space+punct (` ?[^\s\p{L}\p{N}]+...`) alternatives were ruled out. +#[inline(always)] +fn ws_token_end(bytes: &[u8], start: usize) -> usize { + let len = bytes.len(); + let mut p = start; + let mut last_nl_end = 0usize; // 0 = run contains no \r\n + let mut last_char_start = start; + while p < len { + // SAFETY: p < len. + let b = unsafe { *bytes.get_unchecked(p) }; + if b == b'\r' || b == b'\n' { + last_char_start = p; + p += 1; + last_nl_end = p; + } else if is_ascii_ws(b) { + last_char_start = p; + p += 1; + } else if b >= 0x80 { + let (cp, l) = unsafe { decode_cp(bytes, p) }; + if class_of(cp) == CharClass::Whitespace { + last_char_start = p; + p += l; + } else { + break; + } + } else { + break; + } + } + if last_nl_end != 0 { + return last_nl_end; // `\s*[\r\n]+`: through the last newline, even at EOS + } + if p >= len { + return p; // `\s+(?!\S)`: lookahead succeeds at EOS + } + if last_char_start > start { + return last_char_start; // `\s+(?!\S)`: all but the last ws char + } + p // `\s+`: single whitespace char before content +} + +/// Advance past one Qwen2-scheme token starting at `pos`. +/// `pos` must be < `bytes.len()`. +/// +/// `MAX_DIGITS` is the number-run cap: 1 for the Qwen2 pattern's `\p{N}`, +/// 3 for the GLM/dolma2 variant's `\p{N}{1,3}` (the only difference +/// between the two patterns). +#[inline(always)] +fn qwen2_advance(bytes: &[u8], pos: usize) -> usize { + // SAFETY: pos < len per contract. + let b0 = unsafe { *bytes.get_unchecked(pos) }; + + // Hot path 1: ASCII letter — `\p{L}+` with empty prefix + if is_letter(b0) { + return scan_letters_from(bytes, pos + 1); + } + + // Hot path 2: space prefix + if b0 == b' ' { + let Some(&b1) = bytes.get(pos + 1) else { + return pos + 1; // trailing lone space (`\s+(?!\S)` at EOS) + }; + if is_letter(b1) { + return scan_letters_from(bytes, pos + 2); // " word" + } + if b1 < 0x80 { + if is_digit(b1) { + return pos + 1; // numbers never absorb the space + } + if is_ascii_ws(b1) { + return ws_token_end(bytes, pos); + } + // ` ?[^\s\p{L}\p{N}]+[\r\n]*` + let p = scan_other_from(bytes, pos + 2); + return scan_newlines(bytes, p); + } + let (cp, l) = unsafe { decode_cp(bytes, pos + 1) }; + let p1 = pos + 1 + l; + match class_of(cp) { + CharClass::Letter => return scan_letters_from(bytes, p1), + CharClass::Whitespace => return ws_token_end(bytes, pos), + CharClass::Number => return pos + 1, + CharClass::Other => { + let p = scan_other_from(bytes, p1); + return scan_newlines(bytes, p); + } + } + } + + // Non-ASCII + if b0 >= 0x80 { + let (cp, l) = unsafe { decode_cp(bytes, pos) }; + let p0 = pos + l; + let class = class_of(cp); + if class == CharClass::Letter { + return scan_letters_from(bytes, p0); + } + if class == CharClass::Number { + if MAX_DIGITS == 1 { + return p0; // `\p{N}`: exactly one char + } + return scan_numbers_max::(bytes, p0, 1); + } + // Any non-letter/number char except \r\n may prefix a letter run + if let Some(p) = letter_end_at(bytes, p0) { + return scan_letters_from(bytes, p); + } + if class == CharClass::Whitespace { + return ws_token_end(bytes, pos); + } + let p = scan_other_from(bytes, p0); + return scan_newlines(bytes, p); + } + + // ASCII digit: `\p{N}` (one char) or `\p{N}{1,3}` (up to MAX_DIGITS) + if is_digit(b0) { + if MAX_DIGITS == 1 { + return pos + 1; + } + return scan_numbers_max::(bytes, pos + 1, 1); + } + + // Apostrophe: case-insensitive contractions + if b0 == b'\'' { + match bytes.get(pos + 1).map(u8::to_ascii_lowercase) { + Some(b's' | b'd' | b'm' | b't') => return pos + 2, + Some(b'l') if bytes.get(pos + 2).map(u8::to_ascii_lowercase) == Some(b'l') => { + return pos + 3; + } + Some(b'v') if bytes.get(pos + 2).map(u8::to_ascii_lowercase) == Some(b'e') => { + return pos + 3; + } + Some(b'r') if bytes.get(pos + 2).map(u8::to_ascii_lowercase) == Some(b'e') => { + return pos + 3; + } + _ => {} + } + // U+017F LATIN SMALL LETTER LONG S case-folds to 's' under `(?i)` + if bytes.get(pos + 1) == Some(&0xC5) && bytes.get(pos + 2) == Some(&0xBF) { + return pos + 3; + } + // Not a contraction: `'` can still prefix a letter run + if let Some(p) = letter_end_at(bytes, pos + 1) { + return scan_letters_from(bytes, p); + } + let p = scan_other_from(bytes, pos + 1); + return scan_newlines(bytes, p); + } + + // \r and \n are excluded from the letter-run prefix + if b0 == b'\r' || b0 == b'\n' { + return ws_token_end(bytes, pos); + } + + // Other ASCII whitespace (\t, \x0b, \x0c) may prefix a letter run + if is_ascii_ws(b0) { + if let Some(p) = letter_end_at(bytes, pos + 1) { + return scan_letters_from(bytes, p); + } + return ws_token_end(bytes, pos); + } + + // ASCII punctuation/symbol + if let Some(p) = letter_end_at(bytes, pos + 1) { + return scan_letters_from(bytes, p); + } + let p = scan_other_from(bytes, pos + 1); + scan_newlines(bytes, p) +} + +// --------------------------------------------------------------------------- +// Tests: differential against fancy-regex on the original pattern +// --------------------------------------------------------------------------- + +/// Assert a scheme's scanner produces byte-identical spans to `fancy-regex` +/// on the exact `pattern`, over the whole vendored multilingual corpus +/// (~2 MB real prose: English + 8 non-Latin scripts; see +/// `tests/fixtures/PROVENANCE.md`). Compares token-by-token in lockstep so a +/// divergence reports the exact byte offset and surrounding tokens rather +/// than a giant Vec diff. Shared by the o200k-family schemes' tests. +/// +/// Since the scanner only replaces the split stage — byte-level and BPE +/// downstream are untouched — span-identity over a large real corpus is +/// end-to-end token-id equivalence for the change. +#[cfg(test)] +pub(crate) fn assert_corpus_matches_regex( + label: &str, + pattern: &str, + advance: impl Fn(&[u8], usize) -> usize, +) { + let path = concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/fixtures/corpus_multilingual.txt" + ); + let text = std::fs::read_to_string(path).unwrap_or_else(|e| panic!("read corpus: {e}")); + let bytes = text.as_bytes(); + let re = fancy_regex::Regex::new(pattern).unwrap(); + let mut re_iter = re.find_iter(&text); + let mut pos = 0usize; + let mut idx = 0usize; + let mut recent: Vec<(&str, &str)> = Vec::new(); + while pos < bytes.len() { + let end = advance(bytes, pos); + assert!( + end > pos && end <= bytes.len(), + "{label}: no progress at byte {pos}" + ); + let fast = &text[pos..end]; + match re_iter.next() { + Some(m) => { + let m = m.expect("regex match error"); + let re_str = &text[m.start()..m.end()]; + if recent.len() == 8 { + recent.remove(0); + } + recent.push((fast, re_str)); + assert_eq!( + fast, re_str, + "{label}: mismatch at token {idx} (byte {pos}): fast={fast:?} regex={re_str:?}\n recent (fast,regex): {recent:?}" + ); + } + None => panic!("{label}: scanner produced extra token at {idx}: {fast:?}"), + } + pos = end; + idx += 1; + } + assert!( + re_iter.next().is_none(), + "{label}: regex produced tokens past where the scanner stopped" + ); + eprintln!( + "{label}: {idx} tokens match over {} bytes of corpus", + bytes.len() + ); +} + +#[cfg(test)] +mod tests { + use super::*; + + fn regex_spans(pattern: &str, s: &str) -> Vec<(usize, usize)> { + let re = fancy_regex::Regex::new(pattern).unwrap(); + re.find_iter(s) + .map(|m| { + let m = m.unwrap(); + (m.start(), m.end()) + }) + .collect() + } + + fn fast_spans(scheme: FastSplitScheme, s: &str) -> Vec<(usize, usize)> { + let bytes = s.as_bytes(); + let mut spans = Vec::new(); + let mut pos = 0; + while pos < bytes.len() { + let end = scheme.advance(bytes, pos); + assert!( + end > pos && end <= bytes.len(), + "no progress at {pos} in {s:?}" + ); + spans.push((pos, end)); + pos = end; + } + spans + } + + fn assert_parity(scheme: FastSplitScheme, pattern: &str, s: &str) { + assert_eq!( + fast_spans(scheme, s), + regex_spans(pattern, s), + "split mismatch on {s:?}" + ); + } + + #[test] + fn qwen2_edge_cases() { + let cases = [ + "", + "hello", + " hello", + "hello world", + " hello", + " hello", + "\thello", + "\t hello", + " \thello", + "hello\n", + "hello \n", + "hello\r\n\r\n", + "hello \r\n world", + "don't", + "DON'T", + "it's it'S IT'Ll we'VE they're I'd I'm", + "c'mon 'quoted' rock'n'roll", + "'s't'll've're'd'm", + "a1b2c3", + "123", + "1 2 34", + " 5", + "3.14159", + "hello, world! (parens) [brackets]", + " ,,, ", + "!!!\n\n", + "?!\r", + "trailing spaces ", + "trailing tab\t", + "\n\nleading", + "\r\n", + " ", + "\t", + "múltiple ñiño café über", + "日本語のテキスト", + "русский текст", + "한국어 텍스트", + "e\u{301}xpose\u{301}", // combining acute: Mark is not \p{L} + "\u{00A0}nbsp\u{00A0}\u{00A0}", // NBSP is White_Space + "\u{2028}line\u{2029}sep", + "\u{1F600} emoji \u{1F680}\u{1F680}", + "١٢٣ arabic digits", // \p{N} single-char matches + "๑๒๓ thai digits", + "Ⅷ roman numeral", // Nl is \p{N} + "'\u{017F} long s contraction", + "mixed 'ſ fold", + "tab\tbetween\twords", + "\x0b\x0c vertical form", + "space before ünïcode", + "@user #tag $5 100%", + "a'll b'LL c'Ll", + "'''", + "''s", + "x''y", + "-prefix word", + ".hidden", + "\u{3000}ideographic space", + "za\u{017F}ka", // long s inside a letter run + ]; + let extra_digit_cases = [ + "123456", + "12 345 6789", + " 12", + "a123b", + "١٢٣٤ arabic digits", + "๑๒๓๔๕ thai digits", + "1٢3٤5", + "x99999y", + "2024-07-25 12:34:56", + ]; + for scheme in [ + (FastSplitScheme::Qwen2, QWEN2_PATTERN), + (FastSplitScheme::Qwen2N3, QWEN2_N3_PATTERN), + ] { + for case in cases.iter().chain(extra_digit_cases.iter()) { + assert_parity(scheme.0, scheme.1, case); + } + } + } + + #[test] + fn qwen2_fuzz_parity() { + // Deterministic LCG fuzz over a pool that exercises every branch: + // ASCII letters/digits/punct, all ASCII whitespace, CRLF runs, + // contraction heads, non-ASCII letters/digits/whitespace/symbols, + // combining marks, and multi-byte boundary cases. + let pool: Vec = "abcXYZ z09'\"-.,!? \t\r\n\ + \u{00A0}\u{2028}\u{3000}\u{017F}\ + éñÜ日本語һЖ¡½٣๒Ⅷ😀🚀\u{301}\u{10FFFD}" + .chars() + .collect(); + let mut rng: u64 = 0x1234_5678_9ABC_DEF0; + for round in 0..4000 { + rng = rng + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + let len = (rng >> 48) as usize % 64; + let mut s = String::new(); + for _ in 0..len { + rng = rng + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + s.push(pool[(rng >> 33) as usize % pool.len()]); + } + for (scheme, pattern) in [ + (FastSplitScheme::Qwen2, QWEN2_PATTERN), + (FastSplitScheme::Qwen2N3, QWEN2_N3_PATTERN), + ] { + let fast = fast_spans(scheme, &s); + let reference = regex_spans(pattern, &s); + assert_eq!(fast, reference, "fuzz mismatch round {round} on {s:?}"); + } + } + } + + #[test] + fn qwen2_matches_regex_on_corpus() { + super::assert_corpus_matches_regex("qwen2", QWEN2_PATTERN, |b, p| qwen2_advance::<1>(b, p)); + } + + #[test] + fn qwen2n3_matches_regex_on_corpus() { + super::assert_corpus_matches_regex("qwen2n3", QWEN2_N3_PATTERN, |b, p| { + qwen2_advance::<3>(b, p) + }); + } + + #[test] + fn o200k_safe_boundary_excludes_slash_after_newline() { + // `.\n/foo\nbar`: byte 2 is right after `\n` and is `/`. O200k's + // punct tail `[\r\n/]*` lets `.` absorb `\n/`, so byte 2 is an unsafe + // boundary and O200k must skip to the next `\n`-led non-slash point + // (byte 7, after the second `\n`, before `b`). Schemes with a `[\r\n]` + // tail may use the slash boundary. + let bytes = b".\n/foo\nbar"; + assert_eq!( + FastSplitScheme::O200k.find_safe_boundary(bytes, 0, bytes.len()), + Some(7) + ); + for scheme in [ + FastSplitScheme::Qwen2, + FastSplitScheme::Qwen2N3, + FastSplitScheme::Kimi, + ] { + assert_eq!( + scheme.find_safe_boundary(bytes, 0, bytes.len()), + Some(2), + "{scheme:?} has a [\\r\\n] tail; slash boundary is safe" + ); + } + } + + #[test] + fn recognizes_only_known_patterns() { + use super::super::fast_split_o200k::{KIMI_PATTERN, O200K_PATTERN}; + + assert_eq!( + FastSplitScheme::from_pattern(QWEN2_PATTERN), + Some(FastSplitScheme::Qwen2) + ); + assert_eq!( + FastSplitScheme::from_pattern(QWEN2_N3_PATTERN), + Some(FastSplitScheme::Qwen2N3) + ); + assert_eq!( + FastSplitScheme::from_pattern(KIMI_PATTERN), + Some(FastSplitScheme::Kimi) + ); + assert_eq!( + FastSplitScheme::from_pattern(O200K_PATTERN), + Some(FastSplitScheme::O200k) + ); + + // Guard: the Kimi K2/K3 (knext) tiktoken `pat_str`, assembled here + // exactly as `tokenization_kimi.py` builds it, must equal + // `KIMI_PATTERN` and route to the Kimi scanner. knext ships as a + // tiktoken model (no tokenizer.json); its `pat_str` reaches + // `from_pattern` verbatim through the tiktoken -> tokenizer.json + // conversion, so an upstream regex change would otherwise silently + // drop it to the slow PCRE2 path. This literal is the failing + // tripwire for that. + let knext_pat_str = [ + r"[\p{Han}]+", + r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?", + r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?", + r"\p{N}{1,3}", + r" ?[^\s\p{L}\p{N}]+[\r\n]*", + r"\s*[\r\n]+", + r"\s+(?!\S)", + r"\s+", + ] + .join("|"); + assert_eq!(knext_pat_str, KIMI_PATTERN, "Kimi K2/K3 pat_str drifted"); + assert_eq!( + FastSplitScheme::from_pattern(&knext_pat_str), + Some(FastSplitScheme::Kimi) + ); + + // One char off -> no specialization. + assert_eq!( + FastSplitScheme::from_pattern(&QWEN2_PATTERN.replace("'t", "'x")), + None + ); + assert_eq!(FastSplitScheme::from_pattern(""), None); + } +} diff --git a/src/pre_tokenizers/fast_split_o200k.rs b/src/pre_tokenizers/fast_split_o200k.rs new file mode 100644 index 0000000..e83291e --- /dev/null +++ b/src/pre_tokenizers/fast_split_o200k.rs @@ -0,0 +1,783 @@ +//! Hand-specialized scanner for the o200k regex family: o200k_base +//! (GPT-4o, gpt-oss), the Nemotron-3 variant, and the Kimi (moonshotai K2) +//! variant. o200k and Kimi are wired into dispatch (see +//! [`super::fast_split::FastSplitScheme::from_pattern`]). Nemotron-3 is +//! covered by [`advance_pos`]'s const parameters +//! (``) but has no pattern const or `advance` +//! wrapper yet, so it is not routed. +//! +//! Faithfully ported from gigatoken +//! (), MIT License, Copyright (c) +//! 2026 Marcel Rød — its `pretokenize/fast/o200k_family.rs` scalar walker +//! and `pretokenize/unicode.rs` class tables. Shares low-level primitives +//! (`decode_cp`, byte predicates, `scan_numbers_max`) with the sibling +//! [`super::fast_split`] module. +//! +//! The patterns share the shape +//! +//! ```text +//! HAN-RUN?| +//! [^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*[\p{Ll}\p{Lm}\p{Lo}\p{M}]+SUF?| +//! [^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+[\p{Ll}\p{Lm}\p{Lo}\p{M}]*SUF?| +//! \p{N}{1,3} or \p{N}| ?[^\s\p{L}\p{N}]+TAIL|\s*[\r\n]+|\s+(?!\S)|\s+ +//! ``` +//! +//! where `SUF = (?i:'s|'t|'re|'ve|'m|'ll|'d)` exists in o200k and Kimi +//! (`CONTRACTIONS`), the digit group is `\p{N}{1,3}` for o200k/Kimi vs +//! `\p{N}` for Nemotron (`DIGITS3`), the absorbed punct tail is `[\r\n/]*` +//! for o200k/Nemotron vs `[\r\n]*` for Kimi (`SLASH`), and Kimi alone +//! (`HAN`) prepends `[\p{Han}]+` and intersects both letter brackets with +//! `[^\p{Han}]`. See the differential tests at the bottom. + +use std::sync::LazyLock; + +use super::fast_split::{decode_cp, is_ascii_ws, is_digit, is_letter, scan_numbers_max}; + +/// The Kimi (moonshotai Kimi-K2 family) split pattern, as it appears in +/// `tokenizer.json`. +pub(crate) const KIMI_PATTERN: &str = "[\\p{Han}]+|[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}&&[^\\p{Han}]]*[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}&&[^\\p{Han}]]+(?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{Han}]]+[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}&&[^\\p{Han}]]*(?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+"; + +/// The o200k_base split pattern (GPT-4o, gpt-oss), as it appears in +/// `tokenizer.json`. +pub(crate) const O200K_PATTERN: &str = r"[^\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+"; + +/// Advance past one Kimi-scheme token starting at `pos`. +/// `pos` must be `< bytes.len()`; returns `end` with `pos < end <= len`. +#[inline(always)] +pub(super) fn kimi_advance(bytes: &[u8], pos: usize) -> usize { + // Kimi = o200k with contractions, \p{N}{1,3}, no `/` tail, Han runs. + advance_pos::(bytes, pos) +} + +/// Advance past one o200k-scheme token starting at `pos`. +/// `pos` must be `< bytes.len()`; returns `end` with `pos < end <= len`. +#[inline(always)] +pub(super) fn o200k_advance(bytes: &[u8], pos: usize) -> usize { + // o200k = contractions, \p{N}{1,3}, `[\r\n/]*` tail, no Han runs. + advance_pos::(bytes, pos) +} + +// --------------------------------------------------------------------------- +// Character classes +// --------------------------------------------------------------------------- + +/// The o200k letter-run classes: `Upper` is Lu|Lt (only in the first +/// bracket), `Lower` is Ll (only in the second), `Caseless` is Lm|Lo (in +/// both). Marks (`\p{M}`) are their own class: they join letter runs like +/// `Caseless` but, being outside `\p{L}`, also continue `[^\s\p{L}\p{N}]+`. +#[derive(Clone, Copy, PartialEq, Eq)] +#[repr(u8)] +enum O200kCharClass { + Upper = 0, + Lower = 1, + Caseless = 2, + Mark = 3, + Number = 4, + Whitespace = 5, + Other = 6, +} + +/// The Kimi classes: the o200k classes with `\p{Han}` split out. Han runs +/// get their own leading alternative and are excluded from both letter +/// brackets, but the general-category rules are otherwise Han-blind (a Han +/// numeral still counts toward `\p{N}{1,3}`, a Han symbol still spans a +/// `[^\s\p{L}\p{N}]+` run). Each Han variant records the base class it +/// behaves as outside a Han run. +#[derive(Clone, Copy, PartialEq, Eq)] +#[repr(u8)] +enum KimiCharClass { + Upper = 0, + Lower = 1, + Caseless = 2, + Mark = 3, + Number = 4, + Whitespace = 5, + Other = 6, + Han = 7, + HanNumber = 8, + HanOther = 9, +} + +impl KimiCharClass { + #[inline(always)] + fn base(self) -> O200kCharClass { + match self { + KimiCharClass::Upper => O200kCharClass::Upper, + KimiCharClass::Lower => O200kCharClass::Lower, + KimiCharClass::Caseless | KimiCharClass::Han => O200kCharClass::Caseless, + KimiCharClass::Mark => O200kCharClass::Mark, + KimiCharClass::Number | KimiCharClass::HanNumber => O200kCharClass::Number, + KimiCharClass::Whitespace => O200kCharClass::Whitespace, + KimiCharClass::Other | KimiCharClass::HanOther => O200kCharClass::Other, + } + } + + #[inline(always)] + fn is_han(self) -> bool { + self as u8 >= KimiCharClass::Han as u8 + } +} + +// --------------------------------------------------------------------------- +// Packed class tables (4 bits per codepoint, 2 per byte, ~544 KiB each) +// --------------------------------------------------------------------------- + +/// One `O200kCharClass` byte per codepoint — the unpacked form both table +/// builders start from. +fn o200k_classes_unpacked() -> Vec { + use icu_properties::CodePointMapData; + use icu_properties::CodePointSetData; + use icu_properties::props::{GeneralCategory, GeneralCategoryGroup, WhiteSpace}; + + const N: usize = 0x11_0000; + let mut classes = vec![O200kCharClass::Other as u8; N]; + let gc = CodePointMapData::::new(); + for (category, class) in [ + (GeneralCategory::UppercaseLetter, O200kCharClass::Upper), + (GeneralCategory::TitlecaseLetter, O200kCharClass::Upper), + (GeneralCategory::LowercaseLetter, O200kCharClass::Lower), + (GeneralCategory::ModifierLetter, O200kCharClass::Caseless), + (GeneralCategory::OtherLetter, O200kCharClass::Caseless), + ] { + for range in gc.iter_ranges_for_value(category) { + classes[*range.start() as usize..=*range.end() as usize].fill(class as u8); + } + } + for (group, class) in [ + (GeneralCategoryGroup::Mark, O200kCharClass::Mark), + (GeneralCategoryGroup::Number, O200kCharClass::Number), + ] { + for range in gc.iter_ranges_for_group(group) { + classes[*range.start() as usize..=*range.end() as usize].fill(class as u8); + } + } + for range in CodePointSetData::new::().iter_ranges() { + classes[*range.start() as usize..=*range.end() as usize] + .fill(O200kCharClass::Whitespace as u8); + } + classes +} + +fn pack_nibbles(classes: &[u8]) -> Box<[u8]> { + let mut packed = vec![0u8; classes.len() / 2].into_boxed_slice(); + for (i, chunk) in classes.chunks_exact(2).enumerate() { + packed[i] = chunk[0] | (chunk[1] << 4); + } + packed +} + +static KIMI_CLASS_TABLE: LazyLock> = LazyLock::new(build_kimi_class_table); + +fn build_kimi_class_table() -> Box<[u8]> { + use icu_properties::CodePointMapData; + use icu_properties::props::Script; + + let mut classes = o200k_classes_unpacked(); + let script = CodePointMapData::