diff --git a/.github/workflows/atomsplit.yml b/.github/workflows/bitsplit.yml similarity index 72% rename from .github/workflows/atomsplit.yml rename to .github/workflows/bitsplit.yml index c2059f94f..f250a4675 100644 --- a/.github/workflows/atomsplit.yml +++ b/.github/workflows/bitsplit.yml @@ -1,10 +1,10 @@ -name: atomsplit +name: bitsplit on: push: - paths: ["tokenizers/atomsplit/**", "tokenizers/bitmap_gen/**", ".github/workflows/atomsplit.yml"] + paths: ["tokenizers/bitsplit/**", "tokenizers/bitmap_gen/**", ".github/workflows/bitsplit.yml"] pull_request: - paths: ["tokenizers/atomsplit/**", "tokenizers/bitmap_gen/**", ".github/workflows/atomsplit.yml"] + paths: ["tokenizers/bitsplit/**", "tokenizers/bitmap_gen/**", ".github/workflows/bitsplit.yml"] defaults: run: @@ -18,11 +18,11 @@ jobs: - uses: actions/checkout@v4 - uses: dtolnay/rust-toolchain@stable with: { components: rustfmt, clippy } - - run: cargo fmt -p atomsplit -p bitmap_gen -- --check - - run: cargo clippy -p atomsplit --all-targets -- -D warnings - - run: cargo test -p atomsplit # unit + parity gates (cl100k/deepseek/byte_level vs onig) + - run: cargo fmt -p bitsplit -p bitmap_gen -- --check + - run: cargo clippy -p bitsplit --all-targets -- -D warnings + - run: cargo test -p bitsplit # unit + parity gates (cl100k/deepseek/byte_level vs onig) # committed classify tables must match the generator - - run: cargo run -p bitmap_gen && git diff --exit-code atomsplit/src/atom_tables.rs + - run: cargo run -p bitmap_gen && git diff --exit-code bitsplit/src/atom_tables.rs # AVX-512 classify path — GitHub runners often lack AVX-512, so emulate it with Intel SDE and run the # SAME byte-exact test binaries under it. @@ -39,7 +39,7 @@ jobs: echo "$PWD/${SDE_VER}" >> "$GITHUB_PATH" - name: Run tests under SDE (-future = AVX-512) run: | - cargo test -p atomsplit --no-run --message-format=json \ + cargo test -p bitsplit --no-run --message-format=json \ | jq -r 'select(.profile.test == true) | .executable | select(. != null)' \ | while read -r bin; do echo "SDE: $bin"; sde64 -future -- "$bin"; done @@ -55,4 +55,4 @@ jobs: env: RUSTFLAGS: "-C target-feature=+simd128" CARGO_TARGET_WASM32_WASIP1_RUNNER: "wasmtime run --" - run: cargo test -p atomsplit --target wasm32-wasip1 + run: cargo test -p bitsplit --target wasm32-wasip1 diff --git a/tokenizers/Cargo.lock b/tokenizers/Cargo.lock index 9e7c653e3..88ba567cc 100644 --- a/tokenizers/Cargo.lock +++ b/tokenizers/Cargo.lock @@ -67,17 +67,6 @@ version = "1.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" -[[package]] -name = "atomsplit" -version = "0.1.0" -dependencies = [ - "fancy-regex 0.13.0", - "logos", - "memchr", - "onig", - "pcre2", -] - [[package]] name = "autocfg" version = "1.5.1" @@ -96,12 +85,6 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" -[[package]] -name = "beef" -version = "0.5.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3a8241f3ebb85c056b509d4327ad0358fbbba6ffb340bf388f26350aeda225b1" - [[package]] name = "bindgen" version = "0.72.1" @@ -122,30 +105,15 @@ dependencies = [ "syn", ] -[[package]] -name = "bit-set" -version = "0.5.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0700ddab506f33b20a03b13996eccd309a48e5ff77d0d95926aa0210fb4e95f1" -dependencies = [ - "bit-vec 0.6.3", -] - [[package]] name = "bit-set" version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "08807e080ed7f9d5433fa9b275196cfc35414f66a0c79d864dc51a0d825231a3" dependencies = [ - "bit-vec 0.8.0", + "bit-vec", ] -[[package]] -name = "bit-vec" -version = "0.6.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "349f9b6a179ed607305526ca489b34ad0a41aed5f7980fa90eb03160b69598fb" - [[package]] name = "bit-vec" version = "0.8.0" @@ -165,6 +133,14 @@ dependencies = [ "unicode-properties", ] +[[package]] +name = "bitsplit" +version = "0.1.0" +dependencies = [ + "ahash", + "onig", +] + [[package]] name = "bitvec" version = "1.1.1" @@ -233,8 +209,6 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e228eec9be7c17ccb640b59b36a5cd805ea2a564a4c5e162c2f659fea30d3b96" dependencies = [ "find-msvc-tools", - "jobserver", - "libc", "shlex 2.0.1", ] @@ -655,24 +629,13 @@ dependencies = [ "cc", ] -[[package]] -name = "fancy-regex" -version = "0.13.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "531e46835a22af56d1e3b66f04844bed63158bc094a628bec1d321d9b4c44bf2" -dependencies = [ - "bit-set 0.5.3", - "regex-automata", - "regex-syntax", -] - [[package]] name = "fancy-regex" version = "0.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72cf461f865c862bb7dc573f643dd6a2b6842f7c30b07882b56bd148cc2761b8" dependencies = [ - "bit-set 0.8.0", + "bit-set", "regex-automata", "regex-syntax", ] @@ -1180,16 +1143,6 @@ version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" -[[package]] -name = "jobserver" -version = "0.1.35" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1c00acbd29eabad4a2392fa0e921c874934dbbf4194312ad20f04a0ed67a3cb3" -dependencies = [ - "getrandom 0.4.3", - "libc", -] - [[package]] name = "js-sys" version = "0.3.103" @@ -1250,40 +1203,6 @@ version = "0.4.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" -[[package]] -name = "logos" -version = "0.15.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ff472f899b4ec2d99161c51f60ff7075eeb3097069a36050d8037a6325eb8154" -dependencies = [ - "logos-derive", -] - -[[package]] -name = "logos-codegen" -version = "0.15.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "192a3a2b90b0c05b27a0b2c43eecdb7c415e29243acc3f89cc8247a5b693045c" -dependencies = [ - "beef", - "fnv", - "lazy_static", - "proc-macro2", - "quote", - "regex-syntax", - "rustc_version", - "syn", -] - -[[package]] -name = "logos-derive" -version = "0.15.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "605d9697bcd5ef3a42d38efc51541aa3d6a4a25f7ab6d1ed0da5ac632a26b470" -dependencies = [ - "logos-codegen", -] - [[package]] name = "lru-slab" version = "0.1.2" @@ -1476,28 +1395,6 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" -[[package]] -name = "pcre2" -version = "0.2.11" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9e970b0fcce0c7ee6ef662744ff711f21ccd6f11b7cf03cd187a80e89797fc67" -dependencies = [ - "libc", - "log", - "pcre2-sys", -] - -[[package]] -name = "pcre2-sys" -version = "0.2.10" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "18b9073c1a2549bd409bf4a32c94d903bb1a09bf845bc306ae148897fa0760a4" -dependencies = [ - "cc", - "libc", - "pkg-config", -] - [[package]] name = "percent-encoding" version = "2.3.2" @@ -1901,15 +1798,6 @@ version = "2.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94300abf3f1ae2e2b8ffb7b58043de3d399c73fa6f4b73826402a5c457614dbe" -[[package]] -name = "rustc_version" -version = "0.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cfcb3a22ef46e85b45de6ee7e79d063319ebb6594faafcf1c225ea92ab6e9b92" -dependencies = [ - "semver", -] - [[package]] name = "rustix" version = "1.1.4" @@ -1989,12 +1877,6 @@ dependencies = [ "winapi-util", ] -[[package]] -name = "semver" -version = "1.0.28" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8a7852d02fc848982e0c167ef163aaff9cd91dc640ba85e263cb1ce46fae51cd" - [[package]] name = "serde" version = "1.0.228" @@ -2286,13 +2168,13 @@ version = "0.23.2-dev.0" dependencies = [ "ahash", "assert_approx_eq", - "atomsplit", + "bitsplit", "compact_str", "criterion 0.6.0", "daachorse 3.0.2", "dary_heap", "derive_builder", - "fancy-regex 0.17.0", + "fancy-regex", "getrandom 0.3.4", "hf-hub", "indicatif 0.18.5", diff --git a/tokenizers/Cargo.toml b/tokenizers/Cargo.toml index 3759dfdb0..5100c6367 100644 --- a/tokenizers/Cargo.toml +++ b/tokenizers/Cargo.toml @@ -1,6 +1,6 @@ [workspace] resolver = "3" -members = ["bitmap_gen", "atomsplit", "tk-encode", "tk-train"] +members = ["bitmap_gen", "bitsplit", "tk-encode", "tk-train"] [package] authors = [ diff --git a/tokenizers/atomsplit/Cargo.toml b/tokenizers/atomsplit/Cargo.toml deleted file mode 100644 index fb8947f3f..000000000 --- a/tokenizers/atomsplit/Cargo.toml +++ /dev/null @@ -1,48 +0,0 @@ -[package] -name = "atomsplit" -version = "0.1.0" -edition = "2024" -rust-version = "1.89" # AVX-512 VBMI intrinsics (simd_avx_classify) are stable only since 1.89 -authors = [ - "Arthur Zucker ", - "Luc Georges ", - "Simon Brandeis ", -] -homepage = "https://github.com/huggingface/tokenizers" -repository = "https://github.com/huggingface/tokenizers" -documentation = "https://docs.rs/atomsplit/" -license = "Apache-2.0" -keywords = ["tokenizer", "nlp", "pretokenizer", "simd", "unicode"] -categories = ["text-processing", "parsing"] -description = "SIMD Unicode atom classification + FSM pre-tokenization: one classify pass into byte-exact token spans (cl100k, GPT-2/ByteLevel, DeepSeek, Whitespace, Bert, ...)." -exclude = ["benches/data/"] - -[package.metadata.docs.rs] -all-features = true - -[lib] -name = "atomsplit" -path = "src/lib.rs" - -[dependencies] -memchr = "2.8.2" # SIMD search: used for single-byte search in CharDelimiterSplit (1.4–23× vs scalar), or multi-byte string pattern search in `literal`. -# NOTE: classify tables live in src/atom_tables.rs (committed, generated). Regenerate after any atom -# scheme change with `cargo run -p bitmap_gen`. No build script / build-dep — atomsplit builds clean. -# onig is C (Oniguruma) and fancy-regex pulls it in transitively for benches only — neither -# cross-compiles to wasm32 (no wasi libc for the C build), so gate them off wasm. -[target.'cfg(not(target_arch = "wasm32"))'.dev-dependencies] -onig = "6.5.1" # reference regex for the parity tests + benches -fancy-regex = "0.13" # second regex engine in the throughput benches -logos = "0.15" # third: a compile-time DFA lexer-generator (pure Rust) -pcre2 = "0.2" # fourth: PCRE2 with its JIT (C, pregenerated bindings — no libclang) -[[bench]] -name = "regex" # gpt2 · cl100k · o200k · tekken · deepseek, one run per regex family -harness = false - -[[bench]] -name = "classify" -harness = false - -[[bench]] -name = "class_runs" -harness = false diff --git a/tokenizers/atomsplit/README.md b/tokenizers/atomsplit/README.md deleted file mode 100644 index 0802d27a0..000000000 --- a/tokenizers/atomsplit/README.md +++ /dev/null @@ -1,221 +0,0 @@ -# Tag-classify pretokenization — design spec -The key design principle is to run 1 SIMD classification pass over the input, and then run a finite state machine on the produced tags. -SIMD first for the fsm is not great as unrolled regex can be quite complicated and cut often in many cases. The only SIMD you want in the fsm is when you are looking for `*` or `+` patterns. There, SIMD allows you to go fast to the last byte of the category you are looking for. -For simple pretokenizers like whitespace split that emit splits at tag boundary changes, simd can also be used. - -We always have a scalar fallback for both the classification and the finite state machine. -## 1. The generic classify engine (composable lanes) -The key design principle is that no matter the number of atoms (well it has to be <255) the classifier does not change. The classifier operates on byte length. For each byte length we find a smart way to retrieve the class from pre-computed table. A new custom pre-tokenizer would only require us to update the table, never the classifiers. This scales really well as you can combine tags in the FSM to create bigger categories (like white space markers are whitespace and markers) - -| lane | lead byte range | mechanism | shared SIMD ↔ scalar? | -|---------------|----------------------|------------------------------------------------------------------------|------------------------------------| -| ASCII | `0x00–0x7F` (1 B) | 128-entry table, two 64-halves + subtract-trick `OR` | ✓ one table | -| 2-byte | `0xC2–0xDF` (2 B) | 8 groups × (4 sub × 64); **peel** the min group → one 256-lookup | ✓ one table | -| CJK shortcut | `0xE3–0xED` (3 B) | range compares → one tag (current scheme folds all CJK → `Letter`) | SIMD-only optimism (scalar → 3-byte) | -| 3-byte | `0xE0–0xEF` (3 B) | 512 blocks: 425 uniform const / 87 mixed 128-tables; **peel** blocks | ✓ one table | -| cold BMP | any BMP deferred | run-length `(start_cp, tag)`, binary-searched (~1–3 KB) | scalar reader / SIMD `MB`-fixup | -| astral | `0xF0–0xF4` (4 B) | run-length `(start_cp, tag)`, binary-searched | scalar / SIMD stamps `MB` → fixup | -| continuation | `0x80–0xBF` | tagged `Cont` — transparent to every FSM | ✓ | - -The **13 coarse atoms** the default scheme emits (the coarse class is the tag's **low nibble**, so it fits in a `u4`; the whole engine works for any `< 255`): - -``` -0 Letter 1 NumWord 2 NumOther 3 Newline 4 Space 5 WsOther 6 Mark -7 Connector 8 Punct 9 Apostrophe 10 SymOther 11 NumericOther 12 Control - ┌ internal sentinels, never seen by the FSM ┐ - 13 Sentinel 14 MultiByte 15 Cont -``` - -The tag is a full `u8`: the **high nibble** is an optional *refinement* that sub-splits one coarse class for a pretokenizer needing finer granularity — without a second pass. o200k's case split refines `Letter` into `UPPER` (`\p{Lu}\p{Lt}`) / `LOWER` (`\p{Ll}`) / caseless, and `Mark` carries `ALPHA_SYM` for Other_Alphabetic symbols (`\w` but categorically `\p{S}`, e.g. circled letters). Coarse consumers collapse it for free — `in_mask` and the SIMD class path `& 0x0F` off the nibble — so only the FSM that opted in (o200k) ever sees it. - -### 1.1 The only thing you need to know: how UTF-8 lays out a codepoint - -Every table is indexed straight off the raw UTF-8 bytes, so the whole scheme falls out of the encoding. `x`/`y`/`z` are the payload bits of the codepoint; the `0`/`10`/`110`/`1110` prefixes are UTF-8's length tag: - -``` - ├────────────── codepoint bits ──────────────┤ -1 byte 0xxxxxxx U+0000 .. U+007F (ASCII) -2 byte 110xxxyy 10yyyyyy U+0080 .. U+07FF -3 byte 1110xxxx 10yyyyyy 10zzzzzz U+0800 .. U+FFFF (BMP) -4 byte 11110www 10xxxxxx 10yyyyyy 10zzzzzz U+10000 .. U+10FFFF (astral) - └──┬───┘ └──┬───┘ - lead continuation bytes (all start 10……) -``` - -Read a single byte and its **top bits tell you the lane**: - -``` -byte & 0x80 == 0x00 → 0xxxxxxx ASCII (1-byte) -byte & 0xC0 == 0x80 → 10xxxxxx continuation → Cont -byte & 0xE0 == 0xC0 → 110xxxxx 2-byte lead -byte & 0xF0 == 0xE0 → 1110xxxx 3-byte lead -byte & 0xF8 == 0xF0 → 11110xxx 4-byte lead -``` - -### 1.2 How the tables are built (`bitmap_gen`, one source of truth) - -We never hand-write a table. We **synthesize every byte sequence a lane can hold, decode it back to a codepoint, and ask the reference `atom(char)` for its class** — then store that class at the index the raw bytes produce. Because the *decode* and the *index* both come from the same bit layout above, the runtime lookup is exact by construction. One reference function → every table (and the SIMD kernel) derived from it; the generator then re-derives all 1.1 M codepoints and asserts the packed tables read back identically, so a scheme change that breaks a table fails the build. - -``` -for every byte pattern of the lane ──► cp = decode(bytes) ──► atom(char::from(cp)) ──► table[index(bytes)] = tag - (§1.1 layout) (the ONE reference) (same layout → same index) -``` - -### 1.3 Per-lane mechanics - -**ASCII — one 128-entry table, split in two halves.** NEON's table op (`vqtbl`) resolves ≤ 64 entries, so 128 needs two. The subtract trick makes one 8-bit index hit exactly one half (out-of-range lookups return 0, so the `OR` is clean): - -``` -tag(v) = tbl64(ascii_lo, v) OR tbl64(ascii_hi, v - 64) - └ v<64 → lo[v] └ v≥64 → hi[v-64] - └ v≥64 → 0 (oob) └ v<64 → v-64 wraps huge → 0 (oob) - -v=5 : lo[5] | (5-64 wraps → 0) = lo[5] -v=65 : (65 ≥64→0) | hi[1] = hi[1] -``` - -**2-byte — `110xxxyy 10yyyyyy`, 8 groups × 4 sub-groups × 64.** There are only 32 possible leads (`C0..DF`). Split the 5 lead payload bits into `xxx` (group, 8) and `yy` (sub-group, 4); the continuation's `yyyyyy` is the offset (64): - -``` -lead 110 x x x y y cont 10 y y y y y y - └─┬─┘ └┬┘ └──┬───┘ - xxx yy yyyyyy - group subgroup offset - (b0>>2)&7 b0&3 b1&0x3F - -lookup: group_tables[ xxx ][ yy ][ yyyyyy ] ← declared [8][4][64] -in SIMD: the inner [4][64] is 256 contiguous bytes, so ONE 256-lookup does it, - indexed by group_index = (yy<<6) | yyyyyy : - idx 0..63 → sub 0 idx 128..191 → sub 2 - idx 64..127 → sub 1 idx 192..255 → sub 3 (subtract-window OR of 4×vqtbl4) -``` - -**3-byte — `1110xxxx 10yyyyyy 10zzzzzz`, 512 blocks, most of them a single tag.** A *block* = one lead (`E0..EF`, 16) × the high bits of the 2nd byte (`b1>>1`, 32) = 512 blocks, each covering 128 codepoints. Real scripts are homogeneous: `425 / 512` blocks are **one atom** (all of CJK Han is `Letter`, whole symbol ranges are `SymOther`), stored as a single const byte; only `87` "mixed" blocks (e.g. a script interleaved with its punctuation) need an actual 128-entry table. - -``` -1110 x x x x 10 y y y y y y 10 z z z z z z - └──┬──┘ └┬┘└───┬──┘ └───┬────┘ - lead b1-pair (unused here) within-block offset - E0..EF (16) (b1>>1)&0x1F (32) (b1&1)<<6 | z (128) - -block = (lead-0xE0)*32 + ((b1>>1)&0x1F) ── 512 of them -fast3_uni[block] = tag if the whole 128-cp block is one atom (425 blocks — just a const) - = 0xFF otherwise → fast3_mixed[ fast3_slot[block] ] holds a 128-table (87 blocks) -``` - -**CJK shortcut (SIMD only).** The current scheme folds all of CJK to one tag (`CJK_TAG = Atom::Letter`), so the SIMD kernel skips the tables for `E3..ED` and proves "is CJK" with a few range compares (Han `E4..E9`, Hangul `EB..EC`, kana `E3 81..83`, minus a handful of punctuation holes like `・`). It only ever *under*-claims (boundary/hole codepoints fall through to the exact 3-byte tables), so it stays byte-exact. (A scheme that split CJK into distinct tags would drop this shortcut and use the tables.) - -**Cold fallback + astral.** The dense BMP tag stream and the astral range are **run-length encoded** `(start_cp, tag)` and binary-searched — a few KB instead of a 64 KB dense LUT, because tags change rarely across a codepoint range. The SIMD kernel can't see a 4-byte char's 4th byte (it only gathers 3 bytes per lane), so it stamps those lanes `MB` and a per-chunk fixup resolves them via these tables while the chunk is still hot in L1. - -### 1.4 Why "peel" and not "loop the range" - -For 2-byte and 3-byte a chunk may hold lanes from several blocks. Instead of looping every block id between the min and max present (wasting a step per empty gap), we **peel**: `vminvq` the smallest block still unresolved → resolve exactly its lanes with one lookup → mask them out → repeat. Steps = **distinct blocks present** (usually 1 → that's the fast path), independent of how far apart the scripts sit, and no bounds guard. See `src/simd_classify.rs` for the annotated kernel. - -## 2. The FSM: classify is nearly free, so the FSM *is* the cost - -The single SIMD classify pass is ~free — **~0.05 ns/B on ASCII, ~0.3–0.6 on multibyte** (a handful of table lookups per 16 bytes). The pre-tokenizer's real cost is the **FSM** that turns the tag stream into spans: **~1–3 ns/B on dense Latin/code** (many short tokens ⇒ many state transitions) down to **~0.3 on run-heavy CJK**. So all the performance work lives here. The FSM writes spans into a **caller-owned `&mut [Span]` and returns the count** — no `Vec`, no realloc; the buffer is reused across calls. - -### 2.1 A class is a `u16` bitmap - -The classifier hands the FSM a stream of atom ids (`0..12`). A **class** is just a *set of atoms*, and a set of ≤ 16 atoms is one `u16` where **bit `t` = "atom `t` is in this class"**: - -``` -bit : 12 11 10 9 8 7 6 5 4 3 2 1 0 -atom : Ct Nu Sy Ap Pn Cn Mk Ws Sp Nl No Nw Lt (Ct=Control Nu=NumericOther Sy=SymOther - Ap=Apostrophe Pn=Punct Cn=Connector Mk=Mark - Ws=WsOther Sp=Space Nl=Newline No=NumOther - Nw=NumWord Lt=Letter) -WS : . . . . . . . 1 1 1 . . . = 0x0038 (Newline | Space | WsOther) -WORD : . . . . . 1 1 . . . . 1 1 = 0x00C3 (Letter | NumWord | Mark | Connector) -``` - -Building and testing a class are one instruction each — no branches, no table: - -```rust -Atom::WsOther.bit() // = 1 << 5 build a class by OR-ing atoms -in_mask(tag, WS) // = WS & (1 << tag) != 0 test membership: one AND + compare -``` - -A class is any union of atoms — `PUNCT_SYM = Connector|Punct|Apostrophe|SymOther`, `LETTER_MARK = Letter|Mark` — composed *without touching the classifier or re-running classification*. Add a 14th atom and every FSM keeps working; `u16` covers 16 atoms, and the exact same code with a `u32`/`u64` mask covers 32/64 classes for free (the classifier already emits arbitrary `u8` tags). - -### 2.2 The scalar ↔ SIMD duality — SIMD only earns its keep on run-ends - -A pre-tokenizer rule is one of two shapes, and **SIMD only helps the second**: - -1. **Per-char decisions** — cl100k's contraction peek (`'s|'t|…`), the ` ?` / `[^…]?` prefix logic: a branchy little automaton, one char at a time. Branches don't vectorize, so this stays **scalar**. -2. **Run-ends — a `+`/`*` over one class** (`\p{L}+`, `\s+`, `\p{N}+`, a punct run): grab the maximal run. The trick is **skip the invalid, stop only at the valid**. The regex-shaped FSMs (`fsm_cl100k`, `fsm_o200k`, …) do this with the scalar `run_end`: it unrolls **16 tags per chunk** (one bounds check, `get_unchecked` reads), so a 3-char English word finishes in the ≤16-byte scalar tail while a long CJK run skips whole chunks. The *SIMD* form of the same idea lives only in the class family's `class_runs_neon` (below): `vqtbl1` a 16-tag membership LUT, `vminvq == 0xFF` ⇒ all 16 still in-run ⇒ skip the chunk. - -The **dual** of a run-end is a **boundary**: the class-family pre-tokenizers (Whitespace / Bert / Digits / Punctuation) cut at *every* class change, so `class_runs_neon` scans the other way — classify 16 lanes (`vqtbl1`), fill continuation lanes from the left, `movemask` the class-changes, and iterate the set bits to emit one span per segment. It finds **every boundary in a chunk at once**, so short-run text (English words, ~5 chars) is never paid per-char. - -Both fuse in one kernel: `class_runs_neon` first tries the run-end fast path — *whole 16-chunk stays the current class → skip it, no boundary work* — and only movemasks the **mixed** chunks. So it bulk-skips long runs (Digits/Punct/CJK) **and** parallel-boundaries the dense ones, byte-exact with the portable `emit_class_spans` oracle. Two views of the same idea: **run-end skips the invalid to reach the next valid; boundary-extract flags every valid at once.** - -### 2.3 Why it's a finite-state machine - -Each pre-tokenizer is a small automaton over the shared tag alphabet: a few **states** (in a letter run, inside whitespace, at a boundary) and, per incoming tag, a **transition** — extend the run, emit the span, or start a new one. The class masks *are* the transition predicates (`in_mask(tag, LETTER)` = "stay in the letter state"). So cl100k, GPT-2, whitespace-split, deepseek… are all the *same* machine driven by different masks + a few peeked bytes; only the states that are a maximal `+`/`*` run get the SIMD run-end. - -## 3. Measured performance - -Single-thread, **180 KB per language** (uniform size → comparable cache behaviour), **min-of-7 trials**, 14-core Apple Silicon (8 P + 6 E), light background load. `ns/byte`, lower is better. **Every row is byte-exact ✓** against the reference regex (onig for gpt2 / cl100k / o200k; the composed onig×3 `Sequence` for deepseek). Reproduce with `cargo bench --bench regex` (every pre-tokenizer × the four reference engines — the §3.0 chart); the class family's scalar-vs-SIMD boundary extractor is `--bench class_runs`, classify alone is `--bench classify`. - -### 3.0 At a glance — vs every SOTA splitter - -The `regex` bench pits the full pipeline (SIMD classify + scalar FSM) against **four** reference engines — -**onig** and **pcre2** (with JIT), both C; **fancy-regex**, pure-Rust; and **logos**, a compile-time DFA -lexer-generator — for each pre-tokenizer × language. Speedup = engine ÷ our pipeline; **green = we win big, -red = a close race**: - -![pre-tokenization speedup vs onig / fancy / logos / pcre2-JIT](benches/pretok_heatmap.svg) - -We lead on every pre-tokenizer — **~4–60×** vs onig/fancy and **~1.4–10×** vs the JIT / DFA engines — the -only near-ties being **o200k on CJK** (its case-split FSM is heavy there). `n/a` marks a split the engine -can't express: logos has no look-ahead, so deepseek's 3-regex `Sequence` and the punctuation-isolation -splits (punct / bert) have no single-grammar logos form. The GPT FSMs are byte-exact with their regex -(✓); the class-family reference regex is an approximation of the atom mask (`≈` where it diverges), so -those rows are a speed pairing rather than an equality gate. Regenerate: - -```sh -cargo bench --bench regex > bench.txt -python3 benches/heatmap.py bench.txt benches/pretok_heatmap.svg # needs matplotlib + numpy -``` - -### 3.1 cl100k — classify is ~free, the FSM is the cost - -The regex-shaped FSMs are scalar (there is no SIMD cl100k FSM — SIMD lives only in the class family's `class_runs_neon`, §2.2). `pipeline = SIMD classify + scalar FSM`. - -| lang | b/tok | classify | FSM (scalar) | onig | pipeline **vs onig** | -|---|--:|--:|--:|--:|--:| -| English | 4.6 | **0.068** | 1.037 | 36.3 | **32.9×** | -| French | 5.1 | 0.158 | 0.836 | 35.9 | 36.1× | -| Russian | 10.2 | 0.315 | 0.462 | 20.7 | 26.6× | -| Greek | 9.9 | 0.296 | 0.462 | 22.3 | 29.3× | -| Hebrew | 8.3 | 0.302 | 0.481 | 24.2 | 30.9× | -| Arabic | 9.1 | 0.293 | 0.499 | 23.6 | 29.8× | -| Hindi | 5.3 | 0.320 | 0.659 | 39.9 | 40.7× | -| Thai | 11.8 | 0.327 | 0.416 | 23.0 | 30.9× | -| Chinese | 19.3 | 0.601 | 0.289 | 14.4 | 16.2× | -| Japanese | 25.3 | 0.715 | 0.265 | 14.1 | 14.4× | -| Korean | 7.4 | 0.440 | 0.617 | 26.6 | 25.1× | - -Reading it: **classify ≪ FSM** — on dense Latin the FSM is ~5–15× the classify, so that's where the work is (§2). Pure-ASCII English classifies at **0.068** (the ASCII fast path skips whole chunks); accented Latin (French) and multibyte scripts leave that fast path so classify costs more. The FSM is *cheapest* on run-heavy CJK (Chinese/Japanese ~0.27 — long homogeneous runs skip fast via the unrolled `run_end`) and dearest on dense short-token Latin. The full pipeline (SIMD classify + scalar FSM) is **14–41× onig**. - -### 3.2 GPT-2 (ByteLevel), o200k & deepseek — same story, other regexes - -All byte-exact; full-pipeline speedup vs the reference regex (one `cargo bench --bench regex` run covers all four families): - -- **GPT-2 / ByteLevel**: **21–45×** onig (English 34×, Hindi 45×, Chinese 22×). -- **o200k** (GPT-4o — case-aware `[\p{L}\p{M}]+` split, so a heavier FSM): **7–18×** onig (English 17×, Thai 7×, Chinese 10×). -- **deepseek** (a `Sequence` of 3 Isolated splits, collapsed into one FSM): **17–41×** the composed onig×3 (Chinese 41×, Japanese 28×, English 17×). - -### 3.3 Thread scaling (cl100k, ~16 MB doc, newline-partitioned seams) - -Threads spawned once (`std::thread::scope`, no external dep); each chunk's seam sits after a whitespace-run's last `\n`, so no token crosses it — **byte-exact for cl100k/deepseek** (proven: partitioned spans == sequential). MB/s (bytes×iters / wall; illustrative from a prior run — thread scaling has no standalone bench target in the current layout): - -| threads | English | | Chinese | | -|--:|--:|--:|--:|--:| -| 1 | 474 | 1.00× | 818 | 1.00× | -| 2 | 934 | 1.97× | 1580 | 1.93× | -| 4 | 1883 | 3.98× | 2998 | 3.66× | -| 8 | 3739 | 7.90× (99% linear) | 5698 | 6.96× | -| 14 | 4814 | 10.2× | 7229 | 8.83× | - -**~99% linear through the 8 performance cores**; the drop at 14 is the 6 efficiency cores (slower, so "% linear" falls — not contention). A single un-splittable long document uses the overlap-chunk path instead (BPE/pretok locality — see the merge notes). diff --git a/tokenizers/atomsplit/benches/class_runs.rs b/tokenizers/atomsplit/benches/class_runs.rs deleted file mode 100644 index 447ffae58..000000000 --- a/tokenizers/atomsplit/benches/class_runs.rs +++ /dev/null @@ -1,123 +0,0 @@ -//! Scalar vs SIMD fsm for the class family: `emit_class_spans` (scalar run-end core) vs -//! `class_runs_into` (NEON/SIMD128 movemask boundary-extract + homogeneous-chunk early-out). classify is -//! timed separately. `spd = scalar/simd`: >1 → SIMD wins. (cl100k is scalar-only — see benches/cl100k.rs.) -//! -//! Run: cargo bench --bench class_runs -use atomsplit::classify::{classify, mask}; -use atomsplit::fsm::{Span, class_runs_into, emit_class_spans}; -use std::hint::black_box; -use std::time::Instant; - -const CORPORA: &[(&str, &str)] = &[ - ("English", "../data/big.txt"), - ("French", "benches/data/fr.txt"), - ("Russian", "benches/data/ru.txt"), - ("Greek", "benches/data/el.txt"), - ("Arabic", "benches/data/ar.txt"), - ("Hindi", "benches/data/hi.txt"), - ("Thai", "benches/data/th.txt"), - ("Chinese", "benches/data/zh.txt"), - ("Japanese", "../data/unigram_wagahaiwa_nekodearu.txt"), - ("Korean", "benches/data/ko.txt"), -]; - -// scalar (run-end core) vs SIMD (class_runs_into) wrappers per recipe — #[inline(always)] so the timing -// closure inlines the whole chain (not a fn-pointer, which would skew the scalar path slow). -macro_rules! pair { - ($sname:ident, $vname:ident, $d:expr, $i:expr, $a:expr) => { - #[inline(always)] - fn $sname(t: &[u8], tg: &[u8], o: &mut [Span]) -> usize { - emit_class_spans::<$d, $i, $a>(t, tg, o, 0, 0, 0, None) - } - #[inline(always)] - fn $vname(t: &[u8], tg: &[u8], o: &mut [Span]) -> usize { - class_runs_into::<$d, $i, $a>(t, tg, o) - } - }; -} -pair!(s_wss, v_wss, { mask::WS }, 0, 0); -pair!(s_pun, v_pun, 0, { mask::PUNCT }, 0); -pair!(s_dig, v_dig, 0, 0, { mask::NUMERIC }); -pair!(s_ws, v_ws, { mask::WS }, 0, { mask::WORD }); -pair!(s_bert, v_bert, { mask::WS }, { mask::PUNCT }, 0); - -fn ns_per_byte usize>(len: usize, iters: u32, mut f: F) -> f64 { - for _ in 0..3 { - black_box(f()); - } - let mut best = f64::INFINITY; - for _ in 0..7 { - let t = Instant::now(); - let mut acc = 0usize; - for _ in 0..iters { - acc = acc.wrapping_add(f()); - } - black_box(acc); - best = best.min(t.elapsed().as_nanos() as f64 / (iters as usize * len) as f64); - } - best -} - -fn main() { - let manifest = env!("CARGO_MANIFEST_DIR"); - let mut corpora: Vec<(&str, String)> = Vec::new(); - for (label, rel) in CORPORA { - let Ok(s) = std::fs::read_to_string(format!("{manifest}/{rel}")) else { - continue; - }; - if s.trim().is_empty() { - continue; - } - let mut c = s.len().min(180_000); - while c > 0 && !s.is_char_boundary(c) { - c -= 1; - } - corpora.push((label, s[..c].to_string())); - } - - macro_rules! compare { - ($name:expr, $s:ident, $v:ident) => {{ - println!( - "\n== {} ==\n{:<10} {:>7} {:>8} {:>8} {:>8} | {:>6}", - $name, "lang", "bytes", "classify", "scalar", "simd", "spd" - ); - for (label, corpus) in &corpora { - let text = corpus.as_bytes(); - let n = text.len(); - let iters = (4_000_000 / n).clamp(3, 150) as u32; - let mut tags = vec![0u8; n]; - let mut buf = vec![Span::default(); n + 1]; - classify(text, &mut tags); - // parity: scalar == simd - let (ks, kv) = ($s(text, &tags, &mut buf), 0); - let scalar_out: Vec = buf[..ks].to_vec(); - let _ = kv; - let kv = $v(text, &tags, &mut buf); - let parity = scalar_out == buf[..kv]; - let cls = ns_per_byte(n, iters, || { - classify(text, &mut tags); - n - }); - classify(text, &mut tags); - let sc = ns_per_byte(n, iters, || $s(text, &tags, &mut buf)); - let si = ns_per_byte(n, iters, || $v(text, &tags, &mut buf)); - println!( - "{label:<10} {n:>7} {cls:>8.3} {sc:>8.3} {si:>8.3} | {:>5.2}x{}", - sc / si, - if parity { "" } else { " PARITY✗" } - ); - } - }}; - } - - compare!("WhitespaceSplit", s_wss, v_wss); - compare!("Punctuation", s_pun, v_pun); - compare!("Digits", s_dig, v_dig); - compare!("Whitespace \\w", s_ws, v_ws); - compare!("Bert", s_bert, v_bert); - println!( - "\n(ns/byte, lower better. classify = SIMD classify; scalar = emit_class_spans (run-end\n \ - core); simd = class_runs_into (NEON/SIMD128 movemask + early-out). spd = scalar/simd, >1 → SIMD\n \ - wins. cl100k is scalar-only — its perf is in benches/cl100k.rs.)" - ); -} diff --git a/tokenizers/atomsplit/benches/classify.rs b/tokenizers/atomsplit/benches/classify.rs deleted file mode 100644 index da3287fe0..000000000 --- a/tokenizers/atomsplit/benches/classify.rs +++ /dev/null @@ -1,90 +0,0 @@ -//! classify throughput across three MB-fixup branch regimes, to check the per-chunk -//! `if any(out == MB)` fixup branch is predictable: -//! none — no astral (real text) → branch ALWAYS false → should be free. -//! astral — wall of emoji (every chunk has MB) → branch ALWAYS true → predictable, measures fixup cost. -//! sprinkle — BMP text with an emoji every ~24 B → branch flips irregularly → misprediction stress. -//! Run before/after the fixup change and compare. Run: cargo bench --bench classify -use atomsplit::classify::{classify, classify_scalar}; -use std::hint::black_box; -use std::time::Instant; - -fn ns_per_byte(text: &[u8], tags: &mut [u8]) -> f64 { - let iters = (8_000_000 / text.len().max(1)).clamp(20, 400) as u32; - for _ in 0..3 { - classify(text, tags); - black_box(tags[text.len() / 2]); - } - let mut best = f64::INFINITY; - for _ in 0..9 { - let t = Instant::now(); - for _ in 0..iters { - classify(text, tags); - black_box(tags[text.len() / 2]); - } - best = best.min(t.elapsed().as_nanos() as f64 / (iters as usize * text.len()) as f64); - } - best -} - -fn main() { - let manifest = env!("CARGO_MANIFEST_DIR"); - let english = - std::fs::read_to_string(format!("{manifest}/../data/big.txt")).unwrap_or_default(); - let english: String = english.chars().take(120_000).collect(); - - let astral = "😀🎉🚀🔥🌍🐍".repeat(20_000); // pure 4-byte astral → MB in every chunk - - // BMP English with an emoji inserted every ~24 bytes → some chunks have MB, some don't, irregularly. - let mut sprinkle = String::new(); - let mut acc = 0; - for w in english.split_inclusive(' ') { - sprinkle.push_str(w); - acc += w.len(); - if acc >= 24 { - sprinkle.push('🔥'); - acc = 0; - } - } - - println!( - "{:<10} {:>8} {:>10} {:>10}", - "input", "bytes", "clsSIMD", "clsScalar" - ); - for (label, s) in [ - ("none", &english), - ("astral", &astral), - ("sprinkle", &sprinkle), - ] { - if s.is_empty() { - println!("{label:<10} (empty — big.txt missing?)"); - continue; - } - let text = s.as_bytes(); - let mut tags = vec![0u8; text.len()]; - // byte-exactness guard: SIMD == scalar (exercises the astral fixup path) - let mut sc = vec![0u8; text.len()]; - classify(text, &mut tags); - classify_scalar(text, &mut sc); - let ok = if tags == sc { "✓" } else { "✗" }; - let simd = ns_per_byte(text, &mut tags); - let scal = { - let iters = (8_000_000 / text.len().max(1)).clamp(20, 400) as u32; - let mut best = f64::INFINITY; - for _ in 0..5 { - let t = Instant::now(); - for _ in 0..iters { - classify_scalar(text, &mut sc); - black_box(sc[text.len() / 2]); - } - best = - best.min(t.elapsed().as_nanos() as f64 / (iters as usize * text.len()) as f64); - } - best - }; - println!( - "{label:<10} {:>8} {simd:>10.3} {scal:>10.3} {ok}", - text.len() - ); - } - println!("\n(ns/byte, lower better. clsSIMD is the path with the MB-fixup branch.)"); -} diff --git a/tokenizers/atomsplit/benches/data/.gitignore b/tokenizers/atomsplit/benches/data/.gitignore deleted file mode 100644 index 2211df63d..000000000 --- a/tokenizers/atomsplit/benches/data/.gitignore +++ /dev/null @@ -1 +0,0 @@ -*.txt diff --git a/tokenizers/atomsplit/benches/data/fetch.py b/tokenizers/atomsplit/benches/data/fetch.py deleted file mode 100644 index 15370c927..000000000 --- a/tokenizers/atomsplit/benches/data/fetch.py +++ /dev/null @@ -1,17 +0,0 @@ -#!/usr/bin/env python3 -"""Fetch big real Wikipedia article text per language for the cl100k bench (gitignored output). -English/Japanese use the repo's ../data/big.txt + unigram_wagahaiwa_nekodearu.txt instead.""" -import urllib.request, urllib.parse, json, os -ARTS = [("fr","France"),("ru","Россия"),("el","Ελλάδα"),("he","ישראל"),("ar","مصر"), - ("hi","भारत"),("th","ประเทศไทย"),("zh","数学"),("ko","대한민국")] -here = os.path.dirname(os.path.abspath(__file__)) -for lang, title in ARTS: - url = f"https://{lang}.wikipedia.org/w/api.php?" + urllib.parse.urlencode( - {"action":"query","format":"json","prop":"extracts","explaintext":"1","titles":title}) - try: - req = urllib.request.Request(url, headers={"User-Agent":"atomsplit-bench/0.1"}) - text = next(iter(json.load(urllib.request.urlopen(req, timeout=30))["query"]["pages"].values())).get("extract","") - open(f"{here}/{lang}.txt","w").write(text) - print(f"{lang}: {len(text.encode())} bytes") - except Exception as e: - print(f"{lang}: FAILED {e}") diff --git a/tokenizers/atomsplit/benches/heatmap.py b/tokenizers/atomsplit/benches/heatmap.py deleted file mode 100644 index fb51991b1..000000000 --- a/tokenizers/atomsplit/benches/heatmap.py +++ /dev/null @@ -1,74 +0,0 @@ -#!/usr/bin/env python3 -"""Heatmap of the atomsplit `regex` bench: our classify+fsm pipeline speedup vs each SOTA engine, -per pre-tokenizer × language. Green = we win big, red = close race / behind.""" -import re -import sys -import numpy as np -import matplotlib -matplotlib.use("Agg") -import matplotlib.pyplot as plt -from matplotlib.colors import LogNorm - -LOG = sys.argv[1] -OUT = sys.argv[2] - -ENGINES = ["gpt2", "cl100k", "o200k", "deepseek", "ws_split", "whitespace", "digits", "punct", "bert"] -REFS = ["onig", "fancy", "logos", "pcre2"] -# discover languages in first-seen order -langs = [] -# data[pretok][lang] = [vsOnig, vsFncy, vsLogos, vsPcre2] -data = {e: {} for e in ENGINES} - -for line in open(LOG): - toks = [t for t in line.split() if t != "|"] - if len(toks) < 15 or toks[0] not in data: - continue - pre, lang = toks[0], toks[1] - # ...clsSIMD clsScal fsm onig fancy logos pcre2 | vsOnig vsFncy vsLogos vsPcre2 ok - vs = toks[11:15] - def num(x): - x = x.rstrip("x") - return np.nan if x in ("—", "-", "") else float(x) - data[pre][lang] = [num(v) for v in vs] - if lang not in langs: - langs.append(lang) - -cmap = plt.get_cmap("RdYlGn").copy() -cmap.set_bad("#d9d9d9") -norm = LogNorm(vmin=0.8, vmax=60) - -fig, axes = plt.subplots(3, 3, figsize=(15.5, 15), constrained_layout=True) -for ax, pre in zip(axes.flat, ENGINES): - M = np.array([data[pre].get(l, [np.nan] * 4) for l in langs]) # [langs x refs] - ax.imshow(M, aspect="auto", cmap=cmap, norm=norm) - ax.set_title(pre, fontsize=13, fontweight="bold") - ax.set_xticks(range(len(REFS))) - ax.set_xticklabels(REFS, fontsize=9) - ax.set_yticks(range(len(langs))) - ax.set_yticklabels(langs, fontsize=8) - ax.tick_params(length=0) - for i in range(len(langs)): - for j in range(len(REFS)): - v = M[i, j] - if np.isnan(v): - ax.text(j, i, "n/a", ha="center", va="center", fontsize=7, color="#666") - else: - # white text on the dark-red/dark-green extremes, black in the middle - t = norm(v) - col = "white" if (t < 0.18 or t > 0.86) else "black" - ax.text(j, i, f"{v:.0f}×" if v >= 10 else f"{v:.1f}×", - ha="center", va="center", fontsize=7.5, color=col, fontweight="bold") - -sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm) -cb = fig.colorbar(sm, ax=axes, shrink=0.5, aspect=30, pad=0.02, - ticks=[1, 2, 5, 10, 20, 50]) -cb.ax.set_yticklabels(["1× (tie)", "2×", "5×", "10×", "20×", "50×"]) -cb.set_label("speedup = engine ÷ our pipeline (SIMD classify + scalar fsm) · green = we win big", - fontsize=10) - -fig.suptitle("atomsplit pre-tokenization: our classify+fsm pipeline vs SOTA splitters\n" - "(aarch64 / Apple Silicon, release; onig & pcre2-JIT = C, fancy = pure-Rust regex, " - "logos = compile-time DFA lexer; n/a = engine can't express that split)", - fontsize=13, fontweight="bold") -fig.savefig(OUT, format="svg", bbox_inches="tight") -print(f"wrote {OUT}: {len(langs)} langs × {len(REFS)} engines × {len(ENGINES)} pre-tokenizers") diff --git a/tokenizers/atomsplit/benches/pretok_heatmap.svg b/tokenizers/atomsplit/benches/pretok_heatmap.svg deleted file mode 100644 index e09d0782b..000000000 --- a/tokenizers/atomsplit/benches/pretok_heatmap.svg +++ /dev/null @@ -1,8344 +0,0 @@ - - - - - - - - 2026-07-15T17:38:52.195573 - image/svg+xml - - - Matplotlib v3.10.6, https://matplotlib.org/ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/tokenizers/atomsplit/benches/regex.rs b/tokenizers/atomsplit/benches/regex.rs deleted file mode 100644 index cd9552184..000000000 --- a/tokenizers/atomsplit/benches/regex.rs +++ /dev/null @@ -1,429 +0,0 @@ -//! Pre-tokenization on BIG real text (Wikipedia / big.txt), per language, for every pre-tokenizer at -//! once — the GPT regex FSMs (gpt2 · cl100k · o200k · tekken · deepseek) and the class family (WhitespaceSplit · -//! Whitespace · Digits · Punctuation · Bert). Per (engine, corpus): classify (SIMD vs scalar) + fsm in -//! ns/byte, and the full pipeline (SIMD classify + scalar fsm) vs FOUR reference engines — onig (C), -//! fancy-regex (pure Rust), logos (compile-time DFA lexer), pcre2 (C + JIT). GPT regexes come from -//! [`atomsplit::regexes`] (the canonical specs, byte-exact ✓/✗); class-family references are the -//! best-effort equivalent regex (approximate — see the ≈ marker). Single-regex engines have a 1-element -//! chain; deepseek composes 3 `Isolated` Splits. -//! -//! Data: `../data/big.txt` (English) + `../data/unigram_wagahaiwa_nekodearu.txt` (Japanese) ship with -//! the repo; the rest via `benches/data/fetch.py` (gitignored). Missing files are skipped. -//! -//! Run: cargo bench --bench regex -use atomsplit::classify::{classify, classify_scalar, mask}; -use atomsplit::fsm::{ - Span, class_runs_into, fsm_byte_level, fsm_cl100k, fsm_deepseek, fsm_o200k, fsm_tekken, -}; -use atomsplit::regexes; -use fancy_regex::Regex as Fancy; -use logos::Logos; -use onig::Regex; -use pcre2::bytes::Regex as Pcre2; -use std::hint::black_box; -use std::time::Instant; - -type Fsm = fn(&[u8], &[u8], &mut [Span]) -> usize; - -// logos DFA lexers approximating the GPT splits. logos has no look-ahead (`(?!\S)`) nor case-insensitive -// (`(?i:)`), so token boundaries differ slightly — this is a raw-throughput reference (like fancy), not a -// byte-exact oracle (that's onig). deepseek is a 3-split `Sequence`, not one grammar → no logos number. -#[derive(Logos)] -enum LGpt2 { - #[regex(r"'s|'t|'re|'ve|'m|'ll|'d")] - Contraction, - #[regex(r" ?\p{L}+")] - Word, - #[regex(r" ?\p{N}+")] - Num, - #[regex(r" ?[^\s\p{L}\p{N}]+")] - Other, - #[regex(r"\s+")] - Space, -} - -#[derive(Logos)] -enum LCl100k { - // Contraction outranks Word: `[^\r\n\p{L}\p{N}]?\p{L}+` can also match `'s`, but the real regex tries - // the contraction alternative first. - #[regex(r"'s|'t|'re|'ve|'m|'ll|'d", priority = 5)] - Contraction, - #[regex(r"[^\r\n\p{L}\p{N}]?\p{L}+", priority = 4)] - Word, - #[regex(r"\p{N}\p{N}?\p{N}?")] - Num, - #[regex(r" ?[^\s\p{L}\p{N}]+[\r\n]*", priority = 2)] - Other, - #[regex(r"\s+")] - Space, -} - -#[derive(Logos)] -enum LO200k { - #[regex(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}]+('s|'t|'re|'ve|'m|'ll|'d)?", priority = 6)] - LettersA, - #[regex(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}]*('s|'t|'re|'ve|'m|'ll|'d)?", priority = 5)] - LettersB, - #[regex(r"\p{N}\p{N}?\p{N}?")] - Num, - #[regex(r" ?[^\s\p{L}\p{N}]+[\r\n/]*", priority = 2)] - Other, - #[regex(r"\s+")] - Space, -} - -// tekken = LO200k with no contraction suffix and one token per digit. -#[derive(Logos)] -enum LTekken { - #[regex( - 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}]+", - priority = 6 - )] - LettersA, - #[regex( - 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}]*", - priority = 5 - )] - LettersB, - #[regex(r"\p{N}")] - Num, - #[regex(r" ?[^\s\p{L}\p{N}]+[\r\n/]*", priority = 2)] - Other, - #[regex(r"\s+")] - Space, -} - -// class-family logos grammars. The whitespace-dropping ones use `skip` so logos never emits ws tokens. -#[derive(Logos)] -#[logos(skip r"\s+")] -enum LWsSplit { - #[regex(r"\S+")] - Run, -} -#[derive(Logos)] -#[logos(skip r"\s+")] -enum LWhitespace { - #[regex(r"\w+")] - Word, - #[regex(r"[^\w\s]+")] - Sym, -} -#[derive(Logos)] -enum LDigits { - #[regex(r"\p{N}+")] - Num, - #[regex(r"[^\p{N}]+")] - Other, -} - -fn lex_count<'s, T: Logos<'s, Source = str>>(s: &'s str) -> usize -where - T::Extras: Default, -{ - let mut lex = T::lexer(s); - let mut n = 0; - while lex.next().is_some() { - n += 1; - } - n -} - -/// logos throughput (ns/byte) for the engines it can express; `None` for the ones it can't (deepseek's -/// multi-split, and punct/bert whose POSIX-class isolation logos doesn't parse reliably). -fn logos_ns(ename: &str, s: &str, len: usize, iters: u32) -> Option { - let f: fn(&str) -> usize = match ename { - "gpt2" => |s| lex_count::(s), - "cl100k" => |s| lex_count::(s), - "o200k" => |s| lex_count::(s), - "tekken" => |s| lex_count::(s), - "ws_split" => |s| lex_count::(s), - "whitespace" => |s| lex_count::(s), - "digits" => |s| lex_count::(s), - _ => return None, - }; - Some(ns_per_byte(len, iters, || f(s))) -} - -// ── class-family pre-tokenizers: `class_runs_into` recipes + their equivalent reference regex ── -// Each is a monomorphized `Fsm` wrapper (const-generic masks can't go through a fn pointer directly). -fn f_ws_split(t: &[u8], g: &[u8], o: &mut [Span]) -> usize { - class_runs_into::<{ mask::WS }, 0, 0>(t, g, o) -} -fn f_punct(t: &[u8], g: &[u8], o: &mut [Span]) -> usize { - class_runs_into::<0, { mask::PUNCT }, 0>(t, g, o) -} -fn f_digits(t: &[u8], g: &[u8], o: &mut [Span]) -> usize { - class_runs_into::<0, 0, { mask::NUMERIC }>(t, g, o) -} -fn f_whitespace(t: &[u8], g: &[u8], o: &mut [Span]) -> usize { - class_runs_into::<{ mask::WS }, 0, { mask::WORD }>(t, g, o) -} -fn f_bert(t: &[u8], g: &[u8], o: &mut [Span]) -> usize { - class_runs_into::<{ mask::WS }, { mask::PUNCT }, 0>(t, g, o) -} - -// Class-family reference regexes — the exact split each recipe produces under a real regex engine. -// The char-classes map 1:1 to the atomsplit masks: WS=`\s`, WORD=`\w`, NUMERIC=`\p{N}`, -// PUNCT=`[[:punct:]\p{P}]` (ASCII-punctuation ∪ Unicode-P). Isolate = the punct alternative matches a -// single char; the run alternatives use `+`. -const RX_WS_SPLIT: &str = r"\S+"; -const RX_WHITESPACE: &str = r"\w+|[^\w\s]+"; -const RX_DIGITS: &str = r"\p{N}+|[^\p{N}]+"; -const RX_PUNCT: &str = r"[[:punct:]\p{P}]|[^[:punct:]\p{P}]+"; -const RX_BERT: &str = r"[[:punct:]\p{P}]|[^\s[:punct:]\p{P}]+"; - -// (name, native fsm, reference regex chain, keep_gaps, exact). The chain is applied `Isolated`, each -// regex splitting the previous pieces. `keep_gaps=false` for the whitespace-dropping recipes: their -// output is matches only (dropped whitespace is NOT a token), so the reference drops the gaps too. -// `exact`: the GPT FSMs are byte-exact with their regex (parity-tested) → ✓/✗ is a real gate. The -// class-family masks (WORD incl. Other_Alphabetic, PUNCT = P ∪ ASCII-symbols, …) can't be written as a -// regex class, so their reference is an *approximation* — a mismatch is `≈` (still a valid speed pairing), -// not a failure. -const ENGINES: &[(&str, Fsm, &[&str], bool, bool)] = &[ - ("gpt2", fsm_byte_level as Fsm, &[regexes::GPT2], true, true), - ("cl100k", fsm_cl100k as Fsm, &[regexes::CL100K], true, true), - ("o200k", fsm_o200k as Fsm, &[regexes::O200K], true, true), - ("tekken", fsm_tekken as Fsm, &[regexes::TEKKEN], true, true), - ( - "deepseek", - fsm_deepseek as Fsm, - regexes::DEEPSEEK, - true, - true, - ), - ("ws_split", f_ws_split as Fsm, &[RX_WS_SPLIT], false, false), - ( - "whitespace", - f_whitespace as Fsm, - &[RX_WHITESPACE], - false, - false, - ), - ("digits", f_digits as Fsm, &[RX_DIGITS], true, false), - ("punct", f_punct as Fsm, &[RX_PUNCT], true, false), - ("bert", f_bert as Fsm, &[RX_BERT], false, false), -]; - -const CORPORA: &[(&str, &str)] = &[ - ("English", "../data/big.txt"), - ("French", "benches/data/fr.txt"), - ("Russian", "benches/data/ru.txt"), - ("Greek", "benches/data/el.txt"), - ("Hebrew", "benches/data/he.txt"), - ("Arabic", "benches/data/ar.txt"), - ("Hindi", "benches/data/hi.txt"), - ("Thai", "benches/data/th.txt"), - ("Chinese", "benches/data/zh.txt"), - ("Japanese", "../data/unigram_wagahaiwa_nekodearu.txt"), - ("Korean", "benches/data/ko.txt"), -]; - -/// The composed split under `res` — each regex splits the previous pieces. `keep_gaps` keeps the -/// between/after-match gaps as pieces (Isolated split; deepseek); `false` drops them (the whitespace a -/// drop-ws recipe discards). Byte-offset pieces into `text`. -fn onig_pieces(res: &[Regex], text: &str, keep_gaps: bool) -> Vec<(usize, usize)> { - let mut pieces = vec![(0usize, text.len())]; - for re in res { - let mut next = Vec::with_capacity(pieces.len() * 2); - for (s, e) in pieces.drain(..) { - let sub = &text[s..e]; - let mut prev = 0usize; - for (ms, me) in re.find_iter(sub) { - if keep_gaps && ms > prev { - next.push((s + prev, s + ms)); - } - next.push((s + ms, s + me)); - prev = me; - } - if keep_gaps && prev < sub.len() { - next.push((s + prev, e)); - } - } - pieces = next; - } - pieces -} - -/// Same composition under PCRE2 (JIT-compiled), operating on bytes; `find_iter` yields `Result`. -fn pcre2_pieces(res: &[Pcre2], text: &str, keep_gaps: bool) -> Vec<(usize, usize)> { - let bytes = text.as_bytes(); - let mut pieces = vec![(0usize, text.len())]; - for re in res { - let mut next = Vec::with_capacity(pieces.len() * 2); - for (s, e) in pieces.drain(..) { - let sub = &bytes[s..e]; - let mut prev = 0usize; - for m in re.find_iter(sub) { - let Ok(m) = m else { break }; - let (ms, me) = (m.start(), m.end()); - if keep_gaps && ms > prev { - next.push((s + prev, s + ms)); - } - next.push((s + ms, s + me)); - prev = me; - } - if keep_gaps && prev < sub.len() { - next.push((s + prev, e)); - } - } - pieces = next; - } - pieces -} - -/// Same composition under fancy-regex; `find_iter` yields `Result` (a match error ends the pass). -fn fancy_pieces(res: &[Fancy], text: &str, keep_gaps: bool) -> Vec<(usize, usize)> { - let mut pieces = vec![(0usize, text.len())]; - for re in res { - let mut next = Vec::with_capacity(pieces.len() * 2); - for (s, e) in pieces.drain(..) { - let sub = &text[s..e]; - let mut prev = 0usize; - for m in re.find_iter(sub) { - let Ok(m) = m else { break }; - let (ms, me) = (m.start(), m.end()); - if keep_gaps && ms > prev { - next.push((s + prev, s + ms)); - } - next.push((s + ms, s + me)); - prev = me; - } - if keep_gaps && prev < sub.len() { - next.push((s + prev, e)); - } - } - pieces = next; - } - pieces -} - -// MIN over TRIALS timed loops — the fastest trial had the least CPU contention, so it's the truest -// estimate and robust to thermal throttling / background load (which only ever make a trial slower). -fn ns_per_byte usize>(len: usize, iters: u32, mut f: F) -> f64 { - const TRIALS: u32 = 7; - for _ in 0..3 { - black_box(f()); // warm - } - let mut best = f64::INFINITY; - for _ in 0..TRIALS { - let t = Instant::now(); - let mut acc = 0usize; - for _ in 0..iters { - acc = acc.wrapping_add(f()); - } - black_box(acc); - best = best.min(t.elapsed().as_nanos() as f64 / (iters as usize * len) as f64); - } - best -} - -fn main() { - let manifest = env!("CARGO_MANIFEST_DIR"); - println!( - "{:<9} {:<10} {:>7} {:>5} {:>8} {:>8} | {:>8} | {:>8} {:>8} {:>8} {:>8} | {:>7} {:>7} {:>7} {:>7}", - "engine", - "lang", - "bytes", - "b/tok", - "clsSIMD", - "clsScal", - "fsm", - "onig", - "fancy", - "logos", - "pcre2jit", - "vsOnig", - "vsFncy", - "vsLogos", - "vsPcre2" - ); - for &(ename, fsm, rxs, keep_gaps, exact) in ENGINES { - let onig: Vec = rxs.iter().map(|r| Regex::new(r).expect(ename)).collect(); - let fancy: Vec = rxs.iter().map(|r| Fancy::new(r).expect(ename)).collect(); - // PCRE2 with JIT — its speed is the JIT, so bench it there. - let pcre2: Vec = rxs - .iter() - .map(|r| { - pcre2::bytes::RegexBuilder::new() - .utf(true) - .ucp(true) - .jit_if_available(true) - .build(r) - .expect(ename) - }) - .collect(); - for (label, rel) in CORPORA { - let raw = match std::fs::read_to_string(format!("{manifest}/{rel}")) { - Ok(s) if !s.trim().is_empty() => s, - _ => { - println!( - "{ename:<9} {label:<10} (skipped — {rel} missing; run benches/data/fetch.py)" - ); - continue; - } - }; - // UNIFORM cap (char boundary): every language the same byte size → equal cache behaviour, so - // per-byte compute compares across scripts. 180 KB > L1, and all corpora have ≥180 KB. - let mut c = raw.len().min(180_000); - while c > 0 && !raw.is_char_boundary(c) { - c -= 1; - } - let corpus = &raw[..c]; - let text = corpus.as_bytes(); - let n = text.len(); - let iters = (4_000_000 / n).clamp(3, 150) as u32; - - let ref_spans: Vec = onig_pieces(&onig, corpus, keep_gaps) - .iter() - .map(|&(s, e)| Span::new(s as u32, e as u32)) - .collect(); - let mut tags = vec![0u8; n]; - let mut tsc = vec![0u8; n]; - classify(text, &mut tags); - let mut buf = vec![Span::default(); n + 1]; - let k = fsm(text, &tags, &mut buf); - let ok = if buf[..k] == ref_spans[..] { - "✓" - } else if exact { - "✗" // a real regression: the GPT fsm must be byte-exact with its regex - } else { - "≈" // class-family: regex is an approximation of the mask, divergence expected - }; - let btok = n as f64 / k.max(1) as f64; - - let cls_simd = ns_per_byte(n, iters, || { - classify(text, &mut tags); - tags[n / 2] as usize - }); - let cls_scal = ns_per_byte(n, iters, || { - classify_scalar(text, &mut tsc); - tsc[n / 2] as usize - }); - classify(text, &mut tags); - let fsm_ns = ns_per_byte(n, iters, || fsm(text, &tags, &mut buf)); - let onig_ns = ns_per_byte(n, iters, || onig_pieces(&onig, corpus, keep_gaps).len()); - let fancy_ns = ns_per_byte(n, iters, || fancy_pieces(&fancy, corpus, keep_gaps).len()); - let pcre2_ns = ns_per_byte(n, iters, || pcre2_pieces(&pcre2, corpus, keep_gaps).len()); - let logos = logos_ns(ename, corpus, n, iters); - - let pipe = cls_simd + fsm_ns; // SIMD classify + scalar fsm — the full pipeline - let (logos_c, vslogos) = match logos { - Some(l) => (format!("{l:8.2}"), format!("{:6.1}x", l / pipe)), - None => (" —".into(), " —".into()), // deepseek: no single logos grammar - }; - println!( - "{ename:<9} {label:<10} {n:>7} {btok:>5.1} {cls_simd:>8.3} {cls_scal:>8.3} | {fsm_ns:>8.3} | {onig_ns:>8.2} {fancy_ns:>8.2} {logos_c} {pcre2_ns:>8.2} | {:>6.1}x {:>6.1}x {vslogos} {:>6.1}x {ok}", - onig_ns / pipe, - fancy_ns / pipe, - pcre2_ns / pipe - ); - } - } - println!( - "\n(ns/byte, lower better. pipeline = SIMD classify + scalar fsm; vs onig / fancy / logos / \ - pcre2(JIT). GPT FSMs (gpt2/cl100k/o200k/tekken/deepseek): the regex IS the spec, ✓ = byte-exact, \ - ✗ = regression. Class family (ws_split/whitespace/digits/punct/bert): the mask can't be written \ - as a regex class, so the reference is approximate — ✓ where it happens to match, ≈ where it \ - diverges (speed still comparable). logos approximates the grammar (deepseek/punct/bert: n/a).)" - ); -} diff --git a/tokenizers/atomsplit/src/fsm.rs b/tokenizers/atomsplit/src/fsm.rs deleted file mode 100644 index 45c8ca286..000000000 --- a/tokenizers/atomsplit/src/fsm.rs +++ /dev/null @@ -1,387 +0,0 @@ -//! FSM layer: turn the `Atom` tag stream (from [`crate::classify`]) into token spans. Every fsm is -//! NO-PUSH — it writes spans into a caller-preallocated `&mut [Span]` (len ≥ `text.len()`) and returns -//! the token count; no `Vec`, no realloc. Inputs must be well-formed UTF-8 (see the crate-level docs). -//! -//! The class family (WhitespaceSplit / Punctuation / Digits / Whitespace / Bert) goes through -//! [`class_runs_into`]: on aarch64/wasm the SIMD movemask boundary-extractor + homogeneous-chunk -//! early-out (in `simd_fsm`), elsewhere the scalar run-end core ([`emit_class_spans`]). The -//! regex-shaped ones ([`fsm_cl100k`] / [`fsm_o200k`] / [`fsm_tekken`] / [`fsm_deepseek`] / -//! [`fsm_byte_level`]) are scalar jump-tables (only the class family's [`class_runs_into`] has a SIMD -//! path). - -pub(crate) use crate::classify::{Atom, char_len, classify, in_mask, mask}; -// Atom-tag aliases, shared with the per-tokenizer FSM submodules (`fsm/*.rs`) via `use super::*`. -pub(crate) const LET: u8 = Atom::Letter as u8; -pub(crate) const NW: u8 = Atom::NumWord as u8; -pub(crate) const NO: u8 = Atom::NumOther as u8; -pub(crate) const NLN: u8 = Atom::Newline as u8; -pub(crate) const SPC: u8 = Atom::Space as u8; -pub(crate) const WSO: u8 = Atom::WsOther as u8; -pub(crate) const MRK: u8 = Atom::Mark as u8; -pub(crate) const CON: u8 = Atom::Connector as u8; -pub(crate) const PUN: u8 = Atom::Punct as u8; -pub(crate) const APO: u8 = Atom::Apostrophe as u8; -pub(crate) const SYM: u8 = Atom::SymOther as u8; -pub(crate) const NMO: u8 = Atom::NumericOther as u8; -pub(crate) const CTL: u8 = Atom::Control as u8; -pub(crate) const CONT: u8 = Atom::Cont as u8; -pub(crate) const ASM: u8 = Atom::AlphaSymMark as u8; -pub(crate) const ZWJ: u8 = Atom::Zwj as u8; // 0x26 — ZWJ/ZWNJ, tagged in classify so FSMs skip the text peek - -/// Advance over a maximal `m`-membership run (m is a mask); returns the byte index past it. -/// `inline(always)`: it's called once per token (~200K/MB on English) — a real call here doubles fsm cost. -/// -/// logos-style "fast loop": process 16 tags per iteration with ONE bounds check per chunk and unchecked -/// reads (the loop condition proves `i + 16 <= end`), so short runs pay ~1 bounds check instead of one -/// per byte and long runs stay a tight unrolled scan. Byte-identical to the plain `while in_mask` scan. -#[inline(always)] -pub(crate) fn run_end(tags: &[u8], mut i: usize, end: usize, mut m: u16) -> usize { - m |= Atom::Cont.bit(); - debug_assert!(end <= tags.len()); - // SAFETY: `i + 16 <= end <= tags.len()` in the unrolled body, so every `get_unchecked(i + k)` - // (k < 16) is in bounds. The tail is the plain checked scan. - while i + 16 <= end { - for k in 0..16 { - if !in_mask(unsafe { *tags.get_unchecked(i + k) }, m) { - return i + k; - } - } - i += 16; - } - while i < end && in_mask(tags[i], m) { - i += 1; - } - i -} - -/// End of a whitespace token starting at `i`, for the tail shared byte-for-byte by cl100k and o200k -/// (`\s*[\r\n]+ | \s+(?!\S) | \s+`): through the last `\r\n` if any (rule 5), else the whole run at -/// EOF (rule 7), else give the final ws char back to the next token (rule 6). `#[inline]` — hot, -/// called once per whitespace token. (deepseek's tail differs — it also stops before a digit/CJK — and -/// byte_level has no `[\r\n]` rule; both keep their own.) -#[inline] -pub(crate) fn ws_tail(text: &[u8], tags: &[u8], i: usize, end: usize) -> usize { - let re = run_end(tags, i, end, mask::WS); - if let Some(r) = text[i..re].iter().rposition(|&x| x == 0x0A || x == 0x0D) { - i + r + 1 - } else if re == end { - re - } else { - let mut last = re - 1; - while last > i && text[last] & 0xC0 == 0x80 { - last -= 1; - } - if last > i { last } else { re } - } -} - -/// Case-insensitive contraction match at `i` (`'s 't 're 've 'm 'll 'd`), shared by cl100k and o200k: -/// byte length (2 or 3) or 0 if none. Self-guarding — `text[i]` need not be an apostrophe. `#[inline]`. -/// (byte_level's contraction is case-SENSITIVE, so it keeps its own.) -#[inline] -pub(crate) fn contraction(text: &[u8], i: usize) -> usize { - let end = text.len(); - if i >= end || text[i] != 0x27 || i + 1 >= end || text[i + 1] >= 0x80 { - return 0; - } - let lc = text[i + 1] | 0x20; - match lc { - b's' | b't' | b'm' | b'd' => 2, - b'r' | b'v' | b'l' if i + 2 < end && text[i + 2] < 0x80 => { - let l2 = text[i + 2] | 0x20; - usize::from((matches!(lc, b'r' | b'v') && l2 == b'e') || (lc == b'l' && l2 == b'l')) * 3 - } - _ => 0, - } -} - -/// A token span: byte offsets `[start, end)` into the input. `#[repr(C)]` so the FSM output buffer has a -/// stable `[start, end]` layout — the pipeline reuses it with zero conversion, and it can be reinterpreted -/// as bytes / handed across the crate boundary. -#[repr(C)] -#[derive(Clone, Copy, PartialEq, Eq, Debug, Default, Hash, PartialOrd, Ord)] -pub struct Span { - pub start: u32, - pub end: u32, -} - -impl Span { - #[inline] - pub const fn new(start: u32, end: u32) -> Self { - Self { start, end } - } - - /// `[start, end)` as a `usize` range — for slicing the input text. - #[inline] - pub fn range(self) -> core::ops::Range { - self.start as usize..self.end as usize - } -} - -/// Compare against a bare `(start, end)` tuple — convenience for tests/interop. -impl PartialEq<(u32, u32)> for Span { - #[inline] - fn eq(&self, o: &(u32, u32)) -> bool { - self.start == o.0 && self.end == o.1 - } -} - -/// No-`push` class-family pre-tokenizer core: writes spans into the preallocated `out` slice and returns -/// the count. ONE shape covers the whole class family via ``: -/// WhitespaceSplit `<{WS},0,0>` · Punctuation `<0,{PUNCT},0>` · Digits `<0,0,{NUMERIC}>` · -/// Whitespace `<{WS},0,{WORD}>` · Bert `<{WS},{PUNCT},0>`. -/// Class of a char: `DROP`→dropped, `ISOLATE`→own token, `KEEP_A`→run "A", else→run "B" (A/B cut apart). -/// TODO: find a better explanation -#[inline] -#[must_use] -pub fn class_runs_into( - text: &[u8], - tags: &[u8], - out: &mut [Span], -) -> usize { - debug_assert!(out.len() >= text.len() && tags.len() >= text.len()); - #[cfg(target_arch = "aarch64")] - { - crate::simd_fsm::class_runs_neon::(text, tags, out) - } - #[cfg(all(target_arch = "wasm32", target_feature = "simd128"))] - { - crate::simd_fsm::class_runs_wasm::(text, tags, out) - } - #[cfg(not(any( - target_arch = "aarch64", - all(target_arch = "wasm32", target_feature = "simd128") - )))] - { - emit_class_spans::(text, tags, out, 0, 0, 0, None) - } -} - -/// This is the most important function as it's the core of the scalar finite state machine. -/// It allows to emit class spans with different behaviours for tags we want to drop, tags we want -/// to isolate and tags we want to keep. Any other tags are assumed to be keept. -/// -/// This function is used as a fallback to the SIMD fast fsm. It is used for most pre tokenizers -/// but the unrolled regex, which have more complex variations that cannot be expressed with drop, -/// isolate, keep. These 3 generic parameters are u16 bitmap masks over the 16 classes we have and -/// define the behaviour. They are usally one of the [`crate::classify::mask`]. They allow dropping -/// words, isolating whitespace and keeping new line for example. -#[must_use] -#[inline] -pub fn emit_class_spans( - text: &[u8], - tags: &[u8], - out: &mut [Span], - mut write_index: usize, // in the out slice - mut text_pointer: usize, // in the text slice - segment_start: usize, // previous segment_start - segment_class: Option, // previous segment's class -) -> usize { - debug_assert!(out.len() >= text.len() && tags.len() >= text.len()); - let n = text.len(); - // Tie `tags.len() == text.len() == n` so the optimizer drops the interior `tags[i]` / `text[i]` - // bounds checks (callers guarantee len ≥ n; same trick as `cl100k`). Per-byte scanning already - // avoids checks via `run_end`'s unrolled `get_unchecked`; this covers the per-token accesses. - let tags = &tags[..n]; - let text = &text[..n]; - let other = !(DROP | ISOLATE | KEEP_A); // None of the above correspond to a continuation - if let Some(segment_class) = segment_class { - // this will usually be at the tail of a SIMD call. - text_pointer = run_end(tags, text_pointer, n, segment_class); // skip the whole drop run at once - if segment_class != DROP { - out[write_index] = Span { - start: segment_start as u32, - end: text_pointer as u32, - }; - if text_pointer == n { - return write_index + 1; - } - write_index += 1; - } - } - while text_pointer < n { - let t = tags[text_pointer]; - if t == Atom::Cont as u8 { - text_pointer += 1; - continue; - } - // classify the first char. - if in_mask(t, DROP) { - text_pointer = run_end(tags, text_pointer, n, DROP); // skip the whole drop run at once - } else if in_mask(t, ISOLATE) { - let s = text_pointer; - text_pointer += char_len(text[text_pointer]); - out[write_index] = Span { - start: s as u32, - end: text_pointer as u32, - }; // isolate: one char = one token - write_index += 1; - } else { - let s = text_pointer; - text_pointer = if in_mask(t, KEEP_A) { - run_end(tags, text_pointer, n, KEEP_A) - } else { - run_end(tags, text_pointer, n, other) - }; - out[write_index] = Span { - start: s as u32, - end: text_pointer as u32, - }; - write_index += 1; - } - } - write_index -} - -// ── per-tokenizer unrolled FSMs (one file each; shared helpers above via `use super::*`) ── -mod byte_level; -mod cl100k; -mod deepseek; -mod o200k; -pub use byte_level::fsm_byte_level; -pub use cl100k::{fsm_cl100k, fsm_cl100k_cap}; -pub use deepseek::fsm_deepseek; -pub use o200k::{fsm_o200k, fsm_tekken}; - -// ── Composition recipes ──────────────────────────────────────────────────────────────────────── -// Each pre-tokenizer = (classify → fsm shape + params). `tags` and `out` are caller-owned -// scratch, reused across calls — NO per-call alloc, NO push. The class family writes spans into the -// preallocated `out: &mut [Span]` (len ≥ text.len()) via `class_runs_into` and returns the token count. -// In `tk-encode` these delegate from the `pipeline::PreTokenizer` impls (offset conversion happens there). - -/// `WhitespaceSplit` — split on Unicode whitespace and drop it; keeps maximal non-whitespace runs. -pub struct WhitespaceSplit; -impl WhitespaceSplit { - #[inline] - #[must_use] - pub fn pre_tokenize(&self, text: &[u8], tags: &mut [u8], out: &mut [Span]) -> usize { - classify(text, tags); - class_runs_into::<{ mask::WS }, 0, 0>(text, tags, out) - } -} - -/// `Punctuation` — isolate each punctuation char as its own token; non-punct grouped into runs. -pub struct Punctuation; -impl Punctuation { - #[inline] - #[must_use] - pub fn pre_tokenize(&self, text: &[u8], tags: &mut [u8], out: &mut [Span]) -> usize { - classify(text, tags); - class_runs_into::<0, { mask::PUNCT }, 0>(text, tags, out) - } -} - -/// `Digits` — cut numeric runs apart from non-numeric runs (contiguous), keeping both. -pub struct Digits; -impl Digits { - #[inline] - #[must_use] - pub fn pre_tokenize(&self, text: &[u8], tags: &mut [u8], out: &mut [Span]) -> usize { - classify(text, tags); - class_runs_into::<0, 0, { mask::NUMERIC }>(text, tags, out) - } -} - -/// `Whitespace` — the `\w+|[^\w\s]+` pre-tokenizer: drop whitespace, cut word runs from symbol runs. -pub struct Whitespace; -impl Whitespace { - #[inline] - #[must_use] - pub fn pre_tokenize(&self, text: &[u8], tags: &mut [u8], out: &mut [Span]) -> usize { - classify(text, tags); - class_runs_into::<{ mask::WS }, 0, { mask::WORD }>(text, tags, out) - } -} - -/// `Bert` — the BERT basic pre-tokenizer: drop whitespace, isolate punctuation, keep the rest as runs. -pub struct Bert; -impl Bert { - #[inline] - #[must_use] - pub fn pre_tokenize(&self, text: &[u8], tags: &mut [u8], out: &mut [Span]) -> usize { - classify(text, tags); - class_runs_into::<{ mask::WS }, { mask::PUNCT }, 0>(text, tags, out) - } -} - -/// `Cl100k` — the tiktoken cl100k_base / Llama-3 pre-tokenizer (7-rule regex). (o200k is a distinct, -/// case-aware FSM — [`fsm_o200k`].) -pub struct Cl100k; -impl Cl100k { - /// Uses the scalar run-end core: cl100k's letter/ws runs are short on Latin/code (the common case), - /// where a SIMD run-end's setup would lose; only long CJK runs would benefit. - #[inline] - #[must_use] - pub fn pre_tokenize(&self, text: &[u8], tags: &mut [u8], out: &mut [Span]) -> usize { - classify(text, tags); - fsm_cl100k(text, tags, out) - } -} - -/// `DeepSeek` — the DeepSeek-V3/R1 pre-tokenizer (digits{1,3} → CJK-range → big regex, composed). -pub struct DeepSeek; -impl DeepSeek { - #[inline] - #[must_use] - pub fn pre_tokenize(&self, text: &[u8], tags: &mut [u8], out: &mut [Span]) -> usize { - classify(text, tags); - fsm_deepseek(text, tags, out) - } -} - -/// `ByteLevel` — the GPT-2 / Llama / Qwen byte-level pre-tokenizer regex (before byte-mapping). -pub struct ByteLevel; -impl ByteLevel { - #[inline] - #[must_use] - pub fn pre_tokenize(&self, text: &[u8], tags: &mut [u8], out: &mut [Span]) -> usize { - classify(text, tags); - fsm_byte_level(text, tags, out) - } -} - -/// `Split(char, Removed)` — the only pre-tokenizer that keys on a *literal char* rather than an atom -/// class, so it scans bytes directly (no classify pass). UTF-8 is self-synchronizing, so the -/// delimiter's byte pattern only matches on char boundaries. -pub struct CharDelimiterSplit(pub char); -impl CharDelimiterSplit { - /// Split on the literal char (Removed); writes spans into `out` (len ≥ `text.len()`), returns count. - #[inline] - #[must_use] - pub fn pre_tokenize(&self, text: &[u8], _tags: &mut [u8], out: &mut [Span]) -> usize { - debug_assert!(out.len() >= text.len()); - let mut buf = [0u8; 4]; - let delim = self.0.encode_utf8(&mut buf).as_bytes(); - let (n, dl) = (text.len(), delim.len()); - let (mut start, mut i, mut w) = (0usize, 0usize, 0usize); - while i + dl <= n { - // memchr the first delimiter byte, then confirm the full pattern. memchr (already a - // workspace dep) beats a scalar scan 1.4–23× here — the gap widening as the delimiter - // gets rarer over large inputs, since its SIMD skips whole 16/32/64-byte strides. - match memchr::memchr(delim[0], &text[i..n - dl + 1]) { - Some(off) if text[i + off..i + off + dl] == *delim => { - let m = i + off; - if m > start { - out[w] = Span { - start: start as u32, - end: m as u32, - }; // gap before the delimiter (Removed) - w += 1; - } - i = m + dl; - start = i; - } - Some(off) => i += off + 1, // first byte matched mid-pattern; keep scanning - None => break, - } - } - if start < n { - out[w] = Span { - start: start as u32, - end: n as u32, - }; - w += 1; - } - w - } -} diff --git a/tokenizers/atomsplit/src/fsm/byte_level.rs b/tokenizers/atomsplit/src/fsm/byte_level.rs deleted file mode 100644 index 73f84ac9f..000000000 --- a/tokenizers/atomsplit/src/fsm/byte_level.rs +++ /dev/null @@ -1,81 +0,0 @@ -use super::*; - -/// GPT-2 / ByteLevel pretokenization. Regex (same jump-table shape as cl100k, with 3 differences): -/// `'s|'t|'re|'ve|'m|'ll|'d | ?\p{L}+ | ?\p{N}+ | ?[^\s\p{L}\p{N}]+ | \s+(?!\S) | \s+` -/// vs cl100k: (1) contractions are case-SENSITIVE (lowercase only, no `(?i:)`); (2) the ` ?` prefix is -/// a literal SPACE only (not any non-l/n char) and it applies to letters, numbers AND "other"; (3) no -/// `\p{N}{1,3}` cap (numbers are unbounded) and no `\s*[\r\n]`/trailing-`[\r\n]*` rules. -#[must_use] -pub fn fsm_byte_level(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { - debug_assert!(out.len() >= text.len() && tags.len() >= text.len()); - let end = text.len(); - // Tie `tags.len() == end` so the optimizer drops the per-byte bounds check on every interior - // `tags[i]` in this fsm + its `run_end`/`letter_*` scans. (Callers guarantee `tags.len() >= end`.) - let tags = &tags[..end]; - // `\s+(?!\S)|\s+`: the whole run at EOF, else leave the last ws char for the next ` ?`-prefixed run. - let ws = |i: usize| -> usize { - let re = run_end(tags, i, end, mask::WS); - if re == end { - re - } else { - let mut last = re - 1; - while last > i && text[last] & 0xC0 == 0x80 { - last -= 1; - } - if last > i { last } else { re } - } - }; - - let mut i = 0; - let mut w = 0usize; - while i < end { - let start = i; - match tags[i] & 0x0F { - LET => i = run_end(tags, i, end, mask::LETTER), // ` ?\p{L}+` (space taken by the Space arm) - NW | NO => i = run_end(tags, i, end, mask::NUMBER), // ` ?\p{N}+` — UNBOUNDED - MRK | CON | PUN | SYM | NMO | CTL => i = run_end(tags, i, end, mask::NOT_WS_L_N), // ` ?[^…]+` - // `'s|'t|'re|'ve|'m|'ll|'d` (case-sensitive), else `[^\s\p{L}\p{N}]+` (apostrophe ∈ that set) - APO => { - let adv = match (text.get(i + 1), text.get(i + 2)) { - (Some(b's' | b't' | b'm' | b'd'), _) => 2, - (Some(b'r'), Some(b'e')) - | (Some(b'v'), Some(b'e')) - | (Some(b'l'), Some(b'l')) => 3, - _ => 0, - }; - i = if adv > 0 { - i + adv - } else { - run_end(tags, i, end, mask::NOT_WS_L_N) - }; - } - // Space: the ` ?` prefix — attach one space to a following letter / number / "other" run, - // else it's whitespace (rules `\s+(?!\S)|\s+`, which leave one space for the next run). - SPC => { - let a = i + 1; // Space is ASCII (0x20) - i = match tags.get(a).map(|&t| t & 0x0F) { - Some(LET) => run_end(tags, a, end, mask::LETTER), - Some(NW) | Some(NO) => run_end(tags, a, end, mask::NUMBER), - Some(t) if in_mask(t, mask::NOT_WS_L_N) => { - run_end(tags, a, end, mask::NOT_WS_L_N) - } - _ => ws(i), - }; - } - // WsOther / Newline: whitespace only — the ` ?` prefix is a literal 0x20, so tabs/newlines - // never prefix a run. - WSO | NLN => i = ws(i), - // Sentinel / MultiByte / Cont — never a char-start atom; emit one char defensively. - _ => i += char_len(text[i]), - } - // SAFETY: tokens partition the input, so `w < #tokens <= end < out.len()` (out ≥ text.len()+? ; callers size n+1). - unsafe { - *out.get_unchecked_mut(w) = Span { - start: start as u32, - end: i as u32, - } - }; - w += 1; - } - w -} diff --git a/tokenizers/atomsplit/src/fsm/cl100k.rs b/tokenizers/atomsplit/src/fsm/cl100k.rs deleted file mode 100644 index 65976ecef..000000000 --- a/tokenizers/atomsplit/src/fsm/cl100k.rs +++ /dev/null @@ -1,118 +0,0 @@ -use super::*; - -/// TODO: NONE OF THE FOLLOWING HAS BEEN REVIEWED -/// cl100k / Llama-3 pretokenization (7 rules + whitespace-tail, rule-3 cap `\p{N}{1,3}`). Peeks `text` -/// for the ASCII contraction literals. Scalar run-ends. See [`fsm_cl100k_cap`] for the variable-cap -/// family (Qwen2 etc.). ┌── OWNER: shared (scalar) ──┐ -#[must_use] -pub fn fsm_cl100k(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { - fsm_cl100k_cap(text, tags, out, 3) -} - -/// The cl100k family with an explicit rule-3 digit cap: 3 = cl100k / Llama-3, 1 = Qwen2's `\p{N}` (each -/// digit its own token), `usize::MAX` = an unbounded `\p{N}+`. Only the digit rule differs across these. -#[must_use] -pub fn fsm_cl100k_cap(text: &[u8], tags: &[u8], out: &mut [Span], digit_cap: usize) -> usize { - debug_assert!(out.len() >= text.len() && tags.len() >= text.len()); - cl100k(text, tags, out, digit_cap) -} - -fn cl100k(text: &[u8], tags: &[u8], out: &mut [Span], digit_cap: usize) -> usize { - // Leading-atom values, as `const` so the `match` below is a dense jump table (not an if-cascade): - // the dispatch is O(1) and a token never pays for a rule it can't start (e.g. non-number tokens - // never test the number rule — which is what the POC's const-gating removed by hand; here it's free). - let end = text.len(); - // Tie `tags.len() == end` so the optimizer drops the per-byte bounds check on every interior - // `tags[i]` in this fsm + its `run_end`/`letter_*` scans. (Callers guarantee `tags.len() >= end`.) - let tags = &tags[..end]; - let letters = |a: usize| run_end(tags, a, end, mask::LETTER); - // rule 4 body: `[^\s\p{L}\p{N}]+[\r\n]*` from `sp0` (any leading space already consumed). Returns - // the run end, or `sp0` if there is no "other" run there (caller then treats it as whitespace). - let other = |sp0: usize| -> usize { - let mut p = run_end(tags, sp0, end, mask::NOT_WS_L_N); - if p > sp0 { - while p < end && tags[p] == NLN { - p += char_len(text[p]); - } - } - p - }; - // rules 5-7 (`\s*[\r\n] | \s+(?!\S) | \s+`) → the shared `ws_tail`. - let ws = |i: usize| -> usize { ws_tail(text, tags, i, end) }; - - let mut i = 0; - let mut w = 0usize; - while i < end { - let start = i; - let b = text[i]; - match tags[i] & 0x0F { - // rule 2: `\p{L}+` - LET => i = letters(i), - // rule 3: `\p{N}{1,cap}` (cap = 3 cl100k, 1 Qwen2, MAX for `\p{N}+`) - NW | NO => { - let (mut p, mut cnt) = (i, 0); - while p < end && cnt < digit_cap && in_mask(tags[p], mask::NUMBER) { - p += char_len(text[p]); - cnt += 1; - } - i = p; - } - // Space: rule 2 (space prefix + `\p{L}+`) | rule 4 (` ` + "other") | rules 5-7 - SPC => { - let a = i + 1; // Space is ASCII (0x20) - i = if a < end && (tags[a] & 0x0F) == LET { - letters(a) - } else { - let p = other(a); - if p > a { p } else { ws(i) } - }; - } - // WsOther: rule 2 (prefix + `\p{L}+`) | whitespace (never rule 4 — not in NOT_WS_L_N) - WSO => { - let a = i + char_len(b); - i = if a < end && (tags[a] & 0x0F) == LET { - letters(a) - } else { - ws(i) - }; - } - // Newline: whitespace (rule 5 ends at the last newline) - NLN => i = ws(i), - // Apostrophe: rule 1 (contraction) | rule 2 (prefix + `\p{L}+`) | rule 4 - APO => { - let adv = contraction(text, i); // rule 1: `'s 't 're 've 'm 'll 'd` (case-insensitive) - i = if adv > 0 { - i + adv - } else { - let a = i + 1; // Apostrophe is ASCII (0x27) - if a < end && (tags[a] & 0x0F) == LET { - letters(a) - } else { - other(i) - } // c ∈ NOT_WS_L_N ⇒ > i - }; - } - // Mark | Connector | Punct | SymOther | NumericOther | Control (all in NOT_WS_L_N): - // rule 2 (prefix + `\p{L}+`) | rule 4 - MRK | CON | PUN | SYM | NMO | CTL => { - let a = i + char_len(b); - i = if a < end && (tags[a] & 0x0F) == LET { - letters(a) - } else { - other(i) - }; // c ∈ NOT_WS_L_N ⇒ > i - } - // Sentinel / MultiByte / Cont — never a char-start atom; emit one char defensively. - _ => i += char_len(b), - } - // SAFETY: tokens partition the input, so `w < #tokens <= end < out.len()` (out ≥ text.len()+? ; callers size n+1). - unsafe { - *out.get_unchecked_mut(w) = Span { - start: start as u32, - end: i as u32, - } - }; - w += 1; - } - w -} diff --git a/tokenizers/atomsplit/src/fsm/deepseek.rs b/tokenizers/atomsplit/src/fsm/deepseek.rs deleted file mode 100644 index f54afb8fa..000000000 --- a/tokenizers/atomsplit/src/fsm/deepseek.rs +++ /dev/null @@ -1,247 +0,0 @@ -use super::*; - -/// The specific CJK ranges deepseek's Split-2 isolates: Han U+4E00..9FA5 ∪ Hiragana U+3040..309F ∪ -/// Katakana U+30A0..30FF. (All 3-byte, leads E3..E9.) Not "all letters" — only these. -#[inline] -fn ds_is_cjk(cp: u32) -> bool { - (0x4E00..=0x9FA5).contains(&cp) || (0x3040..=0x30FF).contains(&cp) -} -/// Codepoint of a 3-byte UTF-8 char at `text[i]` (only called for leads E2..E9 → always 3-byte). -#[inline] -fn cp3(text: &[u8], i: usize) -> u32 { - ((text[i] as u32 & 0x0F) << 12) - | ((text[i + 1] as u32 & 0x3F) << 6) - | (text[i + 2] as u32 & 0x3F) -} - -/// Any CJK-range char (letter OR punct/sym) at `p` — the full `[一-龥぀-ゟ゠-ヿ]` set Split-2 `[…]+` -/// isolates, INCLUDING CJK punctuation (・ U+30FB, ゠ U+30A0, ゛゜ U+309B/C). Split-3 then re-splits that -/// isolated run into same-kind sub-runs (letters `[\p{L}\p{M}]+` vs punct/sym `[\p{P}\p{S}]+`); because -/// the run is a CLOSED unit, none of it steals a surrounding space or merges with non-CJK punct. -#[inline(always)] -fn ds_is_cjk_at(text: &[u8], p: usize) -> bool { - (0xE3..=0xE9).contains(&text[p]) && ds_is_cjk(cp3(text, p)) -} - -/// deepseek-v3 pretokenization: the `Sequence` of `[N{1,3}]`, `[CJK]+`, `` (all Isolated) -/// collapsed into ONE scalar FSM over the atom stream. Precedence (= the Sequence order): digits → -/// CJK-range runs → the big-regex alts. Because Split-2 isolates CJK *before* the letter rule, the -/// letter run stops at CJK-range codepoints. Peeks bytes for the ASCII `[punct][A-Za-z]+` alt. -/// -/// Byte-exact vs the real composed Sequence (onig ×3, each Isolated) on 10 languages — see -/// `benches/deepseek.rs` (plus Hebrew/Arabic via `tk-encode`'s corpus test). The subtleties the single -/// pass replicates: (1) ws *followed by* a digit/CJK is its own Sequence piece → the whole run is one -/// token (`\s+(?!\S)`); (2) ZWJ/ZWNJ are `\p{Cf}`, not `\p{L}∪\p{M}`, so they end a letter run -/// (`ds_breaks`); (3) Split-2 isolates a maximal CJK-range run and Split-3 re-splits it into same-kind -/// sub-runs — the top-of-loop handler consumes that run as a CLOSED unit (`ds_is_cjk_at`), so CJK punct -/// (・) never steals a surrounding space nor merges with non-CJK punct; (4) chars matching no alt -/// (Control / NumericOther / ZWJ) group into ONE gap piece, and Other_Alphabetic symbols (`ALPHA_SYM`: -/// `\w` but categorically `\p{S}`) take the `[\p{P}\p{S}]` path, not the letter run. -#[must_use] -pub fn fsm_deepseek(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { - debug_assert!(out.len() >= text.len() && tags.len() >= text.len()); - // Leading-atom values as `const` → the `match` is a dense jump table (see `cl100k`). The Split - // precedence (digits → CJK → big-regex alts) is preserved because the atom partition is disjoint. - // `Mark` refined as an Other_Alphabetic symbol (Ⓘ …): coarse `LETTER_MARK`, but categorically `\p{S}` - // — excluded from `[\p{L}\p{M}]`, routed to the `[\p{P}\p{S}]+` run instead (see `punct`). - let end = text.len(); - // Tie `tags.len() == end` so the optimizer drops the per-byte bounds check on every interior - // `tags[i]` in this fsm + its `run_end`/`letter_*` scans. (Callers guarantee `tags.len() >= end`.) - let tags = &tags[..end]; - // maximal `[\p{L}\p{M}]+` run from `a`, stopping at CJK-range chars (Split-2 took those), ZWJ/ZWNJ - // (not `\p{L}∪\p{M}` — see `ds_breaks`), and Other_Alphabetic symbols (`ASM`, categorically `\p{S}`). - // BYTE-wise (`p += 1`, continuation bytes stay in-run, `ds_breaks` only fires at a lead) — the - // `char_len`-per-char form was ~2× slower (see `run_end`'s note). Hot inner loop of the latin path. - let letter_run = |a: usize| -> usize { - let mut p = a; - // ZWJ/ASM are now tags (no text peek); only the CJK-range exclusion still peeks text. - while p < end { - let t = tags[p]; - if t == CONT - || (in_mask(t, mask::LETTER_MARK) && t != ASM && t != ZWJ && !ds_is_cjk_at(text, p)) - { - p += 1; - } else { - break; - } - } - p - }; - // is `text[a]` the start of a deepseek letter/mark char (alt-2 run body / space-prefix target)? - let is_lm = |a: usize| { - a < end - && in_mask(tags[a], mask::LETTER_MARK) - && tags[a] != ASM - && tags[a] != ZWJ - && !ds_is_cjk_at(text, a) - }; - // Split-3 alt-3 tail `[\p{P}\p{S}]+[\r\n]*` from `sp0` (a leading space is already consumed); `sp0` - // if there is no punct/sym run there. STOPS at CJK-range chars — Split-1 isolated those, so a CJK - // punct (・) is never merged into a non-CJK punct run (`!・` → `!`, `・`, not `!・`). - let punct = |sp0: usize| -> usize { - let mut p = sp0; - while p < end - && (in_mask(tags[p], mask::PUNCT_SYM) || tags[p] == ASM) - && !ds_is_cjk_at(text, p) - { - p += char_len(text[p]); - } - if p > sp0 { - while p < end && tags[p] == NLN { - p += char_len(text[p]); - } - } - p - }; - // Split-3 alts d/e/f (whitespace). Unlike cl100k: a ws run FOLLOWED BY a digit/CJK is its own - // Sequence piece (Split-1/2 isolated the next match) → `\s+(?!\S)` takes the WHOLE run; only a - // following letter/punct (same Split-3 piece) leaves the last ws char for its ` ?`/`[^…]?` prefix. - let ws = |i: usize| -> usize { - let re = run_end(tags, i, end, mask::WS); - let next_isolated = re < end && (in_mask(tags[re], mask::NUMBER) || ds_is_cjk_at(text, re)); - if let Some(r) = text[i..re].iter().rposition(|&x| x == 0x0A || x == 0x0D) { - i + r + 1 - } else if re == end || next_isolated { - re // whole ws run is one token - } else { - let mut last = re - 1; - while last > i && text[last] & 0xC0 == 0x80 { - last -= 1; - } - if last > i { last } else { re } - } - }; - - let mut i = 0; - let mut w = 0usize; - while i < end { - let start = i; - let b = text[i]; - // Split-2 isolated a maximal CJK-range run; Split-3 re-splits it into same-kind sub-runs - // (letters `[\p{L}\p{M}]+` vs punct/sym `[\p{P}\p{S}]+`) — a CLOSED unit, handled before the atom - // arms so CJK punct (・) never leaks into alt-3 (stealing a space / merging with non-CJK punct). - if ds_is_cjk_at(text, i) { - let is_letter = in_mask(tags[i], mask::LETTER_MARK); - let mut p = i + 3; // CJK-range chars are all 3-byte (leads E3..E9) - while p < end - && ds_is_cjk_at(text, p) - && in_mask(tags[p], mask::LETTER_MARK) == is_letter - { - p += 3; - } - unsafe { - *out.get_unchecked_mut(w) = Span { - start: start as u32, - end: p as u32, - } - }; - w += 1; - i = p; - continue; - } - // Gap run: maximal Control / NumericOther / ZWJ — none matches a Split-3 alt, so the composed - // Split emits the whole run as ONE unmatched piece. Exception: if it's immediately followed by a - // letter run, the LAST gap char is that run's alt-2 `[^\r\n\p{L}\p{P}\p{S}]?` prefix (splits off). - if matches!(tags[i] & 0x0F, NMO | CTL) || tags[i] == ZWJ { - let (mut p, mut last) = (i, i); - while p < end && (matches!(tags[p] & 0x0F, NMO | CTL) || tags[p] == ZWJ) { - last = p; - p += char_len(text[p]); - } - if is_lm(p) { - if last > i { - unsafe { - *out.get_unchecked_mut(w) = Span { - start: start as u32, - end: last as u32, - } - }; // gap sans prefix char - w += 1; - } - let e = letter_run(p); - unsafe { - *out.get_unchecked_mut(w) = Span { - start: last as u32, - end: e as u32, - } - }; // prefix char + `[\p{L}\p{M}]+` - w += 1; - i = e; - } else { - unsafe { - *out.get_unchecked_mut(w) = Span { - start: start as u32, - end: p as u32, - } - }; // whole gap run is one piece - w += 1; - i = p; - } - continue; - } - match tags[i] & 0x0F { - // Split-1: `\p{N}{1,3}` - NW | NO => { - let (mut p, mut cnt) = (i, 0); - while p < end && cnt < 3 && in_mask(tags[p], mask::NUMBER) { - p += char_len(text[p]); - cnt += 1; - } - i = p; - } - // Split-3 alt-2 `[\p{L}\p{M}]+` (CJK letters + ZWJ/gap chars were consumed above the match). - // An Other_Alphabetic symbol (`ASM`, coarse `Mark`) is categorically `\p{S}` → the alt-3 run. - LET | MRK => { - i = if tags[i] == ASM { - punct(i) - } else { - letter_run(i) - } - } - // Space: alt-2 (space prefix + `[\p{L}\p{M}]+`) | alt-3 (` ` + `[\p{P}\p{S}]+`) | whitespace - SPC => { - let a = i + 1; // Space is ASCII (0x20) - i = if is_lm(a) { - letter_run(a) - } else if a < end && ds_is_cjk_at(text, a) { - ws(i) // next is a Split-1-isolated CJK char → the space is standalone whitespace - } else { - let p = punct(a); - if p > a { p } else { ws(i) } - }; - } - // WsOther: alt-2 (prefix + `[\p{L}\p{M}]+`) | whitespace (not `\p{P}∪\p{S}` → no alt-3) - WSO => { - let a = i + char_len(b); - i = if is_lm(a) { letter_run(a) } else { ws(i) }; - } - // Newline: whitespace - NLN => i = ws(i), - // Connector | Punct | Apostrophe | SymOther (∈ `\p{P}∪\p{S}`): alt-1 `[ascii_punct][A-Za-z]+` - // | alt-3 `[\p{P}\p{S}]+[\r\n]*` - CON | PUN | APO | SYM => { - i = if b.is_ascii_punctuation() && i + 1 < end && text[i + 1].is_ascii_alphabetic() - { - let mut p = i + 1; - while p < end && text[p].is_ascii_alphabetic() { - p += 1; - } - p - } else { - punct(i) // c ∈ PUNCT_SYM ⇒ > i - }; - } - // Sentinel / MultiByte / Cont — never a char-start atom; emit one char defensively. - _ => i += char_len(b), - } - // SAFETY: tokens partition the input, so `w < #tokens <= end < out.len()` (out ≥ text.len()+? ; callers size n+1). - unsafe { - *out.get_unchecked_mut(w) = Span { - start: start as u32, - end: i as u32, - } - }; - w += 1; - } - w -} diff --git a/tokenizers/atomsplit/src/fsm/o200k.rs b/tokenizers/atomsplit/src/fsm/o200k.rs deleted file mode 100644 index 28e15640a..000000000 --- a/tokenizers/atomsplit/src/fsm/o200k.rs +++ /dev/null @@ -1,246 +0,0 @@ -use super::*; - -// ── o200k (GPT-4o): case-aware letter split ────────────────────────────────────────────────────── - -/// o200k class of a letter-run char (all chars in a run are real `[\p{L}\p{M}]`, never ALPHA_SYM/ZWJ): -/// `Atom::UpperLetter` → U (`\p{Lu}\p{Lt}`), `Atom::LowerLetter` → L (`\p{Ll}`), else C (caseless `\p{Lm}\p{Lo} -/// \p{M}`). The two alt char-classes are `[UC]` = "not L" (`!o_is_lower`) and `[LC]` = "not U". -#[inline] -fn o_is_upper(t: u8) -> bool { - t == Atom::UpperLetter as u8 // 0x10: coarse Letter (low nibble 0) + UPPER refine (\p{Lu}∪\p{Lt}) -} -#[inline] -fn o_is_lower(t: u8) -> bool { - t == Atom::LowerLetter as u8 // 0x20: coarse Letter + LOWER refine (\p{Ll}) -} - -/// One o200k letter sub-token from `p` within the run `[.., re)`: alt-1 `[UC]*[LC]+` (tried first) else -/// alt-2 `[UC]+[LC]*` (reached only for an all-U run). Greedy with Perl backtracking — `[UC]*` gives back -/// to the last C so `[LC]+` can take ≥1. Returns the sub-token end, always in `(p, re]`. BYTE-wise (`+=1`, -/// like `run_end`): continuation bytes are tag `Cont`(15) → neither U nor L → transparent to `[UC]`/`[LC]`. -#[inline(always)] -fn o200k_letter_match(tags: &[u8], p: usize, re: usize) -> usize { - // alt-1 `[UC]*`: greedy over "not L" (stops at the next lowercase *lead*), tracking the last C - // char-start so a no-L run needs no separate backtrack pass (`Cont`=15 is neither U nor L, so it's - // "C-like" — the `!= CONT` guard keeps `last_c` on a real char start). - let mut q = p; - let mut last_c = usize::MAX; - while q < re && !o_is_lower(tags[q]) { - if tags[q] != CONT && !o_is_upper(tags[q]) { - last_c = q; - } - q += 1; - } - if q < re { - // tags[q] is L → `[LC]+` from q: greedy over "not U" - let mut e = q; - while e < re && !o_is_upper(tags[e]) { - e += 1; - } - return e; - } - if last_c == usize::MAX { - return re; // no L and no C → all U → alt-2 `[UC]+[LC]*` (empty `[LC]*`) = the whole run - } - // no L: `[UC]*` gives back to the last C, which begins the `[LC]+` - let mut e = last_c; - while e < re && !o_is_upper(tags[e]) { - e += 1; - } - e -} - -/// Emit the o200k case-split of the letter run `[ls, re)` into `out[*w..]`: the first sub-token starts at -/// `pfx` (the optional `[^\r\n\p{L}\p{N}]?` prefix; `pfx == ls` when none), the last absorbs a trailing -/// contraction (`CONTRACTION` — off for tekken). Returns the new cursor (past the contraction). -/// `ls < re` (caller-guaranteed). -#[inline(always)] -fn emit_o200k_letters( - text: &[u8], - tags: &[u8], - pfx: usize, - ls: usize, - re: usize, - out: &mut [Span], - w: &mut usize, -) -> usize { - let (mut p, mut first, mut cursor) = (ls, true, re); - while p < re { - let e = o200k_letter_match(tags, p, re); - let start = if first { pfx } else { p }; - let tok_end = if CONTRACTION && e == re { - e + contraction(text, e) - } else { - e - }; - unsafe { - *out.get_unchecked_mut(*w) = Span { - start: start as u32, - end: tok_end as u32, - } - }; - *w += 1; - first = false; - cursor = tok_end; - p = e; - } - cursor -} - -/// o200k (GPT-4o) pretokenization. Same skeleton as cl100k, but the letter body is `[\p{L}\p{M}]+` -/// split into case sub-runs (`emit_o200k_letters`), the contraction is a *suffix* on letter tokens (not -/// a leading rule), and rule 4 is `[^\s\p{L}\p{N}]+[\r\n/]*`. The letter body excludes ALPHA_SYM symbols -/// and ZWJ/ZWNJ (coarse `Mark` but categorically `\p{S}`/`\p{Cf}`) — they take the prefix / rule-4 path. -/// Unlike deepseek there are no gaps: rule 4's `[^\s\p{L}\p{N}]+` is a catch-all. Scalar; ┌ OWNER: shared ┐ -#[must_use] -pub fn fsm_o200k(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { - o200k::(text, tags, out) -} - -/// Mistral tekken ([`crate::regexes::TEKKEN`]) — the o200k FSM with the contraction suffix off and one -/// token per digit. Every other rule is shared, so both are the same code monomorphized twice. -#[must_use] -pub fn fsm_tekken(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { - o200k::(text, tags, out) -} - -fn o200k( - text: &[u8], - tags: &[u8], - out: &mut [Span], -) -> usize { - debug_assert!(out.len() >= text.len() && tags.len() >= text.len()); - let end = text.len(); - // Tie `tags.len() == end` so the optimizer drops the per-byte bounds check on every interior - // `tags[i]` in this fsm + its `run_end`/`letter_*` scans. (Callers guarantee `tags.len() >= end`.) - let tags = &tags[..end]; - - // Is tag `t` (at byte `p`) a real `[\p{L}\p{M}]` member? Coarse `Letter` (any case) is always in; - // coarse `Mark` is in only as a true `\p{M}` — ALPHA_SYM (`\p{S}`) and ZWJ/ZWNJ (`\p{Cf}`) are `\w` - // but not `[\p{L}\p{M}]`. The `ds_is_zwj` byte-peek is thus paid ONLY for Marks, never for letters. - let member = |t: u8, p: usize| -> bool { - let c = t & 0x0F; - let _ = p; - c == LET || (c == MRK && t != ASM && t != ZWJ) - }; - let is_lm = |a: usize| a < end && member(tags[a], a); - // maximal `[\p{L}\p{M}]+` run from `a` (byte-wise; continuation bytes ride along — see `run_end`). - let letter_end = |a: usize| -> usize { - let mut p = a; - // logos-style fast loop: 16 tags/chunk, one bounds check, unchecked reads. A plain `Letter` - // (low nibble 0, incl Han — o200k keeps all letters) or a `Cont` byte stays in-run with no - // `ds_is_zwj` peek; only a coarse `Mark` lane pays the peek. Byte-exact with the scalar scan. - // SAFETY: `p + 16 <= end <= tags.len()`/`text.len()` in the body. - while p + 16 <= end { - let mut brk = 16; - for k in 0..16 { - let t = unsafe { *tags.get_unchecked(p + k) }; - if t == CONT || t & 0x0F == LET { - continue; - } - if t & 0x0F == MRK && t != ASM && t != ZWJ { - continue; - } - brk = k; - break; - } - if brk < 16 { - return p + brk; - } - p += 16; - } - while p < end && (tags[p] == CONT || member(tags[p], p)) { - p += 1; - } - p - }; - // rule 4 `[^\s\p{L}\p{N}]+[\r\n/]*` from `sp0` (any leading space already consumed); `sp0` if none. - // `/` is in the `+` body too — the trailing class only matters after the `+` stops at a `\r\n`. - let other = |sp0: usize| -> usize { - let mut p = run_end(tags, sp0, end, mask::NOT_WS_L_N); - if p > sp0 { - while p < end && (tags[p] == NLN || text[p] == b'/') { - p += char_len(text[p]); - } - } - p - }; - // rules 5-7 (`\s*[\r\n]+ | \s+(?!\S) | \s+`) → the shared `ws_tail` (identical to cl100k). - let ws = |i: usize| -> usize { ws_tail(text, tags, i, end) }; - // The letter rules: case-split the run starting at `ls`, first sub-token starting at the prefix `pfx`. - let letters = |pfx: usize, ls: usize, out: &mut [Span], w: &mut usize| -> usize { - emit_o200k_letters::(text, tags, pfx, ls, letter_end(ls), out, w) - }; - - let mut i = 0; - let mut w = 0usize; - while i < end { - let start = i; - let b = text[i]; - match tags[i] & 0x0F { - // rule 3: `\p{N}{1,DIGIT_CAP}`, no prefix (3 = o200k, 1 = tekken) - NW | NO => { - let (mut p, mut cnt) = (i, 0); - while p < end && cnt < DIGIT_CAP && in_mask(tags[p], mask::NUMBER) { - p += char_len(text[p]); - cnt += 1; - } - i = p; - } - // letters `[\p{L}\p{M}]+` (case-split) — but ALPHA_SYM/ZWJ (coarse Mark, not `[\p{L}\p{M}]`) - // take the `[^\r\n\p{L}\p{N}]?` prefix / rule-4 path instead. - LET | MRK => { - if tags[i] != ASM && tags[i] != ZWJ { - i = letters(i, i, out, &mut w); - continue; - } - let a = i + char_len(b); - if is_lm(a) { - i = letters(i, a, out, &mut w); - continue; - } - i = other(i); // ∈ NOT_WS_L_N ⇒ > i - } - // Space: ` ?` prefix + letters | ` ?` + rule-4 other | whitespace - SPC => { - let a = i + 1; // Space is ASCII (0x20) - if is_lm(a) { - i = letters(i, a, out, &mut w); - continue; - } - let p = other(a); - i = if p > a { p } else { ws(i) }; - } - // WsOther: prefix + letters | whitespace (∈ `\s` ⇒ never starts rule 4) - WSO => { - let a = i + char_len(b); - if is_lm(a) { - i = letters(i, a, out, &mut w); - continue; - } - i = ws(i); - } - NLN => i = ws(i), - // punct / sym / … (∈ `[^\r\n\p{L}\p{N}]` and `[^\s\p{L}\p{N}]`): prefix + letters | rule-4 other - CON | PUN | APO | SYM | NMO | CTL => { - let a = i + char_len(b); - if is_lm(a) { - i = letters(i, a, out, &mut w); - continue; - } - i = other(i); // ∈ NOT_WS_L_N ⇒ > i - } - // Sentinel / MultiByte / Cont — never a char-start atom; emit one char defensively. - _ => i += char_len(b), - } - // SAFETY: tokens partition the input, so `w < #tokens <= end < out.len()` (out ≥ text.len()+? ; callers size n+1). - unsafe { - *out.get_unchecked_mut(w) = Span { - start: start as u32, - end: i as u32, - } - }; - w += 1; - } - w -} diff --git a/tokenizers/atomsplit/src/lib.rs b/tokenizers/atomsplit/src/lib.rs deleted file mode 100644 index 5f22f52d0..000000000 --- a/tokenizers/atomsplit/src/lib.rs +++ /dev/null @@ -1,33 +0,0 @@ -//! `atomsplit` — SIMD Unicode classification + finite-state pre-tokenization. -//! -//! One SIMD pass ([`classify`]) maps every codepoint to a tiny "atom" alphabet; a family of no-push -//! FSMs ([`fsm`]) turn that atom stream into token spans (byte ranges) — the pre-tokenizer stage that -//! runs before a BPE/WordPiece model. Pre-tokenizers implemented: `WhitespaceSplit`, `Punctuation`, -//! `Digits`, `Whitespace`, `Bert`, `Cl100k`, `DeepSeek`, `ByteLevel`, `CharDelimiterSplit` (o200k and -//! Mistral's tekken are exposed as the [`fsm::fsm_o200k`] / [`fsm::fsm_tekken`] functions rather than -//! recipe structs). -//! -//! For pre-tokenizers that split on single characters (such as the Metaspace `▁` delimiter), -//! we skip the atom classification pass entirely and search for the raw bytes instead. -//! See [`literal`]. -//! -//! Design: every fsm is *no-push* — it writes spans into a caller-preallocated `&mut [fsm::Span]` -//! (length ≥ `text.len()`) and returns the token count; there is no `Vec`/allocation on the hot path. -//! -//! # Preconditions -//! Inputs are `&[u8]` (not `&str`) for zero-copy, but **must be well-formed UTF-8**: a buffer that ends -//! mid-codepoint is a precondition violation and may panic. `tags`/`out` scratch buffers must be -//! `≥ text.len()` (asserted in [`classify`]; documented per-fsm). -mod atom_tables; -pub mod classify; -pub mod fsm; -pub mod literal; -pub mod regexes; -#[cfg(target_arch = "x86_64")] -mod simd_avx_classify; -#[cfg(target_arch = "aarch64")] -mod simd_classify; -mod simd_fsm; -#[cfg(all(target_arch = "wasm32", target_feature = "simd128"))] -mod simd_wasm_classify; -pub mod tables; diff --git a/tokenizers/atomsplit/src/literal.rs b/tokenizers/atomsplit/src/literal.rs deleted file mode 100644 index c826b2dda..000000000 --- a/tokenizers/atomsplit/src/literal.rs +++ /dev/null @@ -1,56 +0,0 @@ -//! Searching for a literal string, for pre-tokenizers that cut on one exact character. -//! -//! The rest of the crate works off atom tags: one SIMD pass gives every character a class, and the -//! FSMs cut where the class changes ([`crate::fsm`]). -//! -//! The atom classification is unnecessary for pre-tokenizers splitting on an exact character or literal string: -//! We can use a simpler byte search looking at 16 bytes at a time. - -use memchr::memmem; -use std::fmt; - -/// The pattern handed to [`Literal::new`] was empty. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct EmptyPattern; - -impl fmt::Display for EmptyPattern { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.write_str("an empty pattern matches everywhere") - } -} - -impl std::error::Error for EmptyPattern {} - -/// A literal string to split on. -/// -/// The finder is boxed because it is large — a few hundred bytes of prefilter state on x86_64 — and -/// callers store it inside enums whose other variants are tiny. -#[derive(Debug, Clone)] -pub struct Literal { - finder: Box>, -} - -impl Literal { - /// # Errors - /// If `pattern` is empty, which would match everywhere. - pub fn new(pattern: &[u8]) -> Result { - if pattern.is_empty() { - return Err(EmptyPattern); - } - Ok(Self { - finder: Box::new(memmem::Finder::new(pattern).into_owned()), - }) - } - - /// The string being searched for. - #[must_use] - pub fn pattern(&self) -> &[u8] { - self.finder.needle() - } - - /// Byte offset of every match, left to right. Matches never overlap, so `"aa"` is found once in - /// `"aaa"` — the same matches a regex engine would report. - pub fn matches<'t>(&'t self, text: &'t [u8]) -> impl Iterator + 't { - self.finder.find_iter(text) - } -} diff --git a/tokenizers/atomsplit/tests/fsm.rs b/tokenizers/atomsplit/tests/fsm.rs deleted file mode 100644 index 1cad95e1b..000000000 --- a/tokenizers/atomsplit/tests/fsm.rs +++ /dev/null @@ -1,107 +0,0 @@ -//! Integration tests for the FSM pre-tokenizers. Kept out of `src/` so the core stays production-only. -use atomsplit::classify::{classify, mask}; -use atomsplit::fsm::{ - CharDelimiterSplit, Span, class_runs_into, emit_class_spans, fsm_byte_level, fsm_cl100k, - fsm_deepseek, fsm_o200k, fsm_tekken, -}; - -/// Run a no-push fsm into a fresh buffer and return the emitted spans. -fn spans(f: impl Fn(&[u8], &[u8], &mut [Span]) -> usize, s: &str) -> Vec { - let mut tags = vec![0u8; s.len()]; - classify(s.as_bytes(), &mut tags); - let mut out = vec![Span::default(); s.len() + 1]; - let k = f(s.as_bytes(), &tags, &mut out); - out.truncate(k); - out -} - -#[test] -fn cl100k_rules() { - // hand-verified against the tiktoken cl100k_base regex - let cl = |s| spans(fsm_cl100k, s); - assert_eq!(cl("Hello world"), vec![(0, 5), (5, 11)]); // "Hello" | " world" - assert_eq!(cl("don't"), vec![(0, 3), (3, 5)]); // "don" | "'t" (contraction) - assert_eq!(cl("a1234"), vec![(0, 1), (1, 4), (4, 5)]); // "a" | "123" | "4" ({1,3} cap) - assert_eq!(cl(" hi"), vec![(0, 1), (1, 4)]); // " " | " hi" - assert_eq!(cl("a, b"), vec![(0, 1), (1, 2), (2, 4)]); // "a" | "," | " b" -} - -/// Mistral's tekken split is o200k's, minus the contraction suffix, with one token per digit. -#[test] -fn tekken_rules() { - let tk = |s| spans(fsm_tekken, s); - assert_eq!(tk("don't"), vec![(0, 3), (3, 5)]); // "don" | "'t" — prefix+letters, not a contraction - assert_eq!(spans(fsm_o200k, "don't"), vec![(0, 5)]); // o200k glues the contraction on - assert_eq!(tk("a1234"), vec![(0, 1), (1, 2), (2, 3), (3, 4), (4, 5)]); // `\p{N}`, one digit each - assert_eq!(tk("XMLHttpRequest"), vec![(0, 7), (7, 14)]); // case split: "XMLHttp" | "Request" - assert_eq!(tk("a/\r\nb"), vec![(0, 1), (1, 4), (4, 5)]); // "/" run + its `[\r\n/]*` tail - assert_eq!(tk("hi ok"), vec![(0, 2), (2, 4), (4, 7)]); // \s+(?!\S) leaves one space -} - -#[test] -fn deepseek_rules() { - let ds = |s| spans(fsm_deepseek, s); - assert_eq!(ds("abc中def"), vec![(0, 3), (3, 6), (6, 9)]); // letters | CJK | letters - assert_eq!(ds("abc123"), vec![(0, 3), (3, 6)]); // letters | digits {1,3} - assert_eq!(ds("_abc"), vec![(0, 4)]); // alt-1: ASCII punct + letters - assert_eq!(ds("hello world"), vec![(0, 5), (5, 11)]); // word | space+word - assert_eq!(ds("!!!"), vec![(0, 3)]); // \p{P}∪\p{S} run -} - -#[test] -fn byte_level_rules() { - let bl = |s| spans(fsm_byte_level, s); - // ` ?\p{L}+`, lowercase contraction, ` ?\p{N}+` UNBOUNDED - assert_eq!(bl("I'm 12345 ok"), vec![(0, 1), (1, 3), (3, 9), (9, 12)]); - assert_eq!(bl("IT'S"), vec![(0, 2), (2, 3), (3, 4)]); // 'S is not a contraction (case-sensitive) - assert_eq!(bl("hi ok"), vec![(0, 2), (2, 4), (4, 7)]); // \s+(?!\S) leaves one space -} - -#[test] -fn char_delimiter_split() { - let mut out = vec![Span::default(); 8]; - // split on '/', Removed → drop delimiters, drop the empty gap between "//" - let k = CharDelimiterSplit('/').pre_tokenize(b"a/bc//d", &mut [], &mut out); - assert_eq!(&out[..k], &[(0, 1), (2, 4), (6, 7)]); -} - -/// Byte-exactness gate for the class family: the NEON boundary extractor (`class_runs_into`) must equal -/// the scalar run-end core (`emit_class_spans`) for every recipe, at every char-aligned truncation length so -/// the < 16-byte NEON tail starts at every offset — including mid-char (chunk loop steps by 16). Corpus -/// mixes ASCII, 2/3-byte scripts, Devanagari letter+matra clusters, consecutive punct, tabs, astral. -#[test] -fn class_runs_into_matches() { - let corpus = "Hello, world!! 123 café × наука 中文。। नरेंद्र मोदी ने ½²¼ ①② 😀a b\t".repeat(30); - let full = corpus.as_bytes(); - let mut tags = vec![0u8; full.len()]; - let mut b1 = vec![Span::default(); full.len()]; - let mut b2 = vec![Span::default(); full.len()]; - - fn eq( - t: &[u8], - tg: &[u8], - x: &mut [Span], - y: &mut [Span], - name: &str, - ) { - let k1 = class_runs_into::(t, tg, x); - let k2 = emit_class_spans::(t, tg, y, 0, 0, 0, None); - assert_eq!(&x[..k1], &y[..k2], "{} @len {}", name, t.len()); - } - let mut sweep = |len: usize| { - let (t, tg) = (&full[..len], &mut tags[..len]); - classify(t, tg); - let (x, y) = (&mut b1[..len], &mut b2[..len]); - eq::<{ mask::WS }, 0, 0>(t, tg, x, y, "WhitespaceSplit"); - eq::<0, { mask::PUNCT }, 0>(t, tg, x, y, "Punctuation"); - eq::<0, 0, { mask::NUMERIC }>(t, tg, x, y, "Digits"); - eq::<{ mask::WS }, 0, { mask::WORD }>(t, tg, x, y, "Whitespace"); - eq::<{ mask::WS }, { mask::PUNCT }, 0>(t, tg, x, y, "Bert"); - }; - sweep(full.len()); - for c in full.len().saturating_sub(64)..full.len() { - if corpus.is_char_boundary(c) { - sweep(c); - } - } -} diff --git a/tokenizers/atomsplit/tests/literal.rs b/tokenizers/atomsplit/tests/literal.rs deleted file mode 100644 index 14ceed9a0..000000000 --- a/tokenizers/atomsplit/tests/literal.rs +++ /dev/null @@ -1,50 +0,0 @@ -//! Tests for the literal search. Kept out of `src/` so the core stays production-only. -use atomsplit::literal::{EmptyPattern, Literal}; - -#[test] -fn finds_every_match_and_nothing_else() { - let literal = Literal::new(b"-").unwrap(); - assert_eq!(literal.matches(b"a-b--c").collect::>(), [1, 3, 4]); - assert_eq!(literal.matches(b"none here").count(), 0); - assert_eq!(literal.matches(b"").count(), 0); - assert_eq!(literal.pattern(), b"-"); -} - -#[test] -fn a_multi_byte_pattern_only_matches_whole() { - // `▁` is U+2581 = E2 96 81. The other characters here share its first byte and nothing else, so a - // search for that byte alone would report all of them. - let literal = Literal::new("▁".as_bytes()).unwrap(); - let text = "a—b“c…d▁e"; - assert_eq!( - literal.matches(text.as_bytes()).collect::>(), - [text.find('▁').unwrap()] - ); -} - -#[test] -fn matches_do_not_overlap() { - let literal = Literal::new(b"aa").unwrap(); - assert_eq!(literal.matches(b"aaa").collect::>(), [0]); - assert_eq!(literal.matches(b"aaaa").collect::>(), [0, 2]); -} - -/// A `Literal` is stored inline in the normalizer and decoder enums, where every other variant is a -/// handful of bytes. `memmem::Finder` itself is a few hundred, so it has to stay behind a pointer. -#[test] -fn a_literal_is_pointer_sized() { - assert_eq!( - size_of::(), - size_of::<*const u8>(), - "a Literal must not carry its finder inline" - ); -} - -#[test] -fn an_empty_pattern_is_rejected() { - assert_eq!(Literal::new(b"").unwrap_err(), EmptyPattern); - assert_eq!( - EmptyPattern.to_string(), - "an empty pattern matches everywhere" - ); -} diff --git a/tokenizers/atomsplit/tests/parity.rs b/tokenizers/atomsplit/tests/parity.rs deleted file mode 100644 index 774e7c002..000000000 --- a/tokenizers/atomsplit/tests/parity.rs +++ /dev/null @@ -1,121 +0,0 @@ -//! Reference-parity gates for the regex-shaped pre-tokenizers. Each of `fsm_cl100k` / `fsm_o200k` / -//! `fsm_tekken` / `fsm_byte_level` / `fsm_deepseek` must be BYTE-EXACT with the real reference — the -//! oniguruma regex, composed exactly as HF applies it (deepseek is a `Sequence` of three Isolated -//! splits) — on two corpora: a multilingual one (ASCII, contractions, digits, punctuation runs, -//! whitespace variants, Latin accents/marks, Cyrillic/Greek/Arabic/Devanagari, Han/Kana/Hangul, ZWJ, -//! astral emoji) and an edge-case one (see [`EDGE`]). -//! -//! This is the byte-exactness gate the hand cases in `fsm.rs` don't provide; the same corpora + gate -//! should run under x86 (Intel SDE) in CI to validate the SIMD paths. -//! -//! Gated off wasm32: the oniguruma reference is a C library that has no wasi libc to build against. -#![cfg(not(target_arch = "wasm32"))] -use atomsplit::classify::classify; -use atomsplit::fsm::{Span, fsm_byte_level, fsm_cl100k, fsm_deepseek, fsm_o200k, fsm_tekken}; -use onig::Regex; -// The oracle regexes are the canonical specs the FSMs implement — single source of truth in atomsplit. -use atomsplit::regexes::{ - CL100K, DEEPSEEK_BIG as DS_BIG, DEEPSEEK_CJK as DS_CJK, DEEPSEEK_NUM as DS_NUM, GPT2, O200K, - TEKKEN, -}; - -const CORPUS: &str = "The quick brown fox. Don't 12345 numbers, \u{00BD}\u{00B2}\u{00BC} \u{2168}! \ - café × naïve — Привет, наука! Ελλάδα 中文分词。ひらがな カタカナ 한글 مرحبا العربية \ - नरेंद्र मोदी x_y a1b2c3 e-mail@host.com 😀👍 hello world\ttabs\nnewlines end "; - -/// Second corpus, aimed at the axes where the o200k-shaped FSMs differ from each other: apostrophes -/// (contraction suffix vs plain prefix+letters), digit-run length (`{1,3}` vs one-per-token), and the -/// `[\r\n/]*` tail after a symbol run. -const EDGE: &str = "IT'S O'Brien can't 'quoted' l'été rock'n'roll\r\n\ - 0 42 999 1000 1234567 v1.2.3 3.14159 1,000,000 \ - https://host/a/b?c=1&d=2 path/to//file /\r\n/ ///x \ - CamelCase XMLHttpRequest IJSSELMEER DžAMBO ŀl a\u{0301}b \ - 日本語1234テスト ½3¼ \u{2168}42\u{2169} #tag @user $9.99 100% \ - end\n\n\nlines\r\n\r\n \t trailing "; - -fn spans(f: impl Fn(&[u8], &[u8], &mut [Span]) -> usize, s: &str) -> Vec { - let mut tags = vec![0u8; s.len()]; - classify(s.as_bytes(), &mut tags); - let mut out = vec![Span::default(); s.len() + 1]; - let k = f(s.as_bytes(), &tags, &mut out); - out.truncate(k); - out -} - -fn onig_spans(re: &Regex, s: &str) -> Vec { - re.find_iter(s) - .map(|(a, b)| Span::new(a as u32, b as u32)) - .collect() -} - -// One Isolated split of text[s..e] by `re`: emit gaps + matches (all pieces), absolute offsets. -fn split_iso(text: &str, s: usize, e: usize, re: &Regex, out: &mut Vec<(usize, usize)>) { - let sub = &text[s..e]; - let mut prev = 0usize; - for (ms, me) in re.find_iter(sub) { - if ms > prev { - out.push((s + prev, s + ms)); - } - out.push((s + ms, s + me)); - prev = me; - } - if prev < sub.len() { - out.push((s + prev, e)); - } -} - -fn deepseek_ref(text: &str) -> Vec { - let (rn, rc, rb) = ( - Regex::new(DS_NUM).unwrap(), - Regex::new(DS_CJK).unwrap(), - Regex::new(DS_BIG).unwrap(), - ); - let mut p1 = Vec::new(); - split_iso(text, 0, text.len(), &rn, &mut p1); - let mut p2 = Vec::new(); - for (s, e) in p1 { - split_iso(text, s, e, &rc, &mut p2); - } - let mut p3 = Vec::new(); - for (s, e) in p2 { - split_iso(text, s, e, &rb, &mut p3); - } - p3.into_iter() - .map(|(s, e)| Span::new(s as u32, e as u32)) - .collect() -} - -/// Both corpora, one regex: the FSM must reproduce the oracle span-for-span. -fn check(fsm: impl Fn(&[u8], &[u8], &mut [Span]) -> usize + Copy, pattern: &str) { - let re = Regex::new(pattern).unwrap(); - for text in [CORPUS, EDGE] { - assert_eq!(spans(fsm, text), onig_spans(&re, text), "{text:?}"); - } -} - -#[test] -fn cl100k_parity() { - check(fsm_cl100k, CL100K); -} - -#[test] -fn o200k_parity() { - check(fsm_o200k, O200K); -} - -#[test] -fn tekken_parity() { - check(fsm_tekken, TEKKEN); -} - -#[test] -fn byte_level_parity() { - check(fsm_byte_level, GPT2); -} - -#[test] -fn deepseek_parity() { - for text in [CORPUS, EDGE] { - assert_eq!(spans(fsm_deepseek, text), deepseek_ref(text), "{text:?}"); - } -} diff --git a/tokenizers/bitmap_gen/Cargo.toml b/tokenizers/bitmap_gen/Cargo.toml index 3500b399a..800c26cc1 100644 --- a/tokenizers/bitmap_gen/Cargo.toml +++ b/tokenizers/bitmap_gen/Cargo.toml @@ -2,7 +2,7 @@ name = "bitmap_gen" version = "0.1.0" edition = "2024" -description = "Dev tool: regenerates atomsplit's committed classify tables (uses unicode-properties). Not for publication." +description = "Dev tool: regenerates bitsplit's committed classify tables (uses unicode-properties). Not for publication." license = "Apache-2.0" repository = "https://github.com/huggingface/tokenizers" authors = ["Arthur Zucker ", "Luc Georges "] diff --git a/tokenizers/bitmap_gen/src/lib.rs b/tokenizers/bitmap_gen/src/lib.rs index 302ff1bb9..224ce82f2 100644 --- a/tokenizers/bitmap_gen/src/lib.rs +++ b/tokenizers/bitmap_gen/src/lib.rs @@ -213,7 +213,7 @@ fn generate_tables(struct_name: &str, kind: &str, classify: &dyn Fn(u32) -> u8) ) .unwrap(); o.push_str("//! `Tables::classify_char`. `bitmap_gen` self-validates all 1.1M codepoints. Spec §1/§7.\n"); - o.push_str("use crate::tables::Tables;\n\n"); + o.push_str("use super::tables::Tables;\n\n"); let emit_u8 = |o: &mut String, name: &str, t: &[u8]| { write!(o, "#[rustfmt::skip]\nstatic {name}: [u8; {}] = [", t.len()).unwrap(); diff --git a/tokenizers/bitmap_gen/src/main.rs b/tokenizers/bitmap_gen/src/main.rs index 0b2990184..e76c05423 100644 --- a/tokenizers/bitmap_gen/src/main.rs +++ b/tokenizers/bitmap_gen/src/main.rs @@ -1,12 +1,9 @@ -//! Regenerate atomsplit's committed classify tables: +//! Regenerate bitsplit's committed classify tables: //! cargo run -p bitmap_gen [-- ] -//! Default `out_path` = ../atomsplit/src/atom_tables.rs. `generate_atom_tables` self-validates every +//! Default `out_path` = ../bitsplit/src/classify/atom_tables.rs. `generate_atom_tables` self-validates every //! codepoint against the reference `atom()`, so an inconsistent scheme change fails HERE, not at ship. fn main() { - let default = concat!( - env!("CARGO_MANIFEST_DIR"), - "/../atomsplit/src/atom_tables.rs" - ); + let default = concat!(env!("CARGO_MANIFEST_DIR"), "/../bitsplit/src/classify/atom_tables.rs"); let out = std::env::args() .nth(1) .unwrap_or_else(|| default.to_string()); diff --git a/tokenizers/bitsplit/Cargo.toml b/tokenizers/bitsplit/Cargo.toml new file mode 100644 index 000000000..e7dc67097 --- /dev/null +++ b/tokenizers/bitsplit/Cargo.toml @@ -0,0 +1,34 @@ +[package] +name = "bitsplit" +version = "0.1.0" +edition = "2024" +rust-version = "1.89" # AVX-512 VBMI intrinsics (simd_avx_classify) are stable only since 1.89 +authors = [ + "Arthur Zucker ", + "Luc Georges ", + "Simon Brandeis ", +] +homepage = "https://github.com/huggingface/tokenizers" +repository = "https://github.com/huggingface/tokenizers" +documentation = "https://docs.rs/bitsplit/" +license = "Apache-2.0" +keywords = ["tokenizer", "nlp", "pretokenizer", "simd", "unicode"] +categories = ["text-processing", "parsing"] +description = "SIMD Unicode atom classification + bitstream pre-tokenization: one classify pass into byte-exact token spans (cl100k, o200k, GPT-2/ByteLevel, DeepSeek, Whitespace, Bert, ...)." +exclude = ["benches/data/"] + +[package.metadata.docs.rs] +all-features = true + +[lib] +name = "bitsplit" +path = "src/lib.rs" + +[dependencies] +ahash = "0.8" +# NOTE: classify tables live in src/classify/atom_tables.rs (committed, generated). Regenerate after any atom +# scheme change with `cargo run -p bitmap_gen`. No build script / build-dep — bitsplit builds clean. + +# onig is C (Oniguruma) and has no wasi libc to build against, so the parity oracle is gated off wasm32. +[target.'cfg(not(target_arch = "wasm32"))'.dev-dependencies] +onig = "6.5.1" diff --git a/tokenizers/bitsplit/mask_splitter_spec.md b/tokenizers/bitsplit/mask_splitter_spec.md new file mode 100644 index 000000000..f6026232b --- /dev/null +++ b/tokenizers/bitsplit/mask_splitter_spec.md @@ -0,0 +1,336 @@ +# bitsplit — mask splitter spec + +A Rust library for **pre-tokenizer splitting as a bitstream program**. Not a general regex engine: +it covers exactly what `tokenizers` needs — the GPT-family regexes, the class-run family, and +literal delimiters — and it does so with one shared operator set so a new grammar is ~150 lines of +boolean algebra rather than a new hand-rolled FSM. + +Lineage: Parabix (Cameron et al., *Boosting the efficiency of text processing on commodity +processors*) for the operator set and carry discipline; *Interleaved Bitstream Execution for +Multi-Pattern Regex Matching on GPUs* (MICRO'25, doi 10.1145/3725843.3756052) for the execution +model — fuse every instruction into ONE block-wise loop instead of one pass per instruction. + +Status: proven in `scratch/bitsplit` for three grammars, all byte-exact against their +`atomsplit::fsm` counterparts. + +--- + +## 1. Why it wins + +An FSM decides one token per unpredictable branch, so its cost tracks bytes-per-token. A bitstream +program decides 64 bytes per register op, branchlessly, so its cost is flat. + +``` + ns/B (english, same run, aarch64) + fsm_deepseek ████████████████████████████████████████████████ 2.93 4.7 B/tok + fsm_cl100k ██████████████████████████████ 1.86 + fsm_byte_level ████████████████████████████ 1.71 + bitsplit ds █████████ 0.58 + bitsplit cl100k████████ 0.50 + bitsplit bl ██████ 0.41 +``` + +The FSM spread across grammars is 71% (1.71 → 2.93); bitsplit's is 20%. **Flat cost across grammars +is the design invariant** — it means the per-grammar layer is thin, which is what makes the library +worth having. + +--- + +## 2. Pipeline + +``` + text ─────────────────────────────────────────────────────────────┐ + │ │ + ▼ │ + classify (atomsplit) one Atom tag per byte │ + │ │ + ▼ ▼ + ┌─────────────────────── per 64-byte block, fused ────────────────────────┐ + │ │ + │ tags ──► LUT ──► dense 3-bit code ──► fill conts ──► 3 bit-planes │ + │ │ │ + │ decode ─────┼──► class │ + │ │ streams │ + │ ▼ │ + │ grammar algebra (Parabix operators) │ + │ │ │ + │ ▼ │ + │ starts │ keep (u64 each) │ + └──────────────────────────────────────────────────────┬───────────────────┘ + │ ~n/8 bytes + ▼ + emit walk ──► spans / sink +``` + +Everything inside the box lives in registers. The **only** intermediate that reaches memory is the +`starts` bitmap. That is the whole point of interleaved execution: the sequential model (one loop +per bitstream instruction, each materialising a full-length stream) is what the MICRO'25 paper +beats, and it is what we must not accidentally rebuild. + +--- + +## 3. The universal intermediate + +Every pre-tokenizer in `tokenizers` reduces to **one bit per byte: "a token starts here"**, plus +optionally "this byte is kept". That single contract is what makes the library abstract without +being a general engine. + +``` + text H e l l o , w o r l d + starts 1 . . . . 1 1 . . . . . + keep 1 1 1 1 1 1 . 1 1 1 1 1 (Removed behaviour drops the space) + └──"Hello"──┘ "," └──"world"──┘ +``` + +Which collapses three families onto one emit: + +| family | how `starts` is produced | +|---|---| +| class runs (WhitespaceSplit, Punctuation, Digits, Whitespace, Bert) | `c & !(c<<1)` per class; DROP/ISOLATE/KEEP_A are 3 bit ops | +| regex grammars (gpt2, cl100k, o200k, tekken, deepseek) | the algebra of §6 | +| literals (Metaspace `▁`, CharDelimiterSplit) | shifted-AND chain of byte compares | + +And `SplitDelimiterBehavior` stops being per-grammar code — it is a post-pass on the mask: + +``` + match . . 1 1 1 . . (a matched delimiter run) + Isolated . . 1 . . 1 . starts at both edges + Removed . . 1 . . 1 . + keep &= !match + MergedWithPrevious drop the leading edge bit + MergedWithNext drop the trailing edge bit + Contiguous merge adjacent same-kind runs +``` + +--- + +## 4. Operator set + +Parabix names, because they are the literature standard. Every operator takes a `&mut Carry` so +cross-block state is explicit and never forgotten. + +| operator | meaning | cost | +|---|---|---| +| `advance(m, n)` | move markers forward n bytes | 1–2 ops | +| `advance_char(m, cont)` | move markers forward one **char** (§5) | ~5 ops | +| `scan_thru(m, c)` | move each marker past the run of `c` it sits in | 1 add | +| `match_star(m, c)` | Kleene closure over a class | 1 add | +| `span_upto(m, e)` | set every bit from marker to end marker | 1 sub | +| `fill_to_last(m, c)` | in each `c`-run, fill from run start through the LAST marker | 2 ops | +| `to_lead(x, cont)` | move each bit back to its char's lead byte | ≤3 steps | +| `run_start(c)` | `c & !(c<<1)` | 2 ops | + +Do all of these in `u128` and read bit 64 as the carry-out. That one habit removes an entire class +of bug: a run reaching bit 63 puts its landing bit at 64 instead of vanishing, and `e - m` still +yields the correct in-block span. + +``` + fill_to_last, worked (finds "start after the LAST newline in a ws run"): + + run of c ┌───────────────────────┐ + c (ws) 1 1 1 1 1 1 1 1 . . + m (nl) . 1 . . 1 . . . . . + scan_thru . . . . . . . . 1 . ← lands past the run + (e-m)|m 1 1 1 1 1 . . . . . ← run start .. last marker + ▲ + the surviving "after last newline" bit is here+1 +``` + +`(e - m) | m` rather than `e - m`: the latter only spans from the *first* marker when a run holds +several. This is the single least obvious identity in the whole library — write a test for it. + +--- + +## 5. The transpose (bytes → bitstreams) + +The measured cost centre: **45–73% of total before optimisation, 36% after.** Optimise it before +touching any grammar. + +``` + tags [Lt][Lt][Pu][Sp][Lt][Lt][Lt][Ap][Lt] ... 16 lanes at a time + │ + │ vqtbl4q (64-entry table, indexed by the RAW tag — + │ refinements 0x10/0x20/0x16/0x26 index straight in) + ▼ + code 0 0 2 4 0 0 0 6 0 ... dense 3-bit code + │ + │ fill continuations from the left: 2 steps, not 3. + │ shift-1 makes every lane right at distance 1, so shifting + │ THAT by 2 covers distances 2 and 3 — the most a 4-byte char needs. + ▼ + code' 0 0 2 4 0 0 0 6 0 ... every byte carries its CHAR's code + │ + │ 3 × (vtst + vand POW + 3 × vpaddq_u8) ← the 64×8 transpose + ▼ + p0 p1 p2 three u64 bit-planes of the code, for 64 bytes + │ + │ decode: ~14 boolean ops + ▼ + l n other nl sp ws apo class streams +``` + +**Filled streams are the load-bearing decision.** Because every byte of a multi-byte char carries +its char's code, "previous char's class" is a plain `<< 1` and no rule anywhere does char-width +arithmetic. Give that up and every rule grows an `advance_char`. + +Three cost rules learned the hard way: + +1. **Reduce plane COUNT, not per-plane cost.** ~13 ops/plane on NEON is the floor; a dense 3-bit + code needing 3 planes beat one-hot class bits needing 7 (0.33 → 0.22 ns/B). +2. **Derive, don't extract.** `ascii = lead & last_byte` (a char that is its own first and last byte + is single-byte, i.e. ASCII) and `ascii_alpha = lm & ascii` (ASCII has no marks). Two `u64` ops + replaced two SIMD extractions. +3. **Gate rare masks per block.** The CJK range test is behind one `vmaxvq`, so Latin text pays ~8 + ops for it instead of ~70. + +Multi-byte-char predicates (the CJK range) belong in **vector space on `vext`-aligned b1/b2 → one +mask**, not as N separate byte-predicate bitstreams recombined with shifts. + +--- + +## 6. A grammar, end to end + +Worked on `"Hi, don't"` under GPT-2 byte-level. Classes: `l` letter, `o` other, `s` space, `a` +apostrophe (a sub-code of other, so it can be flagged). + +``` + index 0 1 2 3 4 5 6 7 8 + text H i , ␣ d o n ' t + class l l o s l l l a l + + l 1 1 . . 1 1 1 . 1 + other . . 1 . . . . 1 . + ws . . . 1 . . . . . + sp . . . 1 . . . . . + + l<<1 . 1 1 . . 1 1 1 . + sp<<1 . . . . 1 . . . . + other<<1 . . . 1 . . . . 1 + + l_start = l & !(l<<1) & !(sp<<1) 1 . . . . . . . 1 + o_start = other & !(other<<1) & !(sp<<1) . . 1 . . . . 1 . + ws_start = ws & !(ws<<1) . . . 1 . . . . . + steal = ws & lb & !(ws>>1) & !eof . . . 1 . . . . . + ───────────────────────────────── + starts 1 . 1 1 . . . 1 1 + flag = starts & apo . . . . . . . 1 . +``` + +Emit walk, with the scalar escape firing at bit 7: + +``` + [0,2) "Hi" [2,3) "," [3,7) " don" ⟨escape⟩ [7,9) "'t" ← bit 8 skipped +``` + +`" don"` shows `steal` doing its job: `\s+(?!\S)` hands the run's last whitespace char to whatever +follows, which is exactly the ` ?` prefix of the next alternative. One rule covers both. + +### Per-grammar checklist + +1. Dense code table `[u8; 64]`, `tag → code`, cont = 7. Merge every class the grammar treats + identically (deepseek: Letter+Mark share a code; gpt2: Mark is "other"). +2. `decode(p0,p1,p2,valid) -> Cls`. Mask with `valid` **only** the class whose code is 0 — past the + block end every plane reads 0. +3. Run starts, one line each. +4. Prefix / steal rules. +5. Bounded repetition, span clearing, absorption. +6. Carries + the edge peek. + +--- + +## 7. Escape hatches + +Not everything belongs in bit algebra. **Flag the bits, resolve them scalar-ly in the emit walk.** + +Contractions (`'s 't 're 've 'm 'll 'd`) are variable-length, case-optional, and outrank every other +alternative. In bit algebra that is miserable; as an escape it is 20 lines. `flag = starts & apo` +costs one word-test per block, so ordinary text pays nothing. + +Two bugs this cost, both worth a regression test: + +- The escape must not re-consume the start bit **at its own end** — the algebra usually has one + there (the letter run resumes) and consuming it emits an empty span. +- Contractions **chain**: `'re've`, `y'all'd've`. Loop until a match fails rather than handing + control back to the algebra, whose start bit you are about to skip. + +Same shape applies to any rare, awkward, variable-length rule. + +--- + +## 8. Rules and pitfalls + +**Prefer backward formulations.** State a rule as "my predecessor did X", never "my successor is +X". Forward needs `advance_char`, which silently drops the marker when the two chars straddle a +block edge — every cl100k failure was this. Backward needs only a shift carry. + +``` + ✗ forward o_prefix = o_start & next_is_letter ; suppress advance_char(o_prefix) + └─ marker leaves the block, lost + + ✓ backward osf = smear(o_start over its char's bytes) + l_start &= !(osf << 1) └─ one carry bit, always exact +``` + +**Carries at a block boundary.** One bit per stream, plus one per open scan: + +``` + block b-1 │ block b + ... 61 62 63 │ 0 1 2 ... + ▲ │ ▲ + └─ last byte's code ──────────┼───┘ seeds the fill and every `p1` + cont word ─────────────────┼───► to_lead underflow patch + open scans (aa/nl runs) ───┼───► carry-in at bit 0 + │ + peek 1 char forward ◄─────────┼─── makes every `n1` exact at the edge +``` + +Other pitfalls, each caught by the fuzzer: + +- Mask shifted-in zeros at u64 edges **before** inverting, or runs truncate silently. +- The first span needs special dispatch (`starts |= 1` on block 0). +- Cont-resolution needs prev-chunk context (`vext`), so starts land on leads. +- Bounded repetition needs a single-byte fast path — three `advance_char`s collapse to one `<< 3` + against `n & (n<<1) & (n<<2)`. Without it, dense digits run **0.58×** the FSM. +- One backward-in-time rule survives (`\s*[\r\n]+`'s "after the LAST newline"): patch the + already-written bitmap. The patch must NOT fire on a bit that a punct tail's `[\r\n]*` already + made a run start — the absorption cut the whitespace run, so a later newline cannot reach past it. + +--- + +## 9. Testing + +Non-negotiable, because every one of these caught a real bug: + +1. Byte-exactness vs the reference FSM on every input, for **every** grammar. +2. Exhaustive 3-char and 4-char sweeps over `[space, \n, a, 1, !, \0, 世, \u{a0}, \t, ']`. +3. Structured fuzz (~200k cases) over a pool hitting every arm: CJK, marks, ZWJ, control, multi-byte + whitespace, Other_Alphabetic symbols, contractions, punct+alpha. +4. **Every corpus re-sliced at 70 offsets.** This is what finds carry bugs; whole-file runs do not. + +## 12. Layout + +``` + src/ + lib.rs Blk, primitives, builders, emit, emit_contr + simd.rs NEON transpose (x86 sibling slots in beside it) + ops.rs the Parabix operator set + Carry + grammar/ + byte_level.rs cl100k.rs o200k.rs deepseek.rs class_runs.rs literal.rs + bin/ + verify.rs byte-exactness oracle vs atomsplit FSMs + bench.rs per-grammar throughput +``` + +Each grammar is independently testable against its reference FSM, which is what keeps the shared +layer honest. + +--- + +## 13. Open + +- **Does the library own classification?** Currently depends on `atomsplit::classify`. Fusing buys + ~0.07 ns/B on English but up to ~0.3 on CJK, where classify (0.57) now exceeds the split (0.24). + Design the block builder so classify can be fused later without changing the operator API. +- **x86 builder.** The scalar path covers it correctly but slowly; the `vpaddq` reduction has a + natural `movemask` equivalent. +- **o200k / tekken.** Case-refined letter runs; the case refinement is already in the Atom tags + (`0x10`/`0x20`), so it is a code-table change, not new machinery. diff --git a/tokenizers/bitsplit/src/classes.rs b/tokenizers/bitsplit/src/classes.rs new file mode 100644 index 000000000..61e954fe9 --- /dev/null +++ b/tokenizers/bitsplit/src/classes.rs @@ -0,0 +1,210 @@ +//! The class-run family: WhitespaceSplit / Punctuation / Digits / Whitespace / Bert, plus +//! CharDelimiterSplit. Not regex-shaped -- these cut where the atom class changes, so one shape +//! (`class_runs_into`) covers all of them. +//! +//! ponytail: still the scalar/NEON run extractor moved over from atomsplit, not a bitstream +//! program. A class-run boundary is `c & !(c << 1)`, so porting it onto the `Blk` streams would +//! delete `simd_classes.rs` outright -- worth doing, but it is not what makes these correct today. + +use crate::Span; +use crate::classify::{Atom, char_len, classify, in_mask, mask}; +use crate::run_end; + +/// No-`push` class-family pre-tokenizer core: writes spans into the preallocated `out` slice and returns +/// the count. ONE shape covers the whole class family via ``: +/// WhitespaceSplit `<{WS},0,0>` · Punctuation `<0,{PUNCT},0>` · Digits `<0,0,{NUMERIC}>` · +/// Whitespace `<{WS},0,{WORD}>` · Bert `<{WS},{PUNCT},0>`. +/// Class of a char: `DROP`→dropped, `ISOLATE`→own token, `KEEP_A`→run "A", else→run "B" (A/B cut apart). +/// TODO: find a better explanation +#[inline] +#[must_use] +pub fn class_runs_into( + text: &[u8], + tags: &[u8], + out: &mut [Span], +) -> usize { + debug_assert!(out.len() >= text.len() && tags.len() >= text.len()); + #[cfg(target_arch = "aarch64")] + { + crate::simd::classes::class_runs_neon::(text, tags, out) + } + #[cfg(all(target_arch = "wasm32", target_feature = "simd128"))] + { + crate::simd::classes::class_runs_wasm::(text, tags, out) + } + #[cfg(not(any( + target_arch = "aarch64", + all(target_arch = "wasm32", target_feature = "simd128") + )))] + { + emit_class_spans::(text, tags, out, 0, 0, 0, None) + } +} + +/// This is the most important function as it's the core of the scalar finite state machine. +/// It allows to emit class spans with different behaviours for tags we want to drop, tags we want +/// to isolate and tags we want to keep. Any other tags are assumed to be keept. +/// +/// This function is used as a fallback to the SIMD fast fsm. It is used for most pre tokenizers +/// but the unrolled regex, which have more complex variations that cannot be expressed with drop, +/// isolate, keep. These 3 generic parameters are u16 bitmap masks over the 16 classes we have and +/// define the behaviour. They are usally one of the [`crate::classify::mask`]. They allow dropping +/// words, isolating whitespace and keeping new line for example. +#[must_use] +#[inline] +pub fn emit_class_spans( + text: &[u8], + tags: &[u8], + out: &mut [Span], + mut write_index: usize, // in the out slice + mut text_pointer: usize, // in the text slice + segment_start: usize, // previous segment_start + segment_class: Option, // previous segment's class +) -> usize { + debug_assert!(out.len() >= text.len() && tags.len() >= text.len()); + let n = text.len(); + // Tie `tags.len() == text.len() == n` so the optimizer drops the interior `tags[i]` / `text[i]` + // bounds checks (callers guarantee len ≥ n; same trick as `cl100k`). Per-byte scanning already + // avoids checks via `run_end`'s unrolled `get_unchecked`; this covers the per-token accesses. + let tags = &tags[..n]; + let text = &text[..n]; + let other = !(DROP | ISOLATE | KEEP_A); // None of the above correspond to a continuation + if let Some(segment_class) = segment_class { + // this will usually be at the tail of a SIMD call. + text_pointer = run_end(tags, text_pointer, n, segment_class); // skip the whole drop run at once + if segment_class != DROP { + out[write_index] = Span { + start: segment_start as u32, + end: text_pointer as u32, + }; + if text_pointer == n { + return write_index + 1; + } + write_index += 1; + } + } + while text_pointer < n { + let t = tags[text_pointer]; + if t == Atom::Cont as u8 { + text_pointer += 1; + continue; + } + // classify the first char. + if in_mask(t, DROP) { + text_pointer = run_end(tags, text_pointer, n, DROP); // skip the whole drop run at once + } else if in_mask(t, ISOLATE) { + let s = text_pointer; + text_pointer += char_len(text[text_pointer]); + out[write_index] = Span { + start: s as u32, + end: text_pointer as u32, + }; // isolate: one char = one token + write_index += 1; + } else { + let s = text_pointer; + text_pointer = if in_mask(t, KEEP_A) { + run_end(tags, text_pointer, n, KEEP_A) + } else { + run_end(tags, text_pointer, n, other) + }; + out[write_index] = Span { + start: s as u32, + end: text_pointer as u32, + }; + write_index += 1; + } + } + write_index +} + +// ── per-tokenizer unrolled FSMs (one file each; shared helpers above via `use super::*`) ── +pub struct WhitespaceSplit; +impl WhitespaceSplit { + #[inline] + #[must_use] + pub fn pre_tokenize(&self, text: &[u8], tags: &mut [u8], out: &mut [Span]) -> usize { + classify(text, tags); + class_runs_into::<{ mask::WS }, 0, 0>(text, tags, out) + } +} + +/// `Punctuation` — isolate each punctuation char as its own token; non-punct grouped into runs. +pub struct Punctuation; +impl Punctuation { + #[inline] + #[must_use] + pub fn pre_tokenize(&self, text: &[u8], tags: &mut [u8], out: &mut [Span]) -> usize { + classify(text, tags); + class_runs_into::<0, { mask::PUNCT }, 0>(text, tags, out) + } +} + +/// `Digits` — cut numeric runs apart from non-numeric runs (contiguous), keeping both. +pub struct Digits; +impl Digits { + #[inline] + #[must_use] + pub fn pre_tokenize(&self, text: &[u8], tags: &mut [u8], out: &mut [Span]) -> usize { + classify(text, tags); + class_runs_into::<0, 0, { mask::NUMERIC }>(text, tags, out) + } +} + +/// `Whitespace` — the `\w+|[^\w\s]+` pre-tokenizer: drop whitespace, cut word runs from symbol runs. +pub struct Whitespace; +impl Whitespace { + #[inline] + #[must_use] + pub fn pre_tokenize(&self, text: &[u8], tags: &mut [u8], out: &mut [Span]) -> usize { + classify(text, tags); + class_runs_into::<{ mask::WS }, 0, { mask::WORD }>(text, tags, out) + } +} + +/// `Bert` — the BERT basic pre-tokenizer: drop whitespace, isolate punctuation, keep the rest as runs. +pub struct Bert; +impl Bert { + #[inline] + #[must_use] + pub fn pre_tokenize(&self, text: &[u8], tags: &mut [u8], out: &mut [Span]) -> usize { + classify(text, tags); + class_runs_into::<{ mask::WS }, { mask::PUNCT }, 0>(text, tags, out) + } +} + + + + +/// `Split(char, Removed)` — the only pre-tokenizer that keys on a *literal char* rather than an atom +/// class, so it scans bytes directly (no classify pass). UTF-8 is self-synchronizing, so the +/// delimiter's byte pattern only matches on char boundaries. +pub struct CharDelimiterSplit(pub char); +impl CharDelimiterSplit { + /// Split on the literal char (`Removed`): writes the gaps between delimiters into `out` + /// (len >= `text.len()`) and returns the count. Straight off [`crate::literal`] — no atom + /// classification, UTF-8 is self-synchronising so the char's bytes only match on boundaries. + #[inline] + #[must_use] + pub fn pre_tokenize(&self, text: &[u8], _tags: &mut [u8], out: &mut [Span]) -> usize { + debug_assert!(out.len() >= text.len()); + let mut buf = [0u8; 4]; + let delim = self.0.encode_utf8(&mut buf).as_bytes(); + let Ok(lit) = crate::literal::Literal::new(delim) else { + return 0; + }; + let (n, dl) = (text.len(), delim.len()); + let (mut start, mut w) = (0usize, 0usize); + for m in lit.matches(text) { + if m > start { + out[w] = Span::new(start as u32, m as u32); + w += 1; + } + start = m + dl; + } + if start < n { + out[w] = Span::new(start as u32, n as u32); + w += 1; + } + w + } +} diff --git a/tokenizers/atomsplit/src/atom_tables.rs b/tokenizers/bitsplit/src/classify/atom_tables.rs similarity index 99% rename from tokenizers/atomsplit/src/atom_tables.rs rename to tokenizers/bitsplit/src/classify/atom_tables.rs index 936f2f2dd..fcdba5840 100644 --- a/tokenizers/atomsplit/src/atom_tables.rs +++ b/tokenizers/bitsplit/src/classify/atom_tables.rs @@ -1,7 +1,7 @@ //! GENERATED by `bitmap_gen` — do NOT edit. Regenerate with `cargo run -p bitmap_gen`. //! Dense atom classify tables shared by the SIMD kernel (`vqtbl`) and the scalar reader //! `Tables::classify_char`. `bitmap_gen` self-validates all 1.1M codepoints. Spec §1/§7. -use crate::tables::Tables; +use super::tables::Tables; #[rustfmt::skip] static GROUP: [[[u8; 64]; 4]; 8] = [[[12,12,12,12,12,12,12,12,12,5,3,5,5,3,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,4,8,8,8,8,8,8,9,8,8,8,8,8,8,8,8,1,1,1,1,1,1,1,1,1,1,8,8,8,8,8,8],[8,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,8,8,8,8,7,8,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,8,8,8,8,12],[12,12,12,12,12,5,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,5,8,10,10,10,10,10,8,10,10,0,8,10,12,10,10,10,10,2,2,10,32,8,8,10,2,0,8,2,2,2,8],[16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,10,16,16,16,16,16,16,16,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,10,32,32,32,32,32,32,32,32]],[[16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,32,16,32,16,32,16,32,16],[32,16,32,16,32,16,32,16,32,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,16,32,16,32,16,32,32],[32,16,16,32,16,32,16,16,32,16,16,16,32,32,16,16,16,16,32,16,16,32,16,16,16,32,32,32,16,16,32,16,16,32,16,32,16,32,16,16,32,16,32,32,16,32,16,16,32,16,16,16,32,16,32,16,16,32,32,0,16,32,32,32],[0,0,0,0,16,16,32,16,16,32,16,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,32,16,16,32,16,32,16,16,16,32,16,32,16,32,16,32]],[[16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,32,32,32,32,32,32,16,16,32,16,16,32],[32,16,32,16,16,16,16,32,16,32,16,32,16,32,16,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32],[32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,0,0,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],[0,0,10,10,10,10,0,0,0,0,0,0,0,0,0,0,0,0,10,10,10,10,10,10,10,10,10,10,10,10,10,10,0,0,0,0,0,10,10,10,10,10,10,10,0,10,0,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10]],[[6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6],[6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,16,32,16,32,0,10,16,32,12,12,0,32,32,32,8,16],[12,12,12,12,10,10,16,8,16,16,16,12,16,12,16,16,32,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,12,16,16,16,16,16,16,16,16,16,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32],[32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,16,32,32,16,16,16,32,32,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,32,32,32,32,16,32,10,16,32,16,16,32,32,16,16,16]],[[16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32],[32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32],[16,32,10,6,6,6,6,6,6,6,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32],[16,16,32,16,32,16,32,16,32,16,32,16,32,16,32,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32]],[[16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,16,32,12,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16],[16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,12,12,0,8,8,8,8,8,8,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32,32],[32,32,32,32,32,32,32,32,32,8,8,12,12,10,10,10,12,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,8,6],[8,6,6,8,6,6,8,6,12,12,12,12,12,12,12,12,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,12,12,12,12,0,0,0,0,8,8,12,12,12,12,12,12,12,12,12,12,12]],[[12,12,12,12,12,12,10,10,10,8,8,10,8,8,10,10,6,6,6,6,6,6,6,6,6,6,6,8,12,8,8,8,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],[0,0,0,0,0,0,0,0,0,0,0,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,1,1,1,1,1,1,1,1,1,1,8,8,8,8,0,0,6,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,8,0,6,6,6,6,6,6,6,12,10,6,6,6,6,6,6,0,0,6,6,10,6,6,6,6,0,0,1,1,1,1,1,1,1,1,1,1,0,0,0,10,10,0]],[[8,8,8,8,8,8,8,8,8,8,8,8,8,8,12,12,0,6,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6],[6,6,6,6,6,6,6,6,6,6,6,12,12,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,6,6,6,6,6,6,6,6,6,6,6,0,12,12,12,12,12,12,12,12,12,12,12,12,12,12],[1,1,1,1,1,1,1,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,6,6,6,6,6,6,6,6,6,0,0,10,8,8,8,0,12,12,6,10,10]]]; diff --git a/tokenizers/atomsplit/src/simd_avx_classify.rs b/tokenizers/bitsplit/src/classify/avx.rs similarity index 99% rename from tokenizers/atomsplit/src/simd_avx_classify.rs rename to tokenizers/bitsplit/src/classify/avx.rs index 68cfef083..c3f595de8 100644 --- a/tokenizers/atomsplit/src/simd_avx_classify.rs +++ b/tokenizers/bitsplit/src/classify/avx.rs @@ -14,8 +14,8 @@ //! x86_64 (SSE4.1 and, if available, AVX-512 VBMI) hardware before trusting it. #![allow(unsafe_op_in_unsafe_fn)] -use crate::atom_tables::ATOM_TABLES; -use crate::classify::{Atom, CONT, MB, char_len, classify_scalar}; +use super::atom_tables::ATOM_TABLES; +use super::{Atom, CONT, MB, char_len, classify_scalar}; use core::arch::x86_64::*; const CJK_TAG: u8 = Atom::Letter as u8; diff --git a/tokenizers/atomsplit/src/classify.rs b/tokenizers/bitsplit/src/classify/mod.rs similarity index 92% rename from tokenizers/atomsplit/src/classify.rs rename to tokenizers/bitsplit/src/classify/mod.rs index 2e4e1f8dc..38fe1d1ed 100644 --- a/tokenizers/atomsplit/src/classify.rs +++ b/tokenizers/bitsplit/src/classify/mod.rs @@ -1,4 +1,16 @@ -use crate::atom_tables::ATOM_TABLES; +//! SIMD Unicode classification: one pass mapping every codepoint to a tiny "atom" alphabet, which +//! is what every grammar in `models/` consumes. The per-arch kernels are siblings; `atom_tables.rs` +//! is generated (`cargo run -p bitmap_gen`) and `tables.rs` is the layout it bakes into. +mod atom_tables; +#[cfg(target_arch = "x86_64")] +mod avx; +#[cfg(target_arch = "aarch64")] +mod neon; +mod tables; +#[cfg(all(target_arch = "wasm32", target_feature = "simd128"))] +mod wasm; + +use atom_tables::ATOM_TABLES; /// The per-codepoint "atom" categories or "tags" that are used by the finite state machine to emit /// spit boundaries. @@ -125,14 +137,14 @@ pub fn classify(text: &[u8], tags: &mut [u8]) { #[cfg(target_arch = "aarch64")] // SAFETY: `tags.len() >= text.len()` (asserted above); NEON vld1q/vst1q are alignment-free. unsafe { - crate::simd_classify::classify_neon(text, tags) + neon::classify_neon(text, tags) } #[cfg(target_arch = "x86_64")] - crate::simd_avx_classify::dispatch(text, tags); + avx::dispatch(text, tags); #[cfg(all(target_arch = "wasm32", target_feature = "simd128"))] // SAFETY: `tags.len() >= text.len()` (asserted above); wasm v128 load/store are alignment-free. unsafe { - crate::simd_wasm_classify::classify_wasm(text, tags) + wasm::classify_wasm(text, tags) } #[cfg(not(any( target_arch = "aarch64", diff --git a/tokenizers/atomsplit/src/simd_classify.rs b/tokenizers/bitsplit/src/classify/neon.rs similarity index 99% rename from tokenizers/atomsplit/src/simd_classify.rs rename to tokenizers/bitsplit/src/classify/neon.rs index ae9ee3939..2b008bf7d 100644 --- a/tokenizers/atomsplit/src/simd_classify.rs +++ b/tokenizers/bitsplit/src/classify/neon.rs @@ -1,5 +1,5 @@ -use crate::atom_tables::ATOM_TABLES; -use crate::classify::{Atom, CONT, MB}; +use super::atom_tables::ATOM_TABLES; +use super::{Atom, CONT, MB}; const CJK_TAG: u8 = Atom::Letter as u8; // ================================================================================================ @@ -98,7 +98,7 @@ unsafe fn any(mask: core::arch::aarch64::uint8x16_t) -> bool { /// TLDR removing the utf8 headers to get the unicode. fn decode(t: &[u8], i: usize) -> u32 { let b = t[i] as u32; - match super::classify::char_len(t[i]) { + match super::char_len(t[i]) { 1 => b, 2 => ((b & 0x1F) << 6) | (t[i + 1] as u32 & 0x3F), 3 => ((b & 0x0F) << 12) | ((t[i + 1] as u32 & 0x3F) << 6) | (t[i + 2] as u32 & 0x3F), @@ -156,7 +156,7 @@ fn decode(t: &[u8], i: usize) -> u32 { #[cfg(target_arch = "aarch64")] #[allow(unsafe_op_in_unsafe_fn, non_snake_case)] pub unsafe fn classify_neon(text: &[u8], tags: &mut [u8]) { - use super::classify::char_len; + use super::char_len; use core::arch::aarch64::*; let n = text.len(); let mut i = 0; diff --git a/tokenizers/atomsplit/src/tables.rs b/tokenizers/bitsplit/src/classify/tables.rs similarity index 100% rename from tokenizers/atomsplit/src/tables.rs rename to tokenizers/bitsplit/src/classify/tables.rs diff --git a/tokenizers/atomsplit/src/simd_wasm_classify.rs b/tokenizers/bitsplit/src/classify/wasm.rs similarity index 99% rename from tokenizers/atomsplit/src/simd_wasm_classify.rs rename to tokenizers/bitsplit/src/classify/wasm.rs index 12b7a6dc5..235175cee 100644 --- a/tokenizers/atomsplit/src/simd_wasm_classify.rs +++ b/tokenizers/bitsplit/src/classify/wasm.rs @@ -19,8 +19,8 @@ //! a SIMD128 wasm engine before trusting it. #![allow(unsafe_op_in_unsafe_fn)] -use crate::atom_tables::ATOM_TABLES; -use crate::classify::{Atom, CONT, MB, char_len}; +use super::atom_tables::ATOM_TABLES; +use super::{Atom, CONT, MB, char_len}; use core::arch::wasm32::*; const CJK_TAG: u8 = Atom::Letter as u8; diff --git a/tokenizers/bitsplit/src/han.rs b/tokenizers/bitsplit/src/han.rs new file mode 100644 index 000000000..f24e76604 --- /dev/null +++ b/tokenizers/bitsplit/src/han.rs @@ -0,0 +1,72 @@ +//! `\p{Han}` (Script=Han) for kimi-k2's leading `[\p{Han}]+` arm. +//! +//! - Han is NOT deepseek's `is_cjk_at` range: that one is Han U+4E00..9FA5 plus kana, this one is +//! the whole script and excludes kana. They are different predicates on purpose. +//! - Every range here is ≥ U+2E80, so a lead byte below 0xE2 exits before any decoding. + +/// Script=Han, sorted and disjoint. Extensions past Ext F are included because onig's `\p{Han}` +/// has them; the parity gate is what actually pins this table to the oracle's Unicode version. +const HAN: &[(u32, u32)] = &[ + (0x2E80, 0x2E99), + (0x2E9B, 0x2EF3), + (0x2F00, 0x2FD5), + (0x3005, 0x3005), + (0x3007, 0x3007), + (0x3021, 0x3029), + (0x3038, 0x303B), + (0x3400, 0x4DBF), + (0x4E00, 0x9FFF), + (0xF900, 0xFA6D), + (0xFA70, 0xFAD9), + (0x20000, 0x2A6DF), + (0x2A700, 0x2B739), + (0x2B740, 0x2B81D), + (0x2B820, 0x2CEA1), + (0x2CEB0, 0x2EBE0), + (0x2EBF0, 0x2EE5D), + (0x2F800, 0x2FA1D), + (0x30000, 0x3134A), + (0x31350, 0x323AF), +]; + +#[inline] +#[must_use] +pub(crate) fn is_han(cp: u32) -> bool { + HAN.binary_search_by(|&(lo, hi)| { + if cp < lo { + core::cmp::Ordering::Greater + } else if cp > hi { + core::cmp::Ordering::Less + } else { + core::cmp::Ordering::Equal + } + }) + .is_ok() +} + +/// Is the char whose lead byte is at `p` in Script=Han? `false` on a truncated tail. +#[inline] +#[must_use] +pub(crate) fn is_han_at(text: &[u8], p: usize) -> bool { + let b = text[p]; + let n = text.len(); + let cp = if (0xE2..=0xEF).contains(&b) { + if p + 2 >= n { + return false; + } + ((b as u32 & 0x0F) << 12) + | ((text[p + 1] as u32 & 0x3F) << 6) + | (text[p + 2] as u32 & 0x3F) + } else if (0xF0..=0xF4).contains(&b) { + if p + 3 >= n { + return false; + } + ((b as u32 & 0x07) << 18) + | ((text[p + 1] as u32 & 0x3F) << 12) + | ((text[p + 2] as u32 & 0x3F) << 6) + | (text[p + 3] as u32 & 0x3F) + } else { + return false; + }; + is_han(cp) +} diff --git a/tokenizers/bitsplit/src/lib.rs b/tokenizers/bitsplit/src/lib.rs new file mode 100644 index 000000000..8e229f2cb --- /dev/null +++ b/tokenizers/bitsplit/src/lib.rs @@ -0,0 +1,573 @@ +//! `bitsplit` — GPT-family pre-tokenization as a **bitstream program**, replacing the scalar FSMs. +//! +//! Follows *Interleaved Bitstream Execution for Multi-Pattern Regex Matching on GPUs* +//! (MICRO'25, doi 10.1145/3725843.3756052). The paper's two ideas that carry over to a CPU: +//! +//! 1. **Bit-parallel regex.** Compile the grammar into character-class *bitstreams* (one bit per +//! input byte) plus boolean ops and carry-propagating adds. 64 input bytes are decided per +//! 64-bit register op, branchlessly — the FSM's per-token unpredictable branch disappears. +//! 2. **Interleaved execution.** Do NOT run one loop per bitstream instruction over the whole +//! input (that is the "sequential" baseline the paper beats: every intermediate stream is +//! materialised and re-read). Instead fuse *all* instructions into ONE block-wise loop, so an +//! intermediate stream lives in a register for the ~100 ops it is needed and dies there. The +//! only thing that reaches memory is the final `starts` bitmap (n/8 bytes). +//! +//! The paper's third contribution (dependency-aware thread-data mapping) is a GPU concern — it +//! resolves cross-block dependencies by recomputing them on other SMs. On one core the blocks are +//! visited in order, so those dependencies are carried in a handful of scalar registers and, for +//! the one genuinely backward-in-time rule, patched into the already-written bitmap. +//! +//! We do not re-derive character classes: [`classify`] already emits one `Atom` tag per byte. The +//! builder folds those 16 atoms into a grammar-specific **dense 3-bit code** and extracts 3 +//! bit-planes of it; every class stream is then a 2–3 op boolean function of the planes. +//! +//! Grammars: [`bitsplit_deepseek`], [`bitsplit_byte_level`] (GPT-2), [`bitsplit_cl100k`]. All three +//! byte-exact with the oniguruma oracle over a block-phase sweep — see `tests/parity.rs`. + +pub mod classes; +pub mod classify; +mod han; +pub mod literal; +pub mod models; +pub mod regexes; +mod simd; + +pub use models::deepseek::bitsplit_deepseek; +pub use models::cl100k::{bitsplit_cl100k, bitsplit_qwen}; +pub use models::gpt2::bitsplit_byte_level; +pub use models::kimi::bitsplit_kimi; +pub use models::o200k::bitsplit_o200k; +pub use models::tekken::bitsplit_tekken; + +/// A token span: byte offsets `[start, end)` into the input. `#[repr(C)]` so the output buffer has a +/// stable `[start, end]` layout — the pipeline reuses it with zero conversion, and it can be +/// reinterpreted as bytes / handed across the crate boundary. +#[repr(C)] +#[derive(Clone, Copy, PartialEq, Eq, Debug, Default, Hash, PartialOrd, Ord)] +pub struct Span { + pub start: u32, + pub end: u32, +} + +impl Span { + #[inline] + #[must_use] + pub const fn new(start: u32, end: u32) -> Self { + Self { start, end } + } + + /// `[start, end)` as a `usize` range — for slicing the input text. + #[inline] + #[must_use] + pub fn range(self) -> core::ops::Range { + self.start as usize..self.end as usize + } +} + +impl PartialEq<(u32, u32)> for Span { + fn eq(&self, other: &(u32, u32)) -> bool { + self.start == other.0 && self.end == other.1 + } +} + +/// Whether this target builds its bitstreams with SIMD. Without it the builder is the portable +/// byte-at-a-time reference, which is slower than the FSM this replaces -- so a caller should keep +/// its FSM path rather than route here. On x86 the kernel needs SSSE3, so this is a runtime check. +#[must_use] +pub fn fast_builder() -> bool { + #[cfg(target_arch = "aarch64")] + { + true + } + #[cfg(target_arch = "x86_64")] + { + has_ssse3() + } + #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))] + { + false + } +} + +#[cfg(target_arch = "x86_64")] +pub(crate) fn has_ssse3() -> bool { + use std::sync::atomic::{AtomicU8, Ordering}; + static CACHED: AtomicU8 = AtomicU8::new(0); + match CACHED.load(Ordering::Relaxed) { + 0 => { + let yes = std::arch::is_x86_feature_detected!("ssse3"); + CACHED.store(1 + u8::from(yes), Ordering::Relaxed); + yes + } + n => n == 2, + } +} + +pub(crate) const CONT: u8 = 15; // Atom::Cont +pub(crate) const CODE_CONT: u8 = 7; // every grammar's dense code for a continuation byte + +/// What a grammar wants in the text-derived `Blk::aux` stream. Tag-derived classes go through the +/// LUT; these three need the raw bytes, so they are the one thing the builder reads `text` for. +pub(crate) const AUX_NONE: u8 = 0; +pub(crate) const AUX_CJK: u8 = 1; // deepseek Split-2: Han U+4E00..9FA5 ∪ kana U+3040..30FF +pub(crate) const AUX_SLASH: u8 = 2; // o200k rule 4's `[\r\n/]*` tail +pub(crate) const AUX_HAN: u8 = 3; // kimi-k2's leading `[\p{Han}]+` arm + +/// The bitstreams for one 64-byte block. `p0`..`p3` are the bit-planes of the **filled** dense +/// code — filled meaning a multi-byte char sets its bits on *all* of its bytes, so "previous char's +/// class" is a plain `<< 1` and no rule does char-width arithmetic. `cont` is the one un-filled +/// stream (it defines the fill); `aux` is text-derived and only built when a grammar asks for it. +/// +/// `p3` exists because o200k needs 9 tag classes (U/L/C/N/NL/SP/WSO/OTHER + cont) and code 7 is +/// reserved — the SIMD kernels find continuation lanes by testing `lut[tag] == 7`, so "other" +/// cannot be the leftover code. It is const-gated off for the 3-plane grammars. +#[derive(Default, Clone, Copy)] +pub(crate) struct Blk { + pub cont: u64, + pub p0: u64, + pub p1: u64, + pub p2: u64, + pub p3: u64, + pub aux: u64, +} + +/// deepseek Split-2's isolated range: Han U+4E00..9FA5 ∪ Hiragana/Katakana U+3040..30FF (all +/// 3-byte, leads E3..E9). Same predicate as `fsm_deepseek`'s `ds_is_cjk_at`. +#[inline] +pub(crate) fn is_cjk_at(text: &[u8], p: usize) -> bool { + let b = text[p]; + if !(0xE3..=0xE9).contains(&b) || p + 2 >= text.len() { + return false; + } + let cp = ((b as u32 & 0x0F) << 12) + | ((text[p + 1] as u32 & 0x3F) << 6) + | (text[p + 2] as u32 & 0x3F); + (0x4E00..=0x9FA5).contains(&cp) || (0x3040..=0x30FF).contains(&cp) +} + +/// The `AUX` predicate at a lead byte. +#[inline] +pub(crate) fn aux_at(text: &[u8], p: usize) -> bool { + match AUX { + AUX_CJK => is_cjk_at(text, p), + AUX_SLASH => text[p] == b'/', + AUX_HAN => crate::han::is_han_at(text, p), + _ => false, + } +} + +/// Build one block. Full blocks go through the NEON kernel; the ragged tail (and every other +/// target) uses the portable byte-at-a-time reference. Both produce the identical `Blk`. +/// `lut` is the grammar's tag → dense code table. +#[inline] +pub(crate) fn build_block( + text: &[u8], + tags: &[u8], + base: usize, + len: usize, + lut: &[u8; 64], + cur_code: u8, + cur_aux: bool, +) -> (Blk, u8) { + #[cfg(target_arch = "aarch64")] + if len == 64 { + // SAFETY: `len == 64` means `base + 64 <= text.len() == tags.len()`. + return unsafe { + crate::simd::neon::build64::(text, tags, base, lut, cur_code, cur_aux) + }; + } + #[cfg(target_arch = "x86_64")] + if len == 64 && has_ssse3() { + // SAFETY: `len == 64` bounds both reads, and SSSE3 is checked above. + return unsafe { + crate::simd::x86::build64::(text, tags, base, lut, cur_code, cur_aux) + }; + } + build_block_scalar::(text, tags, base, len, lut, cur_code, cur_aux) +} + +/// Portable reference builder: one byte at a time. +pub(crate) fn build_block_scalar( + text: &[u8], + tags: &[u8], + base: usize, + len: usize, + lut: &[u8; 64], + mut cur_code: u8, + mut cur_aux: bool, +) -> (Blk, u8) { + let mut b = Blk::default(); + for i in 0..len { + let p = base + i; + let bit = 1u64 << i; + if tags[p] == CONT { + b.cont |= bit; + } else { + cur_code = lut[tags[p] as usize]; + cur_aux = AUX != AUX_NONE && aux_at::(text, p); + } + b.p0 |= bit * u64::from(cur_code & 1 != 0); + b.p1 |= bit * u64::from(cur_code & 2 != 0); + b.p2 |= bit * u64::from(cur_code & 4 != 0); + if P3 { + b.p3 |= bit * u64::from(cur_code & 8 != 0); + } + b.aux |= bit * u64::from(cur_aux); + } + (b, cur_code) +} + +// ── bitstream primitives ──────────────────────────────────────────────────────────────────────── + +/// `ScanThru`: move every marker in `m` forward past the run of `c` it sits in, landing on the +/// first position not in `c`. The classic Parabix carry-propagation trick — one add. Done in `u128` +/// so a run reaching bit 63 puts its landing bit at 64 instead of vanishing: that bit *is* the +/// block's carry-out, and `e - m` still yields the right in-block span. +#[inline] +pub(crate) const fn scanthru(m: u128, c: u128) -> u128 { + m.wrapping_add(c) & !c +} + +/// Advance markers by ONE CHAR: shift one byte, then scan through the continuation bytes. Markers +/// that leave the block are dropped — their state is reconstructed from the scalar carry instead. +#[inline] +pub(crate) const fn adv(m: u64, cont: u64) -> u64 { + scanthru((m as u128) << 1, cont as u128) as u64 +} + +/// Move each bit back to the lead byte of its char. Rules that look at the *next* char are stated +/// at a char's last byte; token starts must sit on leads. ≤3 steps (UTF-8 chars are ≤4 bytes). +/// Returns `(in-block, patch for the previous word)` — a char straddling the block boundary lands +/// its lead in the word we already wrote. +#[inline] +pub(crate) fn to_lead(x: u64, cont: u64, prev_cont: u64) -> (u64, u64) { + let c = ((cont as u128) << 64) | prev_cont as u128; + let mut y = (x as u128) << 64; + for _ in 0..3 { + let m = y & c; + if m == 0 { + break; + } + y = (y & !c) | (m >> 1); + } + ((y >> 64) as u64, y as u64) +} + +/// In each run of `c`, fill from the run start through the LAST marker of `m` (`m ⊆ c`). Used on +/// reversed streams, where it answers "is there a newline at-or-after me, still inside this +/// whitespace run?". `(end - m) | m` rather than `end - m`: the latter only spans from the *first* +/// marker when a run holds several. +#[inline] +pub(crate) fn fill_to_last(m: u64, c: u64) -> u64 { + if m == 0 { + return 0; + } + let (m128, c128) = (m as u128, c as u128); + (scanthru(m128, c128).wrapping_sub(m128) | m128) as u64 +} + +/// Run of `x` that starts at bit 0. +#[inline] +pub(crate) fn lead_run(x: u64, valid: u64) -> u64 { + let z = !x & valid; + if z == 0 { + valid + } else { + (z & z.wrapping_neg()) - 1 + } +} + +/// Run of `x` that ends at bit `len - 1` (0 if that bit is clear). +#[inline] +pub(crate) fn trail_run(x: u64, valid: u64, len: usize) -> u64 { + if x & (1u64 << (len - 1)) == 0 { + return 0; + } + let z = !x & valid; + if z == 0 { + valid + } else { + valid & !((1u64 << (64 - z.leading_zeros())) - 1) + } +} + +/// `\p{N}{1,cap}` — a group boundary every `cap` chars from the run start. The one non-local rule +/// in every tiktoken-family grammar, so it lives here rather than three times over. +/// +/// `cap == 0` or `>= 64` means an unbounded `\p{N}+`: the run is one token and there is nothing to +/// do (and the shifts below would be UB). `dig_since` resumes a run that crossed the block edge — +/// re-masking with `n` at every hop matters, because the carry only says the byte *at* the edge was +/// a digit, not that the run survived it. +#[inline] +pub(crate) fn digit_groups( + cap: usize, + n: u64, + lead: u64, + cont: u64, + prev_is_digit: bool, + dig_run: bool, + dig_since: u32, +) -> u64 { + let mut m = n & lead & !((n << 1) | u64::from(prev_is_digit)); + if cap == 0 || cap >= 64 { + return m; // an unbounded `\p{N}+`: the run is one token + } + let capu = cap as u32; + if dig_run && prev_is_digit { + let mut s = lead & lead.wrapping_neg() & n; // first lead of the block + for _ in 0..((capu - dig_since % capu) % capu) { + s = adv(s, cont) & n & lead; + } + m |= s; + } + let mut groups = m; + if n & cont == 0 { + // Fast path: every digit here is single-byte, so "CAP chars on" is a plain shift against a + // mask asking that the skipped positions were digits too — what the `adv` chain checks. + let mut nk = n; + let mut k = 1; + while k < cap { + nk &= n << k; + k += 1; + } + while m != 0 { + m = (m << cap) & nk; + groups |= m; + } + } else { + while m != 0 { + let mut e = m; + for _ in 0..cap { + e = adv(e, cont) & n & lead; + } + if e == 0 { + break; + } + groups |= e; + m = e; + } + } + groups +} + +// ── scalar helpers, shared by the grammars that need an escape ────────────────────────────── +use crate::classify::{Atom, in_mask, mask}; + + +pub(crate) const LET: u8 = 0x00; // coarse Atom::Letter (low nibble) +pub(crate) const MRK: u8 = 0x06; // coarse Atom::Mark +pub(crate) const ASM: u8 = 0x16; // AlphaSymMark — coarse Mark, categorically \p{S} +pub(crate) const ZWJ: u8 = 0x26; // ZWJ/ZWNJ — coarse Mark, categorically \p{Cf} +pub(crate) const NW: u8 = 0x01; +pub(crate) const NO: u8 = 0x02; +pub(crate) const NLN: u8 = 0x03; +pub(crate) const SPC: u8 = 0x04; +pub(crate) const WSO: u8 = 0x05; + +/// A real `[\p{L}\p{M}]` member. ALPHA_SYM and ZWJ are coarse `Mark` but neither `\p{L}` nor +/// `\p{M}`, so they are NOT letters — that is what keeps them on the rule-4 path. +#[inline] +pub(crate) fn member(t: u8) -> bool { + let c = t & 0x0F; + c == LET || (c == MRK && t != ASM && t != ZWJ) +} + +#[inline] +pub(crate) fn run_end(tags: &[u8], mut i: usize, end: usize, m: u16) -> usize { + let m = m | Atom::Cont.bit(); + while i < end && in_mask(tags[i], m) { + i += 1; + } + i +} + +/// `\s*[\r\n]+ | \s+(?!\S) | \s+`: through the last `\r\n` if any, else the whole run at EOF, else +/// give the final ws char back to whatever follows. +#[inline] +pub(crate) fn ws_tail(text: &[u8], tags: &[u8], i: usize, end: usize) -> usize { + let re = run_end(tags, i, end, mask::WS); + if let Some(r) = text[i..re].iter().rposition(|&x| x == 0x0A || x == 0x0D) { + i + r + 1 + } else if re == end { + re + } else { + let mut last = re - 1; + while last > i && text[last] & 0xC0 == 0x80 { + last -= 1; + } + if last > i { last } else { re } + } +} + +/// One letter sub-token from `p` within `[.., re)`: alt-1 `[UC]*[LC]+` (tried first), else alt-2 +/// `[UC]+[LC]*` for an all-upper run. `[UC]*` gives back to the last C so `[LC]+` can take one. +#[inline] +pub(crate) fn letter_match(tags: &[u8], p: usize, re: usize) -> usize { + let (mut q, mut last_c) = (p, usize::MAX); + while q < re && tags[q] != 0x20 { + if tags[q] != CONT && tags[q] != 0x10 { + last_c = q; + } + q += 1; + } + if q < re { + let mut e = q; + while e < re && tags[e] != 0x10 { + e += 1; + } + return e; + } + if last_c == usize::MAX { + return re; // all upper → alt-2 takes the whole run + } + let mut e = last_c; + while e < re && tags[e] != 0x10 { + e += 1; + } + e +} + + +// ── emit ──────────────────────────────────────────────────────────────────────────────────────── + +/// `starts` bitmap → spans. Each set bit closes the previous token and opens the next; `tzcnt` +/// walks them at ~3 ops per token. +pub(crate) fn emit(starts: &[u64], nblk: usize, n: usize, out: &mut [Span]) -> usize { + let (mut w, mut open) = (0usize, u32::MAX); + for (bi, &word) in starts.iter().enumerate().take(nblk) { + let mut m = word; + while m != 0 { + let pos = (bi * 64 + m.trailing_zeros() as usize) as u32; + if open != u32::MAX { + out[w] = Span::new(open, pos); + w += 1; + } + open = pos; + m &= m - 1; + } + } + if open != u32::MAX { + out[w] = Span::new(open, n as u32); + w += 1; + } + w +} + +/// Emit with a **scalar escape at flagged bits**: contractions (`'s 't 're 've 'm 'll 'd`) are +/// variable-length, case-optional and outrank every other alternative, which makes them miserable +/// in bit algebra and trivial here. `flag` marks the apostrophes that open a token; a block with +/// none takes the plain loop, so the escape costs one test per block on ordinary text. +/// +/// A matched contraction overrides the algebra outright: it emits its own span and skips every +/// start bit inside it, which is how the letter alternative loses the tie (`'sx` → `'s`, `x`). +pub(crate) fn emit_contr( + text: &[u8], + starts: &[u64], + flag: &[u64], + nblk: usize, + n: usize, + ci: bool, + out: &mut [Span], +) -> usize { + let (mut w, mut open, mut skip) = (0usize, u32::MAX, 0usize); + for bi in 0..nblk { + let mut m = starts[bi]; + let f = flag[bi]; + if f == 0 && skip <= bi * 64 { + while m != 0 { + let pos = (bi * 64 + m.trailing_zeros() as usize) as u32; + if open != u32::MAX { + out[w] = Span::new(open, pos); + w += 1; + } + open = pos; + m &= m - 1; + } + continue; + } + while m != 0 { + let j = m.trailing_zeros() as usize; + let pos = bi * 64 + j; + m &= m - 1; + if pos < skip { + continue; + } + if open != u32::MAX { + out[w] = Span::new(open, pos as u32); + w += 1; + } + open = pos as u32; + if f >> j & 1 != 0 { + // Contractions chain (`'re've`, `y'all'd've`): the char after one is a token start + // in its own right, so keep matching until one fails rather than handing control + // back to the bit algebra — whose start bit there we are about to skip. + let mut p = pos; + while p < n { + let l = contr_len(text, p, ci); + if l == 0 { + break; + } + out[w] = Span::new(p as u32, (p + l) as u32); + w += 1; + p += l; + } + if p > pos { + open = if p < n { p as u32 } else { u32::MAX }; + // `+ 1`: the algebra usually also has a start bit exactly at `p` (the letter + // run resumes there) and we have just opened it — consuming it again would + // emit an empty span. `contr_len(p) == 0` here, so nothing is lost. + skip = p + 1; + } + } + } + } + if open != u32::MAX { + out[w] = Span::new(open, n as u32); + w += 1; + } + w +} + + +/// Byte length of the contraction at `i` (2 or 3), or 0. `ci` picks cl100k/o200k's `(?i:)` form +/// over GPT-2's case-sensitive one. +#[inline] +pub(crate) fn contr_len(text: &[u8], i: usize, ci: bool) -> usize { + let n = text.len(); + if i + 1 >= n || text[i] != b'\'' || text[i + 1] >= 0x80 { + return 0; + } + let c1 = if ci { text[i + 1] | 0x20 } else { text[i + 1] }; + match c1 { + b's' | b't' | b'm' | b'd' => 2, + b'r' | b'v' | b'l' if i + 2 < n && text[i + 2] < 0x80 => { + let c2 = if ci { text[i + 2] | 0x20 } else { text[i + 2] }; + usize::from((matches!(c1, b'r' | b'v') && c2 == b'e') || (c1 == b'l' && c2 == b'l')) * 3 + } + _ => 0, + } +} + +/// Builder-only cost probe: runs just the byte→bitstream transpose over every block and folds the +/// streams so nothing is optimised away. The difference against a full grammar is what the +/// bitstream program itself costs. +#[doc(hidden)] +#[must_use] +pub fn build_only(text: &[u8], tags: &[u8]) -> u64 { + let (mut acc, mut code, mut cjk) = (0u64, CODE_CONT, false); + for base in (0..text.len()).step_by(64) { + let len = (text.len() - base).min(64); + let (b, c) = build_block::<{ AUX_CJK }, false>(text, tags, base, len, &models::deepseek::LUT, code, cjk); + code = c; + cjk = b.aux >> (len - 1) & 1 != 0; + acc ^= b.cont ^ b.p0 ^ b.p1 ^ b.p2 ^ b.aux; + } + acc +} + +/// Convenience wrapper: classify + deepseek bitsplit over caller-owned scratch. +#[must_use] +pub fn pre_tokenize(text: &[u8], tags: &mut [u8], starts: &mut [u64], out: &mut [Span]) -> usize { + classify::classify(text, tags); + bitsplit_deepseek(text, tags, starts, out) +} diff --git a/tokenizers/bitsplit/src/literal.rs b/tokenizers/bitsplit/src/literal.rs new file mode 100644 index 000000000..763e96d82 --- /dev/null +++ b/tokenizers/bitsplit/src/literal.rs @@ -0,0 +1,209 @@ +//! Splitting on an exact string — the Metaspace `▁` delimiter (llama2 / gemma), `CharDelimiterSplit`, +//! and any `Split` whose pattern is a plain string rather than a regex. +//! +//! - No atom classification here: a literal does not care what class a byte is, so the tag pass is +//! skipped entirely and this works off the raw bytes. +//! - Same bitstream shape as the grammars: build a `u64` match bitmap per 64-byte block, then walk +//! it with `trailing_zeros`. `shift-and` over the needle's bytes, which is the classic Baeza-Yates +//! bitap specialised to an exact string — one AND per needle byte per 64 input bytes. +//! - Matches never overlap, so `"aa"` is found once in `"aaa"` — the same matches a regex reports. + +use std::fmt; + +/// The pattern handed to [`Literal::new`] was empty. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct EmptyPattern; + +impl fmt::Display for EmptyPattern { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str("an empty pattern matches everywhere") + } +} + +impl std::error::Error for EmptyPattern {} + +/// A literal string to split on. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Literal { + needle: Vec, +} + +impl Literal { + /// # Errors + /// If `pattern` is empty, which would match everywhere. + pub fn new(pattern: &[u8]) -> Result { + if pattern.is_empty() { + return Err(EmptyPattern); + } + Ok(Self { + needle: pattern.to_vec(), + }) + } + + /// The string being searched for. + #[must_use] + pub fn pattern(&self) -> &[u8] { + &self.needle + } + + /// Byte offset of every match, left to right, non-overlapping. + pub fn matches<'t>(&'t self, text: &'t [u8]) -> Matches<'t> { + Matches { + needle: &self.needle, + text, + base: 0, + word: 0, + next_block: 0, + min: 0, + } + } +} + +/// Match positions of `needle` in `text[base..base + 64]`, as a bitmap. A match at bit `i` means +/// `text[base + i..][..needle.len()] == needle`, so the last `needle.len() - 1` bits of the block +/// depend on bytes past it — the shift-and simply reads them from `text`, which is why this takes +/// the whole slice rather than a block. +fn match_bits(text: &[u8], needle: &[u8], base: usize) -> u64 { + let n = text.len(); + // positions in this block that could still fit the needle + let span = n.saturating_sub(base).min(64); + let last = n.saturating_sub(needle.len() - 1); // one past the last possible match start + let room = last.saturating_sub(base).min(span); + if room == 0 { + return 0; + } + let valid = if room == 64 { !0u64 } else { (1u64 << room) - 1 }; + // First byte in SIMD -- this is the whole cost, everything after it runs on the survivors. + let mut m = eq_bits(text, base, needle[0]) & valid; + for (k, &b) in needle.iter().enumerate().skip(1) { + let mut eq = 0u64; + let mut bits = m; + while bits != 0 { + let i = bits.trailing_zeros() as usize; + bits &= bits - 1; + if text[base + i + k] == b { + eq |= 1u64 << i; + } + } + m &= eq; + if m == 0 { + return 0; + } + } + m +} + +/// `text[base..base + 64] == b`, as a bitmap. Falls back to a byte loop on a ragged tail (and on +/// targets with no kernel), which is correct everywhere and only pays on the last block. +#[inline] +fn eq_bits(text: &[u8], base: usize, b: u8) -> u64 { + if base + 64 <= text.len() { + #[cfg(target_arch = "aarch64")] + // SAFETY: bounds checked directly above. + return unsafe { crate::simd::neon::eq64(text, base, b) }; + #[cfg(target_arch = "x86_64")] + if crate::has_ssse3() { + // SAFETY: bounds checked above, SSSE3 checked here. + return unsafe { crate::simd::x86::eq64(text, base, b) }; + } + } + let mut m = 0u64; + for i in 0..(text.len() - base).min(64) { + if text[base + i] == b { + m |= 1u64 << i; + } + } + m +} + +/// Iterator over non-overlapping match offsets. +pub struct Matches<'t> { + needle: &'t [u8], + text: &'t [u8], + base: usize, + word: u64, + next_block: usize, + min: usize, // matches before this would overlap the one just returned +} + +impl Iterator for Matches<'_> { + type Item = usize; + + fn next(&mut self) -> Option { + loop { + while self.word != 0 { + let i = self.word.trailing_zeros() as usize; + self.word &= self.word - 1; + let at = self.base + i; + if at >= self.min { + self.min = at + self.needle.len(); + return Some(at); + } + } + if self.next_block >= self.text.len() { + return None; + } + self.base = self.next_block; + self.next_block += 64; + self.word = match_bits(self.text, self.needle, self.base); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn find(pat: &str, hay: &str) -> Vec { + Literal::new(pat.as_bytes()) + .unwrap() + .matches(hay.as_bytes()) + .collect() + } + + /// The bitmap has to agree with a naive scan on overlap, block edges and multi-byte needles. + #[test] + fn matches_a_naive_scan() { + let cases: &[(&str, String)] = &[ + ("a", "aaaa".into()), + ("aa", "aaaaa".into()), + ("ab", "abab".into()), + ("▁", "▁a▁bb▁".into()), + ("▁▁", "▁▁▁▁".into()), + ("xyz", "xyxyz".into()), + ("a", "b".repeat(200)), + // needle straddling every 64-byte block edge + ("▁w", format!("{}▁w{}", "q".repeat(63), "z".repeat(70))), + ("abc", format!("{}abc", "a".repeat(62))), + ]; + for (pat, hay) in cases { + let (p, h) = (pat.as_bytes(), hay.as_bytes()); + let mut want = Vec::new(); + let mut i = 0; + while i + p.len() <= h.len() { + if &h[i..i + p.len()] == p { + want.push(i); + i += p.len(); // non-overlapping, like a regex + } else { + i += 1; + } + } + assert_eq!(find(pat, hay), want, "{pat:?} in {hay:?}"); + } + } + + /// Every offset of the needle across a 3-block text, so it crosses each edge in every phase. + #[test] + fn finds_the_needle_at_every_offset() { + let needle = "▁w"; + for off in 0..180usize { + let hay = format!("{}{needle}{}", "q".repeat(off), "q".repeat(180 - off)); + assert_eq!(find(needle, &hay), vec![off], "offset {off}"); + } + } + + #[test] + fn an_empty_pattern_is_rejected() { + assert_eq!(Literal::new(b""), Err(EmptyPattern)); + } +} diff --git a/tokenizers/bitsplit/src/models/cl100k.rs b/tokenizers/bitsplit/src/models/cl100k.rs new file mode 100644 index 000000000..9d16e1b2c --- /dev/null +++ b/tokenizers/bitsplit/src/models/cl100k.rs @@ -0,0 +1,269 @@ +//! **cl100k_base** (tiktoken) / Llama-3 -- and byte-for-byte the regex **GLM-4.6** ships: +//! `(?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+` +//! +//! Rule 3's digit cap is the family's only knob, so it is a plain argument: 3 = cl100k / Llama-3 / +//! GLM, 1 = **Qwen** (`\p{N}`), 0 or >= 64 = an unbounded `\p{N}+`. + +use crate::{ + AUX_NONE, CODE_CONT, CONT, Span, build_block, digit_groups, emit_contr, fill_to_last, + lead_run, scanthru, to_lead, trail_run, +}; + +/// Atom tag → dense 3-bit code, shared by both grammars. Unlike deepseek's table, `Mark` is NOT a +/// letter here (`\p{L}` excludes it, so it belongs to the "other" class), and `Apostrophe` gets its +/// own code so the contraction escape can be flagged with one AND — it is still "other" for every +/// run rule, which `decode` restores. +const LUT: [u8; 64] = { + let mut t = [2u8; 64]; // other = [^\s\p{L}\p{N}] + t[0x00] = 0; + t[0x10] = 0; + t[0x20] = 0; // Letter (+ case refinements) + t[0x01] = 1; + t[0x02] = 1; // \p{N} = Nd ∪ Nl ∪ No + t[0x03] = 3; // Newline + t[0x04] = 4; // Space + t[0x05] = 5; // WsOther + t[0x09] = 6; // Apostrophe + t[0x0F] = CODE_CONT; + t +}; + +/// One block's class streams. `other` folds the apostrophe code back in; only `l` needs the +/// `valid` mask, since past the block end every plane reads 0, i.e. code 0. +struct Cls { + l: u64, + n: u64, + other: u64, + nl: u64, + sp: u64, + ws: u64, + apo: u64, +} + +#[inline] +fn decode(p0: u64, p1: u64, p2: u64, valid: u64) -> Cls { + let a = !p2 & !p1; + let nl = !p2 & p1 & p0; + Cls { + l: a & !p0 & valid, + n: a & p0, + other: p1 & !p0 & valid, // codes 2 and 6 — "other" and the apostrophe + nl, + sp: p2 & !p1 & !p0, + ws: (p2 & !p1) | nl, // Space ∪ WsOther ∪ Newline + apo: p2 & p1 & !p0, + } +} + +/// Class bits of a dense code, for the block-edge carry. +const C_L: u8 = 1; +const C_N: u8 = 2; +const C_O: u8 = 4; +const C_NL: u8 = 8; +const C_SP: u8 = 16; +const C_WS: u8 = 32; + +#[inline] +const fn code_bits(code: u8) -> u8 { + match code { + 0 => C_L, + 1 => C_N, + 2 | 6 => C_O, + 3 => C_NL | C_WS, + 4 => C_SP | C_WS, + 5 => C_WS, + _ => 0, // cont never reaches an edge (the fill resolves it) + } +} + +/// cl100k_base / Llama-3 / GLM-4.6 — rule 3 is `\p{N}{1,3}`. +#[must_use] +pub fn bitsplit_cl100k( + text: &[u8], + tags: &[u8], + starts: &mut [u64], + flag: &mut [u64], + out: &mut [Span], +) -> usize { + cl100k(text, tags, starts, flag, out, 3) +} + +/// Qwen2 / Qwen3 — cl100k character-for-character except rule 3 is a bare `\p{N}`. +#[must_use] +pub fn bitsplit_qwen( + text: &[u8], + tags: &[u8], + starts: &mut [u64], + flag: &mut [u64], + out: &mut [Span], +) -> usize { + cl100k(text, tags, starts, flag, out, 1) +} + +fn cl100k( + text: &[u8], + tags: &[u8], + starts: &mut [u64], + flag: &mut [u64], + out: &mut [Span], + digit_cap: usize, +) -> usize { + let ntext = text.len(); + if ntext == 0 { + return 0; + } + let nblk = ntext.div_ceil(64); + assert!( + tags.len() >= ntext && starts.len() >= nblk && flag.len() >= nblk && out.len() >= ntext + ); + let (mut code, mut prev_cont) = (CODE_CONT, 0u64); + let (mut nl_run, mut dig_run, mut dig_since) = (false, false, 0u32); + let mut prev_osf = false; // previous block's last byte belonged to a token-opening "other" char + let mut anl: Option = None; + + for bi in 0..nblk { + let base = bi * 64; + let len = (ntext - base).min(64); + let valid = if len == 64 { !0u64 } else { (1u64 << len) - 1 }; + let last_blk = base + len == ntext; + + let (b, last_code) = build_block::<{ AUX_NONE }, false>(text, tags, base, len, &LUT, code, false); + let c = decode(b.p0, b.p1, b.p2, valid); + + let pb = if base == 0 { 0 } else { code_bits(code) }; + let (nb, nb_lead) = if last_blk { + (0u8, true) + } else { + let q = base + len; + let is_lead = tags[q] != CONT; + ( + if is_lead { + code_bits(LUT[tags[q] as usize]) + } else { + code_bits(last_code) + }, + is_lead, + ) + }; + let has = |v: u8, s: u8| v & s != 0; + let p1 = |x: u64, k: bool| (x << 1) | u64::from(k); + let n1 = |x: u64, k: bool| (x >> 1) | (u64::from(k) << 63); + + let lead = valid & !b.cont; + let lb = ((lead >> 1) & valid) | (u64::from(nb_lead) << (len - 1)); + let eof_bit = if last_blk { 1u64 << (len - 1) } else { 0 }; + + // ── run starts ───────────────────────────────────────────────────────────────────────── + let o_start = c.other & lead & !p1(c.other, has(pb, C_O)) & !p1(c.sp, has(pb, C_SP)); // ` ?[^\s\p{L}\p{N}]+` + let ws_start = c.ws & lead & !p1(c.ws, has(pb, C_WS)); + // `\s+(?!\S)`: the run's last char opens a token — but not a newline, which `\s*[\r\n]+` + // has already swallowed. + let (steal, steal_patch) = to_lead( + c.ws & !c.nl & lb & !eof_bit & !n1(c.ws, has(nb, C_WS)), + b.cont, + prev_cont, + ); + // `\s*[\r\n]+` runs through the run's LAST newline, so a token opens right after it unless + // a further newline still follows inside the run (backward scan → reversal). + let after_nl = p1(c.nl, has(pb, C_NL)) + & c.ws + & lead + & !fill_to_last(c.nl.reverse_bits(), c.ws.reverse_bits()).reverse_bits(); + + // ── `[^\r\n\p{L}\p{N}]?\p{L}+`: the prefix is ANY non-newline non-letter non-digit char, + // not just a space — so a punctuation char that opens a token is swallowed by a following + // letter run (`x!abc` → `x`, `!abc`), while one in mid-run is not (`x!!abc` → `x`, `!!`, + // `abc`), because the greedy other-run already owns it. + // + // Stated backward ("my predecessor opened a token") rather than forward ("my successor is a + // letter"): the forward form needs an `adv`, which silently drops the marker when the two + // chars straddle a block edge. Smearing `o_start` across its char's bytes makes the test a + // plain `p1`, and the smear's only cross-block state is a shift carry. + let mut osf = o_start; + osf |= (osf << 1) & b.cont; + osf |= (osf << 2) & b.cont & (b.cont << 1); + if prev_osf { + osf |= lead_run(b.cont, valid); // a char whose lead sat in the previous block + } + let l_start = c.l + & lead + & !p1(c.l, has(pb, C_L)) + & !p1(c.ws & !c.nl, has(pb, C_WS) && !has(pb, C_NL)) + & !p1(osf, prev_osf); + + // ── rule 3 `\p{N}{1,digit_cap}` + let groups = digit_groups( + digit_cap, + c.n, + lead, + b.cont, + has(pb, C_N), + dig_run, + dig_since, + ); + + // ── the other-run's `[\r\n]*` tail swallows the newlines directly behind it. + let nl_m = ((p1(c.other, has(pb, C_O)) & c.nl & lead) as u128) | u128::from(nl_run); + let nl_e = scanthru(nl_m, c.nl as u128); + let nl_span = nl_e.wrapping_sub(nl_m); + + // ── the one backward-in-time dependency (see deepseek): a newline arriving now retracts an + // "after the last newline" start committed for a run that was still open at the last edge. + if let Some(p) = anl + && has(pb, C_WS) + && c.ws & 1 != 0 + && c.nl & lead_run(c.ws, valid) != 0 + { + starts[p / 64] &= !(1u64 << (p % 64)); + anl = None; + } + + let mut st = groups | l_start | o_start | ws_start | after_nl | steal; + st &= !(nl_span as u64); + st |= nl_e as u64; + st &= lead; + if bi == 0 { + st |= 1; + } + starts[bi] = st; + flag[bi] = st & c.apo; + if bi > 0 { + starts[bi - 1] |= steal_patch; + } + + // ── carries ──────────────────────────────────────────────────────────────────────────── + nl_run = nl_e >> 64 != 0; + let tn = trail_run(c.n, valid, len); + dig_run = tn != 0 && has(nb, C_N); + dig_since = if !dig_run { + 0 + } else { + let g = groups & tn; + let counted = if g == 0 { + dig_since + (c.n & lead & tn).count_ones() + } else { + (c.n & lead & tn & !((1u64 << (63 - g.leading_zeros())) - 1)).count_ones() + }; + if digit_cap == 0 || digit_cap >= 64 { + counted + } else { + counted % digit_cap as u32 + } + }; + let tws = trail_run(c.ws, valid, len); + if tws != 0 && has(nb, C_WS) { + let a = after_nl & tws & !(nl_e as u64); + if a != 0 { + anl = Some(base + 63 - a.leading_zeros() as usize); + } else if !(tws & 1 != 0 && has(pb, C_WS)) { + anl = None; + } + } else { + anl = None; + } + prev_osf = osf >> (len - 1) & 1 != 0; + code = last_code; + prev_cont = b.cont; + } + emit_contr(text, starts, flag, nblk, ntext, true, out) +} diff --git a/tokenizers/bitsplit/src/models/deepseek.rs b/tokenizers/bitsplit/src/models/deepseek.rs new file mode 100644 index 000000000..7685b4d0a --- /dev/null +++ b/tokenizers/bitsplit/src/models/deepseek.rs @@ -0,0 +1,319 @@ +//! DeepSeek-V3/V4 pre-tokenization: the `Sequence` of `\p{N}{1,3}` → `[一-龥぀-ゟ゠-ヿ]+` → +//! the big regex, all `Isolated`, as one bitstream program. Byte-exact with +//! `atomsplit::fsm::fsm_deepseek`. + +use crate::{ + AUX_CJK, CODE_CONT, CONT, Span, adv, build_block, emit, fill_to_last, is_cjk_at, lead_run, scanthru, + to_lead, trail_run, +}; + +/// Atom tag → dense 3-bit code. Letter and Mark share a code because deepseek's letter run is +/// `[\p{L}\p{M}]+`; `AlphaSymMark` (0x16) is categorically `\p{S}` so it takes the punct path, and +/// `Zwj` (0x26) matches no alternative at all so it is a gap char. +pub(crate) const LUT: [u8; 64] = { + let mut t = [6u8; 64]; // gap + t[0x00] = 0; + t[0x10] = 0; + t[0x20] = 0; + t[0x06] = 0; // Letter (+case) | Mark + t[0x01] = 1; + t[0x02] = 1; // \p{N} + t[0x07] = 2; + t[0x08] = 2; + t[0x09] = 2; + t[0x0A] = 2; + t[0x16] = 2; // \p{P} ∪ \p{S} + t[0x03] = 3; + t[0x04] = 4; + t[0x05] = 5; // Newline | Space | WsOther + t[0x0F] = CODE_CONT; + t +}; + +const S_N: u16 = 1 << 0; +const S_LM: u16 = 1 << 1; +const S_PS: u16 = 1 << 2; +const S_WS: u16 = 1 << 3; +const S_NL: u16 = 1 << 4; +const S_SP: u16 = 1 << 5; +const S_CJK: u16 = 1 << 6; +const S_ANY: u16 = S_N | S_LM | S_PS | S_WS | S_CJK; + +/// Dense code → class bits, for the block-edge carries. +const fn code_bits(code: u8) -> u16 { + match code { + 0 => S_LM, + 1 => S_N, + 2 => S_PS, + 3 => S_WS | S_NL, + 4 => S_WS | S_SP, + 5 => S_WS, + _ => 0, // gap | cont + } +} + +/// Filled class bits of the char starting at `p` (a lead byte). +fn bits_at(text: &[u8], tags: &[u8], p: usize) -> u16 { + code_bits(LUT[tags[p] as usize]) | if is_cjk_at(text, p) { S_CJK } else { 0 } +} + +/// Cross-block state: what the paper resolves with selective recomputation across SMs. +struct Carry { + code: u8, // filled dense code of the previous block's last byte (seeds the fill) + cjk: bool, // ...and whether its char is in the Split-2 CJK range + cont: u64, // the previous block's `cont` stream (for `to_lead` underflow) + aa_run: bool, // an alt-1 `[A-Za-z]+` run is still open + nl_run: bool, // a `[\p{P}\p{S}]+[\r\n]*` newline tail is still open + dig_run: bool, // inside a \p{N} run + dig_since: u32, // chars already consumed of the current \p{N}{1,3} group + anl: Option, // a committed "start after the last newline" a later newline may retract +} + +impl Default for Carry { + fn default() -> Self { + // the cont code as the seed: a text opening with a stray continuation byte then classifies + // as gap under both builders instead of as `Letter` (code 0). + Self { + code: CODE_CONT, + cjk: false, + cont: 0, + aa_run: false, + nl_run: false, + dig_run: false, + dig_since: 0, + anl: None, + } + } +} + +/// Pre-tokenize `text` (well-formed UTF-8) with the DeepSeek grammar: writes token spans into `out` +/// and returns the count. `tags` is `atomsplit::classify`'s output (len ≥ `text.len()`), `starts` +/// is scratch for the token-start bitmap (len ≥ `text.len().div_ceil(64)`). +#[must_use] +pub fn bitsplit_deepseek(text: &[u8], tags: &[u8], starts: &mut [u64], out: &mut [Span]) -> usize { + let ntext = text.len(); + if ntext == 0 { + return 0; + } + assert!(tags.len() >= ntext && starts.len() >= ntext.div_ceil(64) && out.len() >= ntext); + let nblk = ntext.div_ceil(64); + let mut cy = Carry::default(); + + for bi in 0..nblk { + let base = bi * 64; + let len = (ntext - base).min(64); + let valid = if len == 64 { !0u64 } else { (1u64 << len) - 1 }; + let last_blk = base + len == ntext; + + let (b, last_code) = build_block::<{ AUX_CJK }, false>(text, tags, base, len, &LUT, cy.code, cy.cjk); + // planes → classes: 3 extractions instead of 7 one-hot masks (see `LUT`). Only `lm` + // needs the `valid` mask — past the block end every plane reads 0, i.e. code 0. + let (pa, pc, pd) = (!b.p2 & !b.p1, !b.p2 & b.p1, b.p2 & !b.p1); + let (s_lm, s_n) = (pa & !b.p0 & valid, pa & b.p0); + let (s_ps, s_nl) = (pc & !b.p0, pc & b.p0); + let (s_sp, s_ws) = (pd & !b.p0, pd | (pc & b.p0)); + let last_cjk = b.aux >> (len - 1) & 1 != 0; + let last_bits = code_bits(last_code) | if last_cjk { S_CJK } else { 0 }; + + // ── edges. The byte before the block is carried; the one after is peeked. One char of + // lookahead is all the grammar needs outside whitespace runs, so peeking it keeps every + // "next char is X" rule exact right at the block boundary — no recomputation needed. + let pb = if base == 0 { + 0 + } else { + code_bits(cy.code) | if cy.cjk { S_CJK } else { 0 } + }; + let (nb, nb_lead, nb_aa) = if last_blk { + (0u16, true, false) + } else { + let q = base + len; + let is_lead = tags[q] != CONT; + let bits = if is_lead { + bits_at(text, tags, q) + } else { + last_bits + }; + (bits, is_lead, text[q].is_ascii_alphabetic()) + }; + let has = |v: u16, s: u16| v & s != 0; + // Split precedence: CJK (Split-2) outranks the big regex, and `fsm_deepseek` tests it ahead + // of the digit arm too — so peel CJK off every other class first. + let cjk = b.aux; + let num = s_n & !cjk; + let lm = s_lm & !cjk; + let ps = s_ps & !cjk; + let gap = valid & !(s_n | s_lm | s_ps | s_ws | cjk); + let cjk_l = cjk & s_lm; // Split-3 re-splits the isolated CJK run into same-kind sub-runs + let cjk_p = cjk & !s_lm; + // `[^\r\n\p{L}\p{P}\p{S}]?` — the letter alternative's optional one-char prefix. Digits and + // CJK are already isolated by Split-1/2, so what remains is non-newline ws + gap chars. + let prefix = (s_ws & !s_nl) | gap; + + // previous-byte / next-byte membership. Streams are filled, so at a lead `p1` reads the + // previous *char*'s class and at a last byte `n1` reads the next char's. + let p1 = |x: u64, c: bool| (x << 1) | u64::from(c); + let n1 = |x: u64, c: bool| (x >> 1) | (u64::from(c) << 63); + let pb_gap = base != 0 && !has(pb, S_ANY); + let pb_cjk = has(pb, S_CJK); + + let lead = valid & !b.cont; + let lb = ((lead >> 1) & valid) | (u64::from(nb_lead) << (len - 1)); + // a char that is both its own first and last byte is single-byte, i.e. ASCII; an ASCII char + // in `\p{L}∪\p{M}` is exactly `[A-Za-z]` (ASCII has no marks). Both streams for 2 ops. + let ascii = lead & lb; + let aa = s_lm & ascii; + + // ── run starts. All purely backward-looking, so the carry alone makes them exact. ─────── + let n_start = num & lead & !p1(num, has(pb, S_N) && !pb_cjk); + let lm_start = lm + & lead + & !p1(lm, has(pb, S_LM) && !pb_cjk) + & !p1(prefix, (has(pb, S_WS) && !has(pb, S_NL)) || pb_gap); + let ws_start = s_ws & lead & !p1(s_ws, has(pb, S_WS)); + let gap_start = gap & lead & !p1(gap, pb_gap); + let ps_start = ps & lead & !p1(ps, has(pb, S_PS) && !pb_cjk) & !p1(s_sp, has(pb, S_SP)); + let cjk_start = (cjk_l & lead & !p1(cjk_l, pb_cjk && has(pb, S_LM))) + | (cjk_p & lead & !p1(cjk_p, pb_cjk && !has(pb, S_LM))); + + // ── Split-1 `\p{N}{1,3}`: the one non-local rule — a group boundary every 3 chars from the + // run start. Marker iteration (the paper's bounded-repetition lowering); ≤21 rounds/block, + // and 0 rounds on text without digit runs. + let mut m = n_start; + if cy.dig_run && has(pb, S_N) { + // resume mid-run: the next group start is (3 - since) chars into the block. Re-mask + // with `num` at every hop — the carry only says the char *at* the edge was a digit, the + // run may well have ended there (`Ⅷ` straddling the edge, then `\t`, then `456`). + let mut s = lead & lead.wrapping_neg() & num; // first lead of the block + for _ in 0..((3 - cy.dig_since % 3) % 3) { + s = adv(s, b.cont) & num & lead; + } + m |= s; + } + let mut groups = m; + if num & b.cont == 0 { + // Fast path: every digit in this block is single-byte, so "3 chars on" is just `<< 3` + // and the three `adv`s collapse into one shift against a precomputed mask (`num3` asks + // that the two skipped positions are digits too, which is what the `adv` chain checked). + // ~3 ops per group instead of ~14 — dense-digit text is otherwise this loop's worst case. + let num3 = num & (num << 1) & (num << 2); + while m != 0 { + m = (m << 3) & num3; + groups |= m; + } + } else { + while m != 0 { + let a = adv(m, b.cont) & num & lead; + let c = adv(adv(a, b.cont) & num & lead, b.cont) & num & lead; + if c == 0 { + break; + } + groups |= c; + m = c; + } + } + + // ── whitespace: `\s*[\r\n]+ | \s+(?!\S) | \s+`. + // (a) the run's first token runs through its LAST newline → a start right after it, unless + // a further newline still follows inside the run (a backward scan, hence the reversal). + let anl = p1(s_nl, has(pb, S_NL)) & s_ws & lead; + let later_nl = fill_to_last(s_nl.reverse_bits(), s_ws.reverse_bits()).reverse_bits(); + let after_nl = anl & !later_nl; + // (b) the run's last char is handed to whatever follows, as its `[^…]?` / ` ?` prefix — + // unless the run ends the input or the next piece is Split-1/2-isolated (`(?!\S)`). + let eof_bit = if last_blk { 1u64 << (len - 1) } else { 0 }; + let steal_lb = s_ws + & !s_nl + & lb + & !eof_bit + & !n1(s_ws, has(nb, S_WS)) + & !n1(s_n | cjk, has(nb, S_N | S_CJK)); + let (steal, steal_patch) = to_lead(steal_lb, b.cont, cy.cont); + // (c) the same one-char give-back out of a gap run (Control / NumericOther / ZWJ match no + // alternative, so the run is one piece minus the char a following letter run claims). + let (gap_steal, gap_patch) = to_lead( + gap & lb & n1(lm, has(nb, S_LM) && !has(nb, S_CJK)), + b.cont, + cy.cont, + ); + + // ── alt-1 `[ascii_punct][A-Za-z]+`: fires only where the scan is actually positioned, i.e. + // at a punct-run start no space swallowed. Its `[A-Za-z]+` run then has no interior starts + // and forces one at its end — which the letter rule alone would suppress (`!c` reads `!ab`, + // `c` after it is a fresh token even though its predecessor is a letter). + let alt1 = ps_start & ascii & n1(aa, nb_aa); + let aa_m = ((alt1 as u128) << 1) | u128::from(cy.aa_run); + let aa_e = scanthru(aa_m, aa as u128); + let aa_span = aa_e.wrapping_sub(aa_m); + // ── a punct run's `[\r\n]*` tail swallows the newlines directly behind it. + let nl_m = + ((p1(ps, has(pb, S_PS) && !pb_cjk) & s_nl & lead) as u128) | u128::from(cy.nl_run); + let nl_e = scanthru(nl_m, s_nl as u128); + let nl_span = nl_e.wrapping_sub(nl_m); + + // ── the one backward-in-time dependency: a newline arriving now retracts the "start after + // the last newline" already committed for a whitespace run that was open at the last edge. + if let Some(p) = cy.anl + && has(pb, S_WS) + && s_ws & 1 != 0 + && s_nl & lead_run(s_ws, valid) != 0 + { + starts[p / 64] &= !(1u64 << (p % 64)); + cy.anl = None; + } + + let mut st = groups + | lm_start + | ws_start + | gap_start + | ps_start + | cjk_start + | after_nl + | steal + | gap_steal; + st &= !(aa_span as u64) & !(nl_span as u64); + st |= aa_e as u64 | nl_e as u64; + st &= lead; + if bi == 0 { + st |= 1; // position 0 always opens a token + } + starts[bi] = st; + if bi > 0 { + starts[bi - 1] |= steal_patch | gap_patch; + } + + // ── carries ──────────────────────────────────────────────────────────────────────────── + cy.aa_run = aa_e >> 64 != 0; + cy.nl_run = nl_e >> 64 != 0; + let tn = trail_run(num, valid, len); + cy.dig_run = tn != 0 && has(nb, S_N) && !has(nb, S_CJK); + cy.dig_since = if !cy.dig_run { + 0 + } else { + let g = groups & tn; + let counted = if g == 0 { + cy.dig_since + (num & lead & tn).count_ones() + } else { + (num & lead & tn & !((1u64 << (63 - g.leading_zeros())) - 1)).count_ones() + }; + counted % 3 + }; + let tws = trail_run(s_ws, valid, len); + if tws != 0 && has(nb, S_WS) { + // ...but not a bit that a punct tail's `[\r\n]*` already made a *run start*: the + // absorption cut the whitespace run, so a later newline cannot reach back past it. + let a = after_nl & tws & !(nl_e as u64); + if a != 0 { + cy.anl = Some(base + 63 - a.leading_zeros() as usize); + } else if !(tws & 1 != 0 && has(pb, S_WS)) { + cy.anl = None; // a fresh run with no newline yet — nothing left to retract + } + } else { + cy.anl = None; + } + cy.code = last_code; + cy.cjk = last_cjk; + cy.cont = b.cont; + } + + emit(starts, nblk, ntext, out) +} diff --git a/tokenizers/bitsplit/src/models/gpt2.rs b/tokenizers/bitsplit/src/models/gpt2.rs new file mode 100644 index 000000000..54567df91 --- /dev/null +++ b/tokenizers/bitsplit/src/models/gpt2.rs @@ -0,0 +1,162 @@ +//! **GPT-2 / ByteLevel**: +//! `'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+` +//! +//! Contractions are an ALTERNATIVE here (`don't` -> `don`, `'t`), and case-sensitive -- unlike +//! cl100k's `(?i:)`. They go to `emit_contr`'s scalar escape: a variable-length, case-optional +//! literal alternation that outranks every other arm is miserable in bit algebra and trivial there. + +use crate::{ + AUX_NONE, CODE_CONT, CONT, Span, build_block, emit_contr, to_lead, +}; + +/// Atom tag → dense 3-bit code, shared by both grammars. Unlike deepseek's table, `Mark` is NOT a +/// letter here (`\p{L}` excludes it, so it belongs to the "other" class), and `Apostrophe` gets its +/// own code so the contraction escape can be flagged with one AND — it is still "other" for every +/// run rule, which `decode` restores. +const LUT: [u8; 64] = { + let mut t = [2u8; 64]; // other = [^\s\p{L}\p{N}] + t[0x00] = 0; + t[0x10] = 0; + t[0x20] = 0; // Letter (+ case refinements) + t[0x01] = 1; + t[0x02] = 1; // \p{N} = Nd ∪ Nl ∪ No + t[0x03] = 3; // Newline + t[0x04] = 4; // Space + t[0x05] = 5; // WsOther + t[0x09] = 6; // Apostrophe + t[0x0F] = CODE_CONT; + t +}; + +/// One block's class streams. `other` folds the apostrophe code back in; only `l` needs the +/// `valid` mask, since past the block end every plane reads 0, i.e. code 0. +struct Cls { + l: u64, + n: u64, + other: u64, + nl: u64, + sp: u64, + ws: u64, + apo: u64, +} + +#[inline] +fn decode(p0: u64, p1: u64, p2: u64, valid: u64) -> Cls { + let a = !p2 & !p1; + let nl = !p2 & p1 & p0; + Cls { + l: a & !p0 & valid, + n: a & p0, + other: p1 & !p0 & valid, // codes 2 and 6 — "other" and the apostrophe + nl, + sp: p2 & !p1 & !p0, + ws: (p2 & !p1) | nl, // Space ∪ WsOther ∪ Newline + apo: p2 & p1 & !p0, + } +} + +/// Class bits of a dense code, for the block-edge carry. +const C_L: u8 = 1; +const C_N: u8 = 2; +const C_O: u8 = 4; +const C_NL: u8 = 8; +const C_SP: u8 = 16; +const C_WS: u8 = 32; + +#[inline] +const fn code_bits(code: u8) -> u8 { + match code { + 0 => C_L, + 1 => C_N, + 2 | 6 => C_O, + 3 => C_NL | C_WS, + 4 => C_SP | C_WS, + 5 => C_WS, + _ => 0, // cont never reaches an edge (the fill resolves it) + } +} + +/// GPT-2 / byte-level pre-tokenization. `starts` and `flag` are scratch bitmaps +/// (len ≥ `text.len().div_ceil(64)`); byte-exact with `atomsplit::fsm::fsm_byte_level`. +#[must_use] +pub fn bitsplit_byte_level( + text: &[u8], + tags: &[u8], + starts: &mut [u64], + flag: &mut [u64], + out: &mut [Span], +) -> usize { + let ntext = text.len(); + if ntext == 0 { + return 0; + } + let nblk = ntext.div_ceil(64); + assert!( + tags.len() >= ntext && starts.len() >= nblk && flag.len() >= nblk && out.len() >= ntext + ); + let (mut code, mut prev_cont) = (CODE_CONT, 0u64); + + for bi in 0..nblk { + let base = bi * 64; + let len = (ntext - base).min(64); + let valid = if len == 64 { !0u64 } else { (1u64 << len) - 1 }; + let last_blk = base + len == ntext; + + let (b, last_code) = build_block::<{ AUX_NONE }, false>(text, tags, base, len, &LUT, code, false); + let c = decode(b.p0, b.p1, b.p2, valid); + + let pb = if base == 0 { 0 } else { code_bits(code) }; + let (nb, nb_lead) = if last_blk { + (0u8, true) + } else { + let q = base + len; + let is_lead = tags[q] != CONT; + ( + if is_lead { + code_bits(LUT[tags[q] as usize]) + } else { + code_bits(last_code) + }, + is_lead, + ) + }; + let has = |v: u8, s: u8| v & s != 0; + let p1 = |x: u64, k: bool| (x << 1) | u64::from(k); + let n1 = |x: u64, k: bool| (x >> 1) | (u64::from(k) << 63); + + let lead = valid & !b.cont; + let lb = ((lead >> 1) & valid) | (u64::from(nb_lead) << (len - 1)); + + // ── every alternative but the contraction is ` ?X+` over a class run, so a token opens at + // each run start — pushed back one char when a literal space sits in front of it (` ?`). + let sp_pfx = p1(c.sp, has(pb, C_SP)); + let l_start = c.l & lead & !p1(c.l, has(pb, C_L)) & !sp_pfx; + let n_start = c.n & lead & !p1(c.n, has(pb, C_N)) & !sp_pfx; + let o_start = c.other & lead & !p1(c.other, has(pb, C_O)) & !sp_pfx; + let ws_start = c.ws & lead & !p1(c.ws, has(pb, C_WS)); + // `\s+(?!\S)` hands the run's LAST whitespace char to whatever follows: as a ` ?` prefix if + // it is a space, else as a token of its own. Either way it opens a token — unless the run + // ends the input, where plain `\s+` takes the lot. GPT-2 has no `[\r\n]` rule, so unlike + // cl100k a newline is ordinary whitespace here and can be the stolen char. + let eof_bit = if last_blk { 1u64 << (len - 1) } else { 0 }; + let (steal, patch) = to_lead( + c.ws & lb & !eof_bit & !n1(c.ws, has(nb, C_WS)), + b.cont, + prev_cont, + ); + + let mut st = (l_start | n_start | o_start | ws_start | steal) & lead; + if bi == 0 { + st |= 1; + } + starts[bi] = st; + flag[bi] = st & c.apo; // apostrophes that open a token → contraction escape + if bi > 0 { + starts[bi - 1] |= patch; + } + code = last_code; + prev_cont = b.cont; + } + emit_contr(text, starts, flag, nblk, ntext, false, out) +} + diff --git a/tokenizers/bitsplit/src/models/kimi.rs b/tokenizers/bitsplit/src/models/kimi.rs new file mode 100644 index 000000000..994f0b35f --- /dev/null +++ b/tokenizers/bitsplit/src/models/kimi.rs @@ -0,0 +1,556 @@ +//! **kimi-k2 / k3** (`moonshotai/Kimi-K2-Instruct`, `tokenization_kimi.py`'s `pat_str`): +//! `[\p{Han}]+ | [^\r\n\p{L}\p{N}]?[\p{Lu}…&&[^\p{Han}]]*[\p{Ll}…&&[^\p{Han}]]+(?i:'s|…)? | …` +//! `| \p{N}{1,3} | ?[^\s\p{L}\p{N}]+[\r\n]* | \s*[\r\n]+ | \s+(?!\S) | \s+` +//! +//! o200k plus a leading `[\p{Han}]+` arm, Han subtracted from both letter classes, and a plain +//! `[\r\n]*` rule-4 tail (o200k has `[\r\n/]*`). Ships tiktoken.model, not a tokenizer.json. +//! +//! Two things are worth knowing before reading the algebra: +//! +//! 1. **The case split is a scalar escape, on purpose.** `[UC]*[LC]+ | [UC]+[LC]*` has no local +//! form: `中Qz` is ONE token (a U after a C is not a boundary) but `ʰABC` is two (it is, when no +//! L follows) — the difference is whether an L appears LATER in the run, so no `p1` decides it. +//! The bit half instead computes a cheap gate: a letter token is escaped only if the block holds +//! an interior upper or a trailing apostrophe, so all-lowercase and Capitalised text never pays. +//! 2. **The contraction is a SUFFIX here, not an alternative.** cl100k emits `'t` as its own token; +//! o200k glues it onto the letter token before it. So this uses `emit_contr_suffix`, which +//! *extends* the open token instead of opening one — the mirror image of `emit_contr`. +//! +use crate::classify::{char_len, in_mask, mask}; +use crate::{ + AUX_HAN, CODE_CONT, CONT, Span, build_block, contr_len, digit_groups, fill_to_last, lead_run, + letter_match, member, run_end, scanthru, to_lead, trail_run, ws_tail, +}; + +/// Atom tag → dense 4-bit code. Unlike cl100k, `\p{M}` IS a letter here (both alt classes list +/// `\p{M}`) — but only a true mark: `AlphaSymMark` (0x16, categorically `\p{S}`) and `Zwj` (0x26, +/// `\p{Cf}`) stay "other", which is what keeps `[\p{L}\p{M}]+` off them. +const LUT: [u8; 64] = { + let mut t = [8u8; 64]; // other = [^\s\p{L}\p{N}] (incl. Connector/Punct/Sym/Control/ASM/ZWJ) + t[0x10] = 0; // \p{Lu} ∪ \p{Lt} + t[0x20] = 1; // \p{Ll} + t[0x00] = 2; // caseless letter (\p{Lm}\p{Lo}) — in BOTH alt classes + t[0x06] = 10; // true \p{M}: a letter for the alts, but ALSO in rule 4's class (see `mark_adj`) + t[0x01] = 3; + t[0x02] = 3; // \p{N} + t[0x03] = 4; // Newline + t[0x04] = 5; // Space + t[0x05] = 6; // WsOther + t[0x09] = 9; // Apostrophe — "other" for every run rule, split out for the contraction flag + t[0x0F] = CODE_CONT; + t +}; + +struct Cls { + u: u64, + l: u64, + c: u64, + n: u64, + nl: u64, + sp: u64, + ws: u64, + oth: u64, + mark: u64, + apo: u64, +} + +#[inline] +fn decode(p0: u64, p1: u64, p2: u64, p3: u64, valid: u64) -> Cls { + let low = !p3; + let a = low & !p2; + let w = low & p2; + Cls { + u: a & !p1 & !p0 & valid, // code 0 — past the block end every plane reads 0, hence `valid` + l: a & !p1 & p0, + c: a & p1 & !p0, + n: a & p1 & p0, + nl: w & !p1 & !p0, + sp: w & !p1 & p0, + ws: w & !(p1 & p0), // codes 4,5,6 — cont (7) excluded + oth: p3 & !p1, // codes 8,9 — the apostrophe is "other" for run purposes + mark: p3 & p1, // code 10 + apo: p3 & p0 & !p1, // code 9 + } +} + +const C_U: u16 = 1 << 0; +const C_L: u16 = 1 << 1; +const C_C: u16 = 1 << 2; +const C_N: u16 = 1 << 3; +const C_NL: u16 = 1 << 4; +const C_SP: u16 = 1 << 5; +const C_WSO: u16 = 1 << 6; +const C_OTH: u16 = 1 << 7; +const C_MARK: u16 = 1 << 8; +const C_HAN: u16 = 1 << 9; +const C_WS: u16 = C_NL | C_SP | C_WSO; +const C_LET: u16 = C_U | C_L | C_C | C_MARK; + +const fn code_bits(code: u8) -> u16 { + match code { + 0 => C_U, + 1 => C_L, + 2 => C_C, + 3 => C_N, + 4 => C_NL, + 5 => C_SP, + 6 => C_WSO, + 8 | 9 => C_OTH, + 10 => C_MARK, + _ => 0, // cont never reaches an edge (the fill resolves it) + } +} + +/// kimi-k2 / k3. +#[must_use] +pub fn bitsplit_kimi( + text: &[u8], + tags: &[u8], + starts: &mut [u64], + flag: &mut [u64], + out: &mut [Span], +) -> usize { + kimi(text, tags, starts, flag, out) +} + +fn kimi( + text: &[u8], + tags: &[u8], + starts: &mut [u64], + flag: &mut [u64], + out: &mut [Span], +) -> usize { + let ntext = text.len(); + if ntext == 0 { + return 0; + } + let nblk = ntext.div_ceil(64); + assert!( + tags.len() >= ntext && starts.len() >= nblk && flag.len() >= nblk && out.len() >= ntext + ); + + let (mut code, mut prev_cont) = (CODE_CONT, 0u64); + let (mut prev_han, mut nl_run, mut prev_osf) = (false, false, false); + let mut prev_absorbed = false; // the block's last byte was eaten by a `[\r\n/]*` tail + let (mut dig_run, mut dig_since) = (false, 0u32); + let mut anl: Option = None; + let mut last_lt: Option = None; + let mut prev_start: Option = None; + + for bi in 0..nblk { + let base = bi * 64; + let len = (ntext - base).min(64); + let valid = if len == 64 { !0u64 } else { (1u64 << len) - 1 }; + let last_blk = base + len == ntext; + + let (b, last_code) = + build_block::<{ AUX_HAN }, true>(text, tags, base, len, &LUT, code, prev_han); + let c = decode(b.p0, b.p1, b.p2, b.p3, valid); + + let han = b.aux & valid; + let slash = 0; // kimi's rule-4 tail is a plain `[\r\n]*` + let last_han = han >> (len - 1) & 1 != 0; + + let pb = if base == 0 { + 0 + } else { + code_bits(code) | if prev_han { C_HAN } else { 0 } + }; + let (nb, nb_lead) = if last_blk { + (0u16, true) + } else { + let q = base + len; + let is_lead = tags[q] != CONT; + let bits = if is_lead { + code_bits(LUT[tags[q] as usize]) + | if crate::aux_at::<{ AUX_HAN }>(text, q) { + C_HAN + } else { + 0 + } + } else { + code_bits(last_code) | if last_han { C_HAN } else { 0 } + }; + (bits, is_lead) + }; + let has = |v: u16, s: u16| v & s != 0; + let p1 = |x: u64, k: bool| (x << 1) | u64::from(k); + let n1 = |x: u64, k: bool| (x >> 1) | (u64::from(k) << 63); + + let lead = valid & !b.cont; + let lb = ((lead >> 1) & valid) | (u64::from(nb_lead) << (len - 1)); + let eof_bit = if last_blk { 1u64 << (len - 1) } else { 0 }; + + // ── Han is peeled off the letter classes first: kimi's alt-1 outranks the letter alts, and + // its letter classes are literally `[…&&[^\p{Han}]]`. With AUX != HAN this is all zero. + let (cu, cl, cc) = (c.u & !han, c.l & !han, c.c & !han); + let letter = cu | cl | cc | c.mark; + let pb_let = has(pb, C_LET) && !has(pb, C_HAN); + + // Rule 4's `[\r\n/]*` tail. `/` is in BOTH the `+` body and the tail, so one tail run can + // collect SEVERAL markers (`!\n/\n`: the `\n` after `!` and the `\n` after `/`). That is + // exactly what `fill_to_last` is for -- `nl_e - nl_m` would span only from the last one. + let tail_cls = c.nl | slash; + let nl_m64 = ((p1(c.oth, has(pb, C_OTH)) & c.nl & lead) | u64::from(nl_run)) & tail_cls; + let nl_m = nl_m64 as u128; + let nl_e = scanthru(nl_m, tail_cls as u128); + let nl_span = fill_to_last(nl_m64, tail_cls) as u128; + + // ── rule 4 ` ?[^\s\p{L}\p{N}]+[\r\n/]*` and rules 5-7 — identical to cl100k. + // A char absorbed by a `[\r\n/]*` tail is NOT part of the `+` body, so an "other" after it + // opens a fresh run (`\u{1f600}\r\n/#` is the tail, then `#` starts again). + let o_prev = c.oth & !(nl_span as u64); + let o_start = c.oth + & lead + & !(nl_span as u64) // a char INSIDE the tail never opens a run + & !p1(o_prev, has(pb, C_OTH) && !prev_absorbed) + & !p1(c.sp, has(pb, C_SP)); + let ws_start = c.ws & lead & !p1(c.ws, has(pb, C_WS)); + let (steal, steal_patch) = to_lead( + c.ws & !c.nl & lb & !eof_bit & !n1(c.ws, has(nb, C_WS)), + b.cont, + prev_cont, + ); + let after_nl = p1(c.nl, has(pb, C_NL)) + & c.ws + & lead + & !fill_to_last(c.nl.reverse_bits(), c.ws.reverse_bits()).reverse_bits(); + + // ── `[^\r\n\p{L}\p{N}]?` before a letter run: any token-opening non-newline non-letter + // non-digit char. Smeared across its char's bytes so the test is a plain `p1` (see cl100k). + let mut osf = o_start; + osf |= (osf << 1) & b.cont; + osf |= (osf << 2) & b.cont & (b.cont << 1); + if prev_osf { + osf |= lead_run(b.cont, valid); + } + let l_start = letter + & lead + & !p1(letter, pb_let) + & !p1(c.ws & !c.nl, has(pb, C_WS) && !has(pb, C_NL)) + & !p1(osf, prev_osf); + // the token holding a letter run may OPEN one char earlier, on the `[^\r\n\p{L}\p{N}]?` + // prefix — that char is the start the escape has to be handed. The class genuinely excludes + // `\r\n` and digits: a newline before a letter run is its own token, never a prefix. + let prefix_cls = c.oth | c.sp | (c.ws & !c.nl); + let l_run = letter & lead & !p1(letter, pb_let); + let next_l_run = nb_lead + && has(nb, C_LET) + && !has(nb, C_HAN) + && letter >> (len - 1) & 1 == 0; + // `n1` only reads the NEXT char at a char's last byte, so mark there and walk back to the + // lead — a 3-byte prefix char (ZWJ, `—`, `½`) otherwise reads its own middle byte. + let (pfx_lead, pfx_patch) = + to_lead(prefix_cls & lb & n1(l_run, next_l_run), b.cont, prev_cont); + let l_start_tok = l_start | pfx_lead; + + // ── the escape gate (see the header). An interior upper, or an apostrophe closing the + // run, means some letter token in this block needs the scalar case/contraction pass. + let interior_u = cu & lead & p1(letter, pb_let); + let apo_after = c.apo & lead & p1(letter, pb_let); + // `\p{M}` is in BOTH the letter classes and rule 4's `[^\s\p{L}\p{N}]`, and which one wins + // depends on whether the punctuation before it STARTED the run (`!\u{301}a` is one token, + // `!!\u{301}a` is two). Adjacency is the gate; the scalar pass resolves it. Real text hits + // this via emoji + variation selector (U+FE0F is \p{Mn}). + let mark_adj = (c.mark & lead & p1(c.oth, has(pb, C_OTH))) + | (c.oth & lead & p1(c.mark, has(pb, C_MARK))); + + // ── kimi's `[\p{Han}]+` + let han_start = han & lead & !p1(han, has(pb, C_HAN)); + + // ── rule 3 `\p{N}{1,DIGIT_CAP}` + let groups = + digit_groups(3, c.n, lead, b.cont, has(pb, C_N), dig_run, dig_since); + + // ── the other-run's `[\r\n/]*` tail (`/` only for the o200k line; kimi has `[\r\n]*`). + // ── the one backward-in-time rule, exactly as in cl100k/deepseek: a newline arriving now + // retracts an "after the last newline" start committed for a run still open at the edge. + if let Some(p) = anl + && has(pb, C_WS) + && c.ws & 1 != 0 + && c.nl & lead_run(c.ws, valid) != 0 + { + starts[p / 64] &= !(1u64 << (p % 64)); + anl = None; + } + + let mut st = groups | l_start | han_start | o_start | ws_start | after_nl | steal; + st &= !(nl_span as u64); + st |= nl_e as u64; + st &= lead; + if bi == 0 { + st |= 1; + } + starts[bi] = st; + // Conservative: flag every letter token in a block that shows either trigger. A false + // positive only costs a rescan of that token, so erring wide is free; missing one is not. + // Only the letter RUNS that actually need the escape, not every letter token in the block: + // one `'s` in a paragraph would otherwise drag every word through the scalar pass. + // `apo_after` sits on the apostrophe, one byte past the run, so shift it back inside. + let needs = interior_u | ((apo_after >> 1) & letter); + let runs_needing = if needs == 0 { + 0 + } else { + fill_to_last(needs.reverse_bits(), letter.reverse_bits()).reverse_bits() + }; + let (pfx_need, _) = to_lead( + prefix_cls & lb & n1(runs_needing & l_run, false), + b.cont, + prev_cont, + ); + let lt_all = st & l_start_tok; // every letter token — what the cross-block patch tracks + let lt = lt_all & (runs_needing | pfx_need); + let trig = (interior_u | apo_after) != 0; + flag[bi] = if trig { lt } else { 0 }; + // mark_adj is rare, so be blunt: escape every token in the block, plus the one still open + // from an earlier block (rule 4's ` ?` means the token can have started on a space). + if mark_adj != 0 { + flag[bi] |= st; + if let Some(p) = prev_start { + flag[p / 64] |= 1u64 << (p % 64); + } + } + if st != 0 { + prev_start = Some(base + 63 - st.leading_zeros() as usize); + } + if trig { + if bi > 0 { + flag[bi - 1] |= starts[bi - 1] & pfx_patch; + } + // the trigger can land in a LATER block than the token it belongs to (`\u{d55c}` ends + // block k, its `'s` opens block k+1), so always re-flag the last letter token seen. + if let Some(p) = last_lt { + flag[p / 64] |= 1u64 << (p % 64); + } + } + if bi > 0 { + starts[bi - 1] |= steal_patch; + } + // a run open at the block edge may meet its trigger in a later block, so flag it now + if lt_all != 0 { + last_lt = Some(base + 63 - lt_all.leading_zeros() as usize); + } + // ...ending on the run's `?` prefix char counts too — the token opened there. + let open_at_edge = (letter | l_start_tok) >> (len - 1) & 1 != 0; + if !last_blk && open_at_edge && let Some(p) = last_lt { + flag[p / 64] |= 1u64 << (p % 64); + } + + // ── carries + nl_run = nl_e >> 64 != 0; + prev_absorbed = nl_span >> (len - 1) & 1 != 0; + let tn = trail_run(c.n, valid, len); + dig_run = tn != 0 && has(nb, C_N); + dig_since = if !dig_run { + 0 + } else { + let g = groups & tn; + let counted = if g == 0 { + dig_since + (c.n & lead & tn).count_ones() + } else { + (c.n & lead & tn & !((1u64 << (63 - g.leading_zeros())) - 1)).count_ones() + }; + counted % 3 + }; + let tws = trail_run(c.ws, valid, len); + if tws != 0 && has(nb, C_WS) { + let a = after_nl & tws & !(nl_e as u64); + if a != 0 { + anl = Some(base + 63 - a.leading_zeros() as usize); + } else if !(tws & 1 != 0 && has(pb, C_WS)) { + anl = None; + } + } else { + anl = None; + } + prev_osf = osf >> (len - 1) & 1 != 0; + prev_han = last_han; + code = last_code; + prev_cont = b.cont; + + } + + emit_kimi(text, tags, starts, flag, nblk, ntext, out) +} + +// ── the scalar escape ─────────────────────────────────────────────────────────────────────────── +// Ported from the FSM this replaces, which is the proven reading of the alternatives. It runs from +// a flagged token start and RESYNCS the moment it lands on an algebra start bit again, so a +// divergent region costs a short scalar walk and nothing downstream. + +/// Emit ONE token starting at `i` (letters emit several) and return the cursor past it. +fn step( + text: &[u8], + tags: &[u8], + i: usize, + end: usize, + out: &mut [Span], + w: &mut usize, +) -> usize { + let han = |p: usize| crate::aux_at::<{ AUX_HAN }>(text, p); + // kimi subtracts Han from both letter classes and isolates it in its own arm + let is_lm = |p: usize| p < end && member(tags[p]) && !han(p); + let letter_end = |mut p: usize| { + while p < end && (tags[p] == CONT || (member(tags[p]) && !han(p))) { + p += 1; + } + p + }; + let other = |sp0: usize| { + let mut p = run_end(tags, sp0, end, mask::NOT_WS_L_N); + if p > sp0 { + while p < end + && (tags[p] == crate::NLN) + { + p += char_len(text[p]); + } + } + p + }; + let emit1 = |a: usize, b: usize, out: &mut [Span], w: &mut usize| { + out[*w] = Span::new(a as u32, b as u32); + *w += 1; + }; + // `[\p{Han}]+` — kimi's alternative 1, ahead of everything else + if han(i) { + let mut p = i; + while p < end && (tags[p] == CONT || han(p)) { + if tags[p] != CONT && !han(p) { + break; + } + p += 1; + } + emit1(i, p, out, w); + return p; + } + // the letter alternatives, with the case split and the optional contraction suffix + let letters = |pfx: usize, ls: usize, out: &mut [Span], w: &mut usize| -> usize { + let re = letter_end(ls); + let (mut p, mut first, mut cursor) = (ls, true, re); + while p < re { + let e = letter_match(tags, p, re); + let start = if first { pfx } else { p }; + let te = if e == re { + e + contr_len(text, e, true) + } else { + e + }; + out[*w] = Span::new(start as u32, te as u32); + *w += 1; + first = false; + cursor = te; + p = e; + } + cursor + }; + + let b = text[i]; + match tags[i] & 0x0F { + crate::NW | crate::NO => { + let (mut p, mut cnt) = (i, 0usize); + while p < end && cnt < 3 && in_mask(tags[p], mask::NUMBER) { + p += char_len(text[p]); + cnt += 1; + } + emit1(i, p, out, w); + p + } + crate::LET | crate::MRK => { + if member(tags[i]) && !han(i) { + return letters(i, i, out, w); + } + let a = i + char_len(b); + if is_lm(a) { + return letters(i, a, out, w); + } + let p = other(i); + emit1(i, p, out, w); + p + } + crate::SPC => { + let a = i + 1; + if is_lm(a) { + return letters(i, a, out, w); + } + let p = other(a); + let e = if p > a { p } else { ws_tail(text, tags, i, end) }; + emit1(i, e, out, w); + e + } + crate::WSO => { + let a = i + char_len(b); + if is_lm(a) { + return letters(i, a, out, w); + } + let e = ws_tail(text, tags, i, end); + emit1(i, e, out, w); + e + } + crate::NLN => { + let e = ws_tail(text, tags, i, end); + emit1(i, e, out, w); + e + } + _ => { + let a = i + char_len(b); + if is_lm(a) { + return letters(i, a, out, w); + } + let p = other(i); + emit1(i, p, out, w); + p + } + } +} + +/// `emit` plus the escape: at a flagged start, hand over to the scalar dispatch and take back +/// control at the first position that is a start bit again. +fn emit_kimi( + text: &[u8], + tags: &[u8], + starts: &[u64], + flag: &[u64], + nblk: usize, + n: usize, + out: &mut [Span], +) -> usize { + let (mut w, mut open, mut skip) = (0usize, usize::MAX, 0usize); + for bi in 0..nblk { + let mut m = starts[bi]; + let f = flag[bi]; + while m != 0 { + let j = m.trailing_zeros() as usize; + let pos = bi * 64 + j; + m &= m - 1; + if pos < skip { + continue; + } + if f >> j & 1 != 0 { + if open != usize::MAX && open < pos { + out[w] = Span::new(open as u32, pos as u32); + w += 1; + } + let mut c = pos; + loop { + c = step(text, tags, c, n, out, &mut w); + if c >= n || starts[c / 64] >> (c % 64) & 1 != 0 { + break; + } + } + open = usize::MAX; + skip = c; + continue; + } + if open != usize::MAX { + out[w] = Span::new(open as u32, pos as u32); + w += 1; + } + open = pos; + } + } + if open != usize::MAX { + out[w] = Span::new(open as u32, n as u32); + w += 1; + } + w +} diff --git a/tokenizers/bitsplit/src/models/mod.rs b/tokenizers/bitsplit/src/models/mod.rs new file mode 100644 index 000000000..d80354a63 --- /dev/null +++ b/tokenizers/bitsplit/src/models/mod.rs @@ -0,0 +1,12 @@ +//! One module per pre-tokenization **regex**, each unrolled against the primitives in +//! [`crate`]. Models that ship the same regex share a module — o200k covers Llama-4, gpt-oss and +//! MiniMax-M2; cl100k covers Llama-3, GLM-4.6 and (at digit cap 1) Qwen. +//! +//! Deliberately not one parameterised grammar: folding these together made every fix a three-way +//! risk and none of them could be read on its own. +pub mod cl100k; +pub mod deepseek; +pub mod gpt2; +pub mod kimi; +pub mod o200k; +pub mod tekken; diff --git a/tokenizers/bitsplit/src/models/o200k.rs b/tokenizers/bitsplit/src/models/o200k.rs new file mode 100644 index 000000000..44dcc6437 --- /dev/null +++ b/tokenizers/bitsplit/src/models/o200k.rs @@ -0,0 +1,527 @@ +//! o200k_base / GPT-4o — byte-for-byte the regex **Llama-4, gpt-oss and MiniMax-M2** also ship: +//! `[^\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|…)?` +//! `|` same with `+`/`*` swapped `| \p{N}{1,3} | ?[^\s\p{L}\p{N}]+[\r\n/]* | \s*[\r\n]+ | \s+(?!\S) | \s+` +//! +//! Same skeleton as cl100k; only the letter half differs. Tekken and kimi are separate files — +//! they are separate regexes, and folding them in here is what made every fix a three-way risk. +//! +//! Two things are worth knowing before reading the algebra: +//! +//! 1. **The case split is a scalar escape, on purpose.** `[UC]*[LC]+ | [UC]+[LC]*` has no local +//! form: `中Qz` is ONE token (a U after a C is not a boundary) but `ʰABC` is two (it is, when no +//! L follows) — the difference is whether an L appears LATER in the run, so no `p1` decides it. +//! The bit half instead computes a cheap gate: a letter token is escaped only if the block holds +//! an interior upper or a trailing apostrophe, so all-lowercase and Capitalised text never pays. +//! 2. **The contraction is a SUFFIX here, not an alternative.** cl100k emits `'t` as its own token; +//! o200k glues it onto the letter token before it. So this uses `emit_contr_suffix`, which +//! *extends* the open token instead of opening one — the mirror image of `emit_contr`. +//! +use crate::classify::{char_len, in_mask, mask}; +use crate::{ + AUX_SLASH, CODE_CONT, CONT, Span, build_block, contr_len, digit_groups, fill_to_last, + lead_run, letter_match, member, run_end, scanthru, to_lead, trail_run, ws_tail, +}; + +/// Atom tag → dense 4-bit code. Unlike cl100k, `\p{M}` IS a letter here (both alt classes list +/// `\p{M}`) — but only a true mark: `AlphaSymMark` (0x16, categorically `\p{S}`) and `Zwj` (0x26, +/// `\p{Cf}`) stay "other", which is what keeps `[\p{L}\p{M}]+` off them. +const LUT: [u8; 64] = { + let mut t = [8u8; 64]; // other = [^\s\p{L}\p{N}] (incl. Connector/Punct/Sym/Control/ASM/ZWJ) + t[0x10] = 0; // \p{Lu} ∪ \p{Lt} + t[0x20] = 1; // \p{Ll} + t[0x00] = 2; // caseless letter (\p{Lm}\p{Lo}) — in BOTH alt classes + t[0x06] = 10; // true \p{M}: a letter for the alts, but ALSO in rule 4's class (see `mark_adj`) + t[0x01] = 3; + t[0x02] = 3; // \p{N} + t[0x03] = 4; // Newline + t[0x04] = 5; // Space + t[0x05] = 6; // WsOther + t[0x09] = 9; // Apostrophe — "other" for every run rule, split out for the contraction flag + t[0x0F] = CODE_CONT; + t +}; + +struct Cls { + u: u64, + l: u64, + c: u64, + n: u64, + nl: u64, + sp: u64, + ws: u64, + oth: u64, + mark: u64, + apo: u64, +} + +#[inline] +fn decode(p0: u64, p1: u64, p2: u64, p3: u64, valid: u64) -> Cls { + let low = !p3; + let a = low & !p2; + let w = low & p2; + Cls { + u: a & !p1 & !p0 & valid, // code 0 — past the block end every plane reads 0, hence `valid` + l: a & !p1 & p0, + c: a & p1 & !p0, + n: a & p1 & p0, + nl: w & !p1 & !p0, + sp: w & !p1 & p0, + ws: w & !(p1 & p0), // codes 4,5,6 — cont (7) excluded + oth: p3 & !p1, // codes 8,9 — the apostrophe is "other" for run purposes + mark: p3 & p1, // code 10 + apo: p3 & p0 & !p1, // code 9 + } +} + +const C_U: u16 = 1 << 0; +const C_L: u16 = 1 << 1; +const C_C: u16 = 1 << 2; +const C_N: u16 = 1 << 3; +const C_NL: u16 = 1 << 4; +const C_SP: u16 = 1 << 5; +const C_WSO: u16 = 1 << 6; +const C_OTH: u16 = 1 << 7; +const C_MARK: u16 = 1 << 8; +const C_WS: u16 = C_NL | C_SP | C_WSO; +const C_LET: u16 = C_U | C_L | C_C | C_MARK; + +const fn code_bits(code: u8) -> u16 { + match code { + 0 => C_U, + 1 => C_L, + 2 => C_C, + 3 => C_N, + 4 => C_NL, + 5 => C_SP, + 6 => C_WSO, + 8 | 9 => C_OTH, + 10 => C_MARK, + _ => 0, // cont never reaches an edge (the fill resolves it) + } +} + +/// o200k_base / GPT-4o — and byte-for-byte the same regex Llama-4, gpt-oss and MiniMax-M2 ship. +#[must_use] +pub fn bitsplit_o200k( + text: &[u8], + tags: &[u8], + starts: &mut [u64], + flag: &mut [u64], + out: &mut [Span], +) -> usize { + o200k(text, tags, starts, flag, out) +} + +fn o200k( + text: &[u8], + tags: &[u8], + starts: &mut [u64], + flag: &mut [u64], + out: &mut [Span], +) -> usize { + let ntext = text.len(); + if ntext == 0 { + return 0; + } + let nblk = ntext.div_ceil(64); + assert!( + tags.len() >= ntext && starts.len() >= nblk && flag.len() >= nblk && out.len() >= ntext + ); + + let (mut code, mut prev_cont) = (CODE_CONT, 0u64); + let (mut nl_run, mut prev_osf) = (false, false); + let mut prev_absorbed = false; // the block's last byte was eaten by a `[\r\n/]*` tail + let (mut dig_run, mut dig_since) = (false, 0u32); + let mut anl: Option = None; + let mut last_lt: Option = None; + let mut prev_start: Option = None; + + for bi in 0..nblk { + let base = bi * 64; + let len = (ntext - base).min(64); + let valid = if len == 64 { !0u64 } else { (1u64 << len) - 1 }; + let last_blk = base + len == ntext; + + let (b, last_code) = + build_block::<{ AUX_SLASH }, true>(text, tags, base, len, &LUT, code, false); + let c = decode(b.p0, b.p1, b.p2, b.p3, valid); + + let slash = b.aux; + + let pb = if base == 0 { + 0 + } else { + code_bits(code) + }; + let (nb, nb_lead) = if last_blk { + (0u16, true) + } else { + let q = base + len; + let is_lead = tags[q] != CONT; + let bits = if is_lead { + code_bits(LUT[tags[q] as usize]) + } else { + code_bits(last_code) + }; + (bits, is_lead) + }; + let has = |v: u16, s: u16| v & s != 0; + let p1 = |x: u64, k: bool| (x << 1) | u64::from(k); + let n1 = |x: u64, k: bool| (x >> 1) | (u64::from(k) << 63); + + let lead = valid & !b.cont; + let lb = ((lead >> 1) & valid) | (u64::from(nb_lead) << (len - 1)); + let eof_bit = if last_blk { 1u64 << (len - 1) } else { 0 }; + + // ── Han is peeled off the letter classes first: kimi's alt-1 outranks the letter alts, and + // its letter classes are literally `[…&&[^\p{Han}]]`. With AUX != HAN this is all zero. + let (cu, cl, cc) = (c.u, c.l, c.c); + let letter = cu | cl | cc | c.mark; + let pb_let = has(pb, C_LET); + + // Rule 4's `[\r\n/]*` tail. `/` is in BOTH the `+` body and the tail, so one tail run can + // collect SEVERAL markers (`!\n/\n`: the `\n` after `!` and the `\n` after `/`). That is + // exactly what `fill_to_last` is for -- `nl_e - nl_m` would span only from the last one. + let tail_cls = c.nl | slash; + let nl_m64 = ((p1(c.oth, has(pb, C_OTH)) & c.nl & lead) | u64::from(nl_run)) & tail_cls; + let nl_m = nl_m64 as u128; + let nl_e = scanthru(nl_m, tail_cls as u128); + let nl_span = fill_to_last(nl_m64, tail_cls) as u128; + + // ── rule 4 ` ?[^\s\p{L}\p{N}]+[\r\n/]*` and rules 5-7 — identical to cl100k. + // A char absorbed by a `[\r\n/]*` tail is NOT part of the `+` body, so an "other" after it + // opens a fresh run (`\u{1f600}\r\n/#` is the tail, then `#` starts again). + let o_prev = c.oth & !(nl_span as u64); + let o_start = c.oth + & lead + & !(nl_span as u64) // a char INSIDE the tail never opens a run + & !p1(o_prev, has(pb, C_OTH) && !prev_absorbed) + & !p1(c.sp, has(pb, C_SP)); + let ws_start = c.ws & lead & !p1(c.ws, has(pb, C_WS)); + let (steal, steal_patch) = to_lead( + c.ws & !c.nl & lb & !eof_bit & !n1(c.ws, has(nb, C_WS)), + b.cont, + prev_cont, + ); + let after_nl = p1(c.nl, has(pb, C_NL)) + & c.ws + & lead + & !fill_to_last(c.nl.reverse_bits(), c.ws.reverse_bits()).reverse_bits(); + + // ── `[^\r\n\p{L}\p{N}]?` before a letter run: any token-opening non-newline non-letter + // non-digit char. Smeared across its char's bytes so the test is a plain `p1` (see cl100k). + let mut osf = o_start; + osf |= (osf << 1) & b.cont; + osf |= (osf << 2) & b.cont & (b.cont << 1); + if prev_osf { + osf |= lead_run(b.cont, valid); + } + let l_start = letter + & lead + & !p1(letter, pb_let) + & !p1(c.ws & !c.nl, has(pb, C_WS) && !has(pb, C_NL)) + & !p1(osf, prev_osf); + // the token holding a letter run may OPEN one char earlier, on the `[^\r\n\p{L}\p{N}]?` + // prefix — that char is the start the escape has to be handed. The class genuinely excludes + // `\r\n` and digits: a newline before a letter run is its own token, never a prefix. + let prefix_cls = c.oth | c.sp | (c.ws & !c.nl); + let l_run = letter & lead & !p1(letter, pb_let); + let next_l_run = nb_lead + && has(nb, C_LET) + && letter >> (len - 1) & 1 == 0; + // `n1` only reads the NEXT char at a char's last byte, so mark there and walk back to the + // lead — a 3-byte prefix char (ZWJ, `—`, `½`) otherwise reads its own middle byte. + let (pfx_lead, pfx_patch) = + to_lead(prefix_cls & lb & n1(l_run, next_l_run), b.cont, prev_cont); + let l_start_tok = l_start | pfx_lead; + + // ── the escape gate (see the header). An interior upper, or an apostrophe closing the + // run, means some letter token in this block needs the scalar case/contraction pass. + let interior_u = cu & lead & p1(letter, pb_let); + let apo_after = c.apo & lead & p1(letter, pb_let); + // `\p{M}` is in BOTH the letter classes and rule 4's `[^\s\p{L}\p{N}]`, and which one wins + // depends on whether the punctuation before it STARTED the run (`!\u{301}a` is one token, + // `!!\u{301}a` is two). Adjacency is the gate; the scalar pass resolves it. Real text hits + // this via emoji + variation selector (U+FE0F is \p{Mn}). + let mark_adj = (c.mark & lead & p1(c.oth, has(pb, C_OTH))) + | (c.oth & lead & p1(c.mark, has(pb, C_MARK))); + + // ── rule 3 `\p{N}{1,DIGIT_CAP}` + let groups = + digit_groups(3, c.n, lead, b.cont, has(pb, C_N), dig_run, dig_since); + + // ── the other-run's `[\r\n/]*` tail (`/` only for the o200k line; kimi has `[\r\n]*`). + // ── the one backward-in-time rule, exactly as in cl100k/deepseek: a newline arriving now + // retracts an "after the last newline" start committed for a run still open at the edge. + if let Some(p) = anl + && has(pb, C_WS) + && c.ws & 1 != 0 + && c.nl & lead_run(c.ws, valid) != 0 + { + starts[p / 64] &= !(1u64 << (p % 64)); + anl = None; + } + + let mut st = groups | l_start | o_start | ws_start | after_nl | steal; + st &= !(nl_span as u64); + st |= nl_e as u64; + st &= lead; + if bi == 0 { + st |= 1; + } + starts[bi] = st; + // Conservative: flag every letter token in a block that shows either trigger. A false + // positive only costs a rescan of that token, so erring wide is free; missing one is not. + // Only the letter RUNS that actually need the escape, not every letter token in the block: + // one `'s` in a paragraph would otherwise drag every word through the scalar pass. + // `apo_after` sits on the apostrophe, one byte past the run, so shift it back inside. + let needs = interior_u | ((apo_after >> 1) & letter); + let runs_needing = if needs == 0 { + 0 + } else { + fill_to_last(needs.reverse_bits(), letter.reverse_bits()).reverse_bits() + }; + let (pfx_need, _) = to_lead( + prefix_cls & lb & n1(runs_needing & l_run, false), + b.cont, + prev_cont, + ); + let lt_all = st & l_start_tok; // every letter token — what the cross-block patch tracks + let lt = lt_all & (runs_needing | pfx_need); + let trig = (interior_u | apo_after) != 0; + flag[bi] = if trig { lt } else { 0 }; + // mark_adj is rare, so be blunt: escape every token in the block, plus the one still open + // from an earlier block (rule 4's ` ?` means the token can have started on a space). + if mark_adj != 0 { + flag[bi] |= st; + if let Some(p) = prev_start { + flag[p / 64] |= 1u64 << (p % 64); + } + } + if st != 0 { + prev_start = Some(base + 63 - st.leading_zeros() as usize); + } + if trig { + if bi > 0 { + flag[bi - 1] |= starts[bi - 1] & pfx_patch; + } + // the trigger can land in a LATER block than the token it belongs to (`\u{d55c}` ends + // block k, its `'s` opens block k+1), so always re-flag the last letter token seen. + if let Some(p) = last_lt { + flag[p / 64] |= 1u64 << (p % 64); + } + } + if bi > 0 { + starts[bi - 1] |= steal_patch; + } + // a run open at the block edge may meet its trigger in a later block, so flag it now + if lt_all != 0 { + last_lt = Some(base + 63 - lt_all.leading_zeros() as usize); + } + // ...ending on the run's `?` prefix char counts too — the token opened there. + let open_at_edge = (letter | l_start_tok) >> (len - 1) & 1 != 0; + if !last_blk && open_at_edge && let Some(p) = last_lt { + flag[p / 64] |= 1u64 << (p % 64); + } + + // ── carries + nl_run = nl_e >> 64 != 0; + prev_absorbed = nl_span >> (len - 1) & 1 != 0; + let tn = trail_run(c.n, valid, len); + dig_run = tn != 0 && has(nb, C_N); + dig_since = if !dig_run { + 0 + } else { + let g = groups & tn; + let counted = if g == 0 { + dig_since + (c.n & lead & tn).count_ones() + } else { + (c.n & lead & tn & !((1u64 << (63 - g.leading_zeros())) - 1)).count_ones() + }; + counted % 3 + }; + let tws = trail_run(c.ws, valid, len); + if tws != 0 && has(nb, C_WS) { + let a = after_nl & tws & !(nl_e as u64); + if a != 0 { + anl = Some(base + 63 - a.leading_zeros() as usize); + } else if !(tws & 1 != 0 && has(pb, C_WS)) { + anl = None; + } + } else { + anl = None; + } + prev_osf = osf >> (len - 1) & 1 != 0; + code = last_code; + prev_cont = b.cont; + + } + + emit_o200k(text, tags, starts, flag, nblk, ntext, out) +} + +// ── the scalar escape ─────────────────────────────────────────────────────────────────────────── +// Ported from the FSM this replaces, which is the proven reading of the alternatives. It runs from +// a flagged token start and RESYNCS the moment it lands on an algebra start bit again, so a +// divergent region costs a short scalar walk and nothing downstream. + +/// Emit ONE token starting at `i` (letters emit several) and return the cursor past it. +fn step( + text: &[u8], + tags: &[u8], + i: usize, + end: usize, + out: &mut [Span], + w: &mut usize, +) -> usize { + let is_lm = |p: usize| p < end && member(tags[p]); + let letter_end = |mut p: usize| { + while p < end && (tags[p] == CONT || member(tags[p])) { + p += 1; + } + p + }; + let other = |sp0: usize| { + let mut p = run_end(tags, sp0, end, mask::NOT_WS_L_N); + if p > sp0 { + while p < end && (tags[p] == crate::NLN || text[p] == b'/') { + p += char_len(text[p]); + } + } + p + }; + let emit1 = |a: usize, b: usize, out: &mut [Span], w: &mut usize| { + out[*w] = Span::new(a as u32, b as u32); + *w += 1; + }; + // the letter alternatives, with the case split and the optional contraction suffix + let letters = |pfx: usize, ls: usize, out: &mut [Span], w: &mut usize| -> usize { + let re = letter_end(ls); + let (mut p, mut first, mut cursor) = (ls, true, re); + while p < re { + let e = letter_match(tags, p, re); + let start = if first { pfx } else { p }; + let te = if e == re { + e + contr_len(text, e, true) + } else { + e + }; + out[*w] = Span::new(start as u32, te as u32); + *w += 1; + first = false; + cursor = te; + p = e; + } + cursor + }; + + let b = text[i]; + match tags[i] & 0x0F { + crate::NW | crate::NO => { + let (mut p, mut cnt) = (i, 0usize); + while p < end && cnt < 3 && in_mask(tags[p], mask::NUMBER) { + p += char_len(text[p]); + cnt += 1; + } + emit1(i, p, out, w); + p + } + crate::LET | crate::MRK => { + if member(tags[i]) { + return letters(i, i, out, w); + } + let a = i + char_len(b); + if is_lm(a) { + return letters(i, a, out, w); + } + let p = other(i); + emit1(i, p, out, w); + p + } + crate::SPC => { + let a = i + 1; + if is_lm(a) { + return letters(i, a, out, w); + } + let p = other(a); + let e = if p > a { p } else { ws_tail(text, tags, i, end) }; + emit1(i, e, out, w); + e + } + crate::WSO => { + let a = i + char_len(b); + if is_lm(a) { + return letters(i, a, out, w); + } + let e = ws_tail(text, tags, i, end); + emit1(i, e, out, w); + e + } + crate::NLN => { + let e = ws_tail(text, tags, i, end); + emit1(i, e, out, w); + e + } + _ => { + let a = i + char_len(b); + if is_lm(a) { + return letters(i, a, out, w); + } + let p = other(i); + emit1(i, p, out, w); + p + } + } +} + +/// `emit` plus the escape: at a flagged start, hand over to the scalar dispatch and take back +/// control at the first position that is a start bit again. +fn emit_o200k( + text: &[u8], + tags: &[u8], + starts: &[u64], + flag: &[u64], + nblk: usize, + n: usize, + out: &mut [Span], +) -> usize { + let (mut w, mut open, mut skip) = (0usize, usize::MAX, 0usize); + for bi in 0..nblk { + let mut m = starts[bi]; + let f = flag[bi]; + while m != 0 { + let j = m.trailing_zeros() as usize; + let pos = bi * 64 + j; + m &= m - 1; + if pos < skip { + continue; + } + if f >> j & 1 != 0 { + if open != usize::MAX && open < pos { + out[w] = Span::new(open as u32, pos as u32); + w += 1; + } + let mut c = pos; + loop { + c = step(text, tags, c, n, out, &mut w); + if c >= n || starts[c / 64] >> (c % 64) & 1 != 0 { + break; + } + } + open = usize::MAX; + skip = c; + continue; + } + if open != usize::MAX { + out[w] = Span::new(open as u32, pos as u32); + w += 1; + } + open = pos; + } + } + if open != usize::MAX { + out[w] = Span::new(open as u32, n as u32); + w += 1; + } + w +} diff --git a/tokenizers/bitsplit/src/models/tekken.rs b/tokenizers/bitsplit/src/models/tekken.rs new file mode 100644 index 000000000..e62ad14ed --- /dev/null +++ b/tokenizers/bitsplit/src/models/tekken.rs @@ -0,0 +1,525 @@ +//! Mistral **tekken** (mistral-small-4 / mistral-4): +//! `[^\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|…)?` +//! `|` same with `+`/`*` swapped `| \p{N} | ?[^\s\p{L}\p{N}]+[\r\n/]* | \s*[\r\n]+ | \s+(?!\S) | \s+` +//! +//! o200k's grammar with two changes: letter tokens take NO contraction suffix, and the digit rule +//! is a bare `\p{N}` — one token per digit. +//! + +//! +//! Two things are worth knowing before reading the algebra: +//! +//! 1. **The case split is a scalar escape, on purpose.** `[UC]*[LC]+ | [UC]+[LC]*` has no local +//! form: `中Qz` is ONE token (a U after a C is not a boundary) but `ʰABC` is two (it is, when no +//! L follows) — the difference is whether an L appears LATER in the run, so no `p1` decides it. +//! The bit half instead computes a cheap gate: a letter token is escaped only if the block holds +//! an interior upper or a trailing apostrophe, so all-lowercase and Capitalised text never pays. +//! 2. **The contraction is a SUFFIX here, not an alternative.** cl100k emits `'t` as its own token; +//! o200k glues it onto the letter token before it. So this uses `emit_contr_suffix`, which +//! *extends* the open token instead of opening one — the mirror image of `emit_contr`. +//! +use crate::classify::{char_len, in_mask, mask}; +use crate::{ + AUX_SLASH, CODE_CONT, CONT, Span, build_block, digit_groups, fill_to_last, + lead_run, letter_match, member, run_end, scanthru, to_lead, trail_run, ws_tail, +}; + +/// Atom tag → dense 4-bit code. Unlike cl100k, `\p{M}` IS a letter here (both alt classes list +/// `\p{M}`) — but only a true mark: `AlphaSymMark` (0x16, categorically `\p{S}`) and `Zwj` (0x26, +/// `\p{Cf}`) stay "other", which is what keeps `[\p{L}\p{M}]+` off them. +const LUT: [u8; 64] = { + let mut t = [8u8; 64]; // other = [^\s\p{L}\p{N}] (incl. Connector/Punct/Sym/Control/ASM/ZWJ) + t[0x10] = 0; // \p{Lu} ∪ \p{Lt} + t[0x20] = 1; // \p{Ll} + t[0x00] = 2; // caseless letter (\p{Lm}\p{Lo}) — in BOTH alt classes + t[0x06] = 10; // true \p{M}: a letter for the alts, but ALSO in rule 4's class (see `mark_adj`) + t[0x01] = 3; + t[0x02] = 3; // \p{N} + t[0x03] = 4; // Newline + t[0x04] = 5; // Space + t[0x05] = 6; // WsOther + t[0x09] = 9; // Apostrophe — "other" for every run rule, split out for the contraction flag + t[0x0F] = CODE_CONT; + t +}; + +struct Cls { + u: u64, + l: u64, + c: u64, + n: u64, + nl: u64, + sp: u64, + ws: u64, + oth: u64, + mark: u64, + apo: u64, +} + +#[inline] +fn decode(p0: u64, p1: u64, p2: u64, p3: u64, valid: u64) -> Cls { + let low = !p3; + let a = low & !p2; + let w = low & p2; + Cls { + u: a & !p1 & !p0 & valid, // code 0 — past the block end every plane reads 0, hence `valid` + l: a & !p1 & p0, + c: a & p1 & !p0, + n: a & p1 & p0, + nl: w & !p1 & !p0, + sp: w & !p1 & p0, + ws: w & !(p1 & p0), // codes 4,5,6 — cont (7) excluded + oth: p3 & !p1, // codes 8,9 — the apostrophe is "other" for run purposes + mark: p3 & p1, // code 10 + apo: p3 & p0 & !p1, // code 9 + } +} + +const C_U: u16 = 1 << 0; +const C_L: u16 = 1 << 1; +const C_C: u16 = 1 << 2; +const C_N: u16 = 1 << 3; +const C_NL: u16 = 1 << 4; +const C_SP: u16 = 1 << 5; +const C_WSO: u16 = 1 << 6; +const C_OTH: u16 = 1 << 7; +const C_MARK: u16 = 1 << 8; +const C_WS: u16 = C_NL | C_SP | C_WSO; +const C_LET: u16 = C_U | C_L | C_C | C_MARK; + +const fn code_bits(code: u8) -> u16 { + match code { + 0 => C_U, + 1 => C_L, + 2 => C_C, + 3 => C_N, + 4 => C_NL, + 5 => C_SP, + 6 => C_WSO, + 8 | 9 => C_OTH, + 10 => C_MARK, + _ => 0, // cont never reaches an edge (the fill resolves it) + } +} + +/// Mistral tekken. +#[must_use] +pub fn bitsplit_tekken( + text: &[u8], + tags: &[u8], + starts: &mut [u64], + flag: &mut [u64], + out: &mut [Span], +) -> usize { + tekken(text, tags, starts, flag, out) +} + +fn tekken( + text: &[u8], + tags: &[u8], + starts: &mut [u64], + flag: &mut [u64], + out: &mut [Span], +) -> usize { + let ntext = text.len(); + if ntext == 0 { + return 0; + } + let nblk = ntext.div_ceil(64); + assert!( + tags.len() >= ntext && starts.len() >= nblk && flag.len() >= nblk && out.len() >= ntext + ); + + let (mut code, mut prev_cont) = (CODE_CONT, 0u64); + let (mut nl_run, mut prev_osf) = (false, false); + let mut prev_absorbed = false; // the block's last byte was eaten by a `[\r\n/]*` tail + let (mut dig_run, mut dig_since) = (false, 0u32); + let mut anl: Option = None; + let mut last_lt: Option = None; + let mut prev_start: Option = None; + + for bi in 0..nblk { + let base = bi * 64; + let len = (ntext - base).min(64); + let valid = if len == 64 { !0u64 } else { (1u64 << len) - 1 }; + let last_blk = base + len == ntext; + + let (b, last_code) = + build_block::<{ AUX_SLASH }, true>(text, tags, base, len, &LUT, code, false); + let c = decode(b.p0, b.p1, b.p2, b.p3, valid); + + let slash = b.aux; + + let pb = if base == 0 { + 0 + } else { + code_bits(code) + }; + let (nb, nb_lead) = if last_blk { + (0u16, true) + } else { + let q = base + len; + let is_lead = tags[q] != CONT; + let bits = if is_lead { + code_bits(LUT[tags[q] as usize]) + } else { + code_bits(last_code) + }; + (bits, is_lead) + }; + let has = |v: u16, s: u16| v & s != 0; + let p1 = |x: u64, k: bool| (x << 1) | u64::from(k); + let n1 = |x: u64, k: bool| (x >> 1) | (u64::from(k) << 63); + + let lead = valid & !b.cont; + let lb = ((lead >> 1) & valid) | (u64::from(nb_lead) << (len - 1)); + let eof_bit = if last_blk { 1u64 << (len - 1) } else { 0 }; + + // ── Han is peeled off the letter classes first: kimi's alt-1 outranks the letter alts, and + // its letter classes are literally `[…&&[^\p{Han}]]`. With AUX != HAN this is all zero. + let (cu, cl, cc) = (c.u, c.l, c.c); + let letter = cu | cl | cc | c.mark; + let pb_let = has(pb, C_LET); + + // Rule 4's `[\r\n/]*` tail. `/` is in BOTH the `+` body and the tail, so one tail run can + // collect SEVERAL markers (`!\n/\n`: the `\n` after `!` and the `\n` after `/`). That is + // exactly what `fill_to_last` is for -- `nl_e - nl_m` would span only from the last one. + let tail_cls = c.nl | slash; + let nl_m64 = ((p1(c.oth, has(pb, C_OTH)) & c.nl & lead) | u64::from(nl_run)) & tail_cls; + let nl_m = nl_m64 as u128; + let nl_e = scanthru(nl_m, tail_cls as u128); + let nl_span = fill_to_last(nl_m64, tail_cls) as u128; + + // ── rule 4 ` ?[^\s\p{L}\p{N}]+[\r\n/]*` and rules 5-7 — identical to cl100k. + // A char absorbed by a `[\r\n/]*` tail is NOT part of the `+` body, so an "other" after it + // opens a fresh run (`\u{1f600}\r\n/#` is the tail, then `#` starts again). + let o_prev = c.oth & !(nl_span as u64); + let o_start = c.oth + & lead + & !(nl_span as u64) // a char INSIDE the tail never opens a run + & !p1(o_prev, has(pb, C_OTH) && !prev_absorbed) + & !p1(c.sp, has(pb, C_SP)); + let ws_start = c.ws & lead & !p1(c.ws, has(pb, C_WS)); + let (steal, steal_patch) = to_lead( + c.ws & !c.nl & lb & !eof_bit & !n1(c.ws, has(nb, C_WS)), + b.cont, + prev_cont, + ); + let after_nl = p1(c.nl, has(pb, C_NL)) + & c.ws + & lead + & !fill_to_last(c.nl.reverse_bits(), c.ws.reverse_bits()).reverse_bits(); + + // ── `[^\r\n\p{L}\p{N}]?` before a letter run: any token-opening non-newline non-letter + // non-digit char. Smeared across its char's bytes so the test is a plain `p1` (see cl100k). + let mut osf = o_start; + osf |= (osf << 1) & b.cont; + osf |= (osf << 2) & b.cont & (b.cont << 1); + if prev_osf { + osf |= lead_run(b.cont, valid); + } + let l_start = letter + & lead + & !p1(letter, pb_let) + & !p1(c.ws & !c.nl, has(pb, C_WS) && !has(pb, C_NL)) + & !p1(osf, prev_osf); + // the token holding a letter run may OPEN one char earlier, on the `[^\r\n\p{L}\p{N}]?` + // prefix — that char is the start the escape has to be handed. The class genuinely excludes + // `\r\n` and digits: a newline before a letter run is its own token, never a prefix. + let prefix_cls = c.oth | c.sp | (c.ws & !c.nl); + let l_run = letter & lead & !p1(letter, pb_let); + let next_l_run = nb_lead + && has(nb, C_LET) + && letter >> (len - 1) & 1 == 0; + // `n1` only reads the NEXT char at a char's last byte, so mark there and walk back to the + // lead — a 3-byte prefix char (ZWJ, `—`, `½`) otherwise reads its own middle byte. + let (pfx_lead, pfx_patch) = + to_lead(prefix_cls & lb & n1(l_run, next_l_run), b.cont, prev_cont); + let l_start_tok = l_start | pfx_lead; + + // ── the escape gate (see the header). An interior upper, or an apostrophe closing the + // run, means some letter token in this block needs the scalar case/contraction pass. + let interior_u = cu & lead & p1(letter, pb_let); + // no contraction suffix here, so an apostrophe after a letter is nothing special + // `\p{M}` is in BOTH the letter classes and rule 4's `[^\s\p{L}\p{N}]`, and which one wins + // depends on whether the punctuation before it STARTED the run (`!\u{301}a` is one token, + // `!!\u{301}a` is two). Adjacency is the gate; the scalar pass resolves it. Real text hits + // this via emoji + variation selector (U+FE0F is \p{Mn}). + let mark_adj = (c.mark & lead & p1(c.oth, has(pb, C_OTH))) + | (c.oth & lead & p1(c.mark, has(pb, C_MARK))); + + // ── rule 3 `\p{N}{1,DIGIT_CAP}` + let groups = + digit_groups(1, c.n, lead, b.cont, has(pb, C_N), dig_run, dig_since); + + // ── the other-run's `[\r\n/]*` tail (`/` only for the o200k line; kimi has `[\r\n]*`). + // ── the one backward-in-time rule, exactly as in cl100k/deepseek: a newline arriving now + // retracts an "after the last newline" start committed for a run still open at the edge. + if let Some(p) = anl + && has(pb, C_WS) + && c.ws & 1 != 0 + && c.nl & lead_run(c.ws, valid) != 0 + { + starts[p / 64] &= !(1u64 << (p % 64)); + anl = None; + } + + let mut st = groups | l_start | o_start | ws_start | after_nl | steal; + st &= !(nl_span as u64); + st |= nl_e as u64; + st &= lead; + if bi == 0 { + st |= 1; + } + starts[bi] = st; + // Conservative: flag every letter token in a block that shows either trigger. A false + // positive only costs a rescan of that token, so erring wide is free; missing one is not. + // Only the letter RUNS that actually need the escape, not every letter token in the block: + // one `'s` in a paragraph would otherwise drag every word through the scalar pass. + // `apo_after` sits on the apostrophe, one byte past the run, so shift it back inside. + let needs = interior_u; + let runs_needing = if needs == 0 { + 0 + } else { + fill_to_last(needs.reverse_bits(), letter.reverse_bits()).reverse_bits() + }; + let (pfx_need, _) = to_lead( + prefix_cls & lb & n1(runs_needing & l_run, false), + b.cont, + prev_cont, + ); + let lt_all = st & l_start_tok; // every letter token — what the cross-block patch tracks + let lt = lt_all & (runs_needing | pfx_need); + let trig = interior_u != 0; + flag[bi] = if trig { lt } else { 0 }; + // mark_adj is rare, so be blunt: escape every token in the block, plus the one still open + // from an earlier block (rule 4's ` ?` means the token can have started on a space). + if mark_adj != 0 { + flag[bi] |= st; + if let Some(p) = prev_start { + flag[p / 64] |= 1u64 << (p % 64); + } + } + if st != 0 { + prev_start = Some(base + 63 - st.leading_zeros() as usize); + } + if trig { + if bi > 0 { + flag[bi - 1] |= starts[bi - 1] & pfx_patch; + } + // the trigger can land in a LATER block than the token it belongs to (`\u{d55c}` ends + // block k, its `'s` opens block k+1), so always re-flag the last letter token seen. + if let Some(p) = last_lt { + flag[p / 64] |= 1u64 << (p % 64); + } + } + if bi > 0 { + starts[bi - 1] |= steal_patch; + } + // a run open at the block edge may meet its trigger in a later block, so flag it now + if lt_all != 0 { + last_lt = Some(base + 63 - lt_all.leading_zeros() as usize); + } + // ...ending on the run's `?` prefix char counts too — the token opened there. + let open_at_edge = (letter | l_start_tok) >> (len - 1) & 1 != 0; + if !last_blk && open_at_edge && let Some(p) = last_lt { + flag[p / 64] |= 1u64 << (p % 64); + } + + // ── carries + nl_run = nl_e >> 64 != 0; + prev_absorbed = nl_span >> (len - 1) & 1 != 0; + let tn = trail_run(c.n, valid, len); + dig_run = tn != 0 && has(nb, C_N); + dig_since = if !dig_run { + 0 + } else { + let g = groups & tn; + let _counted = if g == 0 { + dig_since + (c.n & lead & tn).count_ones() + } else { + (c.n & lead & tn & !((1u64 << (63 - g.leading_zeros())) - 1)).count_ones() + }; + 0 // digit cap 1: every digit is its own group, so nothing carries + }; + let tws = trail_run(c.ws, valid, len); + if tws != 0 && has(nb, C_WS) { + let a = after_nl & tws & !(nl_e as u64); + if a != 0 { + anl = Some(base + 63 - a.leading_zeros() as usize); + } else if !(tws & 1 != 0 && has(pb, C_WS)) { + anl = None; + } + } else { + anl = None; + } + prev_osf = osf >> (len - 1) & 1 != 0; + code = last_code; + prev_cont = b.cont; + + } + + emit_tekken(text, tags, starts, flag, nblk, ntext, out) +} + +// ── the scalar escape ─────────────────────────────────────────────────────────────────────────── +// Ported from the FSM this replaces, which is the proven reading of the alternatives. It runs from +// a flagged token start and RESYNCS the moment it lands on an algebra start bit again, so a +// divergent region costs a short scalar walk and nothing downstream. + +/// Emit ONE token starting at `i` (letters emit several) and return the cursor past it. +fn step( + text: &[u8], + tags: &[u8], + i: usize, + end: usize, + out: &mut [Span], + w: &mut usize, +) -> usize { + let is_lm = |p: usize| p < end && member(tags[p]); + let letter_end = |mut p: usize| { + while p < end && (tags[p] == CONT || member(tags[p])) { + p += 1; + } + p + }; + let other = |sp0: usize| { + let mut p = run_end(tags, sp0, end, mask::NOT_WS_L_N); + if p > sp0 { + while p < end && (tags[p] == crate::NLN || text[p] == b'/') { + p += char_len(text[p]); + } + } + p + }; + let emit1 = |a: usize, b: usize, out: &mut [Span], w: &mut usize| { + out[*w] = Span::new(a as u32, b as u32); + *w += 1; + }; + // the letter alternatives, with the case split and the optional contraction suffix + let letters = |pfx: usize, ls: usize, out: &mut [Span], w: &mut usize| -> usize { + let re = letter_end(ls); + let (mut p, mut first, mut cursor) = (ls, true, re); + while p < re { + let e = letter_match(tags, p, re); + let start = if first { pfx } else { p }; + let te = e; // tekken has no contraction suffix + out[*w] = Span::new(start as u32, te as u32); + *w += 1; + first = false; + cursor = te; + p = e; + } + cursor + }; + + let b = text[i]; + match tags[i] & 0x0F { + crate::NW | crate::NO => { + let (mut p, mut cnt) = (i, 0usize); + while p < end && cnt < 1 && in_mask(tags[p], mask::NUMBER) { + p += char_len(text[p]); + cnt += 1; + } + emit1(i, p, out, w); + p + } + crate::LET | crate::MRK => { + if member(tags[i]) { + return letters(i, i, out, w); + } + let a = i + char_len(b); + if is_lm(a) { + return letters(i, a, out, w); + } + let p = other(i); + emit1(i, p, out, w); + p + } + crate::SPC => { + let a = i + 1; + if is_lm(a) { + return letters(i, a, out, w); + } + let p = other(a); + let e = if p > a { p } else { ws_tail(text, tags, i, end) }; + emit1(i, e, out, w); + e + } + crate::WSO => { + let a = i + char_len(b); + if is_lm(a) { + return letters(i, a, out, w); + } + let e = ws_tail(text, tags, i, end); + emit1(i, e, out, w); + e + } + crate::NLN => { + let e = ws_tail(text, tags, i, end); + emit1(i, e, out, w); + e + } + _ => { + let a = i + char_len(b); + if is_lm(a) { + return letters(i, a, out, w); + } + let p = other(i); + emit1(i, p, out, w); + p + } + } +} + +/// `emit` plus the escape: at a flagged start, hand over to the scalar dispatch and take back +/// control at the first position that is a start bit again. +fn emit_tekken( + text: &[u8], + tags: &[u8], + starts: &[u64], + flag: &[u64], + nblk: usize, + n: usize, + out: &mut [Span], +) -> usize { + let (mut w, mut open, mut skip) = (0usize, usize::MAX, 0usize); + for bi in 0..nblk { + let mut m = starts[bi]; + let f = flag[bi]; + while m != 0 { + let j = m.trailing_zeros() as usize; + let pos = bi * 64 + j; + m &= m - 1; + if pos < skip { + continue; + } + if f >> j & 1 != 0 { + if open != usize::MAX && open < pos { + out[w] = Span::new(open as u32, pos as u32); + w += 1; + } + let mut c = pos; + loop { + c = step(text, tags, c, n, out, &mut w); + if c >= n || starts[c / 64] >> (c % 64) & 1 != 0 { + break; + } + } + open = usize::MAX; + skip = c; + continue; + } + if open != usize::MAX { + out[w] = Span::new(open as u32, pos as u32); + w += 1; + } + open = pos; + } + } + if open != usize::MAX { + out[w] = Span::new(open as u32, n as u32); + w += 1; + } + w +} diff --git a/tokenizers/atomsplit/src/regexes.rs b/tokenizers/bitsplit/src/regexes.rs similarity index 77% rename from tokenizers/atomsplit/src/regexes.rs rename to tokenizers/bitsplit/src/regexes.rs index 875fb892e..3f3a0662d 100644 --- a/tokenizers/atomsplit/src/regexes.rs +++ b/tokenizers/bitsplit/src/regexes.rs @@ -31,3 +31,9 @@ pub const DEEPSEEK_BIG: &str = r##"[!"#$%&'()*+,\-./:;<=>?@\[\\\]^_`{|}~][A-Za-z /// The deepseek chain in application order — convenience for the multi-regex reference. pub const DEEPSEEK: &[&str] = &[DEEPSEEK_NUM, DEEPSEEK_CJK, DEEPSEEK_BIG]; + +/// kimi-k2 / k3 — `moonshotai/Kimi-K2-Instruct`'s `tokenization_kimi.py` `pat_str`. o200k plus a +/// leading `[\p{Han}]+` arm, Han subtracted from both letter classes, and a `[\r\n]*` rule-4 tail +/// (o200k has `[\r\n/]*`). Kimi ships `tiktoken.model` rather than a `tokenizer.json`, so this is +/// the pattern as a converted tokenizer would spell it. Reproduced by [`crate::bitsplit_kimi`]. +pub const KIMI_K2: &str = r"[\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+"; diff --git a/tokenizers/atomsplit/src/simd_fsm.rs b/tokenizers/bitsplit/src/simd/classes.rs similarity index 99% rename from tokenizers/atomsplit/src/simd_fsm.rs rename to tokenizers/bitsplit/src/simd/classes.rs index ed6b172cb..b6700e1c4 100644 --- a/tokenizers/atomsplit/src/simd_fsm.rs +++ b/tokenizers/bitsplit/src/simd/classes.rs @@ -8,7 +8,8 @@ // the non-selected arch some items/imports (char_len, emit_class_spans, ...) are legitimately unused use crate::classify::{Atom, char_len}; -use crate::fsm::{Span, emit_class_spans}; +use crate::Span; +use crate::classes::emit_class_spans; /// Class LookUpTable: tag → 0 drop / 1 isolate / 2 keep-A / 3 keep-B; Cont → 0xFF (fill sentinel). /// This lookup table is built per parameter DROP, ISOLATE and KEEP_A. These as diff --git a/tokenizers/bitsplit/src/simd/mod.rs b/tokenizers/bitsplit/src/simd/mod.rs new file mode 100644 index 000000000..76161133a --- /dev/null +++ b/tokenizers/bitsplit/src/simd/mod.rs @@ -0,0 +1,12 @@ +//! Per-arch kernels. Every one of these is PURE PERF and byte-exact with a portable path that is +//! always compiled, so correctness never depends on a kernel being present: +//! - [`neon`] / [`x86`] build the block's bitstreams; `build_block_scalar` is the reference. +//! - [`classes`] extracts class-run boundaries; `emit_class_spans` is the reference. +//! +//! (`classify` has its own kernels, next to the tables they index.) +#[cfg(target_arch = "aarch64")] +pub(crate) mod neon; +#[cfg(target_arch = "x86_64")] +pub(crate) mod x86; + +pub(crate) mod classes; diff --git a/tokenizers/bitsplit/src/simd/neon.rs b/tokenizers/bitsplit/src/simd/neon.rs new file mode 100644 index 000000000..e685ab7ba --- /dev/null +++ b/tokenizers/bitsplit/src/simd/neon.rs @@ -0,0 +1,229 @@ +//! NEON block builder: 64 tag bytes + 64 text bytes → the 10 class bitstreams of one [`Blk`]. +//! +//! This is the step the paper spends its GPU budget on too — turning bytes into bitstreams. The +//! shape here is a 64×8 bit-matrix transpose done 8 lanes at a time: weight each compare result by +//! its lane's power of two (`POW`) and let three `vpaddq_u8` rounds fold 64 lanes into the 8 bytes +//! of one `u64`. That is ~9 ops per stream per 64 bytes, versus 4 separate 16-bit movemasks. +//! +//! Continuation bytes are resolved **before** extraction (≤3 `vext`+`vbsl`, the same trick +//! `atomsplit::simd_fsm` uses), so every stream comes out *filled* — a multi-byte char sets its bit +//! on all of its bytes. That is what lets the bitstream program read "previous char's class" as a +//! plain `<< 1` with no char-width arithmetic. + +use crate::{AUX_HAN, AUX_NONE, AUX_SLASH, Blk, lead_run}; +use core::arch::aarch64::*; + +const POW: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128]; + +/// 64 lanes of `0x00`/`0xFF` → one `u64`. Three `vpaddq_u8` rounds: lane pairs, then quads, then +/// octets — after which the low 8 bytes of the result are the 8 bytes of the mask, in order. +#[inline(always)] +unsafe fn mm64(v: [uint8x16_t; 4], pow: uint8x16_t) -> u64 { + unsafe { + let ab = vpaddq_u8(vandq_u8(v[0], pow), vandq_u8(v[1], pow)); + let cd = vpaddq_u8(vandq_u8(v[2], pow), vandq_u8(v[3], pow)); + let x = vpaddq_u8(ab, cd); + vgetq_lane_u64::<0>(vreinterpretq_u64_u8(vpaddq_u8(x, x))) + } +} + +/// 64 bytes of `text[base..]` compared against `b`, as a bitmap. Same fold the block builder uses. +/// +/// # Safety +/// `base + 64 <= text.len()`. +#[target_feature(enable = "neon")] +pub(crate) unsafe fn eq64(text: &[u8], base: usize, b: u8) -> u64 { + unsafe { + let pow = vld1q_u8(POW.as_ptr()); + let n = vdupq_n_u8(b); + mm64( + [ + vceqq_u8(vld1q_u8(text.as_ptr().add(base)), n), + vceqq_u8(vld1q_u8(text.as_ptr().add(base + 16)), n), + vceqq_u8(vld1q_u8(text.as_ptr().add(base + 32)), n), + vceqq_u8(vld1q_u8(text.as_ptr().add(base + 48)), n), + ], + pow, + ) + } +} + +/// Build one **full** 64-byte block, folding tags through `lut` into dense codes. +/// Build one **full** 64-byte block. `cur_code` / `cur_cjk` describe the byte just before it, so a +/// block opening mid-char keeps inheriting its lead's class. Returns the block's last filled code. +/// +/// # Safety +/// `base + 64 <= tags.len()` and `base + 64 <= text.len()`. +#[target_feature(enable = "neon")] +pub(crate) unsafe fn build64( + text: &[u8], + tags: &[u8], + base: usize, + lut: &[u8; 64], + cur_code: u8, + cur_aux: bool, +) -> (Blk, u8) { + unsafe { + let pow = vld1q_u8(POW.as_ptr()); + let seven = vdupq_n_u8(7); + let tbl = uint8x16x4_t( + vld1q_u8(lut.as_ptr()), + vld1q_u8(lut.as_ptr().add(16)), + vld1q_u8(lut.as_ptr().add(32)), + vld1q_u8(lut.as_ptr().add(48)), + ); + + // ── tags → dense codes, then fill the continuation lanes from the left so every lane + // carries its char's code. Two steps suffice, not three: after shifting by 1 every lane is + // correct at distance 1, so shifting *that* by 2 covers distances 2 and 3 — which is the + // most a 4-byte char can need. + let mut isc = [vdupq_n_u8(0); 4]; + let mut cd = [vdupq_n_u8(0); 4]; + let mut prev = vdupq_n_u8(cur_code); + for k in 0..4 { + let raw = vqtbl4q_u8(tbl, vld1q_u8(tags.as_ptr().add(base + k * 16))); + isc[k] = vceqq_u8(raw, seven); + let c = vbslq_u8(isc[k], vextq_u8::<15>(prev, raw), raw); + let c = vbslq_u8(vceqq_u8(c, seven), vextq_u8::<14>(prev, c), c); + prev = c; + cd[k] = c; + } + let last_code = vgetq_lane_u8::<15>(prev); + // ── 3 bit-planes of the code. Every stream is a boolean function of these (`decode`), so + // this is 3 extractions where one-hot class bits needed 6 — and `cont` comes from `isc`, + // which the fill needed anyway. + let plane = |bit: u8| { + let d = vdupq_n_u8(bit); + mm64( + [ + vtstq_u8(cd[0], d), + vtstq_u8(cd[1], d), + vtstq_u8(cd[2], d), + vtstq_u8(cd[3], d), + ], + pow, + ) + }; + + let mut b = Blk { + cont: mm64(isc, pow), + p0: plane(1), + p1: plane(2), + p2: plane(4), + p3: if P3 { plane(8) } else { 0 }, + aux: 0, + }; + if AUX == AUX_NONE { + return (b, last_code); + } + // ── text is loaded only for the aux (text-derived) stream. + let ntext = text.len(); + let tv = [ + vld1q_u8(text.as_ptr().add(base)), + vld1q_u8(text.as_ptr().add(base + 16)), + vld1q_u8(text.as_ptr().add(base + 32)), + vld1q_u8(text.as_ptr().add(base + 48)), + ]; + + if AUX == AUX_SLASH { + // single ASCII byte — no fill needed, `/` is its own char + let sl = vdupq_n_u8(b'/'); + b.aux = mm64( + [ + vceqq_u8(tv[0], sl), + vceqq_u8(tv[1], sl), + vceqq_u8(tv[2], sl), + vceqq_u8(tv[3], sl), + ], + pow, + ); + return (b, last_code); + } + if AUX == AUX_HAN { + // ponytail: scalar Han range test; vectorise like the CJK path below if kimi ever + // shows up on a throughput bench. + let lim = ntext.min(base + 64); + let mut leads = 0u64; + for p in base..lim { + if tags[p] != crate::CONT && crate::han::is_han_at(text, p) { + leads |= 1u64 << (p - base); + } + } + b.aux = leads | (leads << 1) | (leads << 2); + if cur_aux { + b.aux |= lead_run(b.cont, !0); + } + return (b, last_code); + } + + // ── the CJK range test lives in the raw bytes, not the tags. It is a 3-byte predicate, so + // it runs in vector space on `vext`-aligned b1/b2 and extracts ONE mask — testing the 9 + // byte predicates as separate bitstreams costs 9 extractions for the same answer. A + // `vmaxvq` over the lead range gates the whole thing away on Latin/code for ~8 ops. + let e3e9 = |v| vcleq_u8(vsubq_u8(v, vdupq_n_u8(0xE3)), vdupq_n_u8(6)); + let any = vmaxvq_u8(vorrq_u8( + vorrq_u8(e3e9(tv[0]), e3e9(tv[1])), + vorrq_u8(e3e9(tv[2]), e3e9(tv[3])), + )); + if any != 0 { + // b1/b2 of a char at lane 15 come from the next chunk — and for the last chunk, from + // the next block, which may not exist. + let tail = { + let mut buf = [0u8; 16]; + let off = base + 64; + let avail = text.len().saturating_sub(off).min(16); + buf[..avail].copy_from_slice(&text[off..off + avail]); + vld1q_u8(buf.as_ptr()) + }; + // Hiragana/Katakana U+3040..30FF is exactly E3 [81-83] xx; Han U+4E00..9FA5 is + // E4 [B8-BF] xx / E5-E8 xx xx / E9 [80-BD] xx / E9 BE [80-A5]. + let cjkv = |v: uint8x16_t, nx: uint8x16_t| { + let (b1, b2) = (vextq_u8::<1>(v, nx), vextq_u8::<2>(v, nx)); + let e9 = vandq_u8( + vceqq_u8(v, vdupq_n_u8(0xE9)), + vorrq_u8( + vcltq_u8(b1, vdupq_n_u8(0xBE)), + vandq_u8( + vceqq_u8(b1, vdupq_n_u8(0xBE)), + vcleq_u8(b2, vdupq_n_u8(0xA5)), + ), + ), + ); + vorrq_u8( + vorrq_u8( + vandq_u8( + vceqq_u8(v, vdupq_n_u8(0xE3)), + vcleq_u8(vsubq_u8(b1, vdupq_n_u8(0x81)), vdupq_n_u8(2)), + ), + vandq_u8( + vceqq_u8(v, vdupq_n_u8(0xE4)), + vcgeq_u8(b1, vdupq_n_u8(0xB8)), + ), + ), + vorrq_u8(vcleq_u8(vsubq_u8(v, vdupq_n_u8(0xE5)), vdupq_n_u8(3)), e9), + ) + }; + let mut leads = mm64( + [ + cjkv(tv[0], tv[1]), + cjkv(tv[1], tv[2]), + cjkv(tv[2], tv[3]), + cjkv(tv[3], tail), + ], + pow, + ); + // `fsm_deepseek` reads 3 bytes unconditionally; refuse to classify a truncated tail. + let lim = ntext.saturating_sub(base + 2); + if lim < 64 { + leads &= (1u64 << lim) - 1; + } + // every CJK char is 3 bytes → fill by two shifts; a char cut by the block edge is + // picked up on the other side by `cur_aux` (its continuation bytes lead that block). + b.aux = leads | (leads << 1) | (leads << 2); + } + if cur_aux { + b.aux |= lead_run(b.cont, !0); + } + (b, last_code) + } +} diff --git a/tokenizers/bitsplit/src/simd/x86.rs b/tokenizers/bitsplit/src/simd/x86.rs new file mode 100644 index 000000000..407f5e481 --- /dev/null +++ b/tokenizers/bitsplit/src/simd/x86.rs @@ -0,0 +1,238 @@ +//! SSSE3 block builder: the x86 twin of `simd.rs`. Same shape — tags through a LUT into dense +//! codes, continuation lanes filled from the left, three bit-planes extracted — with two +//! differences that fall out of the ISA. +//! +//! Plane extraction is cheaper here: `movemask` is native, so a plane is +//! `movemask(slli_epi16(code, 7 - bit))`. Shifting 16-bit lanes puts bit `k` of the low byte at +//! bit 7 and bit `k` of the high byte at bit 15, which is exactly the pair `movemask` reads; the +//! bits that bleed across the byte boundary land where it does not look. +//! +//! The LUT is dearer: `pshufb` indexes 16 entries and the tags run to 0x26, so the table is split +//! into one shuffle per high nibble and selected between. Refinements (`0x10`/`0x20` case, `0x16` +//! AlphaSymMark, `0x26` ZWJ) are the only tags above 0x0F, so three shuffles cover every case. +//! +//! Checked the same way as the NEON kernel: `src/bin/verify.rs` built for `x86_64-apple-darwin` +//! and run under Rosetta, which reports SSSE3, so this path is the one that executes. The scalar +//! builder in `lib.rs` stays the reference it is compared against. + +use crate::{AUX_CJK, AUX_HAN, AUX_NONE, AUX_SLASH, Blk, lead_run}; +use core::arch::x86_64::*; + +/// `x <= k`, unsigned, in the absence of an unsigned byte compare. +#[inline(always)] +unsafe fn le(x: __m128i, k: u8) -> __m128i { + unsafe { _mm_cmpeq_epi8(_mm_min_epu8(x, _mm_set1_epi8(k as i8)), x) } +} + +/// `x >= k`, unsigned. +#[inline(always)] +unsafe fn ge(x: __m128i, k: u8) -> __m128i { + unsafe { _mm_cmpeq_epi8(_mm_max_epu8(x, _mm_set1_epi8(k as i8)), x) } +} + +/// `x - lo <= n`, unsigned, i.e. `x` in `lo ..= lo + n`. +#[inline(always)] +unsafe fn in_range(x: __m128i, lo: u8, n: u8) -> __m128i { + unsafe { le(_mm_sub_epi8(x, _mm_set1_epi8(lo as i8)), n) } +} + +/// `if mask { a } else { b }`, lane-wise. SSE4.1's `blendv` would do it in one, but the and/andnot +/// pair keeps the whole builder at SSSE3. +#[inline(always)] +unsafe fn sel(mask: __m128i, a: __m128i, b: __m128i) -> __m128i { + unsafe { _mm_or_si128(_mm_and_si128(mask, a), _mm_andnot_si128(mask, b)) } +} + +/// 64 bytes of `text[base..]` compared against `b`, as a bitmap. +/// +/// # Safety +/// `base + 64 <= text.len()`; the caller has checked for SSSE3. +#[target_feature(enable = "ssse3")] +pub(crate) unsafe fn eq64(text: &[u8], base: usize, b: u8) -> u64 { + unsafe { + let n = _mm_set1_epi8(b as i8); + let mut m = 0u64; + for k in 0..4 { + let v = _mm_loadu_si128(text.as_ptr().add(base + k * 16).cast()); + m |= (_mm_movemask_epi8(_mm_cmpeq_epi8(v, n)) as u16 as u64) << (16 * k); + } + m + } +} + +/// Build one **full** 64-byte block. `cur_code`/`cur_aux` describe the byte before it. +/// +/// # Safety +/// `base + 64 <= tags.len()` and `base + 64 <= text.len()`; the caller has checked for SSSE3. +#[target_feature(enable = "ssse3")] +pub(crate) unsafe fn build64( + text: &[u8], + tags: &[u8], + base: usize, + lut: &[u8; 64], + cur_code: u8, + cur_aux: bool, +) -> (Blk, u8) { + unsafe { + let low_nibble = _mm_set1_epi8(0x0F); + let seven = _mm_set1_epi8(7); + // one table per high nibble; a tag of 0x30 or above never occurs + let t0 = _mm_loadu_si128(lut.as_ptr().cast()); + let t1 = _mm_loadu_si128(lut.as_ptr().add(16).cast()); + let t2 = _mm_loadu_si128(lut.as_ptr().add(32).cast()); + + let mut cd = [_mm_setzero_si128(); 4]; + let mut isc = [_mm_setzero_si128(); 4]; + let mut prev = _mm_set1_epi8(cur_code as i8); + for k in 0..4 { + let t = _mm_loadu_si128(tags.as_ptr().add(base + k * 16).cast()); + let lo = _mm_and_si128(t, low_nibble); + let hi = _mm_and_si128(_mm_srli_epi16(t, 4), low_nibble); + let raw = sel( + _mm_cmpeq_epi8(hi, _mm_set1_epi8(2)), + _mm_shuffle_epi8(t2, lo), + sel( + _mm_cmpeq_epi8(hi, _mm_set1_epi8(1)), + _mm_shuffle_epi8(t1, lo), + _mm_shuffle_epi8(t0, lo), + ), + ); + isc[k] = _mm_cmpeq_epi8(raw, seven); + // fill the continuation lanes: shift by one, then shift *that* by two, which covers + // the three continuations a 4-byte char can have. + let c = sel(isc[k], _mm_alignr_epi8(raw, prev, 15), raw); + let c = sel(_mm_cmpeq_epi8(c, seven), _mm_alignr_epi8(c, prev, 14), c); + prev = c; + cd[k] = c; + } + let last_code = (_mm_extract_epi16(prev, 7) >> 8) as u8; + + // 64 lanes -> one u64, four native movemasks + let gather = |v: [__m128i; 4]| -> u64 { + let mut m = 0u64; + for (k, chunk) in v.iter().enumerate() { + m |= (_mm_movemask_epi8(*chunk) as u16 as u64) << (16 * k); + } + m + }; + // the shift is a const generic, so one arm per plane: bit k wants `7 - k` + macro_rules! plane { + ($shift:literal) => { + gather([ + _mm_slli_epi16::<$shift>(cd[0]), + _mm_slli_epi16::<$shift>(cd[1]), + _mm_slli_epi16::<$shift>(cd[2]), + _mm_slli_epi16::<$shift>(cd[3]), + ]) + }; + } + + let mut b = Blk { + cont: gather(isc), + p0: plane!(7), + p1: plane!(6), + p2: plane!(5), + p3: if P3 { plane!(4) } else { 0 }, + aux: 0, + }; + if AUX == AUX_NONE { + return (b, last_code); + } + + if AUX == AUX_SLASH { + // single ASCII byte — no fill needed, `/` is its own char + let sl = _mm_set1_epi8(b'/' as i8); + b.aux = gather([ + _mm_cmpeq_epi8(_mm_loadu_si128(text.as_ptr().add(base).cast()), sl), + _mm_cmpeq_epi8(_mm_loadu_si128(text.as_ptr().add(base + 16).cast()), sl), + _mm_cmpeq_epi8(_mm_loadu_si128(text.as_ptr().add(base + 32).cast()), sl), + _mm_cmpeq_epi8(_mm_loadu_si128(text.as_ptr().add(base + 48).cast()), sl), + ]); + return (b, last_code); + } + if AUX == AUX_HAN { + // ponytail: scalar Han range test; vectorise like the CJK path below if kimi ever + // shows up on a throughput bench. + let lim = text.len().min(base + 64); + let mut leads = 0u64; + for p in base..lim { + if tags[p] != crate::CONT && crate::han::is_han_at(text, p) { + leads |= 1u64 << (p - base); + } + } + b.aux = leads | (leads << 1) | (leads << 2); + if cur_aux { + b.aux |= lead_run(b.cont, !0); + } + return (b, last_code); + } + + // ── the CJK range test, on the raw bytes: Hiragana/Katakana U+3040..30FF is E3 [81-83] xx; + // Han U+4E00..9FA5 is E4 [B8-BF] xx / E5-E8 xx xx / E9 [80-BD] xx / E9 BE [80-A5]. + let ntext = text.len(); + let tv = [ + _mm_loadu_si128(text.as_ptr().add(base).cast()), + _mm_loadu_si128(text.as_ptr().add(base + 16).cast()), + _mm_loadu_si128(text.as_ptr().add(base + 32).cast()), + _mm_loadu_si128(text.as_ptr().add(base + 48).cast()), + ]; + let any = tv + .iter() + .any(|v| _mm_movemask_epi8(in_range(*v, 0xE3, 6)) != 0); + if !any { + if cur_aux { + b.aux |= lead_run(b.cont, !0); + } + return (b, last_code); + } + let tail = { + let mut buf = [0u8; 16]; + let off = base + 64; + let avail = ntext.saturating_sub(off).min(16); + buf[..avail].copy_from_slice(&text[off..off + avail]); + _mm_loadu_si128(buf.as_ptr().cast()) + }; + let cjkv = |v: __m128i, nx: __m128i| -> __m128i { + let b1 = _mm_alignr_epi8(nx, v, 1); + let b2 = _mm_alignr_epi8(nx, v, 2); + let e9 = _mm_and_si128( + _mm_cmpeq_epi8(v, _mm_set1_epi8(0xE9u8 as i8)), + _mm_or_si128( + le(b1, 0xBD), + _mm_and_si128( + _mm_cmpeq_epi8(b1, _mm_set1_epi8(0xBEu8 as i8)), + le(b2, 0xA5), + ), + ), + ); + _mm_or_si128( + _mm_or_si128( + _mm_and_si128( + _mm_cmpeq_epi8(v, _mm_set1_epi8(0xE3u8 as i8)), + in_range(b1, 0x81, 2), + ), + _mm_and_si128(_mm_cmpeq_epi8(v, _mm_set1_epi8(0xE4u8 as i8)), ge(b1, 0xB8)), + ), + _mm_or_si128(in_range(v, 0xE5, 3), e9), + ) + }; + let mut leads = gather([ + cjkv(tv[0], tv[1]), + cjkv(tv[1], tv[2]), + cjkv(tv[2], tv[3]), + cjkv(tv[3], tail), + ]); + // `fsm_deepseek` reads three bytes unconditionally; refuse to classify a truncated tail. + let lim = ntext.saturating_sub(base + 2); + if lim < 64 { + leads &= (1u64 << lim) - 1; + } + // every CJK char is 3 bytes, so two shifts fill it; one cut by the block edge is picked up + // on the other side by `cur_aux`. + b.aux = leads | (leads << 1) | (leads << 2); + if cur_aux { + b.aux |= lead_run(b.cont, !0); + } + (b, last_code) + } +} diff --git a/tokenizers/bitsplit/tests/parity.rs b/tokenizers/bitsplit/tests/parity.rs new file mode 100644 index 000000000..2590e5703 --- /dev/null +++ b/tokenizers/bitsplit/tests/parity.rs @@ -0,0 +1,265 @@ +//! Byte-exactness gate for the bitstream splitters. Oracle = oniguruma, composed exactly as HF +//! applies it (deepseek is a `Sequence` of three Isolated splits, not one regex). +//! +//! - The SWEEP is the point. A bitstream program is only interesting where a rule straddles the +//! 64-byte grid: `starts[bi-1] |= patch`, the `anl` retraction, every `pb`/`nb` edge peek. Slicing +//! the corpus at every char boundary walks every construct through every block phase, and the +//! oracle re-runs on the same slice so the expectation is never hand-written. +//! - `find_iter` is the Isolated split only because these regexes match the whole input with no +//! gaps; deepseek's three passes do leave gaps, hence `split_iso`. +#![cfg(not(target_arch = "wasm32"))] + +use bitsplit::Span; +use bitsplit::classify::classify; +use bitsplit::regexes::{ + CL100K, DEEPSEEK_BIG as DS_BIG, DEEPSEEK_CJK as DS_CJK, DEEPSEEK_NUM as DS_NUM, GPT2, KIMI_K2, + O200K, TEKKEN, +}; +use onig::Regex; + +const CORPUS: &str = "The quick brown fox. Don't 12345 numbers, \u{00BD}\u{00B2}\u{00BC} \u{2168}! \ + café × naïve — Привет, наука! Ελλάδα 中文分词。ひらがな カタカナ 한글 مرحبا العربية \ + नरेंद्र मोदी x_y a1b2c3 e-mail@host.com 😀👍 hello world\ttabs\nnewlines end "; + +const EDGE: &str = "IT'S O'Brien can't 'quoted' l'été rock'n'roll\r\n\ + 0 42 999 1000 1234567 v1.2.3 3.14159 1,000,000 \ + https://host/a/b?c=1&d=2 path/to//file /\r\n/ ///x \ + CamelCase XMLHttpRequest IJSSELMEER DžAMBO ŀl a\u{0301}b \ + 日本語1234テスト ½3¼ \u{2168}42\u{2169} #tag @user $9.99 100% \ + end\n\n\nlines\r\n\r\n \t trailing "; + +/// A bitstream splitter, normalised to one shape (`starts` + `flag` scratch; deepseek ignores `flag`). +type Split = fn(&[u8], &[u8], &mut [u64], &mut [u64], &mut [Span]) -> usize; + +fn bs_deepseek(t: &[u8], g: &[u8], s: &mut [u64], _f: &mut [u64], o: &mut [Span]) -> usize { + bitsplit::bitsplit_deepseek(t, g, s, o) +} +fn bs_byte_level(t: &[u8], g: &[u8], s: &mut [u64], f: &mut [u64], o: &mut [Span]) -> usize { + bitsplit::bitsplit_byte_level(t, g, s, f, o) +} +fn bs_cl100k(t: &[u8], g: &[u8], s: &mut [u64], f: &mut [u64], o: &mut [Span]) -> usize { + bitsplit::bitsplit_cl100k(t, g, s, f, o) +} +fn bs_qwen(t: &[u8], g: &[u8], s: &mut [u64], f: &mut [u64], o: &mut [Span]) -> usize { + bitsplit::bitsplit_qwen(t, g, s, f, o) +} +fn bs_o200k(t: &[u8], g: &[u8], s: &mut [u64], f: &mut [u64], o: &mut [Span]) -> usize { + bitsplit::bitsplit_o200k(t, g, s, f, o) +} +fn bs_tekken(t: &[u8], g: &[u8], s: &mut [u64], f: &mut [u64], o: &mut [Span]) -> usize { + bitsplit::bitsplit_tekken(t, g, s, f, o) +} +fn bs_kimi(t: &[u8], g: &[u8], s: &mut [u64], f: &mut [u64], o: &mut [Span]) -> usize { + bitsplit::bitsplit_kimi(t, g, s, f, o) +} + +fn spans(f: Split, s: &str) -> Vec { + if s.is_empty() { + return Vec::new(); + } + let b = s.as_bytes(); + let mut tags = vec![0u8; b.len()]; + classify(b, &mut tags); + let nblk = b.len().div_ceil(64); + let (mut starts, mut flag) = (vec![0u64; nblk], vec![0u64; nblk]); + let mut out = vec![Span::default(); b.len() + 1]; + let k = f(b, &tags, &mut starts, &mut flag, &mut out); + out.truncate(k); + out +} + +fn onig_spans(re: &Regex, s: &str) -> Vec { + re.find_iter(s) + .map(|(a, b)| Span::new(a as u32, b as u32)) + .collect() +} + +/// One Isolated split of `text[s..e]`: gaps + matches, absolute offsets. +fn split_iso(text: &str, s: usize, e: usize, re: &Regex, out: &mut Vec<(usize, usize)>) { + let sub = &text[s..e]; + let mut prev = 0usize; + for (ms, me) in re.find_iter(sub) { + if ms > prev { + out.push((s + prev, s + ms)); + } + out.push((s + ms, s + me)); + prev = me; + } + if prev < sub.len() { + out.push((s + prev, e)); + } +} + +fn deepseek_ref(text: &str) -> Vec { + let (rn, rc, rb) = ( + Regex::new(DS_NUM).unwrap(), + Regex::new(DS_CJK).unwrap(), + Regex::new(DS_BIG).unwrap(), + ); + let mut a = Vec::new(); + split_iso(text, 0, text.len(), &rn, &mut a); + let mut b = Vec::new(); + for (s, e) in a { + split_iso(text, s, e, &rc, &mut b); + } + let mut c = Vec::new(); + for (s, e) in b { + split_iso(text, s, e, &rb, &mut c); + } + c.into_iter() + .map(|(s, e)| Span::new(s as u32, e as u32)) + .collect() +} + +/// The oracle for a grammar, over an arbitrary slice. +enum Oracle { + Whole(Regex), + DeepSeek, +} + +impl Oracle { + fn spans(&self, s: &str) -> Vec { + match self { + Oracle::Whole(re) => onig_spans(re, s), + Oracle::DeepSeek => deepseek_ref(s), + } + } +} + +fn long_corpus() -> String { + let mut s = String::new(); + while s.len() < 4096 { + s.push_str(CORPUS); + s.push_str(EDGE); + } + s +} + +/// Both corpora whole, then the block-phase sweep: every char-boundary prefix and suffix of a +/// >4 KB text, so each rule crosses a block edge in every alignment. +fn check(name: &str, f: Split, oracle: &Oracle) { + for text in [CORPUS, EDGE] { + assert_eq!(spans(f, text), oracle.spans(text), "{name}: whole {text:?}"); + } + + let long = long_corpus(); + let mut checked = 0usize; + + // suffixes: shifts the whole text across the grid + for off in 0..512.min(long.len()) { + if !long.is_char_boundary(off) { + continue; + } + let sub = &long[off..]; + assert_eq!(spans(f, sub), oracle.spans(sub), "{name}: suffix off={off}"); + checked += 1; + } + // prefixes: exercises the last-block / EOF rules (`\s+(?!\S)` vs plain `\s+`) in every phase + for end in (long.len().saturating_sub(512))..=long.len() { + if !long.is_char_boundary(end) { + continue; + } + let sub = &long[..end]; + assert_eq!(spans(f, sub), oracle.spans(sub), "{name}: prefix end={end}"); + checked += 1; + } + assert!(checked > 500, "{name}: sweep too small ({checked})"); +} + +/// Deterministic pseudo-random text from a weighted alphabet — always valid UTF-8 by construction. +/// Aimed at the carry logic: long digit runs, whitespace/newline runs and apostrophes next to +/// letters are what the `\p{N}{1,3}` grouping, the `\s*[\r\n]+` fill and the contraction escape +/// disagree about. +fn fuzz_texts(n: usize) -> Vec { + const ALPHA: &[&str] = &[ + "a", "b", "z", "A", "Q", "é", "ß", "IJ", "0", "1", "9", " ", " ", "\n", "\r\n", "\t", "'", + "'s", "'ll", ".", "!", "/", "#", "_", "中", "文", "ひ", "カ", "한", "م", "😀", "\u{0301}", + "\u{200D}", "½", "\u{2168}", + ]; + let mut st = 0x243F_6A88_85A3_08D3u64; + let mut next = move || { + st ^= st << 13; + st ^= st >> 7; + st ^= st << 17; + st + }; + (0..n) + .map(|_| { + let len = 1 + (next() % 300) as usize; + (0..len) + .map(|_| ALPHA[(next() % ALPHA.len() as u64) as usize]) + .collect() + }) + .collect() +} + +fn check_fuzz(name: &str, f: Split, oracle: &Oracle) { + for (i, t) in fuzz_texts(4000).iter().enumerate() { + assert_eq!(spans(f, t), oracle.spans(t), "{name}: fuzz #{i} {t:?}"); + } +} + +#[test] +fn byte_level_parity() { + let o = Oracle::Whole(Regex::new(GPT2).unwrap()); + check("byte_level", bs_byte_level, &o); + check_fuzz("byte_level", bs_byte_level, &o); +} + +#[test] +fn cl100k_parity() { + let o = Oracle::Whole(Regex::new(CL100K).unwrap()); + check("cl100k", bs_cl100k, &o); + check_fuzz("cl100k", bs_cl100k, &o); +} + +/// Qwen2 / Qwen3: cl100k with a bare `\p{N}`. +#[test] +fn qwen_parity() { + let qwen = CL100K.replace(r"\p{N}{1,3}", r"\p{N}"); + let o = Oracle::Whole(Regex::new(&qwen).unwrap()); + check("qwen", bs_qwen, &o); + check_fuzz("qwen", bs_qwen, &o); +} + +#[test] +fn deepseek_parity() { + check("deepseek", bs_deepseek, &Oracle::DeepSeek); + check_fuzz("deepseek", bs_deepseek, &Oracle::DeepSeek); +} + +/// o200k_base / GPT-4o — and byte-for-byte the regex Llama-4, gpt-oss and MiniMax-M2 ship, so this +/// one test covers four families. +#[test] +fn o200k_parity() { + let o = Oracle::Whole(Regex::new(O200K).unwrap()); + check("o200k", bs_o200k, &o); + check_fuzz("o200k", bs_o200k, &o); +} + +#[test] +fn tekken_parity() { + let o = Oracle::Whole(Regex::new(TEKKEN).unwrap()); + check("tekken", bs_tekken, &o); + check_fuzz("tekken", bs_tekken, &o); +} + +#[test] +fn kimi_parity() { + let o = Oracle::Whole(Regex::new(KIMI_K2).unwrap()); + check("kimi", bs_kimi, &o); + check_fuzz("kimi", bs_kimi, &o); +} + +/// Negative control: the gate above only means something if it can fail. gpt2 and cl100k disagree +/// on plenty (`\p{N}{1,3}`, the `[\r\n]*` tail, the non-space `?` prefix), so crossing them must +/// blow up — if this ever passes, `check` has stopped comparing anything. +#[test] +#[should_panic(expected = "negative-control")] +fn harness_discriminates() { + check( + "negative-control", + bs_byte_level, + &Oracle::Whole(Regex::new(CL100K).unwrap()), + ); +} diff --git a/tokenizers/tk-encode/Cargo.toml b/tokenizers/tk-encode/Cargo.toml index e862ba772..0b7ac82b8 100644 --- a/tokenizers/tk-encode/Cargo.toml +++ b/tokenizers/tk-encode/Cargo.toml @@ -27,7 +27,7 @@ name = "tk_encode" path = "src/lib.rs" [dependencies] -atomsplit = { path = "../atomsplit" } +bitsplit = { path = "../bitsplit" } rand = "0.9" regex = "1.10" rayon = { version = "1.10", optional = true } @@ -60,7 +60,7 @@ memchr = "2.8.2" unicode-normalization = "0.1.25" yada = "0.7.0" libc = "0.2" -wide = "1.6.0" +wide = "1.5.0" # Latest released tokenizers, used as the comparison baseline by the CI benchmark # (examples gated on `bench-baseline`). Optional so production builds never pull it. diff --git a/tokenizers/tk-encode/src/models/bpe/model.rs b/tokenizers/tk-encode/src/models/bpe/model.rs index 4f5080c5c..6b9505569 100644 --- a/tokenizers/tk-encode/src/models/bpe/model.rs +++ b/tokenizers/tk-encode/src/models/bpe/model.rs @@ -8,11 +8,11 @@ use crate::models::bpe::legacy::model::BPE; use crate::models::bpe::merge_hot_cold_queue::{QueueScratch, merge_hot_cold_queue}; use crate::models::bpe::merge_multipass::merge_multipass; use crate::models::bpe::tables::BpeTables; -use crate::pipeline::{self, PipelineToken}; +use crate::pipeline::{self, PipelineToken, Span}; use crate::tokenizer::Result; use crate::utils::byte_level::{self}; -use crate::utils::word_cache::{Lookup, WordCache}; -use crate::vocab::bucket_vocab_store::BucketVocabStore; +use crate::utils::word_cache::{Lookup, MAX_INLINE_IDS, ProbeEmit, WordCache}; +use crate::vocab::bucket_vocab_store::{BucketVocabStore, key_and_hash, key_and_hash_readable}; const GATE_MULTI: u16 = 8; const GATE_ASCII: u16 = 24; @@ -240,13 +240,10 @@ impl PipelineBPE { proven } - /// The id to emit for `sequence` without merging, when the whole pretoken is a vocabulary - /// entry that may be folded. `None` sends the word to the merge engines. #[inline(always)] - fn fold_id(&self, sequence: &str) -> Option { + fn fold_id_keyed(&self, key: u64, hash: u64) -> Option { // One probe; the foldable bit is part of the id that probe already returned. Which entries - // carry it was settled at load -- see `from_bpe`. - let (id, foldable) = self.vocab.get_bytes_foldable(sequence.as_bytes())?; + let (id, foldable) = self.vocab.get_keyed_foldable(key, hash)?; foldable.then_some(id) } @@ -309,10 +306,8 @@ impl pipeline::Model for PipelineBPE { return Ok(()); } - if let Some(id) = self.fold_id(sequence) { - output.push(PipelineToken { id }); - return Ok(()); - } + let bytes = sequence.as_bytes(); + let (key, hash) = key_and_hash(bytes); let BpeScratch { symbols, @@ -320,9 +315,10 @@ impl pipeline::Model for PipelineBPE { word_cache, } = scratch; - // A word seen before costs a probe instead of a merge. + // Cache before fold, for the reason given in `tokenize_spans`: the cache is one load and + // the fold is an MPHF probe, so the fold must not run ahead of it. let insert_at = if let Some(cache) = word_cache.as_mut() { - match cache.lookup(sequence.as_bytes()) { + match cache.lookup_keyed(key, hash) { Lookup::Hit(ids) => { output.extend(ids.iter().map(|&id| PipelineToken { id })); return Ok(()); @@ -333,6 +329,16 @@ impl pipeline::Model for PipelineBPE { None }; + if let Some(id) = self.fold_id_keyed(key, hash) { + output.push(PipelineToken { id }); + if let Some(cache) = word_cache.as_mut() + && let Some(at) = insert_at + { + cache.insert(at, std::iter::once(id)); + } + return Ok(()); + } + let start = output.len(); self.merge_word(sequence, symbols, queue); // the merge engines work in internal ids; `unmap` takes them back to the vocab's own ids @@ -362,15 +368,31 @@ impl pipeline::Model for PipelineBPE { output: &mut Vec, ) -> Result<()> { let BpeScratch { - merge_queue, - skip, - word, + symbols, + queue, word_cache, } = scratch; - // One reservation for the batch. Most pre-tokens are a single token, so the span count is - // a close lower bound on what the batch emits; anything past it grows as usual. - output.reserve(spans.len()); + output.reserve(spans.len() + MAX_INLINE_IDS); + let mut capacity = output.capacity(); + let mut cursor = output.len(); + // A raw write cursor held across the whole chunk. `output.as_mut_ptr()` inside the loop had + // to be *reloaded from memory every span*: the loop also calls `set_len`, `reserve` and + // `extend` on `output`, so the optimiser cannot assume the buffer stayed where it was. That + // reload, the capacity test and the cursor arithmetic were a quarter of encode time in + // `tokenize_spans` itself, with nothing of the model in it. Now the fast path -- a cache hit + // writing inline ids -- touches only these two registers, and `dst` is refreshed solely on + // the paths that can actually move the buffer. This is the other half of gigatoken's + // `probe_emit_chunk`: loop-invariant cursors, refreshed only in the slow path. + // SAFETY: `cursor <= output.len() <= capacity`, so this is inside the allocation. + let mut dst = unsafe { output.as_mut_ptr().add(cursor) }; + + // Unwrapped once for the whole chunk. `word_cache` is an `Option` living in the + // scratch, so every `as_mut()` re-read its discriminant out of memory -- three times per + // span, on a table that is either there for the entire call or not at all. Holding + // `Option<&mut WordCache>` in a local keeps that test in a register, which is the cheap half + // of what gigatoken's `ProbeView` does by carrying the table base and mask across a chunk. + let mut cache_slot = word_cache.as_mut(); for span in spans { // SAFETY: the pre-tokenizer cuts on char boundaries, so a span is always a valid slice @@ -380,33 +402,92 @@ impl pipeline::Model for PipelineBPE { continue; } + if cursor + MAX_INLINE_IDS > capacity { + // SAFETY: `cursor` counts what has been written so far. + unsafe { output.set_len(cursor) }; + output.reserve(spans.len() + MAX_INLINE_IDS); + capacity = output.capacity(); + // `reserve` may have moved the buffer. + // SAFETY: `cursor` is what was written so far, so it is within the new allocation. + dst = unsafe { output.as_mut_ptr().add(cursor) }; + } + + // The cache goes first. It is one direct-mapped load; the fold is an MPHF probe, which + // is a pilot load plus a dependent entry load into the whole vocabulary. Running the + // fold ahead of the cache paid that on every pre-token including the ones the cache + // was about to answer, and on a warm cache that is nearly all of them. + // + // The two still agree on ids: `prove_fold` only sets the bit for an entry that merging + // its own text reproduces, so a folded word and a merged word give the same answer. + // What changes is that a foldable word now gets *inserted*, so its second and later + // occurrences come off the cache instead of re-probing the vocabulary. + let (key, hash) = + key_and_hash_readable(sequence.as_bytes(), chunk.len() - span.start as usize); + let mut placement = None; - if let Some(cache) = word_cache.as_mut() { - match cache.lookup(sequence.as_bytes()) { - Lookup::Hit(ids) => { + if let Some(cache) = cache_slot.as_deref_mut() { + // SAFETY: the capacity check above leaves `MAX_INLINE_IDS` slots past `cursor`, and + let found = unsafe { + cache.probe_emit_keyed(key, hash, dst.cast::()) + }; + match found { + ProbeEmit::Wrote(n) => { + cursor += n; + // SAFETY: the probe wrote `n <= MAX_INLINE_IDS` ids, which the capacity + // check above reserved room for. + dst = unsafe { dst.add(n) }; + continue; + } + ProbeEmit::Hit(ids) => { + // SAFETY: `cursor` counts what has been written so far. + unsafe { output.set_len(cursor) }; output.extend(ids.iter().map(|&id| PipelineToken { id })); + cursor = output.len(); + capacity = output.capacity(); + // SAFETY: `cursor == output.len()`, inside the (possibly moved) allocation. + dst = unsafe { output.as_mut_ptr().add(cursor) }; continue; } - Lookup::Miss(at) => placement = Some(at), + ProbeEmit::Miss(at) => placement = Some(at), } } - let start = output.len(); - if self.ignore_merges - && let Some(id) = self.vocab.get_bytes(sequence.as_bytes()) - { - output.push(PipelineToken { id }); - } else { - self.merge_word(sequence, merge_queue, skip, word); - output.extend(word.get_chars_iter().map(|id| PipelineToken { id })); + + // Cache miss. The fold still answers a word that is its own vocabulary entry in one + // probe, which beats running the merge engine for it. + if let Some(id) = self.fold_id_keyed(key, hash) { + // SAFETY: the check above leaves at least `MAX_INLINE_IDS >= 1` slots past `cursor`. + unsafe { dst.write(PipelineToken { id }) }; + cursor += 1; + // SAFETY: one id written, and the capacity check reserved MAX_INLINE_IDS >= 1. + dst = unsafe { dst.add(1) }; + if let Some(cache) = cache_slot.as_deref_mut() + && let Some(at) = placement + { + cache.insert(at, std::iter::once(id)); + } + continue; } - // The ids come back out of `output` because that is the only place both branches - // above leave them: `ignore_merges` never touches `word`. - if let Some(cache) = word_cache.as_mut() + + // SAFETY: `cursor` counts what the fast paths wrote; the merge below uses `output` + unsafe { output.set_len(cursor) }; + let start = output.len(); + self.merge_word(sequence, symbols, queue); + // the merge engines work in internal ids; `unmap` takes them back to the vocab's own ids + output.extend(symbols.iter().map(|&symbol| PipelineToken { + id: self.tables.unmap.at(symbol as usize), + })); + if let Some(cache) = cache_slot.as_deref_mut() && let Some(at) = placement { cache.insert(at, output[start..].iter().map(|token| token.id)); } + cursor = output.len(); + capacity = output.capacity(); + // SAFETY: `cursor == output.len()`; `extend` above may have moved the buffer. + dst = unsafe { output.as_mut_ptr().add(cursor) }; } + // SAFETY: `cursor` counts every token written, by the fast paths and the slow one alike. + unsafe { output.set_len(cursor) }; Ok(()) } @@ -463,4 +544,38 @@ mod fold_tests { assert_eq!(want, got, "the fold changed the ids for {text:?}"); } } + + #[test] + fn the_batched_path_matches_the_reference() { + let reference = Tokenizer::from_file("../data/gpt2.json").unwrap(); + let pipe = PipelineTokenizer::try_from(&reference).unwrap(); + + let mut text = String::new(); + for i in 0..400 { + text.push_str(" the quick brown fox jumps over the lazy dog"); + text.push_str(" internationalisation unfortunately"); + text.push_str(" def foo(bar): return bar + 1"); + text.push_str(" <|xs0|> <|xs1|> <|endoftext|>"); + text.push_str(" 语言模型 ελληνικά"); + if i % 3 == 0 { + text.push_str(" aaaaaaaaaaaaaaaaaaaaaaaa "); + } + } + + let want: Vec = reference + .encode_fast(text.as_str(), false) + .unwrap() + .get_ids() + .to_vec(); + let got: Vec = pipe + .encode(text.as_str(), false) + .wait() + .unwrap() + .remove(0) + .iter() + .map(|t| t.id) + .collect(); + assert_eq!(want.len(), got.len(), "token count differs"); + assert_eq!(want, got, "the batched path changed the ids"); + } } diff --git a/tokenizers/tk-encode/src/normalizers/replace.rs b/tokenizers/tk-encode/src/normalizers/replace.rs index cc4545b47..635ea5378 100644 --- a/tokenizers/tk-encode/src/normalizers/replace.rs +++ b/tokenizers/tk-encode/src/normalizers/replace.rs @@ -5,7 +5,7 @@ use crate::tokenizer::Decoder; use crate::tokenizer::pattern::Pattern; use crate::tokenizer::{NormalizedString, Normalizer, Result}; use crate::utils::SysRegex; -use atomsplit::literal::Literal; +use bitsplit::literal::Literal; use serde::{Deserialize, Serialize}; /// Represents the different patterns that `Replace` can use diff --git a/tokenizers/tk-encode/src/pre_tokenizers/bert.rs b/tokenizers/tk-encode/src/pre_tokenizers/bert.rs index d2f7e6e8d..fa65cf7cc 100644 --- a/tokenizers/tk-encode/src/pre_tokenizers/bert.rs +++ b/tokenizers/tk-encode/src/pre_tokenizers/bert.rs @@ -21,8 +21,8 @@ impl pipeline::PreTokenizer for BertPreTokenizer { // Bert pre-tokenization = drop whitespace runs, isolate each punctuation char, keep every other // run. One `atomsplit` SIMD classify (bytes → atom tags) + the class-runs FSM, byte-exact with // the legacy `char::is_whitespace` / `is_punc` split above (see the tests). - use atomsplit::classify::{classify, mask}; - use atomsplit::fsm::class_runs_into; + use bitsplit::classify::{classify, mask}; + use bitsplit::classes::class_runs_into; let bytes = text.as_bytes(); let mut tags = vec![0u8; bytes.len()]; classify(bytes, &mut tags); diff --git a/tokenizers/tk-encode/src/pre_tokenizers/byte_level.rs b/tokenizers/tk-encode/src/pre_tokenizers/byte_level.rs index b716136f6..45dedea07 100644 --- a/tokenizers/tk-encode/src/pre_tokenizers/byte_level.rs +++ b/tokenizers/tk-encode/src/pre_tokenizers/byte_level.rs @@ -1,5 +1,5 @@ use crate::utils::byte_level::{BYTES_CHAR_LOOKUP, CHAR_BYTES_LOOKUP, byte_level_transform}; -use crate::utils::{GptFsm, GptFsmPattern}; +use crate::utils::{Grammar, GrammarPattern}; use serde::{Deserialize, Serialize}; use crate::tokenizer::{ @@ -85,7 +85,7 @@ impl PreTokenizer for ByteLevel { if self.use_regex { // GPT-2 byte-level split via the native atomsplit FSM (byte-exact, no regex backend). normalized.split( - GptFsmPattern(GptFsm::Gpt2), + GrammarPattern(Grammar::Gpt2), SplitDelimiterBehavior::Isolated, ) } else { diff --git a/tokenizers/tk-encode/src/pre_tokenizers/delimiter.rs b/tokenizers/tk-encode/src/pre_tokenizers/delimiter.rs index e30f206e6..5f1e9c758 100644 --- a/tokenizers/tk-encode/src/pre_tokenizers/delimiter.rs +++ b/tokenizers/tk-encode/src/pre_tokenizers/delimiter.rs @@ -32,7 +32,7 @@ impl pipeline::PreTokenizer for CharDelimiterSplit { // keeps the runs between, no empty spans. Byte-exact with the char-predicate split. let bytes = text.as_bytes(); let mut spans = vec![pipeline::Span::default(); bytes.len() + 1]; - let n = atomsplit::fsm::CharDelimiterSplit(self.delimiter).pre_tokenize( + let n = bitsplit::classes::CharDelimiterSplit(self.delimiter).pre_tokenize( bytes, &mut [], &mut spans, diff --git a/tokenizers/tk-encode/src/pre_tokenizers/digits.rs b/tokenizers/tk-encode/src/pre_tokenizers/digits.rs index 25cf9f193..38a15cd9f 100644 --- a/tokenizers/tk-encode/src/pre_tokenizers/digits.rs +++ b/tokenizers/tk-encode/src/pre_tokenizers/digits.rs @@ -43,8 +43,8 @@ impl pipeline::PreTokenizer for Digits { fn pre_tokenize(&self, text: &str, out: &mut Vec) -> Result<()> { // isolate each numeric char (`individual_digits`) or keep numeric runs — atomsplit classify + // class-runs FSM. atom `NUMERIC` == `char::is_numeric`, so byte-exact with the scalar path. - use atomsplit::classify::{classify, mask}; - use atomsplit::fsm::class_runs_into; + use bitsplit::classify::{classify, mask}; + use bitsplit::classes::class_runs_into; let bytes = text.as_bytes(); let mut tags = vec![0u8; bytes.len()]; classify(bytes, &mut tags); diff --git a/tokenizers/tk-encode/src/pre_tokenizers/punctuation.rs b/tokenizers/tk-encode/src/pre_tokenizers/punctuation.rs index 52f4e5e5d..b1a466e6c 100644 --- a/tokenizers/tk-encode/src/pre_tokenizers/punctuation.rs +++ b/tokenizers/tk-encode/src/pre_tokenizers/punctuation.rs @@ -44,8 +44,8 @@ impl pipeline::PreTokenizer for Punctuation { // atom `PUNCT` == `is_punc` (ASCII-punct ∪ \p{P}), so Isolated/Removed map to the class-runs FSM // byte-exactly. The merge/contiguous behaviors aren't a class-runs shape → keep the scalar split. if matches!(self.behavior, Isolated | Removed) { - use atomsplit::classify::mask; - use atomsplit::fsm::class_runs_into; + use bitsplit::classify::mask; + use bitsplit::classes::class_runs_into; pipeline::classify_into_spans( text.as_bytes(), |b, t, s| { diff --git a/tokenizers/tk-encode/src/pre_tokenizers/sequence.rs b/tokenizers/tk-encode/src/pre_tokenizers/sequence.rs index b01d4686d..72bee136e 100644 --- a/tokenizers/tk-encode/src/pre_tokenizers/sequence.rs +++ b/tokenizers/tk-encode/src/pre_tokenizers/sequence.rs @@ -107,7 +107,11 @@ impl pipeline::PreTokenizer for PipelineSequence { // deepseek's 3-Split composition → one native FSM pass (also lets the Sequence handle the // trailing byte-map ByteLevel, which the generic child loop can't range-split). if self.is_deepseek() { - pipeline::classify_into_spans(text.as_bytes(), atomsplit::fsm::fsm_deepseek, out); + pipeline::classify_into_spans_bits( + text.as_bytes(), + |t, tags, starts, _flag, out| bitsplit::bitsplit_deepseek(t, tags, starts, out), + out, + ); return Ok(()); } diff --git a/tokenizers/tk-encode/src/pre_tokenizers/split.rs b/tokenizers/tk-encode/src/pre_tokenizers/split.rs index 779c167f1..bddbc0484 100644 --- a/tokenizers/tk-encode/src/pre_tokenizers/split.rs +++ b/tokenizers/tk-encode/src/pre_tokenizers/split.rs @@ -1,6 +1,6 @@ use crate::pipeline; -use crate::utils::{GptFsm, GptFsmPattern, SysRegex, gpt_fsm}; -use atomsplit::literal::Literal; +use crate::utils::{Grammar, GrammarPattern, SysRegex, recognize}; +use bitsplit::literal::Literal; use serde::{Deserialize, Deserializer, Serialize}; use crate::tokenizer::{ @@ -53,7 +53,7 @@ pub struct Split { /// pipeline path when `behavior == Isolated && !invert` (how these regexes always ship). Byte-exact /// with `regex`; `None` falls back to `regex`. #[serde(skip)] - fsm: Option, + fsm: Option, } impl<'de> Deserialize<'de> for Split { @@ -103,7 +103,7 @@ impl Split { let pattern: SplitPattern = pattern.into(); let fsm = match &pattern { SplitPattern::String(_) => None, - SplitPattern::Regex(r) => gpt_fsm(r), + SplitPattern::Regex(r) => recognize(r), }; let search = match &pattern { SplitPattern::String(s) => Search::Literal(Literal::new(s.as_bytes())?), @@ -173,30 +173,23 @@ impl PreTokenizer for Split { .into() })?; pretokenized.split(|_, normalized| { - normalized.split(GptFsmPattern(fsm), SplitDelimiterBehavior::Isolated) + normalized.split(GrammarPattern(fsm), SplitDelimiterBehavior::Isolated) }) } } impl pipeline::PreTokenizer for Split { fn pre_tokenize(&self, text: &str, out: &mut Vec) -> Result<()> { - // A recognized GPT regex (gpt2 / cl100k-Llama-3) in its only real usage — `Isolated`, not - // inverted — routes straight to the native atomsplit FSM. These regexes cover the whole input, - // so `Isolated` == the match list, and the FSM is byte-exact with `regex` (see the tests). - if let Some(fsm) = self + // A recognized GPT regex in its only real usage — `Isolated`, not inverted — routes + // straight to the native bitsplit grammar. These regexes cover the whole input, so + // `Isolated` == the match list, and the grammar is byte-exact with `regex` (see the tests). + if let Some(grammar) = self .fsm .filter(|_| !self.invert && self.behavior == SplitDelimiterBehavior::Isolated) { - pipeline::classify_into_spans( + pipeline::classify_into_spans_bits( text.as_bytes(), - |bytes, tags, spans| match fsm { - GptFsm::Cl100k { digit_cap } => { - atomsplit::fsm::fsm_cl100k_cap(bytes, tags, spans, digit_cap) - } - GptFsm::Gpt2 => atomsplit::fsm::fsm_byte_level(bytes, tags, spans), - GptFsm::O200k => atomsplit::fsm::fsm_o200k(bytes, tags, spans), - GptFsm::Tekken => atomsplit::fsm::fsm_tekken(bytes, tags, spans), - }, + |t, tags, starts, flag, out| grammar.split(t, tags, starts, flag, out), out, ); return Ok(()); @@ -419,7 +412,7 @@ mod tests { fn pipeline_gpt2_uses_fsm_and_matches_legacy() { // The gpt2 pattern is recognized -> the pipeline path routes to the native // atomsplit FSM; its output must equal the legacy fancy-regex path. - let gpt2 = atomsplit::regexes::GPT2; + let gpt2 = bitsplit::regexes::GPT2; let corpus = "The quick brown fox 123!!! double spaces\tand tabs. don't Naïve café. "; let pretok = Split::new(SplitPattern::Regex(gpt2.into()), Isolated, false).unwrap(); assert!(pretok.fsm.is_some(), "gpt2 pattern should be recognized"); @@ -443,7 +436,7 @@ mod tests { fn pipeline_cl100k_llama3_uses_fsm_and_matches_legacy() { // Llama-3's EXACT pre_tokenizer regex (from data/llama-3-tokenizer.json) → recognized → routes // to fsm_cl100k. Output must equal the legacy SysRegex Isolated split, byte-for-byte. - let cl100k = atomsplit::regexes::CL100K; + let cl100k = bitsplit::regexes::CL100K; let corpus = "The quick brown fox 123!!! double spaces\tand tabs. don't Naïve café.\n\n世界 안녕 "; let pretok = Split::new(SplitPattern::Regex(cl100k.into()), Isolated, false).unwrap(); @@ -471,7 +464,7 @@ mod tests { // GPT-4o's EXACT pre_tokenization regex → recognized → routes to fsm_o200k. Output must equal the // legacy SysRegex Isolated split, byte-for-byte. Corpus stresses the case-aware letter split: // camelCase, ALLCAPS→word, McDonald's-style, contractions, accented Ll (é/ß), CJK (caseless). - let o200k = atomsplit::regexes::O200K; + let o200k = bitsplit::regexes::O200K; let corpus = "McDonald's iPhone SQLite HELLOWorld camelCase don't I'll We've 3.14 café Straße 世界 안녕\n\n Mixed CASE end."; let pretok = Split::new(SplitPattern::Regex(o200k.into()), Isolated, false).unwrap(); assert!( @@ -498,12 +491,12 @@ mod tests { // Mistral's tekken regex (mistral-small-4) → recognized → routes to fsm_tekken. Same corpus // shape as o200k, whose grammar it shares: the differences it must get right are apostrophes // (no contraction suffix, so `'s` starts a new token) and one token per digit. - let tekken = atomsplit::regexes::TEKKEN; + let tekken = bitsplit::regexes::TEKKEN; let corpus = "McDonald's iPhone SQLite HELLOWorld camelCase don't I'll We've 3.14159 café Straße 世界 안녕\n\n path/to/file Mixed CASE end."; let pretok = Split::new(SplitPattern::Regex(tekken.into()), Isolated, false).unwrap(); assert_eq!( pretok.fsm, - Some(crate::utils::GptFsm::Tekken), + Some(crate::utils::Grammar::Tekken), "tekken / mistral pattern should route to the native FSM" ); @@ -531,7 +524,7 @@ mod tests { let pretok = Split::new(SplitPattern::Regex(qwen2.into()), Isolated, false).unwrap(); assert_eq!( pretok.fsm, - Some(crate::utils::GptFsm::Cl100k { digit_cap: 1 }), + Some(crate::utils::Grammar::Cl100k { digit_cap: 1 }), "Qwen2 pattern should route to the cl100k FSM with digit cap 1" ); diff --git a/tokenizers/tk-encode/src/pre_tokenizers/whitespace.rs b/tokenizers/tk-encode/src/pre_tokenizers/whitespace.rs index 13f19dda2..20d789eb4 100644 --- a/tokenizers/tk-encode/src/pre_tokenizers/whitespace.rs +++ b/tokenizers/tk-encode/src/pre_tokenizers/whitespace.rs @@ -45,10 +45,10 @@ impl pipeline::PreTokenizer for WhitespaceSplit { fn pre_tokenize(&self, text: &str, out: &mut Vec) -> Result<()> { // drop whitespace runs, keep everything else as runs — atomsplit SIMD classify + class-runs FSM. // atom `WS` == `char::is_whitespace`, so byte-exact with the scalar path. - use atomsplit::classify::mask; + use bitsplit::classify::mask; pipeline::classify_into_spans( text.as_bytes(), - atomsplit::fsm::class_runs_into::<{ mask::WS }, 0, 0>, + bitsplit::classes::class_runs_into::<{ mask::WS }, 0, 0>, out, ); Ok(()) @@ -74,10 +74,10 @@ impl pipeline::PreTokenizer for Whitespace { fn pre_tokenize(&self, text: &str, out: &mut Vec) -> Result<()> { // `\w+|[^\w\s]+`: drop whitespace, cut at the word↔symbol boundary, each run one token — // atomsplit classify + class-runs FSM (`WORD` = `\w`; keep-A = word, keep-B = symbol). - use atomsplit::classify::mask; + use bitsplit::classify::mask; pipeline::classify_into_spans( text.as_bytes(), - atomsplit::fsm::class_runs_into::<{ mask::WS }, 0, { mask::WORD }>, + bitsplit::classes::class_runs_into::<{ mask::WS }, 0, { mask::WORD }>, out, ); Ok(()) diff --git a/tokenizers/tk-encode/src/tokenizer/normalizer.rs b/tokenizers/tk-encode/src/tokenizer/normalizer.rs index 8f1899bdd..937509c80 100644 --- a/tokenizers/tk-encode/src/tokenizer/normalizer.rs +++ b/tokenizers/tk-encode/src/tokenizer/normalizer.rs @@ -1022,7 +1022,7 @@ impl From<&str> for NormalizedString { #[cfg(test)] mod tests { use super::*; - use atomsplit::literal::Literal; + use bitsplit::literal::Literal; use regex::Regex; use unicode_categories::UnicodeCategories; diff --git a/tokenizers/tk-encode/src/tokenizer/pattern.rs b/tokenizers/tk-encode/src/tokenizer/pattern.rs index 5a147f5f4..879e35135 100644 --- a/tokenizers/tk-encode/src/tokenizer/pattern.rs +++ b/tokenizers/tk-encode/src/tokenizer/pattern.rs @@ -1,6 +1,6 @@ use crate::utils::SysRegex; use crate::{Offsets, Result}; -use atomsplit::literal::Literal; +use bitsplit::literal::Literal; use regex::Regex; /// Pattern used to split a NormalizedString diff --git a/tokenizers/tk-encode/src/tokenizer/pipeline.rs b/tokenizers/tk-encode/src/tokenizer/pipeline.rs index 072efaf26..a80c7fdfa 100644 --- a/tokenizers/tk-encode/src/tokenizer/pipeline.rs +++ b/tokenizers/tk-encode/src/tokenizer/pipeline.rs @@ -6,7 +6,7 @@ use std::sync::{Mutex, PoisonError}; use std::vec::IntoIter; use std::{borrow::Cow, convert::TryFrom}; -use atomsplit::classify::classify; +use bitsplit::classify::classify; use crate::models::bpe::{BpeScratch, PipelineBPE}; use crate::models::unigram::{Unigram, UnigramScratch}; @@ -39,7 +39,7 @@ use crate::{ use super::{Result, SplitDelimiterBehavior}; -pub use atomsplit::fsm::Span; +pub use bitsplit::Span; /// We use a thread local scratch for the tags (per byte class) and for the split spans. pub(crate) fn classify_into_spans( @@ -65,6 +65,51 @@ pub(crate) fn classify_into_spans( }); } +/// [`classify_into_spans`] for a splitter that works off bitstreams: it needs two `u64` bitmaps +/// (token starts, and the flags its scalar escapes key on) alongside the tags. +pub(crate) fn classify_into_spans_bits( + bytes: &[u8], + split: impl FnOnce(&[u8], &[u8], &mut [u64], &mut [u64], &mut [Span]) -> usize, + out: &mut Vec, +) { + thread_local! { + static SCRATCH: RefCell<(Vec, Vec, Vec)> = + const { RefCell::new((Vec::new(), Vec::new(), Vec::new())) }; + } + let n = bytes.len(); + if n == 0 { + return; + } + SCRATCH.with(|cell| { + let (tags, starts, flags) = &mut *cell.borrow_mut(); + if tags.len() < n { + tags.resize(n, 0); + } + let words = n.div_ceil(64) + 1; + if starts.len() < words { + starts.resize(words, 0); + flags.resize(words, 0); + } + classify(bytes, &mut tags[..n]); + let base = out.len(); + out.reserve(n + 1); + // SAFETY: as in `classify_into_spans` -- `reserve` covers the splitter's worst case of one + // span per byte plus one, `Span` has no drop glue, and `set_len` counts only what was written. + let k = unsafe { + split( + bytes, + &tags[..n], + &mut starts[..words], + &mut flags[..words], + std::slice::from_raw_parts_mut(out.as_mut_ptr().add(base), n + 1), + ) + }; + debug_assert!(k <= n + 1); + // SAFETY: the splitter wrote `k <= n + 1` spans from `base`. + unsafe { out.set_len(base + k) }; + }); +} + pub trait Normalizer { fn normalize<'a>(&self, input: &'a str) -> Result>; } @@ -994,6 +1039,24 @@ impl PipelineTokenizer { add_special_tokens: bool, ) -> Result> { let mut output = Vec::with_capacity(input.len() / 4); + self.encode_generic_into::(input, add_special_tokens, &mut output)?; + Ok(output) + } + + /// [`Self::encode_generic`] writing into a caller-owned buffer. + /// + /// The allocating form sizes its `Vec` from the input length, which is a guess, and hands back + /// a fresh allocation every call. A caller that encodes many inputs -- a batch, a server loop, + /// a benchmark -- can reserve once and `clear()` between calls instead: fewer allocations, and + /// no first-touch of the token array each time. Measured at ~3% of geomean throughput on + /// tokbench's 29 gpt2 cells (0.9233x -> 0.9536x against gigatoken, same code both sides). + #[doc(hidden)] + pub fn encode_generic_into( + &self, + input: &str, + add_special_tokens: bool, + output: &mut Vec, + ) -> Result<()> { let mut scratch = self.scratch_pool.get(&self.model); let PipelinePostProcessor { prefix, suffix } = &self.post_processor; // Prepend prefix tokens, if any @@ -1036,7 +1099,7 @@ impl PipelineTokenizer { normalized_chunk, &pre_tokens, &mut scratch, - &mut output, + output, )?; } Ok(()) @@ -1052,7 +1115,7 @@ impl PipelineTokenizer { if add_special_tokens && STAGE >= Self::STAGE_POSTPROCESS { output.extend_from_slice(suffix); } - Ok(output) + Ok(()) } } diff --git a/tokenizers/tk-encode/src/utils/byte_level.rs b/tokenizers/tk-encode/src/utils/byte_level.rs index 95a26807a..70d79e938 100644 --- a/tokenizers/tk-encode/src/utils/byte_level.rs +++ b/tokenizers/tk-encode/src/utils/byte_level.rs @@ -4,7 +4,7 @@ use std::sync::LazyLock; // The GPT-2 pre-tokenize regex is the canonical spec in atomsplit (single source of truth); re-export // under the historical name so call sites are unchanged. -pub(crate) use atomsplit::regexes::GPT2 as GPT2_REGEX_STR; +pub(crate) use bitsplit::regexes::GPT2 as GPT2_REGEX_STR; /// Maps each byte to its GPT-2 byte-level unicode character, indexed by the byte value. /// diff --git a/tokenizers/tk-encode/src/utils/mod.rs b/tokenizers/tk-encode/src/utils/mod.rs index 55643b57d..c446f576d 100644 --- a/tokenizers/tk-encode/src/utils/mod.rs +++ b/tokenizers/tk-encode/src/utils/mod.rs @@ -6,7 +6,7 @@ pub(crate) mod word_cache; // Optional system-regex backend, needed only for a *regex* pattern that atomsplit does not cover. // With `fancy-regex` off a stub compiles and those patterns error at load. Everything else works // regardless: the atomsplit-native pre-tokenizers, and any `Split` or `Replace` whose pattern is a -// plain string (searched for directly, see `atomsplit::literal`). +// plain string (searched for directly, see `bitsplit::literal`). #[cfg(feature = "fancy-regex")] mod fancy; #[cfg(feature = "fancy-regex")] @@ -18,7 +18,7 @@ pub use no_regex::SysRegex; // Recognize known GPT pre-tokenization regexes and route them to atomsplit's native (unrolled) FSM. mod unrolled_regex; -pub use unrolled_regex::{GptFsm, GptFsmPattern, gpt_fsm, is_deepseek}; +pub use unrolled_regex::{Grammar, GrammarPattern, recognize, is_deepseek}; pub mod byte_level; pub mod iter; diff --git a/tokenizers/tk-encode/src/utils/no_regex.rs b/tokenizers/tk-encode/src/utils/no_regex.rs index 8645b69c3..2a135e412 100644 --- a/tokenizers/tk-encode/src/utils/no_regex.rs +++ b/tokenizers/tk-encode/src/utils/no_regex.rs @@ -3,7 +3,7 @@ //! The type stays present so `Split` / `Replace` still compile, but construction always fails. Only a //! *regex* pattern ever asks for it: the atomsplit-native pre-tokenizers (GPT-2, cl100k, deepseek, the //! class family, char-delimiter) need no backend, and a plain string pattern is searched for directly -//! (`atomsplit::literal`). A regex atomsplit does not cover errors at load time with a clear message. +//! (`bitsplit::literal`). A regex atomsplit does not cover errors at load time with a clear message. //! Enable `fancy-regex` to get a real backend. use std::error::Error; diff --git a/tokenizers/tk-encode/src/utils/unrolled_regex.rs b/tokenizers/tk-encode/src/utils/unrolled_regex.rs index 78025296b..776387f08 100644 --- a/tokenizers/tk-encode/src/utils/unrolled_regex.rs +++ b/tokenizers/tk-encode/src/utils/unrolled_regex.rs @@ -1,35 +1,60 @@ -//! Recognize a known GPT pre-tokenization regex and route it to the byte-exact native -//! `atomsplit` FSM, so those pre-tokenizers need no system-regex backend. An unrecognized -//! pattern returns `None` and falls back to `SysRegex` (the optional fancy-regex backend). +//! Recognize a known GPT pre-tokenization regex and route it to the byte-exact native `bitsplit` +//! grammar, so those pre-tokenizers need no system-regex backend. An unrecognized pattern returns +//! `None` and falls back to `SysRegex` (the optional fancy-regex backend). -// Canonical GPT pre-tokenization regexes (the look-ahead originals), used as recognition keys — the -// single source of truth lives in `atomsplit::regexes`. `gpt_fsm` maps each to the `atomsplit` FSM that -// reproduces its `Isolated` split byte-for-byte. -use atomsplit::regexes::{GPT2, O200K, TEKKEN}; +// The canonical regexes are the recognition keys; the single source of truth is `bitsplit::regexes`. +use bitsplit::Span; +use bitsplit::regexes::{GPT2, KIMI_K2, O200K, TEKKEN}; // cl100k is recognized structurally (see `cl100k_digit_cap`), so the exact pattern is only a test key. #[cfg(test)] -use atomsplit::regexes::CL100K; +use bitsplit::regexes::CL100K; -/// A recognized GPT pre-tokenization regex that maps to a native `atomsplit` FSM (byte-exact). +/// A recognized GPT pre-tokenization regex and the `bitsplit` grammar that reproduces its +/// `Isolated` split byte-for-byte. One variant per distinct regex; models sharing a regex share a +/// variant (o200k covers Llama-4, gpt-oss and MiniMax-M2). #[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum GptFsm { - /// GPT-2 / ByteLevel regex → `atomsplit::fsm::fsm_byte_level`. +pub enum Grammar { + /// GPT-2 / ByteLevel. Gpt2, - /// cl100k-family regex → `atomsplit::fsm::fsm_cl100k_cap`. `digit_cap` is rule 3's `\p{N}{1,cap}` - /// bound: 3 = cl100k / Llama-3, 1 = Qwen2 (`\p{N}`), `usize::MAX` = an unbounded `\p{N}+`. + /// cl100k family. `digit_cap` is rule 3's `\p{N}{1,cap}` bound: 3 = cl100k / Llama-3 / GLM-4.6, + /// 1 = Qwen (`\p{N}`), `usize::MAX` = an unbounded `\p{N}+`. Cl100k { digit_cap: usize }, - /// o200k / GPT-4o regex → `atomsplit::fsm::fsm_o200k`. + /// o200k / GPT-4o — and byte-for-byte the regex Llama-4, gpt-oss and MiniMax-M2 ship. O200k, - /// Mistral tekken regex → `atomsplit::fsm::fsm_tekken` (o200k's grammar with no contraction - /// suffix and one token per digit). + /// Mistral tekken: o200k with no contraction suffix and one token per digit. Tekken, + /// kimi-k2 / k3: o200k plus a leading `[\p{Han}]+` arm and a plain `[\r\n]*` rule-4 tail. + Kimi, } -/// The cl100k-family template is fixed except rule 3's digit rule. If `pattern` is that template, return -/// the `\p{N}{1,cap}` bound (`\p{N}{1,3}`→3, `\p{N}{1,2}`→2, `\p{N}`→1, `\p{N}+`→`MAX`); else `None`. -/// This is what makes Qwen2 (cl100k with `\p{N}`) unroll without a per-tokenizer exact-string entry. +impl Grammar { + /// Write the token spans into `out`, returning the count. `starts`/`flag` are `u64` scratch + /// bitmaps of length >= `text.len().div_ceil(64)`. + pub fn split( + self, + text: &[u8], + tags: &[u8], + starts: &mut [u64], + flag: &mut [u64], + out: &mut [Span], + ) -> usize { + match self { + Grammar::Gpt2 => bitsplit::bitsplit_byte_level(text, tags, starts, flag, out), + Grammar::Cl100k { digit_cap: 1 } => { + bitsplit::bitsplit_qwen(text, tags, starts, flag, out) + } + Grammar::Cl100k { .. } => bitsplit::bitsplit_cl100k(text, tags, starts, flag, out), + Grammar::O200k => bitsplit::bitsplit_o200k(text, tags, starts, flag, out), + Grammar::Tekken => bitsplit::bitsplit_tekken(text, tags, starts, flag, out), + Grammar::Kimi => bitsplit::bitsplit_kimi(text, tags, starts, flag, out), + } + } +} + +/// The cl100k-family template is fixed except rule 3's digit rule. If `pattern` is that template, +/// return the `\p{N}{1,cap}` bound; else `None`. This is what makes Qwen (cl100k with `\p{N}`) and +/// GLM-4.6 (cl100k verbatim) route without a per-model entry. fn cl100k_digit_cap(pattern: &str) -> Option { - // cl100k rules 1-2 (contraction + word) … … rules 4-7 (other + whitespace). const PRE: &str = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|"; const SUF: &str = r"| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"; match pattern.strip_prefix(PRE)?.strip_suffix(SUF)? { @@ -41,29 +66,24 @@ fn cl100k_digit_cap(pattern: &str) -> Option { } } -/// If `pattern` is a recognized GPT pre-tokenization regex, name the native FSM that reproduces its -/// `Isolated` split byte-for-byte. GPT-2, o200k and tekken are matched exactly; the cl100k family is -/// matched structurally ([`cl100k_digit_cap`]) so digit-cap variants (Qwen2 …) unroll too. An -/// unrecognized pattern → `None` (the SysRegex / fancy-regex path handles it). -pub fn gpt_fsm(pattern: &str) -> Option { - if pattern == GPT2 { - Some(GptFsm::Gpt2) - } else if pattern == O200K { - Some(GptFsm::O200k) - } else if pattern == TEKKEN { - Some(GptFsm::Tekken) - } else { - cl100k_digit_cap(pattern).map(|digit_cap| GptFsm::Cl100k { digit_cap }) +/// If `pattern` is a recognized GPT pre-tokenization regex, name the grammar that reproduces its +/// `Isolated` split byte-for-byte. An unrecognized pattern -> `None` (the SysRegex path handles it). +pub fn recognize(pattern: &str) -> Option { + match pattern { + GPT2 => Some(Grammar::Gpt2), + O200K => Some(Grammar::O200k), + TEKKEN => Some(Grammar::Tekken), + KIMI_K2 => Some(Grammar::Kimi), + _ => cl100k_digit_cap(pattern).map(|digit_cap| Grammar::Cl100k { digit_cap }), } } -/// `Pattern` that runs the native atomsplit FSM for a [`GptFsm`] — the legacy (`NormalizedString`) -/// split path's equivalent of the pipeline's native routing, so GPT pre-tokenizers need no -/// system-regex backend. `Isolated` behaviour keeps every span, and the FSM spans cover the input -/// contiguously, so this is byte-for-byte identical to the original GPT regex `Isolated` split. -pub struct GptFsmPattern(pub GptFsm); +/// `Pattern` that runs the native grammar on the legacy (`NormalizedString`) split path, so GPT +/// pre-tokenizers need no system-regex backend. `Isolated` keeps every span and these grammars +/// cover the input contiguously, so this is byte-for-byte the original regex `Isolated` split. +pub struct GrammarPattern(pub Grammar); -impl crate::tokenizer::pattern::Pattern for GptFsmPattern { +impl crate::tokenizer::pattern::Pattern for GrammarPattern { fn find_matches( &self, inside: &str, @@ -71,28 +91,23 @@ impl crate::tokenizer::pattern::Pattern for GptFsmPattern { if inside.is_empty() { return Ok(vec![((0, 0), false)]); } - use atomsplit::classify::classify; let bytes = inside.as_bytes(); - let mut tags = vec![0u8; bytes.len()]; - classify(bytes, &mut tags); - let mut spans = vec![atomsplit::fsm::Span::default(); bytes.len() + 1]; - let n = match self.0 { - GptFsm::Gpt2 => atomsplit::fsm::fsm_byte_level(bytes, &tags, &mut spans), - GptFsm::Cl100k { digit_cap } => { - atomsplit::fsm::fsm_cl100k_cap(bytes, &tags, &mut spans, digit_cap) - } - GptFsm::O200k => atomsplit::fsm::fsm_o200k(bytes, &tags, &mut spans), - GptFsm::Tekken => atomsplit::fsm::fsm_tekken(bytes, &tags, &mut spans), - }; - Ok(spans[..n] + let n = bytes.len(); + let mut tags = vec![0u8; n]; + bitsplit::classify::classify(bytes, &mut tags); + let words = n.div_ceil(64) + 1; + let (mut starts, mut flag) = (vec![0u64; words], vec![0u64; words]); + let mut spans = vec![Span::default(); n + 1]; + let k = self.0.split(bytes, &tags, &mut starts, &mut flag, &mut spans); + Ok(spans[..k] .iter() .map(|sp| ((sp.start as usize, sp.end as usize), true)) .collect()) } } -// deepseek-v4's pre-tokenizer is a `Sequence` of these three Isolated `Split`s (+ a byte-map -// `ByteLevel`), which `atomsplit::fsm::fsm_deepseek` collapses into one pass. Byte-exact with the +// deepseek-v3/v4's pre-tokenizer is a `Sequence` of these three Isolated `Split`s (+ a byte-map +// `ByteLevel`), which `bitsplit::bitsplit_deepseek` collapses into one pass. Byte-exact with the // shipped tokenizer.json — the big pattern carries LITERAL CR/LF, spliced in via `concat!`. const DS_NUM: &str = r"\p{N}{1,3}"; const DS_CJK: &str = "[\u{4E00}-\u{9FA5}\u{3040}-\u{309F}\u{30A0}-\u{30FF}]+"; @@ -107,7 +122,7 @@ const DS_BIG: &str = concat!( ); /// True iff three `Split` patterns are exactly deepseek's `[\p{N}{1,3}, CJK-range, big-regex]` prefix → -/// `atomsplit::fsm::fsm_deepseek` reproduces the whole composed Isolated split in one pass. +/// `bitsplit::bitsplit_deepseek` reproduces the whole composed Isolated split in one pass. pub fn is_deepseek(p0: &str, p1: &str, p2: &str) -> bool { p0 == DS_NUM && p1 == DS_CJK && p2 == DS_BIG } @@ -119,25 +134,25 @@ mod tests { #[test] fn gpt_fsm_recognizes_family_and_extracts_digit_cap() { // Exact matches for gpt2 / o200k, structural (any digit cap) for the cl100k family. - assert_eq!(gpt_fsm(GPT2), Some(GptFsm::Gpt2)); - assert_eq!(gpt_fsm(O200K), Some(GptFsm::O200k)); - assert_eq!(gpt_fsm(TEKKEN), Some(GptFsm::Tekken)); - assert_eq!(gpt_fsm(CL100K), Some(GptFsm::Cl100k { digit_cap: 3 })); + assert_eq!(recognize(GPT2), Some(Grammar::Gpt2)); + assert_eq!(recognize(O200K), Some(Grammar::O200k)); + assert_eq!(recognize(TEKKEN), Some(Grammar::Tekken)); + assert_eq!(recognize(CL100K), Some(Grammar::Cl100k { digit_cap: 3 })); // Qwen2: cl100k with rule 3 = `\p{N}` → cap 1. let qwen2 = 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+"; - assert_eq!(gpt_fsm(qwen2), Some(GptFsm::Cl100k { digit_cap: 1 })); + assert_eq!(recognize(qwen2), Some(Grammar::Cl100k { digit_cap: 1 })); // Other in-family digit caps. let cap2 = CL100K.replace(r"\p{N}{1,3}", r"\p{N}{1,2}"); - assert_eq!(gpt_fsm(&cap2), Some(GptFsm::Cl100k { digit_cap: 2 })); + assert_eq!(recognize(&cap2), Some(Grammar::Cl100k { digit_cap: 2 })); let unbounded = CL100K.replace(r"\p{N}{1,3}", r"\p{N}+"); assert_eq!( - gpt_fsm(&unbounded), - Some(GptFsm::Cl100k { + recognize(&unbounded), + Some(Grammar::Cl100k { digit_cap: usize::MAX }) ); // Out of family → None (fancy-regex fallback): a foreign digit rule, and a totally unrelated regex. - assert_eq!(gpt_fsm(&CL100K.replace(r"\p{N}{1,3}", r"\p{N}{2,4}")), None); - assert_eq!(gpt_fsm(r"\w+|\s+"), None); + assert_eq!(recognize(&CL100K.replace(r"\p{N}{1,3}", r"\p{N}{2,4}")), None); + assert_eq!(recognize(r"\w+|\s+"), None); } } diff --git a/tokenizers/tk-encode/src/utils/word_cache.rs b/tokenizers/tk-encode/src/utils/word_cache.rs index a92ac3079..9e0d6c02c 100644 --- a/tokenizers/tk-encode/src/utils/word_cache.rs +++ b/tokenizers/tk-encode/src/utils/word_cache.rs @@ -52,28 +52,15 @@ use std::fmt::Debug; -use ahash::RandomState; use std::iter::Iterator; use wide::i8x16; -/// Hashes a word to the 64 bits its home slot and tag are taken from, and to the -/// bottom half of a long word's key ([`LookupKey::new_hash`]). -static PLACEMENT_HASHER: RandomState = RandomState::with_seeds( - 0x243f_6a88_85a3_08d3, - 0x1319_8a2e_0370_7344, - 0xa409_3822_299f_31d0, - 0x082e_fa98_ec4e_6c89, -); - -/// Hashes a long word a second time, to fill the half of its key that -/// [`PLACEMENT_HASHER`] does not reach. The two hashes must be independent, or the -/// key would carry 64 bits of information instead of 127. -static DISCRIMINANT_HASHER: RandomState = RandomState::with_seeds( - 0x4528_21e6_38d0_1377, - 0xbe54_66cf_34e9_0c6c, - 0xc0ac_29b7_c97c_50dd, - 0x3f84_d5b5_b547_0917, -); +use crate::vocab::bucket_vocab_store::key_and_hash; +#[cfg(test)] +use crate::vocab::bucket_vocab_store::INLINE_KEY_BYTES; + + +pub const MAX_INLINE_IDS: usize = 3; /// A table mapping words (`[u8]`) to the token ids they encode to (`[u32]`) pub struct WordCache { @@ -121,12 +108,56 @@ impl<'a> WordCache { /// Looks up a word in the cache. /// On [Lookup::Hit], returns the ids it encodes to. /// On [Lookup::Miss], returns the location in [Self::cached_words] where it should be inserted + #[inline] pub fn lookup(&'a self, word: &[u8]) -> Lookup<'a> { + self.lookup_placed(make_lookup_key(word, self.placement_mask)) + } + + #[inline] + pub fn lookup_keyed(&'a self, key: u64, hash: u64) -> Lookup<'a> { + self.lookup_placed(placement_from(LookupKey(key), hash, self.placement_mask)) + } + + /// + /// + /// + /// # Safety + #[inline] + pub unsafe fn probe_emit(&'a self, word: &[u8], dst: *mut u32) -> ProbeEmit<'a> { + debug_assert!(!word.is_empty(), "probe_emit needs a non-empty word"); + let (key, hash) = key_and_hash(word); + unsafe { self.probe_emit_keyed(key, hash, dst) } + } + + /// + /// # Safety + #[inline] + pub unsafe fn probe_emit_keyed(&'a self, key: u64, hash: u64, dst: *mut u32) -> ProbeEmit<'a> { + let placement = placement_from(LookupKey(key), hash, self.placement_mask); + // SAFETY: `index` is masked with `placement_mask` (`next_pow2 - 1`), and the table is + let slot = unsafe { *self.cached_words.as_ptr().add(placement.index) }; + if slot.key == placement.key && !slot.is_spilled() { + // SAFETY: the caller guarantees room for `MAX_INLINE_IDS`. Lanes past `ids_len` are + unsafe { + for lane in 0..MAX_INLINE_IDS { + dst.add(lane).write(slot.payload[lane]); + } + } + return ProbeEmit::Wrote(slot.ids_len as usize); + } + match self.lookup_placed(placement) { + Lookup::Hit(ids) => ProbeEmit::Hit(ids), + Lookup::Miss(at) => ProbeEmit::Miss(at), + } + } + + #[inline] + fn lookup_placed(&'a self, placement: InsertPlacement) -> Lookup<'a> { let InsertPlacement { key, index: home, tag, - } = make_lookup_key(word, self.placement_mask); + } = placement; let tag_window = self.tag_window(home); let (candidates, first_empty) = tag_window.find_matches_and_first_empty(tag); @@ -163,7 +194,7 @@ impl<'a> WordCache { let len = ids.len(); let InsertPlacement { index, key, tag } = placement; - let word = if len <= 3 { + let word = if len <= MAX_INLINE_IDS { WordCacheSlot::new_self_contained(key, ids) } else { if self.spilled_buffer.len() + len > self.spilled_budget { @@ -264,7 +295,7 @@ impl WordCacheSlot { } pub fn new_self_contained(key: LookupKey, ids: impl ExactSizeIterator) -> Self { - assert!(ids.len() <= 3); + assert!(ids.len() <= MAX_INLINE_IDS); let ids_len = ids.len() as u8; let mut payload = [0u32; 3]; for (slot, id) in payload.iter_mut().zip(ids) { @@ -362,91 +393,36 @@ impl<'a> SelfContained<'a> { /// ``` #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Default)] #[repr(transparent)] -pub struct LookupKey(u128); +pub struct LookupKey(u64); /// The key, home slot and tag of a word. +#[inline] fn make_lookup_key(word: &[u8], placement_mask: u64) -> InsertPlacement { - let placement_hash = PLACEMENT_HASHER.hash_one(word); - let key = if word.len() <= 15 { - LookupKey::new_inline(word) - } else { - LookupKey::new_hash(DISCRIMINANT_HASHER.hash_one(word), placement_hash) - }; + let (key, hash) = key_and_hash(word); + placement_from(LookupKey(key), hash, placement_mask) +} + +#[inline] +fn placement_from(key: LookupKey, hash: u64, placement_mask: u64) -> InsertPlacement { InsertPlacement { key, - index: (placement_hash & placement_mask) as usize, - tag: ((placement_hash >> (64 - 8)) as u8).max(WordCache::EMPTY + 1), + index: (hash & placement_mask) as usize, + tag: ((hash >> (64 - 8)) as u8).max(WordCache::EMPTY + 1), // ^ must be at least 0x01, otherwise can be mistaken for an EMPTY slot } } -impl LookupKey { - pub const TAG_MASK: u128 = 1 << 127; - - /// The key of a word of fifteen bytes or fewer: the word is its own key. - pub fn new_inline(word: &[u8]) -> Self { - let len = word.len(); - assert!(len <= 15); - // yes, this is a bit weird :) - // - // We used to do this: - // ```rust - // payload[..word.len()].copy_from_slice(word); - // payload[15] = word.len() as u8; - // Self(u128::from_le_bytes(payload)) - // ``` - // But that would compile into a memcpy call, probably because the len is only known at runtime. - // memcpy turned out to be quite slow and inefficient. - // - // The head / tail with fixed size compiles into plain register loads which are way faster - let raw = if len >= 8 { - let head = u64::from_le_bytes(word[..8].try_into().unwrap()) as u128; - let tail = u64::from_le_bytes(word[len - 8..].try_into().unwrap()) as u128; - head | tail << (8 * (len - 8)) - } else if len >= 4 { - let head = u32::from_le_bytes(word[..4].try_into().unwrap()) as u128; - let tail = u32::from_le_bytes(word[len - 4..].try_into().unwrap()) as u128; - head | tail << (8 * (len - 4)) - } else if len >= 1 { - let first = word[0] as u128; - let middle = (word[len / 2] as u128) << (8 * (len / 2)); - let last = (word[len - 1] as u128) << (8 * (len - 1)); - first | middle | last - } else { - 0 - }; - Self(raw | (len as u128) << 120) - } - - /// The key of a longer word: 127 bits of hash stand in for the word's bytes. - pub fn new_hash(discriminant: u64, placement: u64) -> Self { - Self(Self::TAG_MASK | (discriminant as u128) << 64 | placement as u128) - } -} - impl std::fmt::Debug for LookupKey { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - let mut debug = f.debug_struct("LookupKey"); - if self.0 & Self::TAG_MASK == 0 { + let len = (self.0 >> 56) as usize; + if len <= 7 { let bytes = self.0.to_le_bytes(); - let len = bytes[15] as usize; - debug - .field("type", &"inline") - .field("word_len", &len) - .field("word", &bytes[..len].escape_ascii().to_string()); + f.debug_tuple("LookupKey") + .field(&String::from_utf8_lossy(&bytes[..len]).into_owned()) + .finish() } else { - debug - .field("type", &"hashed") - .field( - "discriminant", - &format!("{:#x}", ((self.0 & !Self::TAG_MASK) >> 64) as u64), - ) - .field( - "placement", - &format!("{:#x}", (self.0 & u64::MAX as u128) as u64), - ); - } - debug.finish() + write!(f, "LookupKey(hash {:#018x})", self.0) + } } } @@ -461,6 +437,12 @@ pub enum Lookup<'a> { Miss(InsertPlacement), } +pub enum ProbeEmit<'a> { + Wrote(usize), + Hit(&'a [u32]), + Miss(InsertPlacement), +} + struct Window { window: [u8; WordCache::WINDOW_SIZE], offset: usize, @@ -636,7 +618,13 @@ mod tests { let key = |word: &[u8]| make_lookup_key(word, cache.placement_mask).key; assert_ne!(key(b"aaaaaaaaaaaaaa\x7f"), key(b"aaaaaaaaaaaaaa\xff")); assert_ne!(key(b"abcd"), key(b"abcd\0")); - assert_eq!(key(b"aaaaaaaaaaaaaa\xff").0 & LookupKey::TAG_MASK, 0); + let mut seen = std::collections::HashSet::new(); + for len in 1..=INLINE_KEY_BYTES { + for b in 0..=255u8 { + let word: Vec = (0..len).map(|i| b.wrapping_add(i as u8)).collect(); + assert!(seen.insert(key(&word).0), "collision at len={len} b={b}"); + } + } } /// One window shape per row: the needle in various lanes, an empty slot in @@ -716,19 +704,16 @@ mod tests { } /// Every inline length, with a different value in every byte position, so a - /// packing that drops, duplicates or misplaces a byte fails. The reference is - /// the construction the packing must be equivalent to: the bytes copied into - /// a zeroed array, the length written in the top byte. #[test] fn an_inline_key_is_the_words_bytes_with_the_length_on_top() { - for len in 0..=15usize { + for len in 0..=INLINE_KEY_BYTES { let word: Vec = (1..=len as u8).collect(); - let mut padded = [0u8; 16]; + let mut padded = [0u8; 8]; padded[..len].copy_from_slice(&word); - padded[15] = len as u8; + padded[7] = len as u8; assert_eq!( - LookupKey::new_inline(&word), - LookupKey(u128::from_le_bytes(padded)), + key_and_hash(&word).0, + u64::from_le_bytes(padded), "len={len}" ); } @@ -767,7 +752,7 @@ mod tests { assert_ne!(tag, WordCache::EMPTY, "pick a word with a nonzero tag"); cache.quick_lookup[index] = tag; cache.cached_words[index] = - WordCacheSlot::new_self_contained(LookupKey::new_inline(b"decoy"), [7].into_iter()); + WordCacheSlot::new_self_contained(LookupKey(key_and_hash(b"decoy").0), [7].into_iter()); assert_eq!(cache.lookup(b"beta").hit(), None); store(&mut cache, b"beta", &[2]); diff --git a/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs b/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs index 7f62ec0ec..6e18122ce 100644 --- a/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs +++ b/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs @@ -4,16 +4,15 @@ use ahash::RandomState; use ptr_hash::{FastPtrHash, PtrHashParams, hash::NoHash}; use std::fmt; -type Mphf = FastPtrHash; - -// Fixed seeds so a given vocab always hashes identically (the hasher is also stored on the struct, -// so build and query are guaranteed consistent regardless). -const SEEDS: [u64; 4] = [ +static KEY_HASHER: RandomState = RandomState::with_seeds( 0x243F_6A88_85A3_08D3, 0x1319_8A2E_0370_7344, 0xA409_3822_299F_31D0, 0x082E_FA98_EC4E_6C89, -]; +); + +type Mphf = FastPtrHash; + /// Bit 31 of a stored id: the token provably encodes to itself, so a pretoken equal to it can be /// emitted without running the merge loop. See `PipelineBPE::prove_fold`. @@ -26,14 +25,116 @@ const FOLD_BIT: u32 = 1 << 31; /// The id half. 2^31 ids is far past any vocabulary. const VOCAB_ID_MASK: u32 = FOLD_BIT - 1; -#[derive(Clone, Copy, Debug)] +pub(crate) const INLINE_KEY_BYTES: usize = 7; + +#[inline(always)] +fn mix(z: u64) -> u64 { + let z = z.wrapping_mul(0x9E37_79B9_7F4A_7C15); + z ^ (z >> 29) +} + +/// +/// +/// +/// +/// # Safety +/// `inline(always)`, and the long-word arm is a separate `inline(never)` function, because plain +/// `#[inline]` left this outlined: a profile of warm english had 8.3% of self time inside +/// `key_and_hash`, for what should be a masked load, an `or` and a multiply. LLVM weighs the whole +/// body when it decides, and the body used to contain the ahash call for a word too long to pack, so +/// the cheap path paid for the expensive one's size. Splitting them lets the pack inline into the +/// span loop while the hash stays a call -- which is also the shape it wants, since the long arm is +/// already dominated by hashing rather than by call overhead. +#[inline(always)] +pub fn key_and_hash_readable(word: &[u8], readable: usize) -> (u64, u64) { + let len = word.len(); + if len > INLINE_KEY_BYTES || readable < 8 { + return key_and_hash(word); + } + // SAFETY: `readable >= 8` bytes exist from `word.as_ptr()`, and `len <= 7 < 8`. + let raw = unsafe { word.as_ptr().cast::().read_unaligned() }; + // SAFETY: `len <= INLINE_KEY_BYTES == 7`, and both tables have 8 entries. + let (mask, tag) = unsafe { (*KEY_MASK.get_unchecked(len), *LEN_TAG.get_unchecked(len)) }; + let key = (raw & mask) | tag; + debug_assert_eq!(key, key_and_hash(word).0, "masked load must match the stitched pack"); + (key, mix(key)) +} + +/// A word too long to pack into the key: hash it for real. Kept out of line so that the packing +/// path above and in [`key_and_hash`] can be inlined without dragging ahash in with them. Not +/// `#[cold]` on purpose -- for CJK this is the *common* arm, and telling the predictor otherwise +/// would cost more than the call. +#[inline(never)] +fn key_and_hash_long(word: &[u8]) -> (u64, u64) { + let hash = KEY_HASHER.hash_one(word); + (hash, hash) +} + +#[inline(always)] +pub fn key_and_hash(word: &[u8]) -> (u64, u64) { + let len = word.len(); + if len > INLINE_KEY_BYTES { + return key_and_hash_long(word); + } + let raw = if len >= 4 { + let head = u32::from_le_bytes(word[..4].try_into().unwrap()) as u64; + let tail = u32::from_le_bytes(word[len - 4..].try_into().unwrap()) as u64; + head | tail << (8 * (len - 4)) + } else if len >= 1 { + let first = word[0] as u64; + let middle = (word[len / 2] as u64) << (8 * (len / 2)); + let last = (word[len - 1] as u64) << (8 * (len - 1)); + first | middle | last + } else { + 0 + }; + let key = raw | (len as u64) << 56; + (key, mix(key)) +} + +/// +#[derive(Clone, Copy, Debug, Default)] +#[repr(C)] struct Entry { - start: u32, - len: u16, + digest: u32, /// The token id in the low 31 bits, [`FOLD_BIT`] in the top. id: u32, } +const _: () = assert!(size_of::() == 8); + +static KEY_MASK: [u64; 8] = [ + 0x0000_0000_0000_0000, + 0x0000_0000_0000_00FF, + 0x0000_0000_0000_FFFF, + 0x0000_0000_00FF_FFFF, + 0x0000_0000_FFFF_FFFF, + 0x0000_00FF_FFFF_FFFF, + 0x0000_FFFF_FFFF_FFFF, + 0x00FF_FFFF_FFFF_FFFF, +]; +static LEN_TAG: [u64; 8] = [ + 0 << 56, + 1 << 56, + 2 << 56, + 3 << 56, + 4 << 56, + 5 << 56, + 6 << 56, + 7 << 56, +]; + +#[inline(always)] +fn digest_of(hash: u64) -> u32 { + (hash >> 32) as u32 +} + +#[derive(Clone, Copy, Debug, Default)] +struct Span { + start: u32, + len: u16, +} + /// The BucketVocabStore optimizes for space and speed. We don't use a HashMap to prevent duplicating the /// keys. Instead, we just use an `id_to_slot` and `entries` table. When you query bytes, you hash /// on the fly and get an `index` into the `entries` table. When you query an `id`, you fetch in @@ -55,11 +156,10 @@ struct Entry { #[derive(Clone)] pub struct BucketVocabStore { mphf: Mphf, - hasher: RandomState, /// All token bytes, concatenated. Ordered by MPHF slot. bytes: Box<[u8]>, - /// `entries[slot]` -> (offset into `bytes`, length, id). Ordered by MPHF slot. entries: Box<[Entry]>, + spans: Box<[Span]>, /// `id_to_slot[token_id] -> entry_idx` -> index into entries as the entries are not really sorted. id_to_slot: Box<[u32]>, /// Number of real tokens. Cached at build so `len()` is O(1): `entries` is sized to the @@ -102,12 +202,9 @@ impl BucketVocabStore { pub fn build(tokens: Vec<(Vec, u32)>) -> Self { let n = tokens.len(); - let hasher = RandomState::with_seeds(SEEDS[0], SEEDS[1], SEEDS[2], SEEDS[3]); - - // 1. Pre-hash token bytes -> u64 keys using near perfect hash func let keys: Vec = tokens .iter() - .map(|(s, _)| hasher.hash_one(s.as_slice())) + .map(|(s, _)| key_and_hash(s.as_slice()).1) .collect(); // 2. A perfect hash needs distinct keys. Collisions are astronomically unlikely @@ -142,29 +239,24 @@ impl BucketVocabStore { let total: usize = tokens.iter().map(|(s, _)| s.len()).sum(); let max_id = tokens.iter().map(|(_, id)| *id).max().unwrap(); let mut bytes = Vec::with_capacity(total); - let mut entries = vec![ - Entry { - start: 0, - len: 0, - id: 0 - }; - n_slots - ]; + let mut entries = vec![Entry::default(); n_slots]; + let mut spans = vec![Span::default(); n_slots]; let mut id_to_slot = vec![u32::MAX; max_id as usize + 1]; for (s, id) in &tokens { assert!( s.len() <= u16::MAX as usize, "token longer than 65535 bytes" ); - assert!( - *id <= VOCAB_ID_MASK, - "token id {id} needs bit 31, which holds FOLD_BIT" - ); - let slot = mphf.index(&hasher.hash_one(s.as_slice())); + assert!(*id <= VOCAB_ID_MASK, "token id {id} needs bit 31, which holds FOLD_BIT"); + let (key, hash) = key_and_hash(s.as_slice()); + let slot = mphf.index(&hash); entries[slot] = Entry { + digest: digest_of(hash), + id: *id, + }; + spans[slot] = Span { start: bytes.len() as u32, len: s.len() as u16, - id: *id, }; id_to_slot[*id as usize] = slot as u32; bytes.extend_from_slice(s); @@ -172,9 +264,9 @@ impl BucketVocabStore { Self { mphf, - hasher, bytes: bytes.into_boxed_slice(), entries: entries.into_boxed_slice(), + spans: spans.into_boxed_slice(), id_to_slot: id_to_slot.into_boxed_slice(), n, } @@ -185,9 +277,9 @@ impl BucketVocabStore { let empty: [u64; 0] = []; Self { mphf: FastPtrHash::::new(&empty, PtrHashParams::default_fast()), - hasher: RandomState::new(), bytes: Box::new([]), entries: Box::new([]), + spans: Box::new([]), id_to_slot: Box::new([]), n: 0, } @@ -202,34 +294,46 @@ impl BucketVocabStore { if self.entries.is_empty() { return None; } - let slot = self.mphf.index(&self.hasher.hash_one(q)); - + let (key, hash) = key_and_hash(q); + let slot = self.mphf.index(&hash); let e = self.entries[slot]; - let (start, len) = (e.start as usize, e.len as usize); - // Byte equality: confirms `q` really is the token at this slot (perfect hashing only - // guarantees a valid slot for in-vocab keys; this rejects collisions and Out Of Vocab queries). - if len == q.len() && self.bytes[start..start + len] == *q { - Some(e.id & VOCAB_ID_MASK) - } else { - None - } + (e.digest == digest_of(hash)).then_some(e.id & VOCAB_ID_MASK) } /// The id for `q`, together with whether that entry may be folded. One probe and one entry /// load: the flag is a bit of the id the probe already read. #[inline] pub fn get_bytes_foldable(&self, q: &[u8]) -> Option<(u32, bool)> { + let (key, hash) = key_and_hash(q); + self.get_keyed_foldable(key, hash) + } + + #[inline(always)] + pub fn probe_slot(&self, hash: u64) -> usize { + self.mphf.index(&hash) + } + + #[inline(always)] + pub fn entry_at(&self, slot: usize) -> (u32, u32) { + let e = self.entries[slot]; + (e.digest, e.id) + } + + #[inline(always)] + pub fn resolve_foldable(hash: u64, entry: (u32, u32)) -> Option<(u32, bool)> { + let (edigest, eid) = entry; + (edigest == digest_of(hash)).then_some((eid & VOCAB_ID_MASK, eid & FOLD_BIT != 0)) + } + + #[inline] + pub fn get_keyed_foldable(&self, key: u64, hash: u64) -> Option<(u32, bool)> { + let _ = key; if self.entries.is_empty() { return None; } - let slot = self.mphf.index(&self.hasher.hash_one(q)); + let slot = self.mphf.index(&hash); let e = self.entries[slot]; - let (start, len) = (e.start as usize, e.len as usize); - if len == q.len() && self.bytes[start..start + len] == *q { - Some((e.id & VOCAB_ID_MASK, e.id & FOLD_BIT != 0)) - } else { - None - } + (e.digest == digest_of(hash)).then_some((e.id & VOCAB_ID_MASK, e.id & FOLD_BIT != 0)) } /// Records that this token folds to itself. Called once per entry at load, after the proof. @@ -253,9 +357,9 @@ impl BucketVocabStore { if slot == u32::MAX { return None; // id is within range but absent from the vocab } - let e = self.entries[slot as usize]; - let start = e.start as usize; - self.bytes.get(start..start + e.len as usize) + let sp = self.spans[slot as usize]; + let start = sp.start as usize; + self.bytes.get(start..start + sp.len as usize) } #[inline] @@ -285,9 +389,10 @@ impl BucketVocabStore { pub fn content(&self) -> Vec<(String, u32)> { self.entries .iter() - .filter(|e| e.len > 0) + .zip(self.spans.iter()) + .filter(|(_, sp)| sp.len > 0) // Mask: the stored id carries FOLD_BIT, which must never escape this type. - .map(|m| m.id & VOCAB_ID_MASK) + .map(|(e, _)| e.id & VOCAB_ID_MASK) .filter_map(|id| self.id_to_token(id).map(|token| (token, id))) .collect() } @@ -301,9 +406,10 @@ impl BucketVocabStore { pub fn byte_content(&self) -> Vec<(Vec, u32)> { self.entries .iter() - .filter(|e| e.len > 0) + .zip(self.spans.iter()) + .filter(|(_, sp)| sp.len > 0) // Mask: the stored id carries FOLD_BIT, which must never escape this type. - .map(|m| m.id & VOCAB_ID_MASK) + .map(|(e, _)| e.id & VOCAB_ID_MASK) .filter_map(|id| self.id_to_token_bytes(id).map(|token| (token.to_vec(), id))) .collect() } diff --git a/tokenizers/tk-encode/src/vocab/buckets.rs b/tokenizers/tk-encode/src/vocab/buckets.rs index 9a78a8cec..ff46cf50c 100644 --- a/tokenizers/tk-encode/src/vocab/buckets.rs +++ b/tokenizers/tk-encode/src/vocab/buckets.rs @@ -333,14 +333,23 @@ impl Buckets { // needle = the bucket's shared first byte. Assumes a non-empty prefix // (false only if a lone 1-byte token is the sole holder of its first byte); store // the first byte explicitly if that case ever appears. - let needle = self.buckets[0].prefix[0]; - let mut search = 0usize; - while let Some(off) = memchr::memchr(needle, &bytes[search..]) { - let pos = search + off; + // Search for the bucket's whole shared prefix, not just its first byte. + // + // One byte of a long needle is a poor filter on text that is dense in that byte and + // holds no match: `<|endoftext|>` over a corpus full of `<|xs0|>` stops at every `<` + // and dies at the third byte. Restarting `memchr` on `&bytes[pos + 1..]` per + // candidate then pays its prologue again each time. Measured over the same + // one-bucket vocabulary in one process, that cost 0.58 ns/B on a 9.8%-`<` corpus + // that matches nothing, against 0.12 for `memmem`, which picks a *rare* byte of the + // needle instead. Where candidates are already sparse it costs at most 0.015 ns/B. + // + // `nibble_mask_match` already avoids this for two or more buckets and calls it the + // restart penalty. Same positions in the same order, so the leftmost match is + // unchanged and `match_fast` still confirms the length sub-list. + for pos in memchr::memmem::find_iter(bytes, &self.buckets[0].prefix) { if let Some((id, len)) = self.match_fast(bytes, pos, 0) { return Some((id, pos as u32, len)); } - search = pos + 1; } None }