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 ba9632eb1..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" @@ -170,7 +138,7 @@ name = "bitsplit" version = "0.1.0" dependencies = [ "ahash", - "atomsplit", + "onig", ] [[package]] @@ -241,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", ] @@ -663,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", ] @@ -1188,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" @@ -1258,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" @@ -1484,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" @@ -1909,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" @@ -1997,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" @@ -2294,14 +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 fdd4f4894..5100c6367 100644 --- a/tokenizers/Cargo.toml +++ b/tokenizers/Cargo.toml @@ -1,6 +1,6 @@ [workspace] resolver = "3" -members = ["bitmap_gen", "atomsplit", "bitsplit", "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 index ebc3d609e..e7dc67097 100644 --- a/tokenizers/bitsplit/Cargo.toml +++ b/tokenizers/bitsplit/Cargo.toml @@ -2,12 +2,33 @@ 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" -atomsplit = { path = "../atomsplit" } - - - +# 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/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 index 79d619b1e..8e229f2cb 100644 --- a/tokenizers/bitsplit/src/lib.rs +++ b/tokenizers/bitsplit/src/lib.rs @@ -17,24 +17,58 @@ //! 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: `atomsplit::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. +//! 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 their `atomsplit::fsm` counterparts — see `src/bin/verify.rs`. +//! byte-exact with the oniguruma oracle over a block-phase sweep — see `tests/parity.rs`. -pub(crate) use atomsplit::fsm::Span; - -pub mod deepseek; -pub mod gpt; -#[cfg(target_arch = "aarch64")] +pub mod classes; +pub mod classify; +mod han; +pub mod literal; +pub mod models; +pub mod regexes; mod simd; -#[cfg(target_arch = "x86_64")] -mod simd_x86; -pub use deepseek::bitsplit_deepseek; -pub use gpt::{bitsplit_byte_level, bitsplit_cl100k}; +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 @@ -56,7 +90,7 @@ pub fn fast_builder() -> bool { } #[cfg(target_arch = "x86_64")] -fn has_ssse3() -> bool { +pub(crate) fn has_ssse3() -> bool { use std::sync::atomic::{AtomicU8, Ordering}; static CACHED: AtomicU8 = AtomicU8::new(0); match CACHED.load(Ordering::Relaxed) { @@ -72,17 +106,29 @@ fn has_ssse3() -> bool { pub(crate) const CONT: u8 = 15; // Atom::Cont pub(crate) const CODE_CONT: u8 = 7; // every grammar's dense code for a continuation byte -/// The bitstreams for one 64-byte block. `p0`/`p1`/`p2` are the bit-planes of the **filled** dense +/// 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); `cjk` is text-derived and only built when a grammar asks for it. +/// 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 cjk: u64, + pub p3: u64, + pub aux: u64, } /// deepseek Split-2's isolated range: Han U+4E00..9FA5 ∪ Hiragana/Katakana U+3040..30FF (all @@ -99,43 +145,56 @@ pub(crate) fn is_cjk_at(text: &[u8], p: usize) -> bool { (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`. -/// `CJK` asks for the (deepseek-only) range stream; `lut` is the grammar's tag → dense code table. +/// `lut` is the grammar's tag → dense code table. #[inline] -pub(crate) fn build_block( +pub(crate) fn build_block( text: &[u8], tags: &[u8], base: usize, len: usize, lut: &[u8; 64], cur_code: u8, - cur_cjk: bool, + 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::build64::(text, tags, base, lut, cur_code, cur_cjk) }; + 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_cjk) + crate::simd::x86::build64::(text, tags, base, lut, cur_code, cur_aux) }; } - build_block_scalar::(text, tags, base, len, lut, cur_code, cur_cjk) + 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( +pub(crate) fn build_block_scalar( text: &[u8], tags: &[u8], base: usize, len: usize, lut: &[u8; 64], mut cur_code: u8, - mut cur_cjk: bool, + mut cur_aux: bool, ) -> (Blk, u8) { let mut b = Blk::default(); for i in 0..len { @@ -145,12 +204,15 @@ pub(crate) fn build_block_scalar( b.cont |= bit; } else { cur_code = lut[tags[p] as usize]; - cur_cjk = CJK && is_cjk_at(text, p); + 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); - b.cjk |= bit * u64::from(cur_cjk); + if P3 { + b.p3 |= bit * u64::from(cur_code & 8 != 0); + } + b.aux |= bit * u64::from(cur_aux); } (b, cur_code) } @@ -229,6 +291,143 @@ pub(crate) fn trail_run(x: u64, valid: u64, len: usize) -> u64 { } } +/// `\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` @@ -329,10 +528,11 @@ pub(crate) fn emit_contr( 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] -fn contr_len(text: &[u8], i: usize, ci: bool) -> usize { +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; @@ -357,10 +557,10 @@ 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::(text, tags, base, len, &deepseek::LUT, code, cjk); + let (b, c) = build_block::<{ AUX_CJK }, false>(text, tags, base, len, &models::deepseek::LUT, code, cjk); code = c; - cjk = b.cjk >> (len - 1) & 1 != 0; - acc ^= b.cont ^ b.p0 ^ b.p1 ^ b.p2 ^ b.cjk; + cjk = b.aux >> (len - 1) & 1 != 0; + acc ^= b.cont ^ b.p0 ^ b.p1 ^ b.p2 ^ b.aux; } acc } @@ -368,6 +568,6 @@ pub fn build_only(text: &[u8], tags: &[u8]) -> u64 { /// 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 { - atomsplit::classify::classify(text, tags); + 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/gpt.rs b/tokenizers/bitsplit/src/models/cl100k.rs similarity index 62% rename from tokenizers/bitsplit/src/gpt.rs rename to tokenizers/bitsplit/src/models/cl100k.rs index a8b89bc8e..9d16e1b2c 100644 --- a/tokenizers/bitsplit/src/gpt.rs +++ b/tokenizers/bitsplit/src/models/cl100k.rs @@ -1,17 +1,12 @@ -//! The two tiktoken-family grammars as bitstream programs. +//! **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+` //! -//! * [`bitsplit_byte_level`] — GPT-2 / Llama / Qwen: -//! `'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+` -//! * [`bitsplit_cl100k`] — cl100k_base / Llama-3: -//! `(?i:'s|…)|[^\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+` -//! -//! Both share one dense code table, and both hand contractions to the scalar escape in -//! [`crate::emit_contr`] rather than trying to express a variable-length case-optional literal -//! alternation in bit algebra. +//! 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::{ - CODE_CONT, CONT, Span, adv, build_block, emit_contr, fill_to_last, lead_run, scanthru, to_lead, - trail_run, + 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 @@ -81,98 +76,37 @@ const fn code_bits(code: u8) -> u8 { } } -/// 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`. +/// cl100k_base / Llama-3 / GLM-4.6 — rule 3 is `\p{N}{1,3}`. #[must_use] -pub fn bitsplit_byte_level( +pub fn bitsplit_cl100k( 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::(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) + cl100k(text, tags, starts, flag, out, 3) } -/// cl100k_base / Llama-3 pre-tokenization. Byte-exact with `atomsplit::fsm::fsm_cl100k`. +/// Qwen2 / Qwen3 — cl100k character-for-character except rule 3 is a bare `\p{N}`. #[must_use] -pub fn bitsplit_cl100k( +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 { @@ -193,7 +127,7 @@ pub fn bitsplit_cl100k( let valid = if len == 64 { !0u64 } else { (1u64 << len) - 1 }; let last_blk = base + len == ntext; - let (b, last_code) = build_block::(text, tags, base, len, &LUT, code, false); + 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) }; @@ -257,33 +191,16 @@ pub fn bitsplit_cl100k( & !p1(c.ws & !c.nl, has(pb, C_WS) && !has(pb, C_NL)) & !p1(osf, prev_osf); - // ── `\p{N}{1,3}`: a group boundary every 3 chars from the run start. - let mut m = c.n & lead & !p1(c.n, has(pb, C_N)); - if dig_run && has(pb, C_N) { - let mut s = lead & lead.wrapping_neg() & c.n; - for _ in 0..((3 - dig_since % 3) % 3) { - s = adv(s, b.cont) & c.n & lead; - } - m |= s; - } - let mut groups = m; - if c.n & b.cont == 0 { - let n3 = c.n & (c.n << 1) & (c.n << 2); - while m != 0 { - m = (m << 3) & n3; - groups |= m; - } - } else { - while m != 0 { - let a = adv(m, b.cont) & c.n & lead; - let e = adv(adv(a, b.cont) & c.n & lead, b.cont) & c.n & lead; - if e == 0 { - break; - } - groups |= e; - m = e; - } - } + // ── 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); @@ -327,7 +244,11 @@ pub fn bitsplit_cl100k( } else { (c.n & lead & tn & !((1u64 << (63 - g.leading_zeros())) - 1)).count_ones() }; - counted % 3 + 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) { diff --git a/tokenizers/bitsplit/src/deepseek.rs b/tokenizers/bitsplit/src/models/deepseek.rs similarity index 97% rename from tokenizers/bitsplit/src/deepseek.rs rename to tokenizers/bitsplit/src/models/deepseek.rs index 3bafdbe91..7685b4d0a 100644 --- a/tokenizers/bitsplit/src/deepseek.rs +++ b/tokenizers/bitsplit/src/models/deepseek.rs @@ -3,7 +3,7 @@ //! `atomsplit::fsm::fsm_deepseek`. use crate::{ - CODE_CONT, CONT, Span, adv, build_block, emit, fill_to_last, is_cjk_at, lead_run, scanthru, + AUX_CJK, CODE_CONT, CONT, Span, adv, build_block, emit, fill_to_last, is_cjk_at, lead_run, scanthru, to_lead, trail_run, }; @@ -105,14 +105,14 @@ pub fn bitsplit_deepseek(text: &[u8], tags: &[u8], starts: &mut [u64], out: &mut let valid = if len == 64 { !0u64 } else { (1u64 << len) - 1 }; let last_blk = base + len == ntext; - let (b, last_code) = build_block::(text, tags, base, len, &LUT, cy.code, cy.cjk); + 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.cjk >> (len - 1) & 1 != 0; + 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 @@ -138,7 +138,7 @@ pub fn bitsplit_deepseek(text: &[u8], tags: &[u8], starts: &mut [u64], out: &mut 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.cjk; + let cjk = b.aux; let num = s_n & !cjk; let lm = s_lm & !cjk; let ps = s_ps & !cjk; 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.rs b/tokenizers/bitsplit/src/simd/neon.rs similarity index 76% rename from tokenizers/bitsplit/src/simd.rs rename to tokenizers/bitsplit/src/simd/neon.rs index d328dc502..e685ab7ba 100644 --- a/tokenizers/bitsplit/src/simd.rs +++ b/tokenizers/bitsplit/src/simd/neon.rs @@ -10,7 +10,7 @@ //! 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::{Blk, lead_run}; +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]; @@ -27,6 +27,27 @@ unsafe fn mm64(v: [uint8x16_t; 4], pow: uint8x16_t) -> u64 { } } +/// 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. @@ -34,13 +55,13 @@ unsafe fn mm64(v: [uint8x16_t; 4], pow: uint8x16_t) -> u64 { /// # Safety /// `base + 64 <= tags.len()` and `base + 64 <= text.len()`. #[target_feature(enable = "neon")] -pub(crate) unsafe fn build64( +pub(crate) unsafe fn build64( text: &[u8], tags: &[u8], base: usize, lut: &[u8; 64], cur_code: u8, - cur_cjk: bool, + cur_aux: bool, ) -> (Blk, u8) { unsafe { let pow = vld1q_u8(POW.as_ptr()); @@ -89,12 +110,13 @@ pub(crate) unsafe fn build64( p0: plane(1), p1: plane(2), p2: plane(4), - cjk: 0, + p3: if P3 { plane(8) } else { 0 }, + aux: 0, }; - if !CJK { + if AUX == AUX_NONE { return (b, last_code); } - // ── text is loaded only for the (deepseek-only) CJK range test. + // ── text is loaded only for the aux (text-derived) stream. let ntext = text.len(); let tv = [ vld1q_u8(text.as_ptr().add(base)), @@ -103,6 +125,37 @@ pub(crate) unsafe fn build64( 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 @@ -165,11 +218,11 @@ pub(crate) unsafe fn build64( 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_cjk` (its continuation bytes lead that block). - b.cjk = leads | (leads << 1) | (leads << 2); + // picked up on the other side by `cur_aux` (its continuation bytes lead that block). + b.aux = leads | (leads << 1) | (leads << 2); } - if cur_cjk { - b.cjk |= lead_run(b.cont, !0); + 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 similarity index 75% rename from tokenizers/bitsplit/src/simd_x86.rs rename to tokenizers/bitsplit/src/simd/x86.rs index 37bb29838..407f5e481 100644 --- a/tokenizers/bitsplit/src/simd_x86.rs +++ b/tokenizers/bitsplit/src/simd/x86.rs @@ -15,7 +15,7 @@ //! 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::{Blk, lead_run}; +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. @@ -43,18 +43,35 @@ unsafe fn sel(mask: __m128i, a: __m128i, b: __m128i) -> __m128i { unsafe { _mm_or_si128(_mm_and_si128(mask, a), _mm_andnot_si128(mask, b)) } } -/// Build one **full** 64-byte block. `cur_code`/`cur_cjk` describe the byte before it. +/// 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( +pub(crate) unsafe fn build64( text: &[u8], tags: &[u8], base: usize, lut: &[u8; 64], cur_code: u8, - cur_cjk: bool, + cur_aux: bool, ) -> (Blk, u8) { unsafe { let low_nibble = _mm_set1_epi8(0x0F); @@ -115,9 +132,38 @@ pub(crate) unsafe fn build64( p0: plane!(7), p1: plane!(6), p2: plane!(5), - cjk: 0, + p3: if P3 { plane!(4) } else { 0 }, + aux: 0, }; - if !CJK { + 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); } @@ -134,8 +180,8 @@ pub(crate) unsafe fn build64( .iter() .any(|v| _mm_movemask_epi8(in_range(*v, 0xE3, 6)) != 0); if !any { - if cur_cjk { - b.cjk |= lead_run(b.cont, !0); + if cur_aux { + b.aux |= lead_run(b.cont, !0); } return (b, last_code); } @@ -182,10 +228,10 @@ pub(crate) unsafe fn build64( 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_cjk`. - b.cjk = leads | (leads << 1) | (leads << 2); - if cur_cjk { - b.cjk |= lead_run(b.cont, !0); + // 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 7535aa994..bdea45df3 100644 --- a/tokenizers/tk-encode/Cargo.toml +++ b/tokenizers/tk-encode/Cargo.toml @@ -27,7 +27,6 @@ name = "tk_encode" path = "src/lib.rs" [dependencies] -atomsplit = { path = "../atomsplit" } bitsplit = { path = "../bitsplit" } rand = "0.9" regex = "1.10" diff --git a/tokenizers/tk-encode/src/models/bpe/model.rs b/tokenizers/tk-encode/src/models/bpe/model.rs index 17ef63c51..0eb89a8fe 100644 --- a/tokenizers/tk-encode/src/models/bpe/model.rs +++ b/tokenizers/tk-encode/src/models/bpe/model.rs @@ -362,9 +362,8 @@ impl pipeline::Model for PipelineBPE { output: &mut Vec, ) -> Result<()> { let BpeScratch { - merge_queue, - skip, - word, + symbols, + queue, word_cache, } = scratch; @@ -380,6 +379,15 @@ impl pipeline::Model for PipelineBPE { continue; } + // Same order as `tokenize_pipeline`, and it has to stay that way: the fold answers a + // word that is itself a foldable vocabulary entry in one probe, and those words never + // reach the cache. Probing the cache first would populate it with words the fold + // already serves for free, and the two paths would disagree about what it holds. + if let Some(id) = self.fold_id(sequence) { + output.push(PipelineToken { id }); + continue; + } + let mut placement = None; if let Some(cache) = word_cache.as_mut() { match cache.lookup(sequence.as_bytes()) { @@ -390,17 +398,13 @@ impl pipeline::Model for PipelineBPE { Lookup::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 })); - } - // The ids come back out of `output` because that is the only place both branches - // above leave them: `ignore_merges` never touches `word`. + 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) = word_cache.as_mut() && let Some(at) = placement { @@ -410,43 +414,6 @@ impl pipeline::Model for PipelineBPE { Ok(()) } - /// Every pre-token of a chunk in one call. - /// - /// Same work per word as [`Self::tokenize_pipeline`]; what changes is what is *not* repeated. - /// The scratch is destructured once instead of once per word, the output is grown once for the - /// whole batch instead of being capacity-checked on every push, and the virtual call, the - /// slice and the `Result` happen once per chunk rather than once per pre-token. - fn tokenize_spans( - &self, - chunk: &str, - spans: &[Span], - scratch: &mut Self::Scratch, - output: &mut Vec, - ) -> Result<()> { - let BpeScratch { symbols, queue } = 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()); - - for span in spans { - // SAFETY: the pre-tokenizer cuts on char boundaries, so a span is always a valid slice - // of this chunk. Bounds- and UTF-8-checking it again per word measured worth removing. - let sequence = unsafe { chunk.get_unchecked(span.range()) }; - if sequence.is_empty() { - continue; - } - if let Some(id) = self.fold_id(sequence) { - output.push(PipelineToken { id }); - continue; - } - self.merge_word(sequence, symbols, queue); - output.extend(symbols.iter().map(|&symbol| PipelineToken { - id: self.tables.unmap.at(symbol as usize), - })); - } - Ok(()) - } - fn init_scratch(&self) -> Self::Scratch { Self::Scratch { symbols: Vec::with_capacity(64), 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 04a6b4c30..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,53 +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) { - // gpt2 and cl100k-with-the-standard-digit-cap have bitstream splitters, which decide - // 64 bytes per register op instead of one token per unpredictable branch -- but only - // where the bitstream build is SIMD. Its portable builder is slower than the FSM, so - // everything else keeps the FSM. - match fsm { - GptFsm::Gpt2 if bitsplit::fast_builder() => { - pipeline::classify_into_spans_bits( - text.as_bytes(), - bitsplit::bitsplit_byte_level, - out, - ); - return Ok(()); - } - GptFsm::Cl100k { digit_cap: 3 } if bitsplit::fast_builder() => { - pipeline::classify_into_spans_bits( - text.as_bytes(), - bitsplit::bitsplit_cl100k, - out, - ); - return Ok(()); - } - _ => {} - } - 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(()); @@ -442,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"); @@ -466,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(); @@ -494,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!( @@ -521,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" ); @@ -554,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 ef859c16e..55385b574 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( 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); } }