From 059bcb8f60eb344cf6bc88624efb2e8ff01e4cd6 Mon Sep 17 00:00:00 2001 From: SBrandeis <33657802+SBrandeis@users.noreply.github.com> Date: Wed, 29 Jul 2026 16:46:48 +0200 Subject: [PATCH 1/6] new literal module in AtomSplit --- tokenizers/atomsplit/Cargo.toml | 2 +- tokenizers/atomsplit/src/lib.rs | 5 +++ tokenizers/atomsplit/src/literal.rs | 53 +++++++++++++++++++++++++++ tokenizers/atomsplit/tests/literal.rs | 39 ++++++++++++++++++++ 4 files changed, 98 insertions(+), 1 deletion(-) create mode 100644 tokenizers/atomsplit/src/literal.rs create mode 100644 tokenizers/atomsplit/tests/literal.rs diff --git a/tokenizers/atomsplit/Cargo.toml b/tokenizers/atomsplit/Cargo.toml index 0a9de6a86..fb8947f3f 100644 --- a/tokenizers/atomsplit/Cargo.toml +++ b/tokenizers/atomsplit/Cargo.toml @@ -25,7 +25,7 @@ name = "atomsplit" path = "src/lib.rs" [dependencies] -memchr = "2.8.2" # SIMD single-byte search; used only in CharDelimiterSplit (1.4–23× vs scalar). +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 diff --git a/tokenizers/atomsplit/src/lib.rs b/tokenizers/atomsplit/src/lib.rs index ceefc9296..5f22f52d0 100644 --- a/tokenizers/atomsplit/src/lib.rs +++ b/tokenizers/atomsplit/src/lib.rs @@ -7,6 +7,10 @@ //! 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. //! @@ -17,6 +21,7 @@ mod atom_tables; pub mod classify; pub mod fsm; +pub mod literal; pub mod regexes; #[cfg(target_arch = "x86_64")] mod simd_avx_classify; diff --git a/tokenizers/atomsplit/src/literal.rs b/tokenizers/atomsplit/src/literal.rs new file mode 100644 index 000000000..041694cd4 --- /dev/null +++ b/tokenizers/atomsplit/src/literal.rs @@ -0,0 +1,53 @@ +//! 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. +#[derive(Debug, Clone)] +pub struct Literal { + finder: memmem::Finder<'static>, +} + +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: 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/literal.rs b/tokenizers/atomsplit/tests/literal.rs new file mode 100644 index 000000000..80978ba42 --- /dev/null +++ b/tokenizers/atomsplit/tests/literal.rs @@ -0,0 +1,39 @@ +//! 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]); +} + +#[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" + ); +} From f01e1656141c2ecb11b20cbd37730e13c280a024 Mon Sep 17 00:00:00 2001 From: SBrandeis <33657802+SBrandeis@users.noreply.github.com> Date: Wed, 29 Jul 2026 17:15:04 +0200 Subject: [PATCH 2/6] use Literal in Replace normalizer and Split pretok --- tokenizers/tk-encode/Cargo.toml | 9 +- .../tk-encode/src/normalizers/replace.rs | 160 +++++++++++++----- .../tk-encode/src/pre_tokenizers/split.rs | 128 ++++++++++---- .../tk-encode/src/tokenizer/normalizer.rs | 3 +- tokenizers/tk-encode/src/tokenizer/pattern.rs | 73 +++++--- .../tk-encode/src/tokenizer/pipeline.rs | 28 +++ tokenizers/tk-encode/src/utils/mod.rs | 7 +- tokenizers/tk-encode/src/utils/no_regex.rs | 14 +- tokenizers/tk-encode/tests/pipeline_oracle.rs | 9 +- 9 files changed, 310 insertions(+), 121 deletions(-) diff --git a/tokenizers/tk-encode/Cargo.toml b/tokenizers/tk-encode/Cargo.toml index ba64839a5..5f830e3ee 100644 --- a/tokenizers/tk-encode/Cargo.toml +++ b/tokenizers/tk-encode/Cargo.toml @@ -73,10 +73,11 @@ pcre2 = { version = "0.2", optional = true } logos = { version = "0.15", optional = true } # compile-time DFA lexer reference (pure Rust) [features] -# `fancy-regex` is the OPTIONAL system-regex backend, needed ONLY for a `Split` pre-tokenizer with an -# arbitrary (non-GPT) regex or the `Replace` normalizer — the atomsplit-native pre-tokenizers (GPT-2, -# cl100k, deepseek, the class family, char-delimiter) need no backend. Without it a stub compiles and -# those arbitrary-regex paths error at load. Enable with `--features fancy-regex`. +# `fancy-regex` is the OPTIONAL system-regex backend, needed ONLY for a *regex* pattern in a `Split` +# pre-tokenizer or a `Replace` normalizer — 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. Without it a stub compiles and those regex paths error at load. Enable with +# `--features fancy-regex`. default = ["progressbar"] progressbar = ["indicatif"] http = ["hf-hub"] diff --git a/tokenizers/tk-encode/src/normalizers/replace.rs b/tokenizers/tk-encode/src/normalizers/replace.rs index 48fc58180..cc4545b47 100644 --- a/tokenizers/tk-encode/src/normalizers/replace.rs +++ b/tokenizers/tk-encode/src/normalizers/replace.rs @@ -5,6 +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 serde::{Deserialize, Serialize}; /// Represents the different patterns that `Replace` can use @@ -26,7 +27,7 @@ impl From<&str> for ReplacePattern { } } -/// We use this custom deserializer to provide the value for `regex` for `Replace` +/// We use this custom deserializer to build the search for `Replace` #[doc(hidden)] #[derive(Deserialize)] #[serde(tag = "type")] @@ -43,6 +44,27 @@ impl std::convert::TryFrom for Replace { } } +/// How a [`Replace`] looks for its pattern. +#[derive(Debug)] +enum Search { + /// A plain string, scanned for directly — no regex engine involved, so this works in every build. + Literal(Literal), + /// A real regex, which needs the system backend (the `fancy-regex` feature). + Regex(SysRegex), + /// The empty string, which matches nothing. + Nothing, +} + +impl Search { + fn find_matches(&self, inside: &str) -> Result> { + match self { + Self::Literal(literal) => literal.find_matches(inside), + Self::Regex(regex) => regex.find_matches(inside), + Self::Nothing => Ok(vec![((0, inside.len()), false)]), + } + } +} + /// This normalizer will take a `pattern` (for now only a String) /// and replace every occurrence with `content`. #[derive(Debug, Serialize, Deserialize)] @@ -51,7 +73,7 @@ pub struct Replace { pattern: ReplacePattern, pub content: String, #[serde(skip)] - regex: SysRegex, + search: Search, } impl Clone for Replace { @@ -69,45 +91,68 @@ impl PartialEq for Replace { impl Replace { pub fn new, C: Into>(pattern: I, content: C) -> Result { let pattern: ReplacePattern = pattern.into(); - let regex = match &pattern { - ReplacePattern::String(s) => SysRegex::new(®ex::escape(s))?, - ReplacePattern::Regex(r) => SysRegex::new(r)?, + let search = match &pattern { + ReplacePattern::String(s) if s.is_empty() => Search::Nothing, + ReplacePattern::String(s) => Search::Literal(Literal::new(s.as_bytes())?), + ReplacePattern::Regex(r) => Search::Regex(SysRegex::new(r)?), }; Ok(Self { pattern, content: content.into(), - regex, + search, }) } } impl Normalizer for Replace { fn normalize(&self, normalized: &mut NormalizedString) -> Result<()> { - normalized.replace(&self.regex, &self.content) + match &self.search { + Search::Literal(literal) => normalized.replace(literal, &self.content), + Search::Regex(regex) => normalized.replace(regex, &self.content), + Search::Nothing => Ok(()), + } + } +} + +/// Builds the text with every match swapped for `content`. Borrows the input back untouched when +/// nothing matched, which is what keeps the common case free of allocation. +fn replace_matches<'a>( + input: &'a str, + content: &str, + matches: impl Iterator, +) -> Cow<'a, str> { + let mut replaced: Option = None; + let mut last_end = 0; + for (start, end) in matches { + let replaced: &mut String = + replaced.get_or_insert_with(|| String::with_capacity(input.len())); + replaced.push_str(&input[last_end..start]); + replaced.push_str(content); + last_end = end; + } + match replaced { + Some(mut replaced) => { + replaced.push_str(&input[last_end..]); + Cow::Owned(replaced) + } + None => Cow::Borrowed(input), } } impl pipeline::Normalizer for Replace { fn normalize<'a>(&self, input: &'a str) -> Result> { - let iter = self.regex.find_iter(input); - let mut replaced: Option = None; - let mut last_end = 0; - - for (start, end) in iter { - let replaced: &mut String = - replaced.get_or_insert_with(|| String::with_capacity(input.len())); - replaced.push_str(&input[last_end..start]); - replaced.push_str(&self.content); - last_end = end; - } - if let Some(mut replaced) = replaced { - if last_end < input.len() { - replaced.push_str(&input[last_end..]); + Ok(match &self.search { + Search::Literal(literal) => { + let width = literal.pattern().len(); + let matches = literal + .matches(input.as_bytes()) + .map(|start| (start, start + width)); + replace_matches(input, &self.content, matches) } - return Ok(Cow::Owned(replaced)); - } - Ok(input.into()) + Search::Regex(regex) => replace_matches(input, &self.content, regex.find_iter(input)), + Search::Nothing => Cow::Borrowed(input), + }) } } @@ -118,7 +163,7 @@ impl Decoder for Replace { .map(|token| -> Result { let mut new_token = "".to_string(); - for ((start, stop), is_match) in (&self.regex).find_matches(&token)? { + for ((start, stop), is_match) in self.search.find_matches(&token)? { if is_match { new_token.push_str(&self.content); } else { @@ -131,8 +176,7 @@ impl Decoder for Replace { } } -// `Replace` needs a system-regex backend (SysRegex) for every test here. -#[cfg(all(test, feature = "fancy-regex"))] +#[cfg(test)] mod tests { use super::*; @@ -148,6 +192,7 @@ mod tests { } #[test] + #[cfg(feature = "fancy-regex")] // a regex pattern needs a system-regex backend fn test_replace_regex() { let original = "This is a test"; let normalized = "This is a test"; @@ -162,6 +207,7 @@ mod tests { } #[test] + #[cfg(feature = "fancy-regex")] // the regex half of this needs a system-regex backend fn serialization() { let replace = Replace::new("Hello", "Hey").unwrap(); let replace_s = r#"{"type":"Replace","pattern":{"String":"Hello"},"content":"Hey"}"#; @@ -174,6 +220,37 @@ mod tests { assert_eq!(serde_json::from_str::(replace_s).unwrap(), replace); } + /// The goal of the literal path: a plain string pattern builds and runs with no regex backend. + #[test] + fn a_string_pattern_needs_no_backend() { + let replace = Replace::new(" ", "▁").unwrap(); + assert_eq!( + pipeline::Normalizer::normalize(&replace, "a b c").unwrap(), + "a▁b▁▁c" + ); + // Nothing to replace: the input is handed back as it is, with nothing allocated. + assert!(matches!( + pipeline::Normalizer::normalize(&replace, "abc").unwrap(), + Cow::Borrowed("abc") + )); + // An empty pattern would match everywhere, so it matches nowhere instead. + let empty = Replace::new("", "x").unwrap(); + assert_eq!( + pipeline::Normalizer::normalize(&empty, "abc").unwrap(), + "abc" + ); + } + + /// A config spelling its pattern as a string must also *deserialize* with no backend — the + /// regex half of `serialization` above can only run once one is compiled. + #[test] + fn a_string_pattern_deserializes_with_no_backend() { + let replace_s = r#"{"type":"Replace","pattern":{"String":"Hello"},"content":"Hey"}"#; + let replace = Replace::new("Hello", "Hey").unwrap(); + assert_eq!(serde_json::from_str::(replace_s).unwrap(), replace); + assert_eq!(serde_json::to_string(&replace).unwrap(), replace_s); + } + #[test] fn test_replace_decode() { let original = vec!["hello".to_string(), "_hello".to_string()]; @@ -184,26 +261,27 @@ mod tests { ); } - #[test] - fn pipeline_replace_matches_legacy() { - let n = Replace::new("''", "\"").unwrap(); - for input in &["This is a ''test''", "no quotes", ""] { + fn assert_pipeline_matches_legacy(n: &Replace, inputs: &[&str]) { + for input in inputs { let mut ns = NormalizedString::from(*input); - Normalizer::normalize(&n, &mut ns).unwrap(); // legacy oracle + Normalizer::normalize(n, &mut ns).unwrap(); // legacy oracle assert_eq!( ns.get(), - &*pipeline::Normalizer::normalize(&n, input).unwrap() + &*pipeline::Normalizer::normalize(n, input).unwrap() ); } + } + #[test] + fn pipeline_replace_matches_legacy() { + let n = Replace::new("''", "\"").unwrap(); + assert_pipeline_matches_legacy(&n, &["This is a ''test''", "no quotes", ""]); + } + + #[test] + #[cfg(feature = "fancy-regex")] // a regex pattern needs a system-regex backend + fn pipeline_replace_matches_legacy_for_a_regex() { let n = Replace::new(ReplacePattern::Regex(r"\s+".into()), " ").unwrap(); - for input in &["a b c", "single", ""] { - let mut ns = NormalizedString::from(*input); - Normalizer::normalize(&n, &mut ns).unwrap(); // legacy oracle - assert_eq!( - ns.get(), - &*pipeline::Normalizer::normalize(&n, input).unwrap() - ); - } + assert_pipeline_matches_legacy(&n, &["a b c", "single", ""]); } } diff --git a/tokenizers/tk-encode/src/pre_tokenizers/split.rs b/tokenizers/tk-encode/src/pre_tokenizers/split.rs index 61c483559..779c167f1 100644 --- a/tokenizers/tk-encode/src/pre_tokenizers/split.rs +++ b/tokenizers/tk-encode/src/pre_tokenizers/split.rs @@ -1,5 +1,6 @@ use crate::pipeline; use crate::utils::{GptFsm, GptFsmPattern, SysRegex, gpt_fsm}; +use atomsplit::literal::Literal; use serde::{Deserialize, Deserializer, Serialize}; use crate::tokenizer::{ @@ -26,15 +27,26 @@ impl From<&str> for SplitPattern { } } +/// How a [`Split`] looks for its pattern. +#[derive(Debug)] +pub enum Search { + /// A plain string, scanned for directly — no regex engine involved, so this works in every build. + Literal(Literal), + /// A real regex, which needs the system backend (the `fancy-regex` feature). + Regex(SysRegex), + /// No way to search: no backend is compiled and the pattern is a regex. Splitting then only works + /// through the native FSM below, and errors otherwise. + Unavailable, +} + #[derive(Debug, Serialize)] #[serde(tag = "type")] pub struct Split { pub pattern: SplitPattern, - /// System-regex backend for the pattern. `None` only when no backend is compiled *and* the - /// pattern is a recognized GPT regex handled natively by `fsm` — so no backend is needed. - /// With `fancy-regex` enabled this is always `Some`; the default build has no backend (`fsm` only). + /// How the pattern is found. A plain string never needs a backend; a regex does, unless it is one + /// of the GPT patterns the native FSM below covers. #[serde(skip)] - pub regex: Option, + pub search: Search, pub behavior: SplitDelimiterBehavior, pub invert: bool, /// Native `atomsplit` FSM for a recognized GPT regex (gpt2 / cl100k-Llama-3 / o200k), used on the @@ -93,22 +105,21 @@ impl Split { SplitPattern::String(_) => None, SplitPattern::Regex(r) => gpt_fsm(r), }; - // Compile a system-regex backend for the pattern. With `fancy-regex` enabled this succeeds and - // the legacy path is unchanged; with no backend (the default) it's only fatal when atomsplit - // can't cover the pattern (i.e. an arbitrary regex, not a recognized GPT one). - let compiled = match &pattern { - SplitPattern::String(s) => SysRegex::new(®ex::escape(s)), - SplitPattern::Regex(r) => SysRegex::new(r), - }; - let regex = match compiled { - Ok(re) => Some(re), - Err(_) if fsm.is_some() => None, - Err(e) => return Err(e), + let search = match &pattern { + SplitPattern::String(s) => Search::Literal(Literal::new(s.as_bytes())?), + // A regex needs the system backend. Missing it is only fatal when the native FSM cannot + // cover this pattern either — which `pre_tokenize` reports, since a recognized GPT + // pattern in its usual form splits without any backend. + SplitPattern::Regex(r) => match SysRegex::new(r) { + Ok(regex) => Search::Regex(regex), + Err(_) if fsm.is_some() => Search::Unavailable, + Err(e) => return Err(e), + }, }; Ok(Self { pattern, - regex, + search, behavior, invert, fsm, @@ -133,15 +144,27 @@ impl Split { impl PreTokenizer for Split { fn pre_tokenize(&self, pretokenized: &mut PreTokenizedString) -> Result<()> { - if let Some(regex) = &self.regex { - return if self.invert { - pretokenized.split(|_, normalized| normalized.split(Invert(regex), self.behavior)) - } else { - pretokenized.split(|_, normalized| normalized.split(regex, self.behavior)) - }; + match &self.search { + Search::Literal(literal) => { + return if self.invert { + pretokenized + .split(|_, normalized| normalized.split(Invert(literal), self.behavior)) + } else { + pretokenized.split(|_, normalized| normalized.split(literal, self.behavior)) + }; + } + Search::Regex(regex) => { + return if self.invert { + pretokenized + .split(|_, normalized| normalized.split(Invert(regex), self.behavior)) + } else { + pretokenized.split(|_, normalized| normalized.split(regex, self.behavior)) + }; + } + Search::Unavailable => {} } - // No system-regex backend: only a recognized GPT pattern in its canonical usage (Isolated, - // not inverted — how these regexes always ship) can split, via the native atomsplit FSM. + // No way to search for the pattern: only a recognized GPT pattern in its canonical usage + // (Isolated, not inverted — how these regexes always ship) can split, via the native FSM. let fsm = self .fsm .filter(|_| !self.invert && self.behavior == SplitDelimiterBehavior::Isolated) @@ -178,15 +201,18 @@ impl pipeline::PreTokenizer for Split { ); return Ok(()); } - // Not a natively-routed GPT regex: fall back to the system-regex backend. - let regex = self.regex.as_ref().ok_or_else(|| -> crate::tokenizer::Error { - "this `Split` pattern needs a system-regex backend; enable the `fancy-regex` feature" - .into() - })?; - let matches = if self.invert { - Invert(regex).find_matches(text)? - } else { - regex.find_matches(text)? + // Not a natively-routed GPT regex: fall-back to Literal or Regex search + let matches = match (&self.search, self.invert) { + (Search::Literal(literal), false) => literal.find_matches(text)?, + (Search::Literal(literal), true) => Invert(literal).find_matches(text)?, + (Search::Regex(regex), false) => regex.find_matches(text)?, + (Search::Regex(regex), true) => Invert(regex).find_matches(text)?, + (Search::Unavailable, _) => { + return Err( + "this `Split` pattern needs a system-regex backend; enable the `fancy-regex` feature" + .into(), + ); + } }; pipeline::split_matches(out, matches, self.behavior); Ok(()) @@ -523,6 +549,41 @@ mod tests { ); } + /// The goal of the literal path: a plain string pattern splits with no regex backend, on both the + /// legacy and the pipeline path. + #[test] + fn a_string_pattern_needs_no_backend() { + let pretok = Split::new("-", SplitDelimiterBehavior::Removed, false).unwrap(); + let mut legacy = PreTokenizedString::from("a-b--c"); + pretok.pre_tokenize(&mut legacy).unwrap(); + let words: Vec<&str> = legacy + .get_splits(OffsetReferential::Original, OffsetType::Byte) + .iter() + .map(|(word, _, _)| *word) + .collect(); + assert_eq!(words, ["a", "b", "c"]); + assert_eq!( + pipeline_split( + SplitPattern::String("-".into()), + SplitDelimiterBehavior::Removed, + false, + "a-b--c" + ), + [("a", (0, 1)), ("b", (2, 3)), ("c", (5, 6))] + ); + } + + /// A config spelling its pattern as a string must also *deserialize* with no backend — the + /// regex half of `serialization` below can only run once one is compiled. + #[test] + fn a_string_pattern_deserializes_with_no_backend() { + let split_s = + r#"{"type":"Split","pattern":{"String":"Hello"},"behavior":"Removed","invert":true}"#; + let split = Split::new("Hello", SplitDelimiterBehavior::Removed, true).unwrap(); + assert_eq!(serde_json::from_str::(split_s).unwrap(), split); + assert_eq!(serde_json::to_string(&split).unwrap(), split_s); + } + #[cfg(feature = "fancy-regex")] // needs a system-regex backend #[test] fn regex_string() { @@ -548,7 +609,6 @@ mod tests { assert_eq!(pretok_str_for_regex, pretok_str_for_string); } - #[cfg(feature = "fancy-regex")] // needs a system-regex backend #[test] fn invert() { let mut pretok_str = PreTokenizedString::from("Hello Hello Hello"); diff --git a/tokenizers/tk-encode/src/tokenizer/normalizer.rs b/tokenizers/tk-encode/src/tokenizer/normalizer.rs index 5bebd5f7b..8f1899bdd 100644 --- a/tokenizers/tk-encode/src/tokenizer/normalizer.rs +++ b/tokenizers/tk-encode/src/tokenizer/normalizer.rs @@ -1022,6 +1022,7 @@ impl From<&str> for NormalizedString { #[cfg(test)] mod tests { use super::*; + use atomsplit::literal::Literal; use regex::Regex; use unicode_categories::UnicodeCategories; @@ -1479,7 +1480,7 @@ mod tests { // Overlapping let mut s = NormalizedString::from("aaaab"); - s.replace("aaa", "b").unwrap(); + s.replace(&Literal::new(b"aaa").unwrap(), "b").unwrap(); assert_eq!(s.get(), "bab"); // Regex diff --git a/tokenizers/tk-encode/src/tokenizer/pattern.rs b/tokenizers/tk-encode/src/tokenizer/pattern.rs index 5e8821f95..4a1f6054f 100644 --- a/tokenizers/tk-encode/src/tokenizer/pattern.rs +++ b/tokenizers/tk-encode/src/tokenizer/pattern.rs @@ -1,5 +1,6 @@ use crate::utils::SysRegex; use crate::{Offsets, Result}; +use atomsplit::literal::Literal; use regex::Regex; /// Pattern used to split a NormalizedString @@ -19,25 +20,6 @@ impl Pattern for char { } } -impl Pattern for &str { - fn find_matches(&self, inside: &str) -> Result> { - if self.is_empty() { - // If we try to find the matches with an empty string, just don't match anything - return Ok(vec![((0, inside.chars().count()), false)]); - } - - let re = Regex::new(®ex::escape(self))?; - (&re).find_matches(inside) - } -} - -impl Pattern for &String { - fn find_matches(&self, inside: &str) -> Result> { - let s: &str = self; - s.find_matches(inside) - } -} - impl Pattern for &Regex { fn find_matches(&self, inside: &str) -> Result> { if inside.is_empty() { @@ -60,6 +42,30 @@ impl Pattern for &Regex { } } +/// Searching for a plain string, does not need a regex engine: [`Literal`] scans the bytes. +impl Pattern for &Literal { + fn find_matches(&self, inside: &str) -> Result> { + if inside.is_empty() { + return Ok(vec![((0, 0), false)]); + } + + let mut prev = 0; + let mut splits = Vec::with_capacity(inside.len()); + for start in self.matches(inside.as_bytes()) { + let end = start + self.pattern().len(); + if prev != start { + splits.push(((prev, start), false)); + } + splits.push(((start, end), true)); + prev = end; + } + if prev != inside.len() { + splits.push(((prev, inside.len()), false)) + } + Ok(splits) + } +} + impl Pattern for &SysRegex { fn find_matches(&self, inside: &str) -> Result> { if inside.is_empty() { @@ -168,17 +174,28 @@ mod tests { } #[test] - fn str() { - do_test!("aba", "a" => vec![((0, 1), true), ((1, 2), false), ((2, 3), true)]); - do_test!("bbbba", "a" => vec![((0, 4), false), ((4, 5), true)]); - do_test!("aabbb", "a" => vec![((0, 1), true), ((1, 2), true), ((2, 5), false)]); - do_test!("aabbb", "ab" => vec![((0, 1), false), ((1, 3), true), ((3, 5), false)]); - do_test!("aabbab", "ab" => + fn literal() { + let a = Literal::new(b"a").unwrap(); + do_test!("aba", &a => vec![((0, 1), true), ((1, 2), false), ((2, 3), true)]); + do_test!("bbbba", &a => vec![((0, 4), false), ((4, 5), true)]); + do_test!("aabbb", &a => vec![((0, 1), true), ((1, 2), true), ((2, 5), false)]); + do_test!("", &a => vec![((0, 0), false)]); + + let ab = Literal::new(b"ab").unwrap(); + do_test!("aabbb", &ab => vec![((0, 1), false), ((1, 3), true), ((3, 5), false)]); + do_test!("aabbab", &ab => vec![((0, 1), false), ((1, 3), true), ((3, 4), false), ((4, 6), true)] ); - do_test!("", "" => vec![((0, 0), false)]); - do_test!("aaa", "" => vec![((0, 3), false)]); - do_test!("aaa", "b" => vec![((0, 3), false)]); + + let b = Literal::new(b"b").unwrap(); + do_test!("aaa", &b => vec![((0, 3), false)]); + + // Offsets are byte offsets, over the whole input a multi-byte pattern covers. + let metaspace = Literal::new("▁".as_bytes()).unwrap(); + do_test!("a▁b", &metaspace => vec![((0, 1), false), ((1, 4), true), ((4, 5), false)]); + + // An empty pattern would match everywhere, so there is no `Literal` for it. + assert!(Literal::new(b"").is_err()); } #[test] diff --git a/tokenizers/tk-encode/src/tokenizer/pipeline.rs b/tokenizers/tk-encode/src/tokenizer/pipeline.rs index efb12eb67..0eb1aabab 100644 --- a/tokenizers/tk-encode/src/tokenizer/pipeline.rs +++ b/tokenizers/tk-encode/src/tokenizer/pipeline.rs @@ -1004,6 +1004,34 @@ mod tests { } } + /// Test the literal only replace and splits can be run without the fancy-regex feature + #[cfg(not(feature = "fancy-regex"))] + #[test] + fn string_pattern_config_loads_and_encodes_with_no_regex_backend() { + let normalizer: NormalizerWrapper = + serde_json::from_str(r#"{"type":"Replace","pattern":{"String":" "},"content":"▁"}"#) + .unwrap(); + let pre_tokenizer: PreTokenizerWrapper = serde_json::from_str( + r#"{"type":"Split","pattern":{"String":"▁"},"behavior":"MergedWithPrevious","invert":false}"#, + ) + .unwrap(); + + let mut tok = wordlevel_tokenizer(vec![("", 0), ("hello▁", 1), ("world", 2)], None); + tok.with_normalizer(Some(normalizer)).unwrap(); + tok.with_pre_tokenizer(Some(pre_tokenizer)); + + let ids: Vec = PipelineTokenizer::try_from(&tok) + .unwrap() + .encode("hello world", false) + .unwrap() + .iter() + .map(|t| t.id) + .collect(); + // Not the unk id: both the `Replace` and the `Split` really ran on the literal path. + assert_eq!(ids, [1, 2]); + assert_pipeline_matches_reference(&tok, "hello world"); + } + #[test] fn segment_iterator_yields_text_and_specials_in_order() { let input = "aabbcc"; diff --git a/tokenizers/tk-encode/src/utils/mod.rs b/tokenizers/tk-encode/src/utils/mod.rs index 83f39109d..4003f9e09 100644 --- a/tokenizers/tk-encode/src/utils/mod.rs +++ b/tokenizers/tk-encode/src/utils/mod.rs @@ -2,9 +2,10 @@ pub(crate) mod cache; #[cfg(feature = "http")] pub(crate) mod from_pretrained; -// Optional system-regex backend for arbitrary (non-atomsplit) patterns. With `fancy-regex` off a -// stub compiles and arbitrary-regex features error at load — the atomsplit-native pre-tokenizers -// work regardless, so `fancy-regex` is only needed for custom `Split` regexes / `Replace`. +// 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`). #[cfg(feature = "fancy-regex")] mod fancy; #[cfg(feature = "fancy-regex")] diff --git a/tokenizers/tk-encode/src/utils/no_regex.rs b/tokenizers/tk-encode/src/utils/no_regex.rs index 51a8d05a5..8645b69c3 100644 --- a/tokenizers/tk-encode/src/utils/no_regex.rs +++ b/tokenizers/tk-encode/src/utils/no_regex.rs @@ -1,9 +1,10 @@ //! Stub `SysRegex` for builds with **no** system-regex backend (`fancy-regex` off — the default). //! -//! The type stays present so `Split` / `Replace` still compile, but construction always fails: the -//! atomsplit-native pre-tokenizers (GPT-2, cl100k, deepseek, the class family, char-delimiter) need -//! no backend, while a `Split` with an *arbitrary* regex or the `Replace` normalizer error at load -//! time with a clear message. Enable `fancy-regex` to get a real backend. +//! 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. +//! Enable `fancy-regex` to get a real backend. use std::error::Error; #[derive(Debug)] @@ -16,8 +17,9 @@ pub struct SysRegex { impl SysRegex { pub fn new(_regex_str: &str) -> Result> { Err( - "no system-regex backend compiled: enable the `fancy-regex` feature \ - to use a `Split` pre-tokenizer with a custom regex, or the `Replace` normalizer" + "no system-regex backend compiled: enable the `fancy-regex` feature to use a regex \ + pattern in a `Split` pre-tokenizer or a `Replace` normalizer (a plain string pattern \ + needs no backend)" .into(), ) } diff --git a/tokenizers/tk-encode/tests/pipeline_oracle.rs b/tokenizers/tk-encode/tests/pipeline_oracle.rs index 33e8e8372..622208462 100644 --- a/tokenizers/tk-encode/tests/pipeline_oracle.rs +++ b/tokenizers/tk-encode/tests/pipeline_oracle.rs @@ -107,10 +107,6 @@ fn check_model(tok_file: &str) { ); } -// The decode oracle's model set (one per decoder archetype) plus mistral-small-4, whose -// tekken pre-tokenizer is its own encode archetype; whichever files a given checkout has -// get run (bert-wiki + llama-3 ship with `make test`; the rest with `make bench-models`). -// Unsupported models skip, not fail. #[test] fn bert_wiki() { check_model("bert-wiki.json"); @@ -136,6 +132,11 @@ fn llama2() { check_model("llama-2.json"); } +#[test] +fn gemma4() { + check_model("gemma-4.json"); +} + #[test] fn t5_base() { check_model("t5-base.json"); From 1a383917178ecd9aafb9a252d25eb049b77b6ccd Mon Sep 17 00:00:00 2001 From: SBrandeis <33657802+SBrandeis@users.noreply.github.com> Date: Wed, 29 Jul 2026 17:37:22 +0200 Subject: [PATCH 3/6] box the Literal finder `memmem::Finder` carries a few hundred bytes of prefilter state on x86_64 (much less on aarch64, which is why this only showed up in CI): a `Literal` stored inline blew `NormalizerWrapper` and `DecoderWrapper` up to 352 bytes through `Replace`, which `clippy::large_enum_variant` rejects. --- tokenizers/atomsplit/src/literal.rs | 7 +++++-- tokenizers/atomsplit/tests/literal.rs | 11 +++++++++++ 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/tokenizers/atomsplit/src/literal.rs b/tokenizers/atomsplit/src/literal.rs index 041694cd4..c826b2dda 100644 --- a/tokenizers/atomsplit/src/literal.rs +++ b/tokenizers/atomsplit/src/literal.rs @@ -22,9 +22,12 @@ impl fmt::Display for EmptyPattern { 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: memmem::Finder<'static>, + finder: Box>, } impl Literal { @@ -35,7 +38,7 @@ impl Literal { return Err(EmptyPattern); } Ok(Self { - finder: memmem::Finder::new(pattern).into_owned(), + finder: Box::new(memmem::Finder::new(pattern).into_owned()), }) } diff --git a/tokenizers/atomsplit/tests/literal.rs b/tokenizers/atomsplit/tests/literal.rs index 80978ba42..14ceed9a0 100644 --- a/tokenizers/atomsplit/tests/literal.rs +++ b/tokenizers/atomsplit/tests/literal.rs @@ -29,6 +29,17 @@ fn matches_do_not_overlap() { 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); From 179799c88daae7aab7b777dad5b593910b904e74 Mon Sep 17 00:00:00 2001 From: SBrandeis <33657802+SBrandeis@users.noreply.github.com> Date: Wed, 29 Jul 2026 17:37:23 +0200 Subject: [PATCH 4/6] keep the &str / &String patterns, backed by Literal MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Python bindings search with a `&String` pattern — `NormalizedString.replace` and `.split` take a plain `str` — so dropping these impls broke every job that builds the bindings. Reinstating them through `Literal` keeps the public API and still drops the regex engine from the literal path: they used to escape the string and compile a regex on every call. An empty pattern now covers the input by byte length rather than character count, like every other impl (the old count sliced mid-character). --- tokenizers/tk-encode/src/tokenizer/pattern.rs | 34 +++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/tokenizers/tk-encode/src/tokenizer/pattern.rs b/tokenizers/tk-encode/src/tokenizer/pattern.rs index 4a1f6054f..5a147f5f4 100644 --- a/tokenizers/tk-encode/src/tokenizer/pattern.rs +++ b/tokenizers/tk-encode/src/tokenizer/pattern.rs @@ -66,6 +66,25 @@ impl Pattern for &Literal { } } +/// A plain string is one [`Literal`], built here because the pattern only lives for this call. Prefer +/// keeping a [`Literal`] around when the same pattern is searched for repeatedly. +impl Pattern for &str { + fn find_matches(&self, inside: &str) -> Result> { + match Literal::new(self.as_bytes()) { + Ok(literal) => (&literal).find_matches(inside), + // An empty pattern would match everywhere, so it matches nowhere instead. + Err(_) => Ok(vec![((0, inside.len()), false)]), + } + } +} + +impl Pattern for &String { + fn find_matches(&self, inside: &str) -> Result> { + let s: &str = self; + s.find_matches(inside) + } +} + impl Pattern for &SysRegex { fn find_matches(&self, inside: &str) -> Result> { if inside.is_empty() { @@ -198,6 +217,21 @@ mod tests { assert!(Literal::new(b"").is_err()); } + #[test] + fn str() { + do_test!("aba", "a" => vec![((0, 1), true), ((1, 2), false), ((2, 3), true)]); + do_test!("aabbab", "ab" => + vec![((0, 1), false), ((1, 3), true), ((3, 4), false), ((4, 6), true)] + ); + do_test!("aaa", "b" => vec![((0, 3), false)]); + do_test!("", "" => vec![((0, 0), false)]); + // An empty pattern would match everywhere, so it matches nowhere instead. + do_test!("aaa", "" => vec![((0, 3), false)]); + + let owned = String::from("ab"); + do_test!("aabbb", &owned => vec![((0, 1), false), ((1, 3), true), ((3, 5), false)]); + } + #[test] fn functions() { let is_b = |c| c == 'b'; From f4ac843d89c7562345a3fcf63e3b81a50403021b Mon Sep 17 00:00:00 2001 From: SBrandeis <33657802+SBrandeis@users.noreply.github.com> Date: Wed, 29 Jul 2026 18:06:21 +0200 Subject: [PATCH 5/6] feat(atomsplit): run deepseek's pre-tokenizer with no regex backend MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit deepseek's `tokenizer.json` failed to LOAD without the `fancy-regex` feature. Its pre-tokenizer is a `Sequence` of three `Split`s, and `atomsplit` only knew the three as a *chain* (`fsm_deepseek`, the fused pass). `Split::new` accepts a missing backend only for a pattern it recognizes on its own, and no single deepseek pattern was recognized — so deserialization errored out before the `Sequence` ever got the chance to spot the chain. Each split now has its own byte-exact FSM, so each `Split` is recognized alone: fsm_deepseek_num `\p{N}{1,3}` fsm_deepseek_cjk `[一-龥぀-ゟ゠-ヿ]+` fsm_deepseek_big the big regex `fsm_deepseek` keeps fusing all three when the whole chain is present; it and `fsm_deepseek_big` share one body under two compile-time flags saying which earlier splits are folded in, so the grammar lives in one place. Measured performance-neutral on `fsm_deepseek` (median +0.02% over 10 corpora, separate binaries, alternating runs — the microbench's layout noise is several times that). Also: - `DEEPSEEK_BIG` now carries literal CR/LF, matching the string deepseek ships, so `tk-encode` no longer keeps a second copy of the three patterns just to compare against them; `utils::is_deepseek` is gone (recognition is per-split). - `GptFsm::split_into` replaces the duplicated fsm `match` in the legacy and pipeline paths. - `canonicalized_for_pipeline` now rewrites `(invert, Removed)` → `(!invert, Isolated)` only for patterns that match the whole input (`GptFsm::covers_input`). deepseek's leave gaps, where the two forms differ. - The three deepseek `Sequence` tests pointed at a fixture name that no longer exists, so they had been silently skipping; fixed, and un-gated from `fancy-regex` since the splits now build without it (same for the dsv4 pipeline bench). With the feature on they still compare against the real regex engine; with it off they pin fused == chained. Tests: per-split onig parity over both parity corpora + 9 Wikipedia languages, chain-equals-fused over the same, and a `tokenizer.json` load+encode test whose ids must match the legacy path in either build. Co-Authored-By: Claude Opus 5 (1M context) --- tokenizers/atomsplit/src/fsm.rs | 7 +- tokenizers/atomsplit/src/fsm/deepseek.rs | 145 ++++++++++++++---- tokenizers/atomsplit/src/regexes.rs | 19 ++- tokenizers/atomsplit/tests/parity.rs | 96 +++++++++++- .../tk-encode/benches/pipeline_benchmark.rs | 2 +- .../tk-encode/src/pre_tokenizers/sequence.rs | 88 ++++------- .../tk-encode/src/pre_tokenizers/split.rs | 54 +++++-- tokenizers/tk-encode/src/utils/mod.rs | 2 +- .../tk-encode/src/utils/unrolled_regex.rs | 106 ++++++++----- tokenizers/tk-encode/tests/deepseek_native.rs | 91 +++++++++++ 10 files changed, 461 insertions(+), 149 deletions(-) create mode 100644 tokenizers/tk-encode/tests/deepseek_native.rs diff --git a/tokenizers/atomsplit/src/fsm.rs b/tokenizers/atomsplit/src/fsm.rs index 45c8ca286..cf76aa229 100644 --- a/tokenizers/atomsplit/src/fsm.rs +++ b/tokenizers/atomsplit/src/fsm.rs @@ -8,6 +8,11 @@ //! 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). +//! +//! Most fsms stand for ONE regex, so a `Split` carrying that regex routes straight to it. deepseek is +//! the exception: it ships three `Split`s applied in turn, so it gets four entry points — one per split +//! ([`fsm_deepseek_num`] / [`fsm_deepseek_cjk`] / [`fsm_deepseek_big`]) plus [`fsm_deepseek`], which +//! fuses all three into a single pass for a caller that recognizes the whole chain. 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::*`. @@ -240,7 +245,7 @@ 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 deepseek::{fsm_deepseek, fsm_deepseek_big, fsm_deepseek_cjk, fsm_deepseek_num}; pub use o200k::{fsm_o200k, fsm_tekken}; // ── Composition recipes ──────────────────────────────────────────────────────────────────────── diff --git a/tokenizers/atomsplit/src/fsm/deepseek.rs b/tokenizers/atomsplit/src/fsm/deepseek.rs index f54afb8fa..756bd04ab 100644 --- a/tokenizers/atomsplit/src/fsm/deepseek.rs +++ b/tokenizers/atomsplit/src/fsm/deepseek.rs @@ -23,22 +23,88 @@ 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. +/// Does the char tagged `t` match NO alternative of deepseek's big regex? Control / NumericOther / ZWJ +/// never do. `\p{N}` also doesn't — but only when Split-1 is absent (`NUM == false`), since the big +/// regex has no digit rule of its own; there a digit is either a gap char or a letter run's optional +/// `[^\r\n\p{L}\p{P}\p{S}]?` prefix. Consecutive such chars become ONE unmatched piece. +#[inline(always)] +fn ds_is_gap(t: u8) -> bool { + matches!(t & 0x0F, NMO | CTL) || t == ZWJ || (!NUM && in_mask(t, mask::NUMBER)) +} + +/// deepseek Split-1 on its own: `\p{N}{1,3}` under an `Isolated` split — numeric runs cut into +/// ≤3-*char* tokens, every maximal non-numeric run emitted as ONE gap piece. `fsm_deepseek` fuses this +/// in; the standalone form is what a lone `Split` with that pattern needs. +#[must_use] +pub fn fsm_deepseek_num(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { + debug_assert!(out.len() >= text.len() && tags.len() >= text.len()); + let end = text.len(); + let tags = &tags[..end]; + let (mut i, mut w) = (0usize, 0usize); + while i < end { + let start = i; + if in_mask(tags[i], mask::NUMBER) { + let mut cnt = 0; + while i < end && cnt < 3 && in_mask(tags[i], mask::NUMBER) { + i += char_len(text[i]); + cnt += 1; + } + } else { + i = run_end(tags, i, end, !mask::NUMBER); + } + out[w] = Span::new(start as u32, i as u32); + w += 1; + } + w +} + +/// deepseek Split-2 on its own: `[一-龥぀-ゟ゠-ヿ]+` under an `Isolated` split — maximal CJK-range runs, +/// every maximal non-CJK run as ONE gap piece. Unlike the fused [`fsm_deepseek`], the run is NOT cut +/// into letter / punct sub-runs: that cut comes from Split-3 re-splitting the isolated piece. /// -/// 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. +/// The one fsm that reads no tags: a codepoint *range* is not an atom class, so it works off the text. +/// It keeps the shared signature anyway, so callers can dispatch over any fsm without a special case. #[must_use] -pub fn fsm_deepseek(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { +pub fn fsm_deepseek_cjk(text: &[u8], _tags: &[u8], out: &mut [Span]) -> usize { + debug_assert!(out.len() >= text.len()); + let end = text.len(); + let (mut i, mut w) = (0usize, 0usize); + while i < end { + let start = i; + if ds_is_cjk_at(text, i) { + while i < end && ds_is_cjk_at(text, i) { + i += 3; // CJK-range chars are all 3-byte + } + } else { + while i < end && !ds_is_cjk_at(text, i) { + i += char_len(text[i]); + } + } + out[w] = Span::new(start as u32, i as u32); + w += 1; + } + w +} + +/// deepseek pretokenization over the atom stream, as one scalar pass. The two flags say which of the +/// `Sequence`'s earlier Splits are fused in — they are pure compile-time switches, so each instantiation +/// is the straight-line FSM for its grammar: +/// * `` = [`fsm_deepseek`], the whole `Sequence`. Precedence (= Sequence order): digits → +/// CJK-range runs → the big-regex alts, so the letter run stops at CJK codepoints and a `\p{N}{1,3}` +/// token wins over them all. +/// * `` = [`fsm_deepseek_big`], Split-3's regex alone: digits match no alt (they become +/// gap chars, or the leading `[^\r\n\p{L}\p{P}\p{S}]?` of a following letter run) and CJK is just +/// `\p{L}`, free to join an adjacent letter run. +/// +/// Peeks bytes for the ASCII `[punct][A-Za-z]+` alt. The subtleties either instantiation replicates: +/// (1) a ws run *followed by* a fused-Split match is its own Sequence piece → `\s+(?!\S)` takes the +/// WHOLE run; (2) ZWJ/ZWNJ are `\p{Cf}`, not `\p{L}∪\p{M}`, so they end a letter run; (3) a CJK run is +/// consumed as a CLOSED unit, so CJK punct (・) never steals a surrounding space nor merges with +/// non-CJK punct; (4) chars matching no alt group into ONE gap piece, and Other_Alphabetic symbols +/// (`ASM`: `\w` but categorically `\p{S}`) take the `[\p{P}\p{S}]` path, not the letter run. +#[must_use] +#[inline] // so each wrapper below gets its own flat copy, as if hand-written, rather than a shared call +fn fsm_ds(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. @@ -48,6 +114,10 @@ pub fn fsm_deepseek(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { // 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]; + // The CJK test below is spelled out at each use site as `CJK && ds_is_cjk_at(…)` rather than hoisted + // into one local closure: it sits in `letter_run`'s per-byte loop, and a closure captured by another + // closure stopped inlining there — worth ~7% on the multilingual bench. + // // 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 @@ -58,7 +128,10 @@ pub fn fsm_deepseek(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { 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)) + || (in_mask(t, mask::LETTER_MARK) + && t != ASM + && t != ZWJ + && !(CJK && ds_is_cjk_at(text, p))) { p += 1; } else { @@ -73,16 +146,16 @@ pub fn fsm_deepseek(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { && in_mask(tags[a], mask::LETTER_MARK) && tags[a] != ASM && tags[a] != ZWJ - && !ds_is_cjk_at(text, a) + && !(CJK && 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 + // if there is no punct/sym run there. STOPS at CJK-range chars — Split-2 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) + && !(CJK && ds_is_cjk_at(text, p)) { p += char_len(text[p]); } @@ -98,7 +171,8 @@ pub fn fsm_deepseek(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { // 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)); + let next_isolated = re < end + && ((NUM && in_mask(tags[re], mask::NUMBER)) || (CJK && 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 { @@ -120,7 +194,7 @@ pub fn fsm_deepseek(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { // 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) { + if CJK && 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 @@ -139,12 +213,12 @@ pub fn fsm_deepseek(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { 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 { + // A maximal run of chars that match no alt (see `ds_is_gap`) is ONE unmatched piece. Exception: + // if a letter run follows immediately, the LAST gap char is that run's alt-2 + // `[^\r\n\p{L}\p{P}\p{S}]?` prefix, so it splits off the gap and joins the letters. + if ds_is_gap::(tags[i]) { let (mut p, mut last) = (i, i); - while p < end && (matches!(tags[p] & 0x0F, NMO | CTL) || tags[p] == ZWJ) { + while p < end && ds_is_gap::(tags[p]) { last = p; p += char_len(text[p]); } @@ -180,7 +254,7 @@ pub fn fsm_deepseek(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { continue; } match tags[i] & 0x0F { - // Split-1: `\p{N}{1,3}` + // Split-1: `\p{N}{1,3}` — reachable only with `NUM`; otherwise `ds_is_gap` took these above. NW | NO => { let (mut p, mut cnt) = (i, 0); while p < end && cnt < 3 && in_mask(tags[p], mask::NUMBER) { @@ -203,8 +277,8 @@ pub fn fsm_deepseek(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { 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 if CJK && a < end && ds_is_cjk_at(text, a) { + ws(i) // next is a Split-2-isolated CJK char → the space is standalone whitespace } else { let p = punct(a); if p > a { p } else { ws(i) } @@ -245,3 +319,18 @@ pub fn fsm_deepseek(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { } w } + +/// deepseek pretokenization: the `Sequence` of `[\p{N}{1,3}]`, `[CJK]+`, `` (all `Isolated`) +/// as ONE pass — see [`fsm_ds`]. Byte-exact vs the real composed Sequence (onig ×3) on both parity +/// corpora and 9 Wikipedia languages (`tests/parity.rs`, `benches/regex.rs`). +#[must_use] +pub fn fsm_deepseek(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { + fsm_ds::(text, tags, out) +} + +/// deepseek Split-3 on its own: the big regex under an `Isolated` split — see [`fsm_ds`]. Without the +/// earlier Splits, digits are unmatched gap chars and CJK is plain `\p{L}`. +#[must_use] +pub fn fsm_deepseek_big(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { + fsm_ds::(text, tags, out) +} diff --git a/tokenizers/atomsplit/src/regexes.rs b/tokenizers/atomsplit/src/regexes.rs index 875fb892e..971acce34 100644 --- a/tokenizers/atomsplit/src/regexes.rs +++ b/tokenizers/atomsplit/src/regexes.rs @@ -23,11 +23,24 @@ pub const O200K: &str = r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*[\p{ /// [`crate::fsm::fsm_tekken`]. pub const TEKKEN: &str = r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*[\p{Ll}\p{Lm}\p{Lo}\p{M}]+|[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+[\p{Ll}\p{Lm}\p{Lo}\p{M}]*|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n/]*|\s*[\r\n]+|\s+(?!\S)|\s+"; -/// deepseek-v3 `Sequence`: `NUM` → `CJK` → `BIG`, each `Isolated`. Reproduced by -/// [`crate::fsm::fsm_deepseek`] as one pass. +/// deepseek `Sequence`: `NUM` → `CJK` → `BIG`, each `Isolated`. Fused into one pass by +/// [`crate::fsm::fsm_deepseek`]; each also stands alone ([`crate::fsm::fsm_deepseek_num`] / +/// [`crate::fsm::fsm_deepseek_cjk`] / [`crate::fsm::fsm_deepseek_big`]). +/// +/// These are byte-for-byte the strings deepseek's `tokenizer.json` ships, so a loader can recognize a +/// `Split` by string equality — hence `BIG` carries LITERAL CR/LF (spliced in with `concat!`) rather +/// than the `\r` / `\n` escapes an equivalent regex could use. pub const DEEPSEEK_NUM: &str = r"\p{N}{1,3}"; pub const DEEPSEEK_CJK: &str = r"[一-龥぀-ゟ゠-ヿ]+"; -pub const DEEPSEEK_BIG: &str = r##"[!"#$%&'()*+,\-./:;<=>?@\[\\\]^_`{|}~][A-Za-z]+|[^\r\n\p{L}\p{P}\p{S}]?[\p{L}\p{M}]+| ?[\p{P}\p{S}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"##; +pub const DEEPSEEK_BIG: &str = concat!( + r##"[!"#$%&'()*+,\-./:;<=>?@\[\\\]^_`{|}~][A-Za-z]+|[^"##, + "\r\n", + r##"\p{L}\p{P}\p{S}]?[\p{L}\p{M}]+| ?[\p{P}\p{S}]+["##, + "\r\n", + r##"]*|\s*["##, + "\r\n", + r##"]+|\s+(?!\S)|\s+"##, +); /// The deepseek chain in application order — convenience for the multi-regex reference. pub const DEEPSEEK: &[&str] = &[DEEPSEEK_NUM, DEEPSEEK_CJK, DEEPSEEK_BIG]; diff --git a/tokenizers/atomsplit/tests/parity.rs b/tokenizers/atomsplit/tests/parity.rs index 774e7c002..40b5f5366 100644 --- a/tokenizers/atomsplit/tests/parity.rs +++ b/tokenizers/atomsplit/tests/parity.rs @@ -11,7 +11,10 @@ //! 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 atomsplit::fsm::{ + Span, fsm_byte_level, fsm_cl100k, fsm_deepseek, fsm_deepseek_big, fsm_deepseek_cjk, + fsm_deepseek_num, 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::{ @@ -119,3 +122,94 @@ fn deepseek_parity() { assert_eq!(spans(fsm_deepseek, text), deepseek_ref(text), "{text:?}"); } } + +/// Wikipedia-per-language corpora (`benches/data/fetch.py`). Gitignored, so absent is a skip: the two +/// inline corpora are always checked, these widen the alphabet (RTL + format marks, Thai, Devanagari, +/// Greek, Han/Kana). +fn wikipedia_corpora() -> Vec { + let got: Vec = ["fr", "ru", "el", "he", "ar", "hi", "th", "zh", "ko"] + .iter() + .filter_map(|l| { + std::fs::read_to_string(format!( + "{}/benches/data/{l}.txt", + env!("CARGO_MANIFEST_DIR") + )) + .ok() + }) + .collect(); + eprintln!("wikipedia corpora: {}/9 (fetch.py to widen)", got.len()); + got +} + +/// One regex as a SINGLE Isolated split — matches *and* the gaps between them, which is what each of +/// deepseek's three `Split`s emits on its own (the fused [`fsm_deepseek`] never sees a gap piece, +/// since the three together cover the input). Line-by-line so a divergence names a short text. +fn check_iso(fsm: impl Fn(&[u8], &[u8], &mut [Span]) -> usize + Copy, pattern: &str) { + let re = Regex::new(pattern).unwrap(); + let corpora = wikipedia_corpora(); + for text in [CORPUS, EDGE] + .into_iter() + .chain(corpora.iter().map(String::as_str)) + { + for line in text.lines() { + let mut want = Vec::new(); + split_iso(line, 0, line.len(), &re, &mut want); + let want: Vec = want + .into_iter() + .map(|(s, e)| Span::new(s as u32, e as u32)) + .collect(); + assert_eq!(spans(fsm, line), want, "{line:?}"); + } + } +} + +#[test] +fn deepseek_num_parity() { + check_iso(fsm_deepseek_num, DS_NUM); +} + +#[test] +fn deepseek_cjk_parity() { + check_iso(fsm_deepseek_cjk, DS_CJK); +} + +#[test] +fn deepseek_big_parity() { + check_iso(fsm_deepseek_big, DS_BIG); +} + +/// The three standalone FSMs chained the way the `Sequence` chains its `Split`s (each one re-splitting +/// every piece of the previous, seeing ONLY that piece's text) must equal the fused [`fsm_deepseek`]. +/// This is what makes the fusion an optimization rather than a second implementation — and it covers +/// the multilingual bench corpora, not just the two inline ones. +#[test] +fn deepseek_split_chain_equals_fused() { + let chain = |text: &str| -> Vec { + let mut pieces = vec![Span::new(0, text.len() as u32)]; + for fsm in [ + fsm_deepseek_num as fn(&[u8], &[u8], &mut [Span]) -> usize, + fsm_deepseek_cjk, + fsm_deepseek_big, + ] { + pieces = pieces + .iter() + .flat_map(|p| { + spans(fsm, &text[p.range()]) + .into_iter() + .map(|s| Span::new(p.start + s.start, p.start + s.end)) + .collect::>() + }) + .collect(); + } + pieces + }; + let corpora = wikipedia_corpora(); + for text in [CORPUS, EDGE] + .into_iter() + .chain(corpora.iter().map(String::as_str)) + { + for line in text.lines() { + assert_eq!(spans(fsm_deepseek, line), chain(line), "{line:?}"); + } + } +} diff --git a/tokenizers/tk-encode/benches/pipeline_benchmark.rs b/tokenizers/tk-encode/benches/pipeline_benchmark.rs index 54e0feca9..f2eddbeed 100644 --- a/tokenizers/tk-encode/benches/pipeline_benchmark.rs +++ b/tokenizers/tk-encode/benches/pipeline_benchmark.rs @@ -20,7 +20,7 @@ use tk_encode::pipeline::PipelineTokenizer; // pre-tokenizer unroll (`fsm_deepseek`) + BPE end-to-end. const TOKENIZERS: &[(&str, &str)] = &[ ("bert", "../data/bert-wiki.json"), - ("dsv4", "../data/deepseek-v4-flash-base-tokenizer.json"), + ("dsv4", "../data/deepseek-v4.json"), ]; const CORPORA: &[(&str, &str)] = &[ diff --git a/tokenizers/tk-encode/src/pre_tokenizers/sequence.rs b/tokenizers/tk-encode/src/pre_tokenizers/sequence.rs index b01d4686d..b48400089 100644 --- a/tokenizers/tk-encode/src/pre_tokenizers/sequence.rs +++ b/tokenizers/tk-encode/src/pre_tokenizers/sequence.rs @@ -58,25 +58,20 @@ impl PipelineSequence { Self { pre_tokenizers } } - /// Same recognition as [`crate::utils::is_deepseek`], on the converted children: the first three are - /// Isolated, non-inverted `Split`s carrying deepseek's `[\p{N}{1,3}, CJK, big]` regexes (the trailing - /// byte-map `ByteLevel` converts to `PipelinePreTokenizer::None`). Routes the whole split to one - /// `fsm_deepseek` pass. + /// True iff the first three children are Isolated, non-inverted `Split`s carrying deepseek's + /// `[\p{N}{1,3}, CJK, big]` regexes, in that order (the trailing byte-map `ByteLevel` converts to + /// `PipelinePreTokenizer::None`). Each child already recognized its own pattern at construction; + /// spotting the chain lets us replace three passes with one fused `fsm_deepseek`. fn is_deepseek(&self) -> bool { - use crate::pre_tokenizers::split::SplitPattern; use crate::tokenizer::SplitDelimiterBehavior::Isolated; - let regex = |i: usize| match self.pre_tokenizers.get(i) { - Some(PipelinePreTokenizer::Split(s)) if s.behavior == Isolated && !s.invert => { - match &s.pattern { - SplitPattern::Regex(r) => Some(r.as_str()), - SplitPattern::String(_) => None, - } - } + use crate::utils::GptFsm::{DeepSeekBig, DeepSeekCjk, DeepSeekNum}; + let fsm = |i: usize| match self.pre_tokenizers.get(i) { + Some(PipelinePreTokenizer::Split(s)) if s.behavior == Isolated && !s.invert => s.fsm, _ => None, }; matches!( - (regex(0), regex(1), regex(2)), - (Some(a), Some(b), Some(c)) if crate::utils::is_deepseek(a, b, c) + (fsm(0), fsm(1), fsm(2)), + (Some(DeepSeekNum), Some(DeepSeekCjk), Some(DeepSeekBig)) ) } } @@ -266,15 +261,16 @@ mod tests { } } - #[cfg(feature = "fancy-regex")] // deepseek `Split`s need a backend at construction (legacy baseline) - #[test] - fn pipeline_deepseek_uses_fsm_and_matches_legacy() { - // Load deepseek-v4's real pre_tokenizer, rebuild a Sequence of just its 3 Splits (drop the - // trailing byte-map ByteLevel), and prove: (1) the exact fixture patterns are recognized, - // (2) the fsm_deepseek pipeline output == the 3-regex-split legacy output, byte-for-byte. - let path = "../data/deepseek-v4-flash-base-tokenizer.json"; + /// deepseek-v4's real pre_tokenizer as a `Sequence` of just its 3 `Split`s (dropping the trailing + /// byte-map ByteLevel), plus the pipeline form. `None` when the fixture isn't downloaded. + /// + /// Constructing these needs NO system-regex backend — all three patterns are recognized — so what + /// `legacy_pretokenize` compares against depends on the build: fancy-regex (the reference) when the + /// feature is on, the three standalone atomsplit FSMs when it's off. Both must equal the fused pass. + fn deepseek_seq() -> Option<(Sequence, PipelineSequence)> { + let path = "../data/deepseek-v4.json"; if !std::path::Path::new(path).exists() { - return; // fixture not downloaded in this environment + return None; } let v: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(path).unwrap()).unwrap(); @@ -292,7 +288,14 @@ mod tests { pipe.is_deepseek(), "deepseek's exact 3-Split sequence must be recognized" ); + Some((seq, pipe)) + } + #[test] + fn pipeline_deepseek_uses_fsm_and_matches_legacy() { + let Some((seq, pipe)) = deepseek_seq() else { + return; + }; for text in [ "中文 with 123 numbers!! and ケーキ don't", "hello 世界\n\n表 x", @@ -307,27 +310,14 @@ mod tests { } } - // CJK-range PUNCTUATION (・ U+30FB, ゠, ゛゜) sits inside Split-1's `[一-龥぀-ゟ゠-ヿ]` range, so - // Split-1 isolates it (`fsm_deepseek` handles a CJK-range run as a closed unit) — a preceding space + // CJK-range PUNCTUATION (・ U+30FB, ゠, ゛゜) sits inside Split-2's `[一-龥぀-ゟ゠-ヿ]` range, so + // Split-2 isolates it (`fsm_deepseek` handles a CJK-range run as a closed unit) — a preceding space // stays separate and it never merges with adjacent non-CJK punct. - #[cfg(feature = "fancy-regex")] // deepseek `Split`s need a backend at construction (legacy baseline) #[test] fn pipeline_deepseek_cjk_punct_whitespace_edge() { - let path = "../data/deepseek-v4-flash-base-tokenizer.json"; - if !std::path::Path::new(path).exists() { + let Some((seq, pipe)) = deepseek_seq() else { return; - } - let v: serde_json::Value = - serde_json::from_str(&std::fs::read_to_string(path).unwrap()).unwrap(); - let splits: Vec = v["pre_tokenizer"]["pretokenizers"] - .as_array() - .unwrap() - .iter() - .filter(|c| c["type"] == "Split") - .map(|c| serde_json::from_value(c.clone()).unwrap()) - .collect(); - let seq = Sequence::new(splits); - let pipe: PipelineSequence = seq.clone().try_into().unwrap(); + }; let text = "hello 世界\n\n表 ・ x"; // standalone ・ with surrounding spaces assert_eq!( pipeline_pretokenize(&pipe, text), @@ -335,29 +325,15 @@ mod tests { ); } - // fsm_deepseek == the 3-Split onig Sequence over multilingual Wikipedia corpora — the broad byte-exact + // fsm_deepseek == the 3-Split Sequence over multilingual Wikipedia corpora — the broad byte-exact // guard. `he.txt` is why it exists: Hebrew mixes format controls (RLM, `\p{Cf}`) and Other_Alphabetic // symbols (Ⓘ, `\p{S}` but is_alphabetic), which stress the *gap* grouping (consecutive unmatched chars // = one piece) and the `ALPHA_SYM` Mark refinement (a `\w` char that is NOT `[\p{L}\p{M}]`). - #[cfg(feature = "fancy-regex")] // deepseek `Split`s need a backend at construction (legacy baseline) #[test] fn pipeline_deepseek_matches_legacy_on_corpora() { - let path = "../data/deepseek-v4-flash-base-tokenizer.json"; - if !std::path::Path::new(path).exists() { + let Some((seq, pipe)) = deepseek_seq() else { return; - } - let v: serde_json::Value = - serde_json::from_str(&std::fs::read_to_string(path).unwrap()).unwrap(); - let splits: Vec = v["pre_tokenizer"]["pretokenizers"] - .as_array() - .unwrap() - .iter() - .filter(|c| c["type"] == "Split") - .map(|c| serde_json::from_value(c.clone()).unwrap()) - .collect(); - let seq = Sequence::new(splits); - let pipe: PipelineSequence = seq.clone().try_into().unwrap(); - assert!(pipe.is_deepseek()); + }; // `he`/`ar` (RTL, RLM/format-mark + Other_Alphabetic-symbol heavy) are the cases the atomsplit // deepseek bench doesn't cover; the other 8 languages are byte-exact-gated there. for lang in ["he", "ar"] { diff --git a/tokenizers/tk-encode/src/pre_tokenizers/split.rs b/tokenizers/tk-encode/src/pre_tokenizers/split.rs index 779c167f1..fa6f6c26a 100644 --- a/tokenizers/tk-encode/src/pre_tokenizers/split.rs +++ b/tokenizers/tk-encode/src/pre_tokenizers/split.rs @@ -49,11 +49,11 @@ pub struct Split { pub search: Search, pub behavior: SplitDelimiterBehavior, pub invert: bool, - /// Native `atomsplit` FSM for a recognized GPT regex (gpt2 / cl100k-Llama-3 / o200k), used on the - /// pipeline path when `behavior == Isolated && !invert` (how these regexes always ship). Byte-exact - /// with `regex`; `None` falls back to `regex`. + /// Native `atomsplit` FSM for a recognized GPT regex (gpt2 / cl100k-Llama-3 / o200k / tekken / + /// deepseek's three), used when `behavior == Isolated && !invert` (how these regexes always ship). + /// Byte-exact with `regex`; `None` falls back to `regex`. #[serde(skip)] - fsm: Option, + pub(crate) fsm: Option, } impl<'de> Deserialize<'de> for Split { @@ -132,9 +132,13 @@ impl Split { /// Isolated)`, the form the native FSM fast path requires (the inverted match /// set is the gaps, and these patterns leave no gaps). Rewrite to it so /// cl100k/o200k route to `fsm_cl100k`/`fsm_o200k` instead of the SysRegex fallback. + /// + /// Only whole-covering patterns qualify ([`GptFsm::covers_input`]): where a pattern *does* leave + /// gaps, `Removed` on the inverted match set drops those gaps while `Isolated` keeps them, so the + /// two forms differ and the rewrite is unsound. pub(crate) fn canonicalized_for_pipeline(self) -> Result { use crate::tokenizer::SplitDelimiterBehavior::{Isolated, Removed}; - if self.fsm.is_some() && self.invert && self.behavior == Removed { + if self.fsm.is_some_and(GptFsm::covers_input) && self.invert && self.behavior == Removed { Split::new(self.pattern, Isolated, false) } else { Ok(self) @@ -180,23 +184,17 @@ impl PreTokenizer for Split { 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). + // A recognized GPT regex (gpt2 / cl100k-Llama-3 / …) in its only real usage — `Isolated`, not + // inverted — routes straight to the native atomsplit FSM, which is byte-exact with `regex` (see + // the tests). The FSM emits the gap pieces too, so this is the full `Isolated` split even for + // deepseek's patterns, which (unlike the GPT ones) do not cover the whole input on their own. if let Some(fsm) = self .fsm .filter(|_| !self.invert && self.behavior == SplitDelimiterBehavior::Isolated) { pipeline::classify_into_spans( 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), - }, + |bytes, tags, spans| fsm.split_into(bytes, tags, spans), out, ); return Ok(()); @@ -584,6 +582,30 @@ mod tests { assert_eq!(serde_json::to_string(&split).unwrap(), split_s); } + #[test] + fn canonicalization_skips_patterns_that_leave_gaps() { + // `(invert, Removed)` → `(!invert, Isolated)` only holds when the pattern matches the whole + // input. deepseek's do not, so rewriting one would turn "drop the unmatched text" into "keep + // it" — the rewrite must leave them alone even though they have a native FSM. + let cl100k = Split::new( + SplitPattern::Regex(atomsplit::regexes::CL100K.into()), + Removed, + true, + ) + .unwrap(); + let canon = cl100k.canonicalized_for_pipeline().unwrap(); + assert_eq!((canon.behavior, canon.invert), (Isolated, false)); + + let ds = Split::new( + SplitPattern::Regex(atomsplit::regexes::DEEPSEEK_NUM.into()), + Removed, + true, + ) + .unwrap(); + let canon = ds.canonicalized_for_pipeline().unwrap(); + assert_eq!((canon.behavior, canon.invert), (Removed, true)); + } + #[cfg(feature = "fancy-regex")] // needs a system-regex backend #[test] fn regex_string() { diff --git a/tokenizers/tk-encode/src/utils/mod.rs b/tokenizers/tk-encode/src/utils/mod.rs index 4003f9e09..bfd89c7a8 100644 --- a/tokenizers/tk-encode/src/utils/mod.rs +++ b/tokenizers/tk-encode/src/utils/mod.rs @@ -17,7 +17,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::{GptFsm, GptFsmPattern, gpt_fsm}; pub mod byte_level; pub mod iter; diff --git a/tokenizers/tk-encode/src/utils/unrolled_regex.rs b/tokenizers/tk-encode/src/utils/unrolled_regex.rs index 78025296b..544ee082b 100644 --- a/tokenizers/tk-encode/src/utils/unrolled_regex.rs +++ b/tokenizers/tk-encode/src/utils/unrolled_regex.rs @@ -5,7 +5,8 @@ // 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}; +use atomsplit::fsm::Span; +use atomsplit::regexes::{DEEPSEEK_BIG, DEEPSEEK_CJK, DEEPSEEK_NUM, GPT2, 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; @@ -23,6 +24,42 @@ pub enum GptFsm { /// Mistral tekken regex → `atomsplit::fsm::fsm_tekken` (o200k's grammar with no contraction /// suffix and one token per digit). Tekken, + /// deepseek Split-1 `\p{N}{1,3}` → `atomsplit::fsm::fsm_deepseek_num`. + DeepSeekNum, + /// deepseek Split-2 `[一-龥぀-ゟ゠-ヿ]+` → `atomsplit::fsm::fsm_deepseek_cjk`. + DeepSeekCjk, + /// deepseek Split-3 (the big regex) → `atomsplit::fsm::fsm_deepseek_big`. + DeepSeekBig, +} + +impl GptFsm { + /// Does every byte of the input land in some match? True for the GPT regexes, whose final `\s+` + /// alternative makes them total. False for deepseek's three, which each match only part of the text + /// (they are composed, and only together do they cover it) — so those leave *gap* pieces, and any + /// rewrite that treats "matched" and "kept" as the same set is invalid for them. + pub fn covers_input(self) -> bool { + !matches!( + self, + Self::DeepSeekNum | Self::DeepSeekCjk | Self::DeepSeekBig + ) + } + + /// Run this pattern's `atomsplit` FSM: `tags` from `classify`, spans into `out` (len ≥ `text.len()`), + /// token count returned. The single dispatch point — both the legacy [`GptFsmPattern`] and the + /// pipeline `Split` route through it. + #[inline] + pub fn split_into(self, text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { + use atomsplit::fsm::*; + match self { + Self::Gpt2 => fsm_byte_level(text, tags, out), + Self::Cl100k { digit_cap } => fsm_cl100k_cap(text, tags, out, digit_cap), + Self::O200k => fsm_o200k(text, tags, out), + Self::Tekken => fsm_tekken(text, tags, out), + Self::DeepSeekNum => fsm_deepseek_num(text, tags, out), + Self::DeepSeekCjk => fsm_deepseek_cjk(text, tags, out), + Self::DeepSeekBig => fsm_deepseek_big(text, tags, out), + } + } } /// The cl100k-family template is fixed except rule 3's digit rule. If `pattern` is that template, return @@ -42,18 +79,18 @@ 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). +/// `Isolated` split byte-for-byte. GPT-2, o200k, tekken and deepseek's three 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 }) + match pattern { + GPT2 => Some(GptFsm::Gpt2), + O200K => Some(GptFsm::O200k), + TEKKEN => Some(GptFsm::Tekken), + DEEPSEEK_NUM => Some(GptFsm::DeepSeekNum), + DEEPSEEK_CJK => Some(GptFsm::DeepSeekCjk), + DEEPSEEK_BIG => Some(GptFsm::DeepSeekBig), + _ => cl100k_digit_cap(pattern).map(|digit_cap| GptFsm::Cl100k { digit_cap }), } } @@ -75,15 +112,8 @@ impl crate::tokenizer::pattern::Pattern for GptFsmPattern { 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), - }; + let mut spans = vec![Span::default(); bytes.len() + 1]; + let n = self.0.split_into(bytes, &tags, &mut spans); Ok(spans[..n] .iter() .map(|sp| ((sp.start as usize, sp.end as usize), true)) @@ -91,27 +121,6 @@ impl crate::tokenizer::pattern::Pattern for GptFsmPattern { } } -// 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 -// 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}]+"; -const DS_BIG: &str = concat!( - r##"[!"#$%&'()*+,\-./:;<=>?@\[\\\]^_`{|}~][A-Za-z]+|[^"##, - "\r\n", - r##"\p{L}\p{P}\p{S}]?[\p{L}\p{M}]+| ?[\p{P}\p{S}]+["##, - "\r\n", - r##"]*|\s*["##, - "\r\n", - r##"]+|\s+(?!\S)|\s+"##, -); - -/// 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. -pub fn is_deepseek(p0: &str, p1: &str, p2: &str) -> bool { - p0 == DS_NUM && p1 == DS_CJK && p2 == DS_BIG -} - #[cfg(test)] mod tests { use super::*; @@ -140,4 +149,17 @@ mod tests { 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); } + + #[test] + fn gpt_fsm_recognizes_deepseeks_three_splits() { + // The `Sequence` fuses these into one `fsm_deepseek` pass (see `PipelineSequence`), but each + // must be recognized on its own — that is what lets a deepseek `tokenizer.json` LOAD with no + // system-regex backend, since `Split::new` only tolerates a missing backend for a known pattern. + assert_eq!(gpt_fsm(DEEPSEEK_NUM), Some(GptFsm::DeepSeekNum)); + assert_eq!(gpt_fsm(DEEPSEEK_CJK), Some(GptFsm::DeepSeekCjk)); + assert_eq!(gpt_fsm(DEEPSEEK_BIG), Some(GptFsm::DeepSeekBig)); + // The big pattern ships with LITERAL CR/LF; the `\r`/`\n`-escaped form is an equivalent regex + // but a different string, so it is (correctly) not recognized by string equality. + assert_eq!(gpt_fsm(&DEEPSEEK_BIG.replace('\r', r"\r")), None); + } } diff --git a/tokenizers/tk-encode/tests/deepseek_native.rs b/tokenizers/tk-encode/tests/deepseek_native.rs new file mode 100644 index 000000000..48444cfd3 --- /dev/null +++ b/tokenizers/tk-encode/tests/deepseek_native.rs @@ -0,0 +1,91 @@ +//! deepseek's pre-tokenizer is a `Sequence` of three `Split`s whose regexes `atomsplit` reproduces +//! natively, so a deepseek `tokenizer.json` must **load and encode with no system-regex backend** — +//! `fancy-regex` off, which is the default build. Recognition happens per-`Split` (that is what lets +//! `Split::new` accept a missing backend), and the `Sequence` then fuses the recognized chain into one +//! `fsm_deepseek` pass. +//! +//! The pipeline (fused) and legacy (three chained splits) paths must agree on ids. That comparison is +//! meaningful in BOTH builds, and checks a different thing in each: with `fancy-regex` the legacy side +//! is oniguruma — the reference; without it, the legacy side is the three standalone FSMs, so the test +//! pins fused == chained end-to-end (`atomsplit`'s `parity.rs` gates both against oniguruma). + +mod common; + +use std::convert::TryFrom; +use std::path::Path; + +use common::{DATA, FIXTURES, WINDOWS, window}; +use tk_encode::Tokenizer; +use tk_encode::pipeline::PipelineTokenizer; + +const MODEL: &str = "deepseek-v4.json"; + +const PROBE: &str = "中文 with 123 numbers!! and ケーキ don't\n\n 純粋なCJK日本語テキスト x"; + +fn load() -> Option<(Tokenizer, PipelineTokenizer)> { + let path = Path::new(DATA).join(MODEL); + if !path.exists() { + eprintln!("skip: {MODEL} absent (fetch with `make bench-models`)"); + return None; + } + // The load itself is the assertion: with no backend compiled, an unrecognized `Split` regex is a + // hard error here, so this only succeeds because all three deepseek patterns are recognized. + let tree = + Tokenizer::from_file(&path).expect("deepseek must load with no system-regex backend"); + let pipeline = PipelineTokenizer::try_from(&tree).expect("pipeline must build"); + Some((tree, pipeline)) +} + +fn ids(tree: &Tokenizer, pipeline: &PipelineTokenizer, text: &str) -> (Vec, Vec) { + let legacy = tree + .encode_fast(text, false) + .expect("legacy encode") + .get_ids() + .to_vec(); + let fused = pipeline + .encode(text, false) + .expect("pipeline encode") + .iter() + .map(|t| t.id) + .collect(); + (legacy, fused) +} + +#[test] +fn deepseek_encodes_without_a_regex_backend() { + let Some((tree, pipeline)) = load() else { + return; + }; + let (legacy, fused) = ids(&tree, &pipeline, PROBE); + assert!(!fused.is_empty(), "encode produced no tokens"); + assert_eq!(fused, legacy); +} + +#[test] +fn deepseek_fused_matches_chained_splits_on_fixtures() { + let Some((tree, pipeline)) = load() else { + return; + }; + let mut checked = 0; + for &(group, stem) in FIXTURES { + let fixture = Path::new(DATA) + .join("fixtures") + .join(group) + .join(format!("{stem}.txt")); + let Ok(text) = std::fs::read_to_string(&fixture) else { + continue; // fixture not fetched (`make fixtures`) + }; + let mut start = 0; + for &w in WINDOWS { + let chunk = window(&text, start, w); + start += w; + if chunk.is_empty() { + continue; + } + let (legacy, fused) = ids(&tree, &pipeline, chunk); + assert_eq!(fused, legacy, "diverged on {group}/{stem} window {w}"); + checked += 1; + } + } + eprintln!("windows checked: {checked}"); +} From 780b71dfa5a3cbc5ce53299d1508d451e271a7ef Mon Sep 17 00:00:00 2001 From: SBrandeis <33657802+SBrandeis@users.noreply.github.com> Date: Wed, 29 Jul 2026 18:54:11 +0200 Subject: [PATCH 6/6] bench: measure the pipeline in the configuration it ships MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The CI benchmark built `fixture_bench` with `bench-baseline`, which enabled `fancy-regex`. So the throughput, memory and model-coverage numbers described a build nobody ships, while the binary size next to them was measured on a different binary (`binsize_pipeline`, built without it). The backend is not a rounding error: on a minimal encode program it is +1.47 MB stripped, +71%. `fancy-regex` was in that feature only to supply one of the *reference regex engines* the report times our split against. Those move to their own binary: bench-baseline = the released crate, to compare against -> fixture_bench bench-engines = onig + PCRE2 + fancy-regex + logos -> pretok_engines `fixture_bench` now builds with no regex backend, and the shared work list — corpus loading, the model manifest, the timing helpers, `--shard` — lives in `examples/bench_common/` so a number from one binary is comparable to the other. The report is unchanged: each shard runs both binaries and the `report` job folds the engine numbers back onto the row they describe as `pretok_vs_regex`, which is where the renderer already looks. Verified end to end locally — merge attaches 198 rows (9 models x 22 fixtures) and the renderer emits the same tables. All 10 manifest models load with no backend, which is what makes this possible: deepseek needed per-split FSMs (previous commit) and the SentencePiece pair (gemma-4, llama-2) needed literal patterns, whose `Split(" ")` / `Replace(" ", "▁")` are not regexes at all. `tests/no_regex_backend.rs` pins that, so adding a model that needs an engine fails a test instead of quietly turning into an error card in the report. Ran the benchmark this way over all 22 fixtures: 9/10 models benched with ids_match true against the released crate; t5-base is the pre-existing Metaspace pipeline gap, unrelated to the backend. Co-Authored-By: Claude Opus 5 (1M context) --- .github/workflows/pipeline-bench.yml | 57 ++- tokenizers/atomsplit/src/fsm/deepseek.rs | 6 +- tokenizers/tk-encode/Cargo.toml | 27 +- .../tk-encode/examples/bench_common/mod.rs | 182 ++++++++ .../tk-encode/examples/fixture_bench.rs | 421 +----------------- .../tk-encode/examples/pretok_engines.rs | 259 +++++++++++ .../tk-encode/tests/no_regex_backend.rs | 65 +++ 7 files changed, 597 insertions(+), 420 deletions(-) create mode 100644 tokenizers/tk-encode/examples/bench_common/mod.rs create mode 100644 tokenizers/tk-encode/examples/pretok_engines.rs create mode 100644 tokenizers/tk-encode/tests/no_regex_backend.rs diff --git a/.github/workflows/pipeline-bench.yml b/.github/workflows/pipeline-bench.yml index a92f62726..7698b545b 100644 --- a/.github/workflows/pipeline-bench.yml +++ b/.github/workflows/pipeline-bench.yml @@ -11,6 +11,12 @@ name: Pipeline Benchmark # release dep is behind tk-encode's `bench-baseline` feature, so production # builds never pull it. # +# What gets measured is tk-encode as it ships: `bench-baseline` brings in the released crate to +# compare against but leaves `fancy-regex` off, so the throughput numbers and the binary size +# describe the same build. The reference regex engines we time our split against are a separate +# binary (`pretok_engines`, feature `bench-engines`) precisely because one of them IS that backend — +# its per-fixture numbers are merged back into the report in the `report` job. +# # A `build` job compiles the bench + binsize example binaries ONCE (each with its # exact per-example feature set) and uploads them; the work is then fanned out across # a matrix: each `bench` shard downloads that binary and benches a slice of the models @@ -103,8 +109,9 @@ jobs: - name: Setup sccache uses: mozilla-actions/sccache-action@v0.0.9 - # `onig_sys` (a `bench-baseline` reference-regex dep) generates its bindings with bindgen, - # which needs libclang. pcre2/fancy-regex don't. The runner image doesn't ship it. + # `onig_sys` generates its bindings with bindgen, which needs libclang. It arrives two ways: + # inside the released crate (`bench-baseline`) and as a reference engine (`bench-engines`). + # pcre2/fancy-regex don't need it. The runner image doesn't ship it. - name: Install libclang (onig_sys → bindgen) run: sudo apt-get update && sudo apt-get install -y --no-install-recommends libclang-dev @@ -112,14 +119,18 @@ jobs: # them would corrupt what's measured: # • fixture_bench + binsize_baseline WITH `bench-baseline` → link the released crate # (fixture_bench = the throughput/memory bench; binsize_baseline = "cost of the - # released lib" size number). - # • binsize_pipeline WITHOUT it → the real shipping config whose size we report (none - # of the benchmark-only reference-regex deps the feature drags in). - # These are byte-for-byte the same cargo invocations the shards / report ran before. + # released lib" size number). `bench-baseline` does NOT enable `fancy-regex`, so + # fixture_bench measures tk-encode in the same configuration it ships in — the one + # binsize_pipeline reports a size for. + # • pretok_engines WITH `bench-engines` → the only binary allowed the reference regex + # engines, since one of them IS the `fancy-regex` backend. + # • binsize_pipeline with NO features → the real shipping config whose size we report. - name: Build bench + binsize binaries run: | cargo build --release -p tk-encode --features tk-encode/bench-baseline \ --example fixture_bench --example binsize_baseline + cargo build --release -p tk-encode --features tk-encode/bench-engines \ + --example pretok_engines cargo build --release -p tk-encode --example binsize_pipeline - name: Upload bench binaries @@ -128,6 +139,7 @@ jobs: name: bench-binaries path: | tokenizers/target/release/examples/fixture_bench + tokenizers/target/release/examples/pretok_engines tokenizers/target/release/examples/binsize_baseline tokenizers/target/release/examples/binsize_pipeline retention-days: 3 @@ -222,11 +234,21 @@ jobs: > "pipeline_bench_${{ matrix.shard }}.json" cat "pipeline_bench_${{ matrix.shard }}.json" + # Separate binary because it links the reference regex engines (see the build job). Same + # shard, so its rows line up with the benchmark's; the report job merges the two. + - name: Compare pre-tokenize against regex engines (shard ${{ matrix.shard }}) + run: | + chmod +x target/release/examples/pretok_engines + ./target/release/examples/pretok_engines --shard ${{ matrix.shard }} ${{ env.SHARDS }} \ + > "pretok_engines_${{ matrix.shard }}.json" + - name: Upload shard partial uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: partial-${{ matrix.shard }} - path: tokenizers/pipeline_bench_${{ matrix.shard }}.json + path: | + tokenizers/pipeline_bench_${{ matrix.shard }}.json + tokenizers/pretok_engines_${{ matrix.shard }}.json retention-days: 3 # Fan-in: concatenate the shard partials (in shard order = manifest order), @@ -272,12 +294,16 @@ jobs: merge-multiple: true path: tokenizers/partials + # Two partials per shard: the benchmark itself, and the regex-engine comparison that had to be + # built separately (see the build job). The engine numbers are folded back onto the row they + # describe as `pretok_vs_regex`, which is where the renderer looks for them. - name: Merge shard partials run: | python3 - <<'PY' import glob, json + shard_no = lambda f: int(f.rsplit('_', 1)[1].split('.')[0]) parts = sorted(glob.glob('partials/**/pipeline_bench_*.json', recursive=True), - key=lambda f: int(f.rsplit('_', 1)[1].split('.')[0])) + key=shard_no) merged = None for f in parts: d = json.load(open(f)) @@ -286,8 +312,21 @@ jobs: merged["models"].extend(d["models"]) if merged is None: raise SystemExit("no shard partials found") + + engines = {} + for f in sorted(glob.glob('partials/**/pretok_engines_*.json', recursive=True), + key=shard_no): + engines.update(json.load(open(f))) + attached = 0 + for m in merged["models"]: + per_fixture = engines.get(m["model"], {}) + for row in m["results"]: + if row["fixture"] in per_fixture: + row["pretok_vs_regex"] = per_fixture[row["fixture"]] + attached += 1 json.dump(merged, open('pipeline_bench.json', 'w')) - print(f"merged {len(parts)} shard(s) -> {len(merged['models'])} models") + print(f"merged {len(parts)} shard(s) -> {len(merged['models'])} models, " + f"{attached} rows with engine numbers") PY head -c 300 pipeline_bench.json; echo diff --git a/tokenizers/atomsplit/src/fsm/deepseek.rs b/tokenizers/atomsplit/src/fsm/deepseek.rs index 756bd04ab..d7b6394b0 100644 --- a/tokenizers/atomsplit/src/fsm/deepseek.rs +++ b/tokenizers/atomsplit/src/fsm/deepseek.rs @@ -114,9 +114,9 @@ fn fsm_ds(text: &[u8], tags: &[u8], out: &mut // 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]; - // The CJK test below is spelled out at each use site as `CJK && ds_is_cjk_at(…)` rather than hoisted - // into one local closure: it sits in `letter_run`'s per-byte loop, and a closure captured by another - // closure stopped inlining there — worth ~7% on the multilingual bench. + // The CJK test below is spelled out at each use site as `CJK && ds_is_cjk_at(…)` rather than + // hoisted into one local closure: it sits in `letter_run`'s per-byte loop, where a closure + // captured by another closure did not inline. // // 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}`). diff --git a/tokenizers/tk-encode/Cargo.toml b/tokenizers/tk-encode/Cargo.toml index 5f830e3ee..3caa70d82 100644 --- a/tokenizers/tk-encode/Cargo.toml +++ b/tokenizers/tk-encode/Cargo.toml @@ -64,10 +64,10 @@ libc = "0.2" # Latest released tokenizers, used as the comparison baseline by the CI benchmark # (examples gated on `bench-baseline`). Optional so production builds never pull it. tokenizers-release = { package = "tokenizers", version = "=0.23.1", optional = true } -# Reference regex engines `fixture_bench` times the classify+fsm pre-tokenize against: Oniguruma and -# PCRE2 (both C) alongside the pure-Rust `fancy-regex` (optional, above). All three are pulled in by -# `bench-baseline` — optional + behind that feature so the C builds happen ONLY for the CI benchmark, -# never for `cargo test`/production. +# Reference regex engines the `pretok_engines` example times our classify+fsm split against: +# Oniguruma and PCRE2 (both C) alongside the pure-Rust `fancy-regex` (optional, above). Behind +# `bench-engines` so the C libraries are built ONLY for that comparison, never for +# `cargo test`/production, and never for the throughput benchmark. onig = { version = "6.5.1", optional = true } pcre2 = { version = "0.2", optional = true } logos = { version = "0.15", optional = true } # compile-time DFA lexer reference (pure Rust) @@ -83,13 +83,14 @@ progressbar = ["indicatif"] http = ["hf-hub"] unstable_wasm = ["fancy-regex", "getrandom/wasm_js"] rustls-tls = ["hf-hub?/rustls-tls"] -bench-baseline = [ - "dep:tokenizers-release", - "dep:onig", - "dep:pcre2", - "dep:logos", - "fancy-regex", -] +# The released crate as the comparison baseline + correctness oracle. Deliberately does NOT enable +# `fancy-regex`: the throughput/memory benchmark then measures the pipeline in its real shipping +# configuration, and the same build is what the binary-size numbers describe. +bench-baseline = ["dep:tokenizers-release"] +# Reference regex engines for the `pretok_engines` comparison only (onig + PCRE2 are C libraries, +# `fancy-regex` is the pure-Rust one). Separate from `bench-baseline` so timing our split against +# them cannot drag a regex backend into the benchmarked pipeline. +bench-engines = ["dep:onig", "dep:pcre2", "dep:logos", "fancy-regex"] [dev-dependencies] criterion = "0.6" @@ -106,6 +107,10 @@ harness = false name = "fixture_bench" required-features = ["bench-baseline"] +[[example]] +name = "pretok_engines" +required-features = ["bench-engines"] + [[example]] name = "binsize_baseline" required-features = ["bench-baseline"] diff --git a/tokenizers/tk-encode/examples/bench_common/mod.rs b/tokenizers/tk-encode/examples/bench_common/mod.rs new file mode 100644 index 000000000..84d70b50b --- /dev/null +++ b/tokenizers/tk-encode/examples/bench_common/mod.rs @@ -0,0 +1,182 @@ +//! Inputs and timing shared by the two benchmark binaries. +//! +//! Both walk the same work list — every model in `examples/bench_models.json` against every corpus +//! in `data/fixtures/` — but measure different things, so they are separate binaries built with +//! different features: `fixture_bench` compares the pipeline against the released crate, while +//! `pretok_engines` compares our split against real regex engines. Keeping the corpus loading, the +//! model list and the timing here means a number from one is directly comparable to a number from +//! the other. +//! +//! Both binaries are built by CI and take `--shard `, which selects the i-th of `n` slices of +//! the model list so the work fans out over parallel runners. + +// Each binary uses a subset of this module; the rest is not dead code, just unused over there. +#![allow(dead_code)] + +use std::hint::black_box; +use std::path::{Path, PathBuf}; +use std::time::Instant; + +use serde_json::Value; + +pub const DATA_DIR: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../data"); +pub const MANIFEST: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/examples/bench_models.json"); +/// Input size per encode call. ~10 kB is large enough that per-call overhead is amortized, so the +/// number reflects steady-state throughput (`benches/pipeline_benchmark.rs` sweeps the sizes). +pub const CHUNK_BYTES: usize = 10 * 1024; +pub const MAX_CHUNKS: usize = 100; +/// Timed passes per measurement; the median is reported. +pub const REPS: usize = 5; + +/// One corpus, already cut into encode-sized chunks. +pub struct Fixture { + pub group: &'static str, + pub name: String, + pub chunks: Vec, + pub bytes: usize, +} + +/// Cuts `text` on line boundaries into chunks of at least [`CHUNK_BYTES`], so a chunk never splits a +/// line and every model sees the same inputs. +pub fn make_chunks(text: &str) -> Vec { + let mut chunks = Vec::new(); + let mut cur = String::new(); + for line in text.lines().filter(|l| !l.trim().is_empty()) { + if !cur.is_empty() { + cur.push('\n'); + } + cur.push_str(line); + if cur.len() >= CHUNK_BYTES { + chunks.push(std::mem::take(&mut cur)); + if chunks.len() == MAX_CHUNKS { + return chunks; + } + } + } + if !cur.is_empty() { + chunks.push(cur); + } + chunks +} + +/// The sorted `.txt` fixtures under `data/fixtures/{lang,modalities}`, tagged with group. +pub fn fixture_paths() -> Vec<(&'static str, PathBuf)> { + let mut out = Vec::new(); + for group in ["lang", "modalities"] { + let dir = Path::new(DATA_DIR).join("fixtures").join(group); + let mut paths: Vec<_> = std::fs::read_dir(&dir) + .unwrap_or_else(|e| panic!("{}: {e} — run `make fixtures` first", dir.display())) + .map(|e| e.unwrap().path()) + .filter(|p| p.extension().is_some_and(|x| x == "txt")) + .collect(); + paths.sort(); + out.extend(paths.into_iter().map(|p| (group, p))); + } + out +} + +/// Every corpus in `data/fixtures/{lang,modalities}`, read and chunked once. +pub fn load_fixtures() -> Vec { + fixture_paths() + .into_iter() + .map(|(group, path)| { + let name = path.file_stem().unwrap().to_str().unwrap().to_string(); + let chunks = make_chunks(&std::fs::read_to_string(&path).unwrap()); + let bytes = chunks.iter().map(String::len).sum(); + Fixture { + group, + name, + chunks, + bytes, + } + }) + .collect() +} + +/// The models this process is responsible for, read from [`MANIFEST`]. +/// +/// `args` is the raw command line. Without `--shard i n` the slice is the whole manifest. +pub fn shard(args: &[String]) -> Vec { + let (i, n): (usize, usize) = match (args.get(1).map(String::as_str), args.get(2), args.get(3)) { + (Some("--shard"), Some(i), Some(n)) => { + (i.parse().unwrap(), n.parse::().unwrap().max(1)) + } + _ => (0, 1), + }; + let full: Vec = + serde_json::from_str(&std::fs::read_to_string(MANIFEST).unwrap()).unwrap(); + let (lo, hi) = (i * full.len() / n, (i + 1) * full.len() / n); + eprintln!("shard {i}/{n}: models {lo}..{hi} of {}", full.len()); + full[lo.min(full.len())..hi.min(full.len())].to_vec() +} + +// ── timing ────────────────────────────────────────────────────────────────── + +pub fn median_secs(mut samples: Vec) -> f64 { + samples.sort_by(|a, b| a.partial_cmp(b).unwrap()); + samples[samples.len() / 2] +} + +/// Median ns/byte of a warmed-up `run` over `len` bytes. `run` returns a value that is `black_box`'d +/// so the work cannot be optimized away; the `black_box` belongs here, in the benchmark, never in the +/// library. +pub fn timed_ns(len: usize, mut run: impl FnMut() -> usize) -> f64 { + if len == 0 { + return 0.0; + } + run(); // warm-up + let mut s = Vec::with_capacity(REPS); + for _ in 0..REPS { + let t = Instant::now(); + black_box(run()); + s.push(t.elapsed().as_secs_f64()); + } + median_secs(s) * 1e9 / len as f64 +} + +// ── the model list ────────────────────────────────────────────────────────── + +/// Path to a manifest entry's `tokenizer.json`, defaulting to `.json`. +pub fn model_path(entry: &Value) -> PathBuf { + let name = entry["name"].as_str().unwrap(); + let file = entry + .get("file") + .and_then(Value::as_str) + .map(str::to_string) + .unwrap_or_else(|| format!("{name}.json")); + Path::new(DATA_DIR).join(file) +} + +fn split_regex(p: &Value) -> Option { + (p["type"].as_str() == Some("Split")) + .then(|| p["pattern"]["Regex"].as_str().map(str::to_string)) + .flatten() +} + +/// The ordered Split regexes a model's pre-tokenizer applies (deepseek → 3; a lone `Split` → 1; a +/// byte-map `ByteLevel` with no Split → GPT-2's implicit regex, the canonical spec in atomsplit). +/// Empty → the model has no regex to compare a regex engine against (Bert, Metaspace, +/// WhitespaceSplit, a literal-string `Split`, …). +pub fn pretok_regexes(path: &Path) -> Vec { + let v: Value = std::fs::read_to_string(path) + .ok() + .and_then(|s| serde_json::from_str(&s).ok()) + .unwrap_or(Value::Null); + let pt = &v["pre_tokenizer"]; + match pt["type"].as_str() { + Some("Split") => split_regex(pt).into_iter().collect(), + Some("ByteLevel") => vec![atomsplit::regexes::GPT2.to_string()], + Some("Sequence") => { + let arr = pt["pretokenizers"].as_array().cloned().unwrap_or_default(); + let res: Vec = arr.iter().filter_map(split_regex).collect(); + if !res.is_empty() { + res + } else if arr.iter().any(|p| p["type"] == "ByteLevel") { + vec![atomsplit::regexes::GPT2.to_string()] + } else { + vec![] + } + } + _ => vec![], + } +} diff --git a/tokenizers/tk-encode/examples/fixture_bench.rs b/tokenizers/tk-encode/examples/fixture_bench.rs index 9ac6d89bf..e8dab00fb 100644 --- a/tokenizers/tk-encode/examples/fixture_bench.rs +++ b/tokenizers/tk-encode/examples/fixture_bench.rs @@ -51,11 +51,12 @@ use std::convert::TryFrom; use std::hint::black_box; -use std::path::{Path, PathBuf}; +use std::path::Path; use std::process::Command; use std::time::Instant; -use logos::Logos; +mod bench_common; + use rayon::ThreadPoolBuilder; use rayon::iter::{IntoParallelRefIterator, ParallelIterator}; use serde_json::{Value, json}; @@ -63,13 +64,13 @@ use tk_encode::pipeline::{Model, PipelineTokenizer}; use tk_encode::{AddedToken, ModelWrapper, Tokenizer}; use tokenizers_release::{AddedToken as BaselineAddedToken, Tokenizer as BaselineTokenizer}; -const DATA_DIR: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../data"); -const MANIFEST: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/examples/bench_models.json"); +use bench_common::{ + CHUNK_BYTES, Fixture, REPS, fixture_paths, load_fixtures, make_chunks, median_secs, model_path, + shard, +}; + // Keep in sync with the `tokenizers-release` pin in Cargo.toml. const BASELINE_VERSION: &str = "0.23.1"; -const CHUNK_BYTES: usize = 10 * 1024; -const MAX_CHUNKS: usize = 100; -const REPS: usize = 5; const PROBE: &str = "The quick brown fox jumps 123."; // Chunks per fixture for the memory children's encode pass: enough to warm the // lazy structures and grow the output buffers, small enough to stay cheap. @@ -120,75 +121,8 @@ fn inject_added_tokens_baseline(tok: &mut BaselineTokenizer) { // ── fixtures ──────────────────────────────────────────────────────────────── -struct Fixture { - group: &'static str, - name: String, - chunks: Vec, - bytes: usize, -} - -fn make_chunks(text: &str) -> Vec { - let mut chunks = Vec::new(); - let mut cur = String::new(); - for line in text.lines().filter(|l| !l.trim().is_empty()) { - if !cur.is_empty() { - cur.push('\n'); - } - cur.push_str(line); - if cur.len() >= CHUNK_BYTES { - chunks.push(std::mem::take(&mut cur)); - if chunks.len() == MAX_CHUNKS { - return chunks; - } - } - } - if !cur.is_empty() { - chunks.push(cur); - } - chunks -} - -/// The sorted `.txt` fixtures under `data/fixtures/{lang,modalities}`, tagged with group. -fn fixture_paths() -> Vec<(&'static str, PathBuf)> { - let mut out = Vec::new(); - for group in ["lang", "modalities"] { - let dir = Path::new(DATA_DIR).join("fixtures").join(group); - let mut paths: Vec<_> = std::fs::read_dir(&dir) - .unwrap_or_else(|e| panic!("{}: {e} — run `make fixtures` first", dir.display())) - .map(|e| e.unwrap().path()) - .filter(|p| p.extension().is_some_and(|x| x == "txt")) - .collect(); - paths.sort(); - out.extend(paths.into_iter().map(|p| (group, p))); - } - out -} - -/// Every corpus in `data/fixtures/{lang,modalities}`, read and chunked once. -fn load_fixtures() -> Vec { - fixture_paths() - .into_iter() - .map(|(group, path)| { - let name = path.file_stem().unwrap().to_str().unwrap().to_string(); - let chunks = make_chunks(&std::fs::read_to_string(&path).unwrap()); - let bytes = chunks.iter().map(String::len).sum(); - Fixture { - group, - name, - chunks, - bytes, - } - }) - .collect() -} - // ── timing helpers ────────────────────────────────────────────────────────── -fn median_secs(mut samples: Vec) -> f64 { - samples.sort_by(|a, b| a.partial_cmp(b).unwrap()); - samples[samples.len() / 2] -} - fn time_pass(encode: &dyn Fn(&str) -> usize, chunks: &[String]) -> f64 { let start = Instant::now(); let mut n = 0usize; @@ -199,22 +133,6 @@ fn time_pass(encode: &dyn Fn(&str) -> usize, chunks: &[String]) -> f64 { start.elapsed().as_secs_f64() } -/// Median ns/byte of a warmed-up `run` over `len` bytes (`run` returns a value that's -/// `black_box`'d so the work isn't optimized away). -fn timed_ns(len: usize, mut run: impl FnMut() -> usize) -> f64 { - if len == 0 { - return 0.0; - } - run(); // warm-up - let mut s = Vec::with_capacity(REPS); - for _ in 0..REPS { - let t = Instant::now(); - black_box(run()); - s.push(t.elapsed().as_secs_f64()); - } - median_secs(s) * 1e9 / len as f64 -} - // ── phase 1: throughput ───────────────────────────────────────────────────── /// Warm single-thread throughput + the id gates for one fixture. @@ -333,9 +251,9 @@ fn stage_secs(pipeline: &PipelineTokenizer, chunks: &[String]) median_secs(samples) } -/// Stage decomposition + regex-engine references for one fixture: the -/// `stage_ns_per_byte` and `pretok_vs_regex` objects of its report row. -fn bench_stages(pipeline: &PipelineTokenizer, f: &Fixture, regexes: &[String]) -> (Value, Value) { +/// The `stage_ns_per_byte` object of one fixture's report row. How our split compares to a regex +/// engine is measured by the `pretok_engines` binary instead, which is built with those engines. +fn bench_stages(pipeline: &PipelineTokenizer, f: &Fixture) -> Value { let t_frame = stage_secs::<{ PipelineTokenizer::STAGE_FRAME }>(pipeline, &f.chunks); let t_norm = stage_secs::<{ PipelineTokenizer::STAGE_NORMALIZE }>(pipeline, &f.chunks); let t_split = stage_secs::<{ PipelineTokenizer::STAGE_SPLIT }>(pipeline, &f.chunks); @@ -358,243 +276,16 @@ fn bench_stages(pipeline: &PipelineTokenizer, f: &Fixture, regexes: &[String]) - f.name ); - // pre_tokenize (= classify SIMD + fsm) vs classify-scalar and vs real regex engines - // over the same corpus, so the report shows the split beating a regex engine both - // WITH and WITHOUT SIMD. `scalar_pipe` = pre_tokenize + (cls_scalar − cls_simd): - // fsm is the scalar jump-table in both pipes, SIMD/scalar is the classify pass only. - let corpus: String = f.chunks.concat(); - let cls_simd = classify_ns(corpus.as_bytes(), false); - let cls_scalar = classify_ns(corpus.as_bytes(), true); - let onig_ns = regex_reference_ns::(&corpus, regexes); - let fancy_ns = regex_reference_ns::(&corpus, regexes); - let pcre2_ns = regex_reference_ns::(&corpus, regexes); - let logos_ns = logos_reference_ns(regexes, &corpus); - if [onig_ns, fancy_ns, pcre2_ns, logos_ns] - .iter() - .any(Option::is_some) - { - let scalar_pipe = ns_split + (cls_scalar - cls_simd).max(0.0); - let vs = |r: Option| { - r.map_or("—".into(), |v| { - format!( - "{:.1}×/{:.1}×", - v / ns_split.max(1e-9), - v / scalar_pipe.max(1e-9) - ) - }) - }; - eprintln!( - " {} pre-tok: SIMD-cls {ns_split:.2} / scalar-cls {scalar_pipe:.2} ns/B · vs onig {} · vs fancy {} · vs pcre2 {} · vs logos {}", - f.name, - vs(onig_ns), - vs(fancy_ns), - vs(pcre2_ns), - vs(logos_ns) - ); - } - - ( - json!({ - "added_split": ns_added, - "normalize": ns_norm, - "pre_tokenize": ns_split, - "model": ns_model, - "post": ns_post, - "total": nspb(t_post), - }), - json!({ - "cls_simd": cls_simd, - "cls_scalar": cls_scalar, - "onig": onig_ns, - "fancy": fancy_ns, - "pcre2": pcre2_ns, - "logos": logos_ns, - }), - ) -} - -/// Median ns/byte to classify `bytes` once via the SIMD or scalar path. -fn classify_ns(bytes: &[u8], scalar: bool) -> f64 { - let mut tags = vec![0u8; bytes.len()]; - timed_ns(bytes.len(), || { - if scalar { - atomsplit::classify::classify_scalar(bytes, &mut tags); - } else { - atomsplit::classify::classify(bytes, &mut tags); - } - tags[bytes.len() / 2] as usize + json!({ + "added_split": ns_added, + "normalize": ns_norm, + "pre_tokenize": ns_split, + "model": ns_model, + "post": ns_post, + "total": nspb(t_post), }) } -// ── regex-engine references ───────────────────────────────────────────────── -// The pipeline's `pre_tokenize` stage is `classify (SIMD) + fsm`; these reference -// numbers time the model's own pre-tokenizer regex(es) — the split a regex-based -// tokenizer actually pays for — under three real engines. Each engine only needs to -// enumerate matches; the composed Isolated split chain is shared. - -/// A regex engine timed through the composed split chain. -trait SplitEngine: Sized { - fn compile(pattern: &str) -> Option; - /// Call `on_match(start, end)` for every match in `hay`, in order. - fn for_each_match(&self, hay: &str, on_match: impl FnMut(usize, usize)); -} - -/// Oniguruma (C) — what the reference tokenizer itself uses. -impl SplitEngine for onig::Regex { - fn compile(pattern: &str) -> Option { - onig::Regex::new(pattern).ok() - } - fn for_each_match(&self, hay: &str, mut on_match: impl FnMut(usize, usize)) { - for (s, e) in self.find_iter(hay) { - on_match(s, e); - } - } -} - -/// fancy-regex (pure Rust). `find_iter` yields `Result`; a match error -/// aborts that piece's pass (rare, backtrack-limit) and it is left un-split. -impl SplitEngine for fancy_regex::Regex { - fn compile(pattern: &str) -> Option { - fancy_regex::Regex::new(pattern).ok() - } - fn for_each_match(&self, hay: &str, mut on_match: impl FnMut(usize, usize)) { - for m in self.find_iter(hay) { - let Ok(m) = m else { break }; - on_match(m.start(), m.end()); - } - } -} - -/// PCRE2 (C) — built with `utf(true).ucp(true)` so `\p{L}`/`\p{N}`/`\s` are -/// Unicode-aware and byte offsets land on char boundaries, matching the other -/// engines, and **JIT-compiled** so PCRE2 is benched at its best. -impl SplitEngine for pcre2::bytes::Regex { - fn compile(pattern: &str) -> Option { - pcre2::bytes::RegexBuilder::new() - .utf(true) - .ucp(true) - .jit_if_available(true) - .build(pattern) - .ok() - } - fn for_each_match(&self, hay: &str, mut on_match: impl FnMut(usize, usize)) { - for m in self.find_iter(hay.as_bytes()) { - let Ok(m) = m else { break }; - on_match(m.start(), m.end()); - } - } -} - -/// ns/byte for the composed Isolated split chain under engine `E` — each regex splits -/// the previous pieces (gaps + matches), exactly how the reference tokenizer applies a -/// `Sequence` of Splits. `None` when the model has no regex pre-tokenizer, or the -/// engine rejects a pattern. -fn regex_reference_ns(text: &str, patterns: &[String]) -> Option { - if patterns.is_empty() || text.is_empty() { - return None; - } - let engines: Vec = patterns - .iter() - .map(|p| E::compile(p)) - .collect::>()?; - Some(timed_ns(text.len(), || { - let mut pieces = vec![(0usize, text.len())]; - for re in &engines { - let mut next = Vec::with_capacity(pieces.len() * 2); - for (s, e) in pieces.drain(..) { - let mut prev = 0usize; - re.for_each_match(&text[s..e], |ms, me| { - if ms > prev { - next.push((s + prev, s + ms)); - } - next.push((s + ms, s + me)); - prev = me; - }); - if prev < e - s { - next.push((s + prev, e)); - } - } - pieces = next; - } - pieces.len() - })) -} - -// logos DFA lexers approximating the GPT splits (no look-ahead / case-insensitive → -// boundaries differ slightly; a raw-throughput reference like fancy, not a byte-exact -// oracle). Only families logos can express get a number; deepseek / variants / -// non-regex pretoks report null. -#[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 { - #[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, -} - -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) when the model's pre-tokenizer is a single regex logos can -/// express (matched against the canonical gpt2/cl100k/o200k specs); `None` otherwise. -fn logos_reference_ns(regexes: &[String], text: &str) -> Option { - if text.is_empty() || regexes.len() != 1 { - return None; - } - let r = regexes[0].as_str(); - let f: fn(&str) -> usize = if r == atomsplit::regexes::GPT2 { - |s| lex_count::(s) - } else if r == atomsplit::regexes::CL100K { - |s| lex_count::(s) - } else if r == atomsplit::regexes::O200K { - |s| lex_count::(s) - } else { - return None; - }; - Some(timed_ns(text.len(), || f(text))) -} - // ── phase 3: multi-thread scaling + memory ────────────────────────────────── /// Thread counts for the scaling sweep: 1 (the single-thread anchor) + 2/4/8 + the device max, @@ -973,16 +664,6 @@ fn measure_memory(model: &Path, baseline_ok: bool) -> Value { /// Local path to a manifest entry's config: `data/`, else `data/.json`. /// Both come from the test-data dataset (see the Makefile `bench-models` target). -fn model_path(entry: &Value) -> PathBuf { - let name = entry["name"].as_str().unwrap(); - let file = entry - .get("file") - .and_then(Value::as_str) - .map(str::to_string) - .unwrap_or_else(|| format!("{name}.json")); - Path::new(DATA_DIR).join(file) -} - fn model_kind(tok: &Tokenizer) -> &'static str { match tok.get_model() { ModelWrapper::BPE(_) => "BPE", @@ -1015,39 +696,6 @@ fn pretok_label(path: &Path) -> String { } } -fn split_regex(p: &Value) -> Option { - (p["type"].as_str() == Some("Split")) - .then(|| p["pattern"]["Regex"].as_str().map(str::to_string)) - .flatten() -} - -/// The ordered Split regexes a model's pre-tokenizer applies (deepseek → 3; a lone `Split` → 1; a -/// byte-map `ByteLevel` with no Split → GPT-2's implicit regex, the canonical spec in atomsplit). -/// Empty → no regex reference (Bert, Metaspace, WhitespaceSplit, …) → engines report null. -fn pretok_regexes(path: &Path) -> Vec { - let v: Value = std::fs::read_to_string(path) - .ok() - .and_then(|s| serde_json::from_str(&s).ok()) - .unwrap_or(Value::Null); - let pt = &v["pre_tokenizer"]; - match pt["type"].as_str() { - Some("Split") => split_regex(pt).into_iter().collect(), - Some("ByteLevel") => vec![atomsplit::regexes::GPT2.to_string()], - Some("Sequence") => { - let arr = pt["pretokenizers"].as_array().cloned().unwrap_or_default(); - let res: Vec = arr.iter().filter_map(split_regex).collect(); - if !res.is_empty() { - res - } else if arr.iter().any(|p| p["type"] == "ByteLevel") { - vec![atomsplit::regexes::GPT2.to_string()] - } else { - vec![] - } - } - _ => vec![], - } -} - fn main() { let args: Vec = std::env::args().collect(); if args.get(1).map(String::as_str) == Some("--memory") { @@ -1055,34 +703,14 @@ fn main() { return; } - // Optional model sharding for CI matrix fan-out: `--shard ` benches only the i-th of `n` - // contiguous manifest chunks, so the models split across parallel isolated runners and the partial - // JSONs are concatenated downstream. Without it, `(0, 1)` = the whole manifest, unchanged. - let (shard, nshards): (usize, usize) = - match (args.get(1).map(String::as_str), args.get(2), args.get(3)) { - (Some("--shard"), Some(i), Some(n)) => { - (i.parse().unwrap(), n.parse::().unwrap().max(1)) - } - _ => (0, 1), - }; - let full: Vec = - serde_json::from_str(&std::fs::read_to_string(MANIFEST).unwrap()).unwrap(); - let (lo, hi) = ( - shard * full.len() / nshards, - (shard + 1) * full.len() / nshards, - ); - let manifest = &full[lo.min(full.len())..hi.min(full.len())]; - eprintln!( - "shard {shard}/{nshards}: models {lo}..{hi} of {}", - full.len() - ); + let manifest = shard(&args); let fixtures = load_fixtures(); // The whole corpus, flattened once: the multi-thread sweep runs over all fixtures so // thread-spawn/scheduling overhead is amortized and the scaling curve is stable. let all_chunks: Vec = fixtures.iter().flat_map(|f| f.chunks.clone()).collect(); let mut models: Vec = Vec::new(); - for entry in manifest { + for entry in &manifest { let name = entry["name"].as_str().unwrap().to_string(); let repo = entry.get("repo").and_then(Value::as_str).unwrap_or(""); let desc = entry.get("desc").and_then(Value::as_str).unwrap_or(""); @@ -1156,12 +784,11 @@ fn main() { .iter() .map(|f| bench_throughput(baseline.as_ref(), &tok, f)) .collect(); - let regexes = pretok_regexes(&path); for (row, f) in rows.iter_mut().zip(&fixtures) { - let (stages, pretok) = bench_stages(&pipeline, f, ®exes); - let row = row.as_object_mut().unwrap(); - row.insert("stage_ns_per_byte".into(), stages); - row.insert("pretok_vs_regex".into(), pretok); + let stages = bench_stages(&pipeline, f); + row.as_object_mut() + .unwrap() + .insert("stage_ns_per_byte".into(), stages); } // Decode: probe once whether the pipeline can decode yet (a loud stub // today → the pipeline decode series is `null`, rendered "pending"). The diff --git a/tokenizers/tk-encode/examples/pretok_engines.rs b/tokenizers/tk-encode/examples/pretok_engines.rs new file mode 100644 index 000000000..73bb72f7e --- /dev/null +++ b/tokenizers/tk-encode/examples/pretok_engines.rs @@ -0,0 +1,259 @@ +//! How fast is `classify + fsm` against a real regex engine? +//! +//! Most tokenizers pre-tokenize by running a regex over the text. This crate instead classifies every +//! character into an "atom" class in one SIMD pass and then cuts with a small state machine +//! (`atomsplit`), which produces the same pieces without a regex engine. This binary puts the two +//! side by side: for every model in the manifest it takes that model's own pre-tokenizer regex and +//! times it under three engines — oniguruma and PCRE2 (both C, PCRE2 JIT-compiled) and fancy-regex +//! (pure Rust) — plus a logos DFA lexer where the grammar can be expressed, over the same corpora +//! `fixture_bench` uses. +//! +//! It also reports the classify pass alone, SIMD and scalar, so the comparison can be read both with +//! and without SIMD: the state machine is the same scalar jump table either way. +//! +//! This lives apart from `fixture_bench` because the engines are extra dependencies, two of them C +//! libraries, and one of them is the `fancy-regex` backend itself. Pulling them into the throughput +//! benchmark would mean benchmarking a build nobody ships. Built with `--features bench-engines`. +//! +//! Emits `{model: {fixture: {cls_simd, cls_scalar, onig, fancy, pcre2, logos}}}` on stdout, which CI +//! merges into the benchmark report as each row's `pretok_vs_regex`. + +mod bench_common; + +use bench_common::{load_fixtures, model_path, pretok_regexes, shard, timed_ns}; +use logos::Logos; +use serde_json::{Map, Value, json}; + +fn main() { + let args: Vec = std::env::args().collect(); + let manifest = shard(&args); + let fixtures = load_fixtures(); + + let mut models = Map::new(); + for entry in &manifest { + let name = entry["name"].as_str().unwrap().to_string(); + let regexes = pretok_regexes(&model_path(entry)); + eprintln!("== {name} ({} regex(es)) ==", regexes.len()); + let mut rows = Map::new(); + for f in &fixtures { + let corpus: String = f.chunks.concat(); + let cls_simd = classify_ns(corpus.as_bytes(), false); + let cls_scalar = classify_ns(corpus.as_bytes(), true); + let onig = regex_reference_ns::(&corpus, ®exes); + let fancy = regex_reference_ns::(&corpus, ®exes); + let pcre2 = regex_reference_ns::(&corpus, ®exes); + let logos = logos_reference_ns(®exes, &corpus); + eprintln!( + " {} cls {cls_simd:.2}/{cls_scalar:.2} · onig {} · fancy {} · pcre2 {} · logos {}", + f.name, + ns_or_dash(onig), + ns_or_dash(fancy), + ns_or_dash(pcre2), + ns_or_dash(logos), + ); + rows.insert( + f.name.clone(), + json!({ + "cls_simd": cls_simd, + "cls_scalar": cls_scalar, + "onig": onig, + "fancy": fancy, + "pcre2": pcre2, + "logos": logos, + }), + ); + } + models.insert(name, Value::Object(rows)); + } + println!( + "{}", + serde_json::to_string_pretty(&Value::Object(models)).unwrap() + ); +} + +/// ns/byte for the log line, or `—` for a model this engine has no number for. +fn ns_or_dash(v: Option) -> String { + v.map_or("—".into(), |v| format!("{v:.2}")) +} + +/// Median ns/byte to classify `bytes` once via the SIMD or scalar path. +fn classify_ns(bytes: &[u8], scalar: bool) -> f64 { + let mut tags = vec![0u8; bytes.len()]; + timed_ns(bytes.len(), || { + if scalar { + atomsplit::classify::classify_scalar(bytes, &mut tags); + } else { + atomsplit::classify::classify(bytes, &mut tags); + } + tags[bytes.len() / 2] as usize + }) +} + +// ── regex-engine references ───────────────────────────────────────────────── +// The pipeline's `pre_tokenize` stage is `classify (SIMD) + fsm`; these reference +// numbers time the model's own pre-tokenizer regex(es) — the split a regex-based +// tokenizer actually pays for — under three real engines. Each engine only needs to +// enumerate matches; the composed Isolated split chain is shared. + +/// A regex engine timed through the composed split chain. +trait SplitEngine: Sized { + fn compile(pattern: &str) -> Option; + /// Call `on_match(start, end)` for every match in `hay`, in order. + fn for_each_match(&self, hay: &str, on_match: impl FnMut(usize, usize)); +} + +/// Oniguruma (C) — what the reference tokenizer itself uses. +impl SplitEngine for onig::Regex { + fn compile(pattern: &str) -> Option { + onig::Regex::new(pattern).ok() + } + fn for_each_match(&self, hay: &str, mut on_match: impl FnMut(usize, usize)) { + for (s, e) in self.find_iter(hay) { + on_match(s, e); + } + } +} + +/// fancy-regex (pure Rust). `find_iter` yields `Result`; a match error +/// aborts that piece's pass (rare, backtrack-limit) and it is left un-split. +impl SplitEngine for fancy_regex::Regex { + fn compile(pattern: &str) -> Option { + fancy_regex::Regex::new(pattern).ok() + } + fn for_each_match(&self, hay: &str, mut on_match: impl FnMut(usize, usize)) { + for m in self.find_iter(hay) { + let Ok(m) = m else { break }; + on_match(m.start(), m.end()); + } + } +} + +/// PCRE2 (C) — built with `utf(true).ucp(true)` so `\p{L}`/`\p{N}`/`\s` are +/// Unicode-aware and byte offsets land on char boundaries, matching the other +/// engines, and **JIT-compiled** so PCRE2 is benched at its best. +impl SplitEngine for pcre2::bytes::Regex { + fn compile(pattern: &str) -> Option { + pcre2::bytes::RegexBuilder::new() + .utf(true) + .ucp(true) + .jit_if_available(true) + .build(pattern) + .ok() + } + fn for_each_match(&self, hay: &str, mut on_match: impl FnMut(usize, usize)) { + for m in self.find_iter(hay.as_bytes()) { + let Ok(m) = m else { break }; + on_match(m.start(), m.end()); + } + } +} + +/// ns/byte for the composed Isolated split chain under engine `E` — each regex splits +/// the previous pieces (gaps + matches), exactly how the reference tokenizer applies a +/// `Sequence` of Splits. `None` when the model has no regex pre-tokenizer, or the +/// engine rejects a pattern. +fn regex_reference_ns(text: &str, patterns: &[String]) -> Option { + if patterns.is_empty() || text.is_empty() { + return None; + } + let engines: Vec = patterns + .iter() + .map(|p| E::compile(p)) + .collect::>()?; + Some(timed_ns(text.len(), || { + let mut pieces = vec![(0usize, text.len())]; + for re in &engines { + let mut next = Vec::with_capacity(pieces.len() * 2); + for (s, e) in pieces.drain(..) { + let mut prev = 0usize; + re.for_each_match(&text[s..e], |ms, me| { + if ms > prev { + next.push((s + prev, s + ms)); + } + next.push((s + ms, s + me)); + prev = me; + }); + if prev < e - s { + next.push((s + prev, e)); + } + } + pieces = next; + } + pieces.len() + })) +} + +// logos DFA lexers approximating the GPT splits (no look-ahead / case-insensitive → +// boundaries differ slightly; a raw-throughput reference like fancy, not a byte-exact +// oracle). Only families logos can express get a number; deepseek / variants / +// non-regex pretoks report null. +#[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 { + #[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, +} + +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) when the model's pre-tokenizer is a single regex logos can +/// express (matched against the canonical gpt2/cl100k/o200k specs); `None` otherwise. +fn logos_reference_ns(regexes: &[String], text: &str) -> Option { + if text.is_empty() || regexes.len() != 1 { + return None; + } + let r = regexes[0].as_str(); + let f: fn(&str) -> usize = if r == atomsplit::regexes::GPT2 { + |s| lex_count::(s) + } else if r == atomsplit::regexes::CL100K { + |s| lex_count::(s) + } else if r == atomsplit::regexes::O200K { + |s| lex_count::(s) + } else { + return None; + }; + Some(timed_ns(text.len(), || f(text))) +} diff --git a/tokenizers/tk-encode/tests/no_regex_backend.rs b/tokenizers/tk-encode/tests/no_regex_backend.rs new file mode 100644 index 000000000..0efbeccd6 --- /dev/null +++ b/tokenizers/tk-encode/tests/no_regex_backend.rs @@ -0,0 +1,65 @@ +//! Every model the CI benchmark reports on must load with no regex backend. +//! +//! `fancy-regex` is optional. Without it there is no engine for a user-supplied regex, so a config +//! needing one fails to load — and the benchmark renders that model as an error card instead of +//! benching it. The benchmark is built without the feature on purpose (see `bench-baseline` in +//! Cargo.toml), so that the throughput and memory numbers describe the configuration tk-encode +//! actually ships, the same one the binary-size numbers describe. +//! +//! This test walks the benchmark's own model list and fails if any entry stops loading, which is what +//! adding a model with, say, a `Replace` normalizer holding a real regex would do. It is the guard on +//! that build choice. Run in the default build to mean anything; with `fancy-regex` on it passes +//! trivially, since then every pattern has an engine. + +use std::convert::TryFrom; +use std::path::Path; + +use serde_json::Value; +use tk_encode::Tokenizer; +use tk_encode::pipeline::PipelineTokenizer; + +const MANIFEST: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/examples/bench_models.json"); +const DATA: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../data"); +const PROBE: &str = "The quick brown fox jumps 123 中文 don't!"; + +#[test] +fn every_benchmarked_model_loads_without_a_regex_backend() { + let manifest: Vec = + serde_json::from_str(&std::fs::read_to_string(MANIFEST).unwrap()).unwrap(); + assert!(!manifest.is_empty(), "empty model manifest"); + + let (mut checked, mut failed) = (0, Vec::new()); + for entry in &manifest { + let name = entry["name"].as_str().unwrap(); + let file = entry["file"] + .as_str() + .map_or_else(|| format!("{name}.json"), str::to_string); + let path = Path::new(DATA).join(&file); + if !path.exists() { + eprintln!("skip {name}: {file} not fetched (`make bench-models`)"); + continue; + } + checked += 1; + match Tokenizer::from_file(&path) { + // A model the pipeline cannot build or encode yet is a separate, tracked gap (the + // benchmark reports those as roadmap cards); only the *load* is pinned here. + Ok(tok) => { + if let Ok(pipeline) = PipelineTokenizer::try_from(&tok) + && let Ok(tokens) = pipeline.encode(PROBE, false) + { + assert!(!tokens.is_empty(), "{name}: encoded to nothing"); + } + } + Err(e) => failed.push(format!("{name}: {e}")), + } + } + assert!( + checked > 0, + "no model tokenizers present — run `make bench-models`" + ); + assert!( + failed.is_empty(), + "{checked} models checked, these need a regex backend to load:\n {}", + failed.join("\n ") + ); +}