From a03f6e5729d0cf7a5738ca0a19a5cc4e22e60ca7 Mon Sep 17 00:00:00 2001 From: itayh Date: Tue, 16 Jun 2026 16:34:24 +0300 Subject: [PATCH 1/4] normalizers: Added support for prepend and replace normalizers (encountered on RedHatAI/gemma-4-31B-it-FP8-Dynamic) --- src/json_structs.rs | 14 +++- src/normalizers.rs | 16 +++++ src/normalizers/prepend.rs | 61 ++++++++++++++++++ src/normalizers/replace.rs | 102 ++++++++++++++++++++++++++++++ tests/gemma_replace_normalizer.rs | 93 +++++++++++++++++++++++++++ 5 files changed, 285 insertions(+), 1 deletion(-) create mode 100644 src/normalizers/prepend.rs create mode 100644 src/normalizers/replace.rs create mode 100644 tests/gemma_replace_normalizer.rs diff --git a/src/json_structs.rs b/src/json_structs.rs index a9aaa5e..6fc5504 100644 --- a/src/json_structs.rs +++ b/src/json_structs.rs @@ -86,7 +86,7 @@ pub enum NormalizerConfig { prepend: String, }, Replace { - pattern: Value, + pattern: ReplacePattern, #[serde(default)] content: String, }, @@ -104,6 +104,18 @@ pub enum NormalizerConfig { Other(Value), } +/// Pattern for a `Replace` normalizer: either a literal string or a regex. +/// +/// Serialized in `tokenizer.json` as an externally-tagged object, e.g. +/// `{"String": " "}` or `{"Regex": "\\s+"}`. +#[derive(Clone, Debug, Deserialize)] +pub enum ReplacePattern { + /// Literal substring, matched verbatim. + String(String), + /// Regular expression source. + Regex(String), +} + /// Pre-tokenizer pipeline step from `tokenizer.json`. #[derive(Clone, Debug, Deserialize, EnumDiscriminants)] #[strum_discriminants(name(PreTokenizerKind), derive(strum::Display, Hash, Ord, PartialOrd))] diff --git a/src/normalizers.rs b/src/normalizers.rs index 5a1c8b6..9bace63 100644 --- a/src/normalizers.rs +++ b/src/normalizers.rs @@ -1,8 +1,12 @@ mod nfc; +mod prepend; +mod replace; use std::borrow::Cow; pub use self::nfc::Nfc; +pub use self::prepend::Prepend; +pub use self::replace::Replace; use crate::json_structs::{NormalizerConfig, NormalizerKind}; /// Errors from constructing a normalizer. @@ -10,12 +14,18 @@ use crate::json_structs::{NormalizerConfig, NormalizerKind}; pub enum Error { #[error("unsupported normalizer type: {0}")] Unsupported(String), + + /// A `Replace` normalizer's regex pattern failed to compile. + #[error(transparent)] + Replace(#[from] replace::Error), } /// A compiled normalizer ready for use. #[derive(Debug)] pub enum Normalizer { Nfc(Nfc), + Prepend(Prepend), + Replace(Replace), Sequence(Vec), } @@ -24,6 +34,10 @@ impl Normalizer { pub fn from_config(config: NormalizerConfig) -> Result { match config { NormalizerConfig::Nfc => Ok(Self::Nfc(Nfc)), + NormalizerConfig::Prepend { prepend } => Ok(Self::Prepend(Prepend::new(prepend))), + NormalizerConfig::Replace { pattern, content } => { + Ok(Self::Replace(Replace::new(pattern, content)?)) + } NormalizerConfig::Sequence { normalizers } => { let steps = normalizers .into_iter() @@ -46,6 +60,8 @@ impl Normalizer { pub fn normalize<'a>(&self, input: &'a str) -> Cow<'a, str> { match self { Self::Nfc(nfc) => nfc.normalize(input), + Self::Prepend(prepend) => prepend.normalize(input), + Self::Replace(replace) => replace.normalize(input), Self::Sequence(steps) => { let mut current = Cow::Borrowed(input); for step in steps { diff --git a/src/normalizers/prepend.rs b/src/normalizers/prepend.rs new file mode 100644 index 0000000..84042ad --- /dev/null +++ b/src/normalizers/prepend.rs @@ -0,0 +1,61 @@ +use std::borrow::Cow; + +/// `Prepend` normalizer: prepends a fixed string to the (non-empty) input. +/// +/// Mirrors HuggingFace `tokenizers.normalizers.Prepend`. SentencePiece-style +/// tokenizers (e.g. gemma) use `Prepend("▁")` so the first word receives the +/// same metaspace prefix as interior words. HF prepends only when the input is +/// non-empty; an empty input is left untouched. +#[derive(Debug)] +pub struct Prepend { + prepend: String, +} + +impl Prepend { + /// Build a `Prepend` normalizer from its prefix string. + pub fn new(prepend: String) -> Self { + Self { prepend } + } + + /// Apply the prepend to `input`. + /// + /// Returns `Cow::Borrowed` when the prefix is empty or the input is empty, + /// avoiding allocation. + pub fn normalize<'a>(&self, input: &'a str) -> Cow<'a, str> { + if self.prepend.is_empty() || input.is_empty() { + Cow::Borrowed(input) + } else { + let mut out = String::with_capacity(self.prepend.len() + input.len()); + out.push_str(&self.prepend); + out.push_str(input); + Cow::Owned(out) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn prepends_metaspace() { + let p = Prepend::new("\u{2581}".to_string()); + assert_eq!(p.normalize("hello"), "\u{2581}hello"); + } + + #[test] + fn empty_input_unchanged() { + let p = Prepend::new("\u{2581}".to_string()); + let out = p.normalize(""); + assert_eq!(out, ""); + assert!(matches!(out, Cow::Borrowed(_))); + } + + #[test] + fn empty_prefix_unchanged() { + let p = Prepend::new(String::new()); + let out = p.normalize("hello"); + assert_eq!(out, "hello"); + assert!(matches!(out, Cow::Borrowed(_))); + } +} diff --git a/src/normalizers/replace.rs b/src/normalizers/replace.rs new file mode 100644 index 0000000..cd399ff --- /dev/null +++ b/src/normalizers/replace.rs @@ -0,0 +1,102 @@ +use std::borrow::Cow; + +use fancy_regex::Regex; + +use crate::json_structs::ReplacePattern; + +/// Errors from constructing a [`Replace`] normalizer. +#[derive(Debug, thiserror::Error)] +pub enum Error { + /// The regex pattern failed to compile. + #[error("replace normalizer regex error: {0}")] + Regex(#[from] fancy_regex::Error), +} + +/// `Replace` normalizer: replaces every match of `pattern` with `content`. +/// +/// Mirrors HuggingFace `tokenizers.normalizers.Replace`. The pattern is either +/// a literal string (matched verbatim) or a regular expression. For +/// SentencePiece-style tokenizers (e.g. gemma) this is what maps literal spaces +/// to the metaspace marker `▁` (U+2581), e.g. `Replace(" ", "▁")`. +#[derive(Debug)] +pub enum Replace { + /// Literal-substring replacement. Fast path: no regex engine needed. + Literal { pattern: String, content: String }, + /// Regex replacement. + Regex { regex: Regex, content: String }, +} + +impl Replace { + /// Build a `Replace` normalizer from its parsed JSON pattern + content. + pub fn new(pattern: ReplacePattern, content: String) -> Result { + match pattern { + ReplacePattern::String(pattern) => Ok(Self::Literal { pattern, content }), + ReplacePattern::Regex(source) => Ok(Self::Regex { + regex: Regex::new(&source)?, + content, + }), + } + } + + /// Apply the replacement to `input`. + /// + /// Returns `Cow::Borrowed` when nothing matched, avoiding allocation on the + /// common no-op path. + pub fn normalize<'a>(&self, input: &'a str) -> Cow<'a, str> { + match self { + Self::Literal { pattern, content } => { + // An empty literal pattern would match infinitely in HF's model + // it is a no-op; mirror that by returning the input untouched. + if pattern.is_empty() || !input.contains(pattern.as_str()) { + Cow::Borrowed(input) + } else { + Cow::Owned(input.replace(pattern.as_str(), content)) + } + } + Self::Regex { regex, content } => regex.replace_all(input, content.as_str()), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn lit(pattern: &str, content: &str) -> Replace { + Replace::new(ReplacePattern::String(pattern.to_string()), content.to_string()).unwrap() + } + + #[test] + fn literal_space_to_metaspace() { + let r = lit(" ", "\u{2581}"); + assert_eq!(r.normalize("hello world"), "hello\u{2581}world"); + assert_eq!(r.normalize(" world"), "\u{2581}world"); + assert_eq!(r.normalize("a b c"), "a\u{2581}b\u{2581}c"); + } + + #[test] + fn literal_no_match_borrows() { + let r = lit(" ", "\u{2581}"); + let out = r.normalize("noscript"); + assert_eq!(out, "noscript"); + assert!(matches!(out, Cow::Borrowed(_))); + } + + #[test] + fn empty_pattern_is_noop() { + let r = lit("", "\u{2581}"); + let out = r.normalize("hello"); + assert_eq!(out, "hello"); + assert!(matches!(out, Cow::Borrowed(_))); + } + + #[test] + fn regex_collapses_whitespace() { + let r = Replace::new( + ReplacePattern::Regex(r"\s+".to_string()), + "\u{2581}".to_string(), + ) + .unwrap(); + assert_eq!(r.normalize("a b\t c"), "a\u{2581}b\u{2581}c"); + } +} diff --git a/tests/gemma_replace_normalizer.rs b/tests/gemma_replace_normalizer.rs new file mode 100644 index 0000000..a570dc9 --- /dev/null +++ b/tests/gemma_replace_normalizer.rs @@ -0,0 +1,93 @@ +//! End-to-end validation that fastokens' normalizer converter reproduces the +//! gemma (SentencePiece) normalizer chain exactly, against the reference +//! HuggingFace `tokenizers` crate. +//! +//! gemma's `tokenizer.json` normalizer is a `Sequence` of `Prepend("▁")` then +//! `Replace(" ", "▁")` — this is what maps whitespace to the metaspace marker +//! `▁` (U+2581) that the vocab is keyed on. If `Replace` is dropped or +//! mishandled, leading/interior spaces produce wrong token IDs. +//! +//! The reference oracle is the real `tokenizers` crate (a dev-dependency): we +//! deserialize the *same JSON* into both fastokens' `Normalizer` and the +//! `tokenizers` `NormalizerWrapper`, then assert byte-for-byte equal output on +//! a corpus of edge cases (leading space, interior spaces, runs of spaces, +//! unicode, empty string). + +use fastokens::normalizers::Normalizer; +use tokenizers::tokenizer::{NormalizedString, Normalizer as _}; +use tokenizers::NormalizerWrapper; + +/// The gemma normalizer chain, verbatim from a gemma `tokenizer.json`. +const GEMMA_NORMALIZER_JSON: &str = r#"{ + "type": "Sequence", + "normalizers": [ + { "type": "Prepend", "prepend": "▁" }, + { "type": "Replace", "pattern": { "String": " " }, "content": "▁" } + ] +}"#; + +/// Reference normalization via the HuggingFace `tokenizers` crate. +fn reference(wrapper: &NormalizerWrapper, input: &str) -> String { + let mut ns = NormalizedString::from(input); + wrapper.normalize(&mut ns).expect("reference normalize"); + ns.get().to_string() +} + +#[test] +fn gemma_chain_matches_reference() { + let ft: Normalizer = Normalizer::from_config( + serde_json::from_str(GEMMA_NORMALIZER_JSON).expect("parse fastokens config"), + ) + .expect("build fastokens normalizer (Replace must be supported)"); + + let hf: NormalizerWrapper = + serde_json::from_str(GEMMA_NORMALIZER_JSON).expect("parse reference normalizer"); + + let cases = [ + "hello world", + " world", // leading space -> ▁world + " double space", // runs of spaces + "a b c", // interior spaces + "trailing ", + "no_space", + "", // empty + "café résumé", // unicode + space + "tabs\tand\nnewlines", // non-space whitespace must be left alone + "mix ▁ already", // pre-existing metaspace marker + ]; + + for case in cases { + let got = ft.normalize(case); + let want = reference(&hf, case); + assert_eq!( + got.as_ref(), + want, + "gemma normalizer mismatch for input {case:?}: fastokens={got:?} reference={want:?}" + ); + } +} + +/// A bare `Replace` with a String pattern must match the reference, including +/// the leading-space case that broke gemma. +#[test] +fn replace_string_pattern_matches_reference() { + let json = r#"{ "type": "Replace", "pattern": { "String": " " }, "content": "▁" }"#; + let ft = Normalizer::from_config(serde_json::from_str(json).unwrap()).unwrap(); + let hf: NormalizerWrapper = serde_json::from_str(json).unwrap(); + + for case in [" world", "a b", " ", "none", ""] { + assert_eq!(ft.normalize(case).as_ref(), reference(&hf, case), "input {case:?}"); + } +} + +/// The `Regex` pattern variant must also be supported and match the reference. +#[test] +fn replace_regex_pattern_matches_reference() { + let json = r#"{ "type": "Replace", "pattern": { "Regex": " {2,}" }, "content": "_" }"#; + let ft = Normalizer::from_config(serde_json::from_str(json).unwrap()).unwrap(); + let hf: NormalizerWrapper = serde_json::from_str(json).unwrap(); + + for case in ["a b", "a b", "a b c", "none"] { + assert_eq!(ft.normalize(case).as_ref(), reference(&hf, case), "input {case:?}"); + } +} From 76bd7f87f5f5b50891722ef7ae9c9b6218e78892 Mon Sep 17 00:00:00 2001 From: itayh Date: Tue, 16 Jun 2026 16:34:24 +0300 Subject: [PATCH 2/4] ci: disable free-threaded (py3.13t) wheel builds in linux matrix maturin/pyo3 here only supports free-threaded on Python 3.14+, so the -i python3.13t step failed and aborted the x86_64 job before the wheel artifact uploaded. We don't need free-threaded wheels (target is regular CPython 3.12; the cp39-abi3 wheel covers it). --- .github/workflows/build-maturin.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/build-maturin.yml b/.github/workflows/build-maturin.yml index db63613..cece401 100644 --- a/.github/workflows/build-maturin.yml +++ b/.github/workflows/build-maturin.yml @@ -32,11 +32,11 @@ jobs: - runner: ubuntu-22.04 target: x86_64 manylinux: auto - ft: true + ft: false - runner: ubuntu-22.04 target: x86 manylinux: auto - ft: true + ft: false - runner: ubuntu-22.04 target: aarch64 manylinux: manylinux_2_28 @@ -48,7 +48,7 @@ jobs: - runner: ubuntu-22.04 target: ppc64le manylinux: manylinux_2_28 - ft: true + ft: false steps: - uses: actions/checkout@v6 - uses: actions/setup-python@v6 From b111a04d98379ec41f50f612b580e078ace99ac3 Mon Sep 17 00:00:00 2001 From: itayh Date: Tue, 16 Jun 2026 17:28:23 +0300 Subject: [PATCH 3/4] decoders: support Replace/ByteFallback/Fuse/Strip for SentencePiece decode chain MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit gemma's tokenizer.json decoder is a Sequence of Replace("▁"->" "), ByteFallback, Fuse, Strip(" ", 1, 0). DecoderConfig already parsed these, but the runtime Decoder only implemented ByteLevel/Sequence and stubbed Fuse as a no-op, so tokenizer construction failed with 'unsupported decoder type: Replace' (encountered on RedHatAI/gemma-4-31B-it-FP8-Dynamic under vLLM). - Implement Replace (reusing normalizers::Replace, the inverse direction), ByteFallback (run-buffered UTF-8 reassembly; one U+FFFD per byte token on an invalid run), real Fuse (concat to one token), and Strip. - Fix Fuse: was identity; gemma's Strip runs after Fuse and expects a single fused token, so a no-op Fuse would strip a leading space from every piece. - Decoder Replace.pattern: Value -> ReplacePattern to reuse the normalizer's pattern engine; re-export normalizers::ReplaceError for the build path. - Add tests/gemma_replace_decoder.rs: byte-for-byte parity with the HF tokenizers crate on the verbatim gemma decoder chain (metaspace, multi-byte fallback, emoji, invalid UTF-8, strip-after-fuse, empty input). Co-Authored-By: Claude Opus 4.8 (1M context) --- src/decoders.rs | 125 ++++++++++++++++++++++++++++++++- src/json_structs.rs | 2 +- src/normalizers.rs | 2 +- tests/gemma_replace_decoder.rs | 80 +++++++++++++++++++++ 4 files changed, 206 insertions(+), 3 deletions(-) create mode 100644 tests/gemma_replace_decoder.rs diff --git a/src/decoders.rs b/src/decoders.rs index 5bf7e46..5ced38e 100644 --- a/src/decoders.rs +++ b/src/decoders.rs @@ -1,6 +1,7 @@ mod byte_level; use crate::json_structs::{DecoderConfig, DecoderKind}; +use crate::normalizers::{Replace, ReplaceError}; pub use self::byte_level::ByteLevelDecoder; @@ -9,12 +10,34 @@ pub use self::byte_level::ByteLevelDecoder; pub enum Error { #[error("unsupported decoder type: {0}")] Unsupported(String), + + /// A `Replace` decoder's regex pattern failed to compile. + #[error(transparent)] + Replace(#[from] ReplaceError), } /// A compiled decoder ready for use. #[derive(Debug)] pub enum Decoder { ByteLevel(ByteLevelDecoder), + /// Per-token substring/regex replacement. For SentencePiece-style + /// tokenizers (gemma, llama) this maps the metaspace marker `▁` back to a + /// space — the inverse of the `Replace` normalizer (`Replace("▁", " ")`). + Replace(Replace), + /// Reassembles `<0xHH>` byte tokens into UTF-8 text. Runs of byte tokens + /// are accumulated and decoded together; an invalid UTF-8 run becomes one + /// replacement char (`U+FFFD`) per byte token, matching HuggingFace. + ByteFallback, + /// Concatenates every token into a single token. + Fuse, + /// Strips up to `start` leading and `stop` trailing occurrences of + /// `content` from each token (gemma uses `Strip(" ", 1, 0)` to drop the + /// leading metaspace space the SentencePiece pre-tokenizer added). + Strip { + content: char, + start: usize, + stop: usize, + }, Sequence(Vec), } @@ -23,7 +46,20 @@ impl Decoder { pub fn from_config(config: DecoderConfig) -> Result { match config { DecoderConfig::ByteLevel => Ok(Self::ByteLevel(ByteLevelDecoder)), - DecoderConfig::Fuse => Ok(Self::Sequence(vec![])), // identity/no-op + DecoderConfig::Replace { pattern, content } => { + Ok(Self::Replace(Replace::new(pattern, content)?)) + } + DecoderConfig::ByteFallback => Ok(Self::ByteFallback), + DecoderConfig::Fuse => Ok(Self::Fuse), + DecoderConfig::Strip { + content, + start, + stop, + } => Ok(Self::Strip { + content, + start, + stop, + }), DecoderConfig::Sequence { decoders } => { let steps = decoders .into_iter() @@ -50,6 +86,20 @@ impl Decoder { pub fn decode_chain(&self, tokens: Vec) -> Result, Error> { match self { Self::ByteLevel(bl) => Ok(bl.decode_chain(tokens)), + Self::Replace(r) => Ok(tokens + .into_iter() + .map(|t| r.normalize(&t).into_owned()) + .collect()), + Self::ByteFallback => Ok(decode_byte_fallback(tokens)), + Self::Fuse => Ok(vec![tokens.concat()]), + Self::Strip { + content, + start, + stop, + } => Ok(tokens + .into_iter() + .map(|t| strip_token(&t, *content, *start, *stop)) + .collect()), Self::Sequence(steps) => { let mut current = tokens; for step in steps { @@ -66,3 +116,76 @@ impl Decoder { Ok(result.join("")) } } + +/// `ByteFallback` decode: reassemble `<0xHH>` byte tokens into UTF-8 text. +/// +/// Mirrors HuggingFace's `ByteFallback` decoder exactly: a maximal run of +/// `<0xHH>` tokens is buffered and, when the run ends, decoded as UTF-8. A run +/// that is not valid UTF-8 is replaced by one `U+FFFD` per byte token in the +/// run (not one per invalid byte). +fn decode_byte_fallback(tokens: Vec) -> Vec { + let mut out: Vec = Vec::with_capacity(tokens.len()); + let mut pending: Vec = Vec::new(); + + let flush = |pending: &mut Vec, out: &mut Vec| { + if pending.is_empty() { + return; + } + match String::from_utf8(pending.clone()) { + Ok(s) => out.push(s), + Err(_) => { + for _ in 0..pending.len() { + out.push("\u{fffd}".to_string()); + } + } + } + pending.clear(); + }; + + for token in tokens { + let byte = if token.len() == 6 && token.starts_with("<0x") && token.ends_with('>') { + u8::from_str_radix(&token[3..5], 16).ok() + } else { + None + }; + match byte { + Some(b) => pending.push(b), + None => { + flush(&mut pending, &mut out); + out.push(token); + } + } + } + flush(&mut pending, &mut out); + out +} + +/// `Strip` decode for a single token: remove up to `start` leading and up to +/// `stop` trailing occurrences of `content`. Mirrors HuggingFace's `Strip` +/// decoder (counts are caps, and stop scanning at the first non-matching char). +fn strip_token(token: &str, content: char, start: usize, stop: usize) -> String { + let chars: Vec = token.chars().collect(); + + let mut start_cut = 0; + for (i, &c) in chars.iter().enumerate().take(start) { + if c == content { + start_cut = i + 1; + } else { + break; + } + } + + let mut stop_cut = chars.len(); + for (i, &c) in chars.iter().rev().enumerate().take(stop) { + if c == content { + stop_cut = chars.len() - (i + 1); + } else { + break; + } + } + + if start_cut >= stop_cut { + return String::new(); + } + chars[start_cut..stop_cut].iter().collect() +} diff --git a/src/json_structs.rs b/src/json_structs.rs index 6fc5504..2f15050 100644 --- a/src/json_structs.rs +++ b/src/json_structs.rs @@ -274,7 +274,7 @@ pub enum DecoderConfig { stop: usize, }, Replace { - pattern: Value, + pattern: ReplacePattern, #[serde(default)] content: String, }, diff --git a/src/normalizers.rs b/src/normalizers.rs index 9bace63..a2653e8 100644 --- a/src/normalizers.rs +++ b/src/normalizers.rs @@ -6,7 +6,7 @@ use std::borrow::Cow; pub use self::nfc::Nfc; pub use self::prepend::Prepend; -pub use self::replace::Replace; +pub use self::replace::{Error as ReplaceError, Replace}; use crate::json_structs::{NormalizerConfig, NormalizerKind}; /// Errors from constructing a normalizer. diff --git a/tests/gemma_replace_decoder.rs b/tests/gemma_replace_decoder.rs new file mode 100644 index 0000000..e6b63c9 --- /dev/null +++ b/tests/gemma_replace_decoder.rs @@ -0,0 +1,80 @@ +//! End-to-end validation that fastokens' decoder converter reproduces the +//! gemma (SentencePiece) decoder chain exactly, against the reference +//! HuggingFace `tokenizers` crate. +//! +//! gemma's `tokenizer.json` decoder is a `Sequence` of +//! `Replace("▁", " ")` → `ByteFallback` → `Fuse` → `Strip(" ", 1, 0)`: +//! it maps the metaspace marker `▁` (U+2581) back to spaces, reassembles +//! `<0xHH>` byte tokens into UTF-8, fuses the pieces, and drops the single +//! leading space SentencePiece prepends. If `Replace` (or `ByteFallback`/ +//! `Fuse`/`Strip`) is dropped or mishandled, `decode` produces wrong text — +//! and tokenizer construction outright fails with +//! `unsupported decoder type: Replace`. +//! +//! The reference oracle is the real `tokenizers` crate (a dev-dependency): we +//! deserialize the *same JSON* into both fastokens' `Decoder` and the +//! `tokenizers` `DecoderWrapper`, then assert byte-for-byte equal output on a +//! corpus of edge cases (metaspace, multi-byte fallback, invalid UTF-8, emoji, +//! word continuations, empty input). + +use fastokens::decoders::Decoder; +use tokenizers::tokenizer::Decoder as _; +use tokenizers::DecoderWrapper; + +/// The gemma decoder chain, verbatim from a gemma `tokenizer.json`. +const GEMMA_DECODER_JSON: &str = r#"{ + "type": "Sequence", + "decoders": [ + { "type": "Replace", "pattern": { "String": "▁" }, "content": " " }, + { "type": "ByteFallback" }, + { "type": "Fuse" }, + { "type": "Strip", "content": " ", "start": 1, "stop": 0 } + ] +}"#; + +fn cases() -> Vec> { + let to_vec = |xs: &[&str]| xs.iter().map(|s| s.to_string()).collect::>(); + vec![ + to_vec(&["▁Hello", "▁world"]), // simple metaspace words + to_vec(&["▁a", "▁b", "▁c"]), // interior spaces + to_vec(&["lo"]), // continuation (no leading ▁) + to_vec(&["▁caf", "<0xC3>", "<0xA9>"]), // café via byte fallback + to_vec(&["▁", "<0xF0>", "<0x9F>", "<0x98>", "<0x80>"]), // 😀 emoji + to_vec(&["<0xFF>"]), // invalid UTF-8 -> U+FFFD + to_vec(&["▁ok", "<0xFF>", "<0xFE>", "▁end"]), // invalid run between text + to_vec(&["▁mix", "<0xC3>", "▁bad"]), // lone byte (invalid) mid-text + to_vec(&["▁"]), // bare metaspace -> "" after strip + Vec::new(), // empty token list + ] +} + +#[test] +fn gemma_decoder_chain_matches_reference() { + let ft: Decoder = Decoder::from_config( + serde_json::from_str(GEMMA_DECODER_JSON).expect("parse fastokens decoder config"), + ) + .expect("build fastokens decoder (Replace/ByteFallback/Fuse/Strip must be supported)"); + + let hf: DecoderWrapper = + serde_json::from_str(GEMMA_DECODER_JSON).expect("parse reference decoder"); + + for tokens in cases() { + let ft_chain = ft + .decode_chain(tokens.clone()) + .expect("fastokens decode_chain"); + let hf_chain = hf + .decode_chain(tokens.clone()) + .expect("reference decode_chain"); + assert_eq!( + ft_chain, hf_chain, + "decode_chain mismatch on tokens {tokens:?}" + ); + + // The joined result is what `Tokenizer::decode` ultimately returns. + assert_eq!( + ft.decode(tokens.clone()).expect("fastokens decode"), + hf_chain.join(""), + "joined decode mismatch on tokens {tokens:?}" + ); + } +} From 8c5e1d3dfcac6f78b623affab68fc9a6a5adc412 Mon Sep 17 00:00:00 2001 From: itayh Date: Wed, 17 Jun 2026 10:27:14 +0300 Subject: [PATCH 4/4] python: release the GIL during encode / encode_batch The PyTokenizer encode hot path held the GIL for the entire Rust tokenization. Under a high-concurrency serving frontend (e.g. vLLM tokenizing hundreds of requests at once) this serializes all encodes on the GIL and blocks the caller's event loop, inflating time-to-first-token while decode is unaffected. A controlled same-node A/B on gemma-4-31B (VLLM_USE_FASTOKENS on vs off) showed ~+130% TTFT p50 / -14% throughput with the GIL held, entirely in the encode/TTFT portion (ITL/TPOT flat). The read lock was already designed for GIL-released concurrent reads (see the TokenizerState doc comment), so wrap the tokenization in py.allow_threads(): acquire the read lock and build the PyEncoding inside the GIL-released closure, and do only the Py::new allocation under the GIL. encode_batch gets the same treatment so rayon parallelizes without pinning the GIL for the whole batch. (Also: tidy the relocated pad_to_multiple_of math to base.div_ceil(m).) Co-Authored-By: Claude Opus 4.8 (1M context) --- python/src/lib.rs | 94 +++++++++++++++++++++++++++++------------------ 1 file changed, 59 insertions(+), 35 deletions(-) diff --git a/python/src/lib.rs b/python/src/lib.rs index 45ff6ab..6f59e0b 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -633,15 +633,26 @@ impl PyTokenizer { add_special_tokens: bool, py: Python<'_>, ) -> PyResult> { - let state = self.read(); - let mut ids = state - .inner - .encode_with_special_tokens(input, add_special_tokens) - .map_err(|e| PyValueError::new_err(e.to_string()))?; - state.do_truncate(&mut ids); - let target = state.single_pad_target(ids.len()); + // Release the GIL for the tokenization work. The read lock is built + // for GIL-released concurrent reads (see `TokenizerState` docs), so + // concurrent `encode` calls — e.g. a serving frontend tokenizing many + // requests at once — run in parallel instead of serializing on the + // GIL (and blocking the caller's event loop). Only the final Python + // object allocation needs the GIL. + let encoding = py + .allow_threads(|| -> Result { + let state = self.read(); + let mut ids = state + .inner + .encode_with_special_tokens(input, add_special_tokens) + .map_err(|e| e.to_string())?; + state.do_truncate(&mut ids); + let target = state.single_pad_target(ids.len()); + Ok(build_encoding(ids, state.pad.as_ref(), target)) + }) + .map_err(PyValueError::new_err)?; - Py::new(py, build_encoding(ids, state.pad.as_ref(), target)) + Py::new(py, encoding) } /// Encode a batch of inputs in parallel. @@ -657,36 +668,49 @@ impl PyTokenizer { ) -> PyResult>> { use rayon::prelude::*; - let state = self.read(); - let mut batch: Vec> = inputs - .par_iter() - .map(|s| { - state - .inner - .encode_with_special_tokens(s.as_str(), add_special_tokens) - .map_err(|e| PyValueError::new_err(e.to_string())) - }) - .collect::>>()?; - - for ids in &mut batch { - state.do_truncate(ids); - } + // Release the GIL for the batch tokenization. rayon parallelizes the + // encode across cores; holding the GIL would otherwise block every + // other Python thread for the whole batch. Only the Python object + // allocations need the GIL, so they run after. + let encodings = py + .allow_threads(|| -> Result, String> { + let state = self.read(); + let mut batch: Vec> = inputs + .par_iter() + .map(|s| { + state + .inner + .encode_with_special_tokens(s.as_str(), add_special_tokens) + .map_err(|e| e.to_string()) + }) + .collect::, _>>()?; + + for ids in &mut batch { + state.do_truncate(ids); + } - let pad_target: Option = state.pad.as_ref().map(|p| { - let max_len = batch.iter().map(|ids| ids.len()).max().unwrap_or(0); - let base = p.length.unwrap_or(max_len).max(max_len); - match p.pad_to_multiple_of { - Some(m) if m > 0 => (base + m - 1) / m * m, - _ => base, - } - }); + let pad_target: Option = state.pad.as_ref().map(|p| { + let max_len = batch.iter().map(|ids| ids.len()).max().unwrap_or(0); + let base = p.length.unwrap_or(max_len).max(max_len); + match p.pad_to_multiple_of { + Some(m) if m > 0 => base.div_ceil(m) * m, + _ => base, + } + }); + + Ok(batch + .into_iter() + .map(|ids| { + let target = pad_target.unwrap_or(ids.len()); + build_encoding(ids, state.pad.as_ref(), target) + }) + .collect()) + }) + .map_err(PyValueError::new_err)?; - batch + encodings .into_iter() - .map(|ids| { - let target = pad_target.unwrap_or(ids.len()); - Py::new(py, build_encoding(ids, state.pad.as_ref(), target)) - }) + .map(|enc| Py::new(py, enc)) .collect() }