toklen 0.2.0

#![doc = include_str!("../README.md")]
#![allow(clippy::upper_case_acronyms)]

//---------------------------------------------------------------
// Modules

mod added_tokens;
mod byte_level;
mod config;
mod models;
mod normalizers;
mod pre_tokenized;
mod pre_tokenizers;

//---------------------------------------------------------------
// Exports

use std::borrow::Cow;

use self::added_tokens::{AddedTokens, Segment};
use self::config::TokenizerJson;
use self::models::Model;
use self::normalizers::Normalizer;
use self::pre_tokenized::{PreTokenizedString, PtSplit};
use self::pre_tokenizers::PreTokenizer;

/// Errors that can occur when constructing a [`Tokenizer`].
#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error("failed to parse JSON: {0}")]
    Json(#[from] serde_json::Error),

    #[error("normalizer error: {0}")]
    Normalizer(#[from] normalizers::Error),

    #[error("pre-tokenizer error: {0}")]
    PreTokenizer(#[from] pre_tokenizers::Error),

    #[error("model error: {0}")]
    Model(String),
}

/// A token-counting tokenizer backed by `tokenizer.json`.
pub struct Tokenizer {
    added_tokens: Option<AddedTokens>,
    normalizer: Option<Normalizer>,
    pre_tokenizer: Option<PreTokenizer>,
    model: Model,
}

impl Tokenizer {
    /// Build the pipeline steps from a parsed JSON config.
    #[rustfmt::skip]
    fn build(json: TokenizerJson) -> Result<Self, Error> {
        let added_tokens = AddedTokens::from_configs(&json.added_tokens).map_err(Error::Model)?;
        let normalizer = json.normalizer.map(Normalizer::from_config).transpose()?;
        let pre_tokenizer = json.pre_tokenizer.map(PreTokenizer::from_config).transpose()?;
        let model = Model::from_config(json.model).map_err(Error::Model)?;
        Ok(Self { added_tokens, normalizer, pre_tokenizer, model })
    }

    /// Create a tokenizer from `tokenizer.json` content.
    ///
    /// Accepts anything that can be referenced as bytes: `&str`, `&[u8]`,
    /// `String`, `Vec<u8>`, etc.
    #[inline]
    pub fn from_json(json: impl AsRef<[u8]>) -> Result<Self, Error> {
        let json: TokenizerJson = serde_json::from_slice(json.as_ref())?;
        Self::build(json)
    }

    /// Count the number of tokens that `input` would produce when tokenized.
    ///
    /// Returns `Ok(count)` on success, or `Err(estimate)` with a fallback
    /// estimate (`input.len() / 4`) on failure.
    ///
    /// Inputs longer than `i32::MAX` bytes are rejected immediately with an
    /// estimate — the internal pipeline uses 32-bit offsets for compactness.
    pub fn encode_len(&self, input: &str) -> Result<usize, usize> {
        if input.is_empty() {
            return Ok(0);
        }
        if input.len() >= i32::MAX as usize {
            return Err(input.len() >> 2);
        }
        match self.count_tokens_internal(input) {
            Ok(count) => Ok(count),
            Err(_) => Err(input.len() >> 2),
        }
    }

    /// Internal counting pipeline.
    #[rustfmt::skip]
    fn count_tokens_internal(&self, input: &str) -> Result<usize, String> {
        let mut pts = self.build_pre_tokenized(input);

        // Normal path: full pre-tokenize then count each split.
        if let Some(pt) = &self.pre_tokenizer {
            pt.pre_tokenize(&mut pts).map_err(|e| e.to_string())?;
        }

        let mut count = pts.added_token_count();

        for text in pts.texts() {
            count += self.model.count_tokens(text)?;
        }

        Ok(count)
    }

    /// Build a [`PreTokenizedString`]: split on added tokens, then normalize.
    #[rustfmt::skip]
    fn build_pre_tokenized(&self, input: &str) -> PreTokenizedString {
        // Fast path: no added tokens → skip allocation.
        let segments = match &self.added_tokens {
            Some(at) => at.split(input),
            None => return self.normalize_into_pts(input),
        };

        // Fast path: exactly one Text segment, normalization is borrowed.
        if segments.len() == 1 && let Segment::Text(text) = segments[0] {
            return self.normalize_into_pts(text);
        }

        // Slow path: multiple segments with mixed added-token / text.
        let mut buffer = String::with_capacity(input.len());
        let mut splits = Vec::new();

        for seg in &segments {
            match seg {
                Segment::Text("") => continue,
                Segment::Text(text) => {
                    let start = buffer.len();
                    buffer.push_str(&self.normalize(text));
                    splits.push(PtSplit::from_range(start..buffer.len()));
                }
                Segment::Token(id) => {
                    splits.push(PtSplit::from_token(buffer.len(), *id));
                }
            }
        }

        PreTokenizedString::new(buffer, splits)
    }

    /// Normalize `text` and wrap into a [`PreTokenizedString`].
    #[inline]
    fn normalize_into_pts(&self, text: &str) -> PreTokenizedString {
        match self.normalize(text) {
            Cow::Borrowed("") => PreTokenizedString::EMPTY,
            Cow::Owned(s) if s.is_empty() => PreTokenizedString::EMPTY,
            Cow::Borrowed(s) => PreTokenizedString {
                splits: vec![PtSplit::from_range(0..s.len())],
                buffer: s.to_owned(),
            },
            Cow::Owned(s) => PreTokenizedString {
                splits: vec![PtSplit::from_range(0..s.len())],
                buffer: s,
            },
        }
    }

    /// Apply normalizer to `text`, returning `Cow::Borrowed` when unchanged.
    #[inline]
    fn normalize<'a>(&self, text: &'a str) -> Cow<'a, str> {
        match &self.normalizer {
            Some(n) => n.normalize(text),
            None => Cow::Borrowed(text),
        }
    }
}