rustytorch/crates/production/rtx-inference/src/tokenizer.rs

//! Serving-boundary tokenization.
//!
//! Historically the HTTP/streaming serving layers tokenized text by mapping
//! each UTF-8 byte to a token id (`bytes().map(i32::from)`) and reversed it
//! with `String::from_utf8_lossy`. That byte-level scheme is lossless and
//! reversible but not vocabulary-aware, so it wastes context and produces
//! nonsense token counts for anything downstream that expects real
//! sub-word tokens.
//!
//! [`ServingTokenizer`] makes tokenization pluggable: load a real
//! `tokenizers`-crate vocabulary (`tokenizer.json`) via
//! [`ServingTokenizer::from_file`], or fall back to the original byte-level
//! behavior via [`ServingTokenizer::byte_level`] (the default).

use std::path::Path;

/// Tokenizer used at the serving boundary (HTTP inference, streaming).
///
/// Defaults to byte-level encoding for backward compatibility; construct a
/// [`ServingTokenizer::Vocab`] from a `tokenizer.json` file to use a real
/// vocabulary tokenizer instead.
#[derive(Debug, Clone)]
pub enum ServingTokenizer {
    /// Real vocabulary tokenizer backed by the `tokenizers` crate.
    Vocab(Box<tokenizers::Tokenizer>),
    /// Byte-level fallback: each UTF-8 byte is one token id. Lossless and
    /// reversible, but not vocabulary-aware.
    ByteLevel,
}

impl Default for ServingTokenizer {
    fn default() -> Self {
        Self::ByteLevel
    }
}

impl ServingTokenizer {
    /// Byte-level tokenizer: each UTF-8 byte is one token id.
    #[must_use]
    pub fn byte_level() -> Self {
        Self::ByteLevel
    }

    /// Load a real vocabulary tokenizer from a `tokenizer.json` file.
    ///
    /// # Errors
    ///
    /// Returns an error if the file cannot be read or does not describe a
    /// valid `tokenizers` configuration.
    pub fn from_file(path: impl AsRef<Path>) -> Result<Self, String> {
        let tokenizer = tokenizers::Tokenizer::from_file(path.as_ref())
            .map_err(|e| format!("failed to load tokenizer from {:?}: {e}", path.as_ref()))?;
        Ok(Self::Vocab(Box::new(tokenizer)))
    }

    /// Encode text into token ids.
    ///
    /// The `ByteLevel` variant maps each UTF-8 byte to its numeric value,
    /// matching the historical serving-boundary behavior exactly.
    #[must_use]
    pub fn encode(&self, text: &str) -> Vec<i32> {
        match self {
            Self::ByteLevel => text.bytes().map(i32::from).collect(),
            Self::Vocab(tokenizer) => match tokenizer.encode(text, false) {
                Ok(encoding) => encoding.get_ids().iter().map(|&id| id as i32).collect(),
                Err(_) => Vec::new(),
            },
        }
    }

    /// Decode token ids back into text.
    ///
    /// The `ByteLevel` variant clamps each id to a `u8` and reconstructs the
    /// string with `String::from_utf8_lossy`, matching the historical
    /// serving-boundary behavior exactly.
    #[must_use]
    pub fn decode(&self, ids: &[i32]) -> String {
        match self {
            Self::ByteLevel => {
                let bytes: Vec<u8> = ids.iter().map(|&t| t.clamp(0, 255) as u8).collect();
                String::from_utf8_lossy(&bytes).into_owned()
            }
            Self::Vocab(tokenizer) => {
                let ids: Vec<u32> = ids.iter().map(|&id| id.max(0) as u32).collect();
                tokenizer.decode(&ids, true).unwrap_or_default()
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn byte_level_round_trip_ascii() {
        let tok = ServingTokenizer::byte_level();
        let ids = tok.encode("Hello, world!");
        assert_eq!(tok.decode(&ids), "Hello, world!");
    }

    #[test]
    fn byte_level_round_trip_multibyte_utf8() {
        let tok = ServingTokenizer::byte_level();
        let text = "héllo wörld — 日本語";
        let ids = tok.encode(text);
        assert_eq!(tok.decode(&ids), text);
    }

    #[test]
    fn byte_level_default_matches_byte_level_constructor() {
        let default_tok = ServingTokenizer::default();
        assert!(matches!(default_tok, ServingTokenizer::ByteLevel));
    }

    /// Builds a tiny in-memory WordLevel tokenizer over a 5-word vocab and
    /// verifies encode/decode round-trips for a sentence made purely of
    /// vocab words.
    #[test]
    fn vocab_tokenizer_round_trip() {
        use std::collections::HashMap;
        use tokenizers::models::wordlevel::WordLevel;
        use tokenizers::pre_tokenizers::whitespace::Whitespace;

        let mut vocab: HashMap<String, u32> = HashMap::new();
        vocab.insert("[UNK]".to_string(), 0);
        vocab.insert("the".to_string(), 1);
        vocab.insert("quick".to_string(), 2);
        vocab.insert("brown".to_string(), 3);
        vocab.insert("fox".to_string(), 4);

        let model = WordLevel::builder()
            .vocab(vocab)
            .unk_token("[UNK]".to_string())
            .build()
            .expect("valid wordlevel model");

        let mut tokenizer = tokenizers::Tokenizer::new(model);
        tokenizer.with_pre_tokenizer(Some(Whitespace));

        let serving = ServingTokenizer::Vocab(Box::new(tokenizer));

        let text = "the quick brown fox";
        let ids = serving.encode(text);
        assert_eq!(ids, vec![1, 2, 3, 4]);
        assert_eq!(serving.decode(&ids), text);
    }
}