aprender-serve 0.69.1

//! Tokenizer for text encoding and decoding
//!
//! Implements tokenization for transformer models:
//! - BPE (Byte Pair Encoding) - Used by GPT models
//! - Vocabulary management
//! - Special token handling
//!
//! ## Example
//!
//! ```rust,ignore
//! use realizar::Tokenizer;
//!
//! let tokenizer = Tokenizer::from_vocab(vocab);
//! let token_ids = tokenizer.encode("Hello, world!")?;
//! let text = tokenizer.decode(&token_ids)?;
//! ```

use std::collections::HashMap;

use crate::error::{RealizarError, Result};

/// Vocabulary mapping between tokens and IDs
#[derive(Debug, Clone)]
pub struct Vocabulary {
    /// Token to ID mapping
    token_to_id: HashMap<String, u32>,
    /// ID to token mapping
    id_to_token: HashMap<u32, String>,
}

impl Vocabulary {
    /// Create a new vocabulary from token list
    ///
    /// # Arguments
    ///
    /// * `tokens` - List of tokens in order (index = token ID)
    ///
    /// # Errors
    ///
    /// Returns error if tokens list is empty or contains duplicates
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let vocab = Vocabulary::from_tokens(vec![
    ///     "<unk>".to_string(),
    ///     "hello".to_string(),
    ///     "world".to_string(),
    /// ])?;
    /// ```
    pub fn from_tokens(tokens: Vec<String>) -> Result<Self> {
        if tokens.is_empty() {
            return Err(RealizarError::UnsupportedOperation {
                operation: "create_vocabulary".to_string(),
                reason: "Vocabulary cannot be empty".to_string(),
            });
        }

        let mut token_to_id = HashMap::new();
        let mut id_to_token = HashMap::new();

        for (id, token) in tokens.into_iter().enumerate() {
            let id = u32::try_from(id).map_err(|_| RealizarError::UnsupportedOperation {
                operation: "convert_token_id".to_string(),
                reason: format!("Token ID {id} exceeds u32 limit"),
            })?;

            if token_to_id.contains_key(&token) {
                return Err(RealizarError::UnsupportedOperation {
                    operation: "create_vocabulary".to_string(),
                    reason: format!("Duplicate token: {token}"),
                });
            }

            token_to_id.insert(token.clone(), id);
            id_to_token.insert(id, token);
        }

        Ok(Self {
            token_to_id,
            id_to_token,
        })
    }

    /// Get token ID for a token
    ///
    /// Returns `None` if token not in vocabulary
    #[must_use]
    pub fn get_id(&self, token: &str) -> Option<u32> {
        self.token_to_id.get(token).copied()
    }

    /// Get token for a token ID
    ///
    /// Returns `None` if ID not in vocabulary
    #[must_use]
    pub fn get_token(&self, id: u32) -> Option<&str> {
        self.id_to_token.get(&id).map(String::as_str)
    }

    /// Get vocabulary size
    #[must_use]
    pub fn size(&self) -> usize {
        self.token_to_id.len()
    }
}

/// Tokenizer for encoding and decoding text
#[derive(Debug, Clone)]
pub struct Tokenizer {
    /// Vocabulary
    vocab: Vocabulary,
    /// Unknown token ID
    unk_token_id: u32,
}

/// BPE (Byte Pair Encoding) tokenizer
///
/// Implements subword tokenization using byte pair encoding algorithm.
/// Used by GPT-2, GPT-3, and many other models.
///
/// ## vLLM-Style Property Caching (PMAT-805)
///
/// Properties like vocab_size and max_token_id are cached at construction
/// to avoid repeated HashMap operations during hot inference paths.
/// This is especially important for large vocabularies (e.g., Qwen2's 151K tokens).
#[derive(Debug, Clone)]
pub struct BPETokenizer {
    /// Token to ID mapping
    token_to_id: HashMap<String, u32>,
    /// ID to token mapping
    id_to_token: HashMap<u32, String>,
    /// The model's unknown token, or `None` when it has none (#3609).
    ///
    /// An unknown token belongs to a model family, not to BPE. LLaMA-style vocabularies
    /// declare `<unk>`; byte-level ones (GPT-2, Qwen) need none, because every byte has a
    /// token, and Qwen3.5 declares none. When this is `None`, construction has proved
    /// that every byte has a token, so `encode` never needs a stand-in.
    unk_token_id: Option<u32>,
    /// Cached vocabulary size (PMAT-805: vLLM pattern)
    vocab_size: usize,
    /// Cached maximum token ID (PMAT-805: vLLM pattern)
    max_token_id: u32,
    /// GH-88: Optional BPE merge rules for proper encoding.
    /// When present, `encode()` uses merge-based BPE instead of greedy longest-match.
    /// Required for HuggingFace vocabularies (SafeTensors/APR imports).
    merge_rules: Vec<(String, String)>,
    /// GH-88: Special tokens for atomic tokenization (not split by BPE).
    special_tokens: HashMap<String, u32>,
}

impl BPETokenizer {
    /// Create a new BPE tokenizer
    ///
    /// # Arguments
    ///
    /// * `vocab` - List of tokens (index = token ID)
    /// * `merges` - List of merge pairs in priority order
    /// * `unk_token` - The model's unknown token if it has one (`"<unk>"` or
    ///   `Some("<unk>")`), or `None` if it has none. Having one is a property of the
    ///   model, not a requirement of this type (#3609).
    ///
    /// # Errors
    ///
    /// Returns an error if the vocabulary is empty, or if `unk_token` names a token the
    /// vocabulary does not contain. With no unknown token, also returns an error naming
    /// the first byte that has neither a `<0xNN>` token nor a byte-level glyph token,
    /// because `encode` would have nothing to emit for it. Absent means absent: no
    /// stand-in is synthesised and no input is dropped.
    pub fn new<'a>(
        vocab: Vec<String>,
        _merges: Vec<(String, String)>,
        unk_token: impl Into<Option<&'a str>>,
    ) -> Result<Self> {
        if vocab.is_empty() {
            return Err(RealizarError::UnsupportedOperation {
                operation: "create_bpe_tokenizer".to_string(),
                reason: "Vocabulary cannot be empty".to_string(),
            });
        }

        let mut token_to_id = HashMap::new();
        let mut id_to_token = HashMap::new();

        for (id, token) in vocab.into_iter().enumerate() {
            let id = u32::try_from(id).map_err(|_| RealizarError::UnsupportedOperation {
                operation: "convert_token_id".to_string(),
                reason: format!("Token ID {id} exceeds u32 limit"),
            })?;
            token_to_id.insert(token.clone(), id);
            id_to_token.insert(id, token);
        }

        let unk_token_id =
            match unk_token.into() {
                Some(unk) => Some(*token_to_id.get(unk).ok_or_else(|| {
                    RealizarError::UnsupportedOperation {
                        operation: "create_bpe_tokenizer".to_string(),
                        reason: format!(
                        "unknown token '{unk}' was named, but the vocabulary does not contain it"
                    ),
                    }
                })?),
                None => {
                    if let Some(byte) =
                        (0u8..=255).find(|&b| byte_token_id(&token_to_id, b).is_none())
                    {
                        return Err(RealizarError::UnsupportedOperation {
                            operation: "create_bpe_tokenizer".to_string(),
                            reason: format!(
                            "the model has no unknown token, and byte 0x{byte:02X} has neither a \
                             '<0x{byte:02X}>' token nor a byte-level glyph token, so encode would \
                             have nothing to emit for it"
                        ),
                        });
                    }
                    None
                },
            };

        // PMAT-805: Cache vocabulary properties at construction (vLLM pattern)
        // This avoids repeated HashMap operations during inference
        let vocab_size = token_to_id.len();
        let max_token_id = id_to_token.keys().copied().max().unwrap_or(0);

        Ok(Self {
            token_to_id,
            id_to_token,
            unk_token_id,
            vocab_size,
            max_token_id,
            merge_rules: Vec::new(),
            special_tokens: HashMap::new(),
        })
    }

    /// GH-88: Create a BPE tokenizer with merge rules for proper encoding.
    ///
    /// HuggingFace vocabularies require merge-based BPE (not greedy longest-match).
    /// When merge rules are provided, `encode()` delegates to `bpe_encode()` from
    /// `apr::tokenizer` which handles special tokens atomically and applies merges.
    pub fn with_merges<'a>(
        vocab: Vec<String>,
        merges: Vec<(String, String)>,
        unk_token: impl Into<Option<&'a str>>,
    ) -> Result<Self> {
        let mut tokenizer = Self::new(vocab, vec![], unk_token)?;
        // #3609: the merge path (`bpe_encode`) maps each byte to its byte-level glyph and
        // nothing else, so with no unknown token every glyph must be in the vocabulary,
        // or that byte would be dropped from the input.
        if tokenizer.unk_token_id.is_none() {
            if let Some(byte) =
                (0u8..=255).find(|&b| byte_glyph_id(&tokenizer.token_to_id, b).is_none())
            {
                return Err(RealizarError::UnsupportedOperation {
                    operation: "create_bpe_tokenizer".to_string(),
                    reason: format!(
                        "the model has no unknown token, and byte 0x{byte:02X} has no byte-level \
                         glyph token, which merge-based encoding needs"
                    ),
                });
            }
        }
        // Extract special tokens from vocabulary for atomic tokenization
        let special_tokens: HashMap<String, u32> = tokenizer
            .token_to_id
            .iter()
            .filter(|(k, _)| {
                (k.starts_with("<|") && k.ends_with("|>"))
                    || (k.starts_with("<")
                        && k.ends_with(">")
                        && k.len() > 2
                        && !k.starts_with("<0x"))
            })
            .map(|(k, v)| (k.clone(), *v))
            .collect();
        tokenizer.merge_rules = merges;
        tokenizer.special_tokens = special_tokens;
        Ok(tokenizer)
    }

    /// Get maximum token ID (cached, O(1))
    ///
    /// PMAT-805: Returns pre-computed max token ID for fast access
    #[inline]
    #[must_use]
    pub fn max_token_id(&self) -> u32 {
        self.max_token_id
    }

    /// Encode text to token IDs.
    ///
    /// GH-88: When merge rules are present (HuggingFace tokenizers), uses proper
    /// BPE encoding with special token support. Otherwise falls back to greedy
    /// longest-match (sufficient for GGUF vocabularies).
    ///
    /// # Arguments
    ///
    /// * `text` - Input text to encode
    ///
    /// # Returns
    ///
    /// Vector of token IDs
    #[must_use]
    pub fn encode(&self, text: &str) -> Vec<u32> {
        contract_pre_encode!();
        if text.is_empty() {
            return Vec::new();
        }

        // GH-88: Use proper BPE when merge rules are available
        if !self.merge_rules.is_empty() {
            return crate::apr::tokenizer::bpe_encode(
                text,
                &self.token_to_id,
                &self.merge_rules,
                &self.special_tokens,
            );
        }

        // Convert to GPT-2 encoding: space -> Ġ, newline -> Ċ
        let processed: String = text
            .chars()
            .map(|c| match c {
                ' ' => 'Ġ',  // U+0120
                '\n' => 'Ċ', // U+010A
                '\r' => 'Ḃ', // U+1E02
                _ => c,
            })
            .collect();

        let mut tokens = Vec::new();
        let mut remaining = processed.as_str();

        while !remaining.is_empty() {
            // Greedy longest match
            let mut best_len = 0;
            let mut best_id = None;

            // Collect character byte offsets for proper slicing
            let char_indices: Vec<usize> = remaining
                .char_indices()
                .map(|(i, _)| i)
                .chain(std::iter::once(remaining.len()))
                .collect();

            // Try all prefixes up to 32 chars
            for char_count in 1..=char_indices.len().saturating_sub(1).min(32) {
                let byte_end = char_indices[char_count];
                let prefix = &remaining[..byte_end];
                if let Some(&id) = self.token_to_id.get(prefix) {
                    best_len = byte_end;
                    best_id = Some(id);
                }
            }

            if let Some(id) = best_id {
                tokens.push(id);
                remaining = &remaining[best_len..];
            } else {
                // No match - try byte tokens like <0x48>
                let ch = remaining.chars().next().expect("non-empty");
                let ch_len = ch.len_utf8();

                for byte in remaining[..ch_len].bytes() {
                    tokens.push(self.byte_fallback_id(byte));
                }
                remaining = &remaining[ch_len..];
            }
        }

        contract_post_encode!(&tokens);
        tokens
    }

    /// Decode token IDs to text
    ///
    /// # Arguments
    ///
    /// * `token_ids` - Token IDs to decode
    ///
    /// # Errors
    ///
    /// Returns error if any token ID is invalid
    pub fn decode(&self, token_ids: &[u32]) -> Result<String> {
        contract_pre_decode!();
        let mut bytes: Vec<u8> = Vec::new();

        for &id in token_ids {
            let token =
                self.id_to_token
                    .get(&id)
                    .ok_or_else(|| RealizarError::UnsupportedOperation {
                        operation: "decode_bpe_token".to_string(),
                        reason: format!("Invalid token ID: {id}"),
                    })?;

            // Skip special tokens
            if token.starts_with("<|") && token.ends_with("|>") {
                continue;
            }
            if token == "<s>" || token == "</s>" || token == "<unk>" || token == "<pad>" {
                continue;
            }

            // Handle byte tokens like <0xE6>
            if token.starts_with("<0x") && token.ends_with('>') && token.len() == 6 {
                if let Ok(byte_val) = u8::from_str_radix(
                    token
                        .get(3..5)
                        .expect("byte token <0xNN> has len 6, indices 3..5 always valid"),
                    16,
                ) {
                    bytes.push(byte_val);
                    continue;
                }
            }

            // Decode GPT-2 style byte-level BPE
            for c in token.chars() {
                match c {
                    'Ġ' => bytes.push(b' '),  // U+0120 -> space
                    'Ċ' => bytes.push(b'\n'), // U+010A -> newline
                    'ċ' => bytes.push(b'\n'), // lowercase variant
                    'Ḃ' => bytes.push(b'\r'), // U+1E02 -> carriage return
                    '▁' => bytes.push(b' '),  // U+2581 SentencePiece -> space
                    _ => {
                        // Try GPT-2 unicode-to-byte mapping
                        if let Some(byte) = Self::gpt2_char_to_byte(c) {
                            bytes.push(byte);
                        } else {
                            // Regular UTF-8 character
                            let mut buf = [0u8; 4];
                            let encoded = c.encode_utf8(&mut buf);
                            bytes.extend_from_slice(encoded.as_bytes());
                        }
                    },
                }
            }
        }

        // Decode as UTF-8, replacing invalid sequences
        let result = String::from_utf8_lossy(&bytes).into_owned();
        contract_post_decode!(&result);
        Ok(result)
    }

    /// Convert a GPT-2 byte-level-BPE unicode character back to its original byte.
    ///
    /// PMAT-837: the prior body treated the U+0100.. region as a LINEAR `code - 0x100`
    /// offset and only accepted `0..=32 | 127..=160 | 173`. GPT-2's byte_encoder above
    /// U+0120 is a STAIRCASE (U+0121→0x7F, U+0122..=U+0142→0x80..=0xA0, U+0143→0xAD), and
    /// the Latin-1 self-mapped bytes (cp ≤ 0xFF) were omitted entirely — so 129/256 byte
    /// values returned None and `decode` re-encoded the codepoint as multi-byte UTF-8,
    /// turning every non-ASCII generation into mojibake. Delegate to the correct in-crate
    /// inverse map (the GGUF path's `gpt2_unicode_to_byte`).
    fn gpt2_char_to_byte(c: char) -> Option<u8> {
        crate::gguf::utils::gpt2_unicode_to_byte(c)
    }

    /// Get vocabulary size (cached, O(1))
    ///
    /// PMAT-805: Returns pre-computed vocab size for fast access.
    /// Avoids repeated HashMap::len() calls during inference.
    #[inline]
    #[must_use]
    pub fn vocab_size(&self) -> usize {
        self.vocab_size
    }

    /// Get token ID for a token
    #[must_use]
    pub fn get_token_id(&self, token: &str) -> Option<u32> {
        self.token_to_id.get(token).copied()
    }

    /// The unknown token's ID, or `None` when the model has none (#3609).
    #[must_use]
    pub fn unk_token_id(&self) -> Option<u32> {
        self.unk_token_id
    }

    /// The token for one byte that no vocabulary entry covers: its `<0xNN>` token, else the
    /// model's unknown token, else its byte-level glyph (#3609). With no unknown token,
    /// `new` proved that every byte has a `<0xNN>` or a glyph token, so this cannot miss.
    fn byte_fallback_id(&self, byte: u8) -> u32 {
        if let Some(&id) = self.token_to_id.get(&format!("<0x{byte:02X}>")) {
            return id;
        }
        if let Some(unk) = self.unk_token_id {
            return unk;
        }
        byte_glyph_id(&self.token_to_id, byte).expect(
            "BPETokenizer::new proved every byte encodable without an unknown token (#3609)",
        )
    }

    /// Get token for a token ID
    #[must_use]
    pub fn get_token(&self, id: u32) -> Option<&str> {
        self.id_to_token.get(&id).map(String::as_str)
    }

    /// Returns true if `id` is a registered special token (e.g. BOS/EOS/PAD).
    ///
    /// PMAT-803: model-backed embeddings mean-pool over the *content* tokens, so the
    /// embedding handler uses this to skip special tokens during pooling. Returns
    /// false when no special tokens were registered (vocab-only tokenizers).
    #[must_use]
    pub fn is_special_token(&self, id: u32) -> bool {
        self.special_tokens.values().any(|&v| v == id)
    }
}

/// The token for one raw byte that no vocabulary token covers: SentencePiece's `<0xNN>`
/// byte-fallback token, else the GPT-2 byte-level glyph (#3609).
fn byte_token_id(token_to_id: &HashMap<String, u32>, byte: u8) -> Option<u32> {
    token_to_id
        .get(&format!("<0x{byte:02X}>"))
        .copied()
        .or_else(|| byte_glyph_id(token_to_id, byte))
}

/// The GPT-2 byte-level glyph token for one byte (`é`'s 0xC3 → `Ã`), if the vocabulary has it.
fn byte_glyph_id(token_to_id: &HashMap<String, u32>, byte: u8) -> Option<u32> {
    let glyph = crate::gguf::utils::gpt2_byte_to_unicode(byte);
    token_to_id
        .get(glyph.encode_utf8(&mut [0u8; 4]) as &str)
        .copied()
}

/// The unknown token to pass to [`BPETokenizer::new`] when the caller has only a token list.
///
/// That is the vocabulary's own `<unk>` entry if it has one, and `None` if it does not.
/// Nothing is synthesised: a vocabulary without `<unk>` (Qwen3.5, GPT-2) gets no unknown
/// token, and `BPETokenizer::new` then proves that every byte has a token (#3609).
///
/// This finds the token by NAME. It is the interim for callers that receive a `Vec<String>`
/// and nothing else. The model's DECLARATION (`tokenizer.ggml.unknown_token_id`, or
/// `tokenizer.json`'s `unk_token`) is the design, and threading it through every loader
/// is #3675.
#[must_use]
pub fn vocabulary_unk_token(vocab: &[String]) -> Option<&'static str> {
    vocab.iter().any(|t| t == "<unk>").then_some("<unk>")
}

/// Viterbi algorithm state: `(best_score, best_token)` at each position
type ViterbiState = (Vec<f32>, Vec<Option<String>>);

/// `SentencePiece` tokenizer (Unigram model)
///
/// Implements subword tokenization using unigram language model.
/// Used by `LLaMA`, T5, ALBERT, and many other models.
///
/// Unlike BPE which uses greedy merges, `SentencePiece` finds the
/// most likely segmentation using token scores (log probabilities).
#[derive(Debug, Clone)]
pub struct SentencePieceTokenizer {
    /// Token to ID mapping
    token_to_id: HashMap<String, u32>,
    /// ID to token mapping
    id_to_token: HashMap<u32, String>,
    /// Token scores (log probabilities)
    scores: HashMap<String, f32>,
    /// Unknown token ID
    unk_token_id: u32,
}

include!("tokenizer_sentence_piece.rs");
include!("tokenizer_sentencepiece_encode.rs");
include!("tokenizer_04.rs");
include!("tokenizer_contract_tests.rs");