normalizers: support Replace and Prepend (fixes gemma-4-31B tokenizer load)

.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

python/src/lib.rs

@@ -633,15 +633,26 @@ impl PyTokenizer {
         add_special_tokens: bool,
         py: Python<'_>,
     ) -> PyResult<Py<PyEncoding>> {
-        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<PyEncoding, String> {
+                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<Vec<Py<PyEncoding>>> {
         use rayon::prelude::*;
 
-        let state = self.read();
-        let mut batch: Vec<Vec<u32>> = 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::<PyResult<Vec<_>>>()?;
-
-        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<Vec<PyEncoding>, String> {
+                let state = self.read();
+                let mut batch: Vec<Vec<u32>> = inputs
+                    .par_iter()
+                    .map(|s| {
+                        state
+                            .inner
+                            .encode_with_special_tokens(s.as_str(), add_special_tokens)
+                            .map_err(|e| e.to_string())
+                    })
+                    .collect::<Result<Vec<_>, _>>()?;
+
+                for ids in &mut batch {
+                    state.do_truncate(ids);
+                }
 
-        let pad_target: Option<usize> = 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<usize> = 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()
     }
 

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<Decoder>),
 }
 
@@ -23,7 +46,20 @@ impl Decoder {
     pub fn from_config(config: DecoderConfig) -> Result<Self, Error> {
         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<String>) -> Result<Vec<String>, 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<String>) -> Vec<String> {
+    let mut out: Vec<String> = Vec::with_capacity(tokens.len());
+    let mut pending: Vec<u8> = Vec::new();
+
+    let flush = |pending: &mut Vec<u8>, out: &mut Vec<String>| {
+        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<char> = 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()
+}

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))]
@@ -262,7 +274,7 @@ pub enum DecoderConfig {
         stop: usize,
     },
     Replace {
-        pattern: Value,
+        pattern: ReplacePattern,
         #[serde(default)]
         content: String,
     },

src/normalizers.rs

@@ -1,21 +1,31 @@
 mod nfc;
+mod prepend;
+mod replace;
 
 use std::borrow::Cow;
 
 pub use self::nfc::Nfc;
+pub use self::prepend::Prepend;
+pub use self::replace::{Error as ReplaceError, Replace};
 use crate::json_structs::{NormalizerConfig, NormalizerKind};
 
 /// Errors from constructing a normalizer.
 #[derive(Debug, thiserror::Error)]
 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<Normalizer>),
 }
 
@@ -24,6 +34,10 @@ impl Normalizer {
     pub fn from_config(config: NormalizerConfig) -> Result<Self, Error> {
         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 {