toklen 0.1.0

#![doc = include_str!("../README.md")]

use std::borrow::Cow;

use self::added_tokens::{AddedTokens, Segment};
use self::bpe::Model;
use self::config::TokenizerJson;
use self::normalizer::Normalizer;
use self::pre_tokenized::{PreTokenizedString, Split};
use self::pre_tokenizer::PreTokenizer;

mod added_tokens;
mod bpe;
mod byte_level;
pub mod config;
mod normalizer;
mod pre_tokenized;
mod pre_tokenizer;
mod split;

/// 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] normalizer::Error),

    #[error("pre-tokenizer error: {0}")]
    PreTokenizer(#[from] pre_tokenizer::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,
    /// When the pre-tokenizer is `Sequence([Split, ByteLevel(bulk)])`,
    /// we store a Split-only pre-tokenizer and fuse ByteLevel into BPE.
    split_only: Option<PreTokenizer>,
}

impl Tokenizer {
    /// Build the pipeline steps from a parsed JSON config.
    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 model = Model::from_config(json.model).map_err(Error::Model)?;
        let pre_tokenizer = json
            .pre_tokenizer
            .map(PreTokenizer::from_config)
            .transpose()?;

        let split_only = {
            if let Some(pret) = pre_tokenizer.as_ref()
                && let PreTokenizer::Sequence(steps) = pret
                && steps.len() == 2
            {
                let is_split = matches!(&steps[0], PreTokenizer::Split(_));
                let is_bulk_bl =
                    matches!(&steps[1], PreTokenizer::ByteLevel(bl) if bl.is_bulk_only());
                (is_split && is_bulk_bl).then(|| steps[0].clone())
            } else {
                None
            }
        };

        Ok(Self {
            added_tokens,
            normalizer,
            pre_tokenizer,
            model,
            split_only,
        })
    }

    /// Create a tokenizer from `tokenizer.json` content.
    ///
    /// Accepts anything that can be referenced as bytes: `&str`, `&[u8]`,
    /// `String`, `Vec<u8>`, etc.
    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 `u32::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() >= u32::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.
    fn count_tokens_internal(&self, input: &str) -> Result<usize, String> {
        let mut pts = self.build_pre_tokenized(input);

        // Fused path: Split only, then batch-tokenize with inline ByteLevel.
        if let Some(split) = &self.split_only {
            split.pre_tokenize(&mut pts).map_err(|e| e.to_string())?;
            return pts.count_tokens_batched(|buf, splits, count| {
                self.model.count_tokens_batch_fused(buf, splits, count)
            });
        }

        // 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())?;
        }

        pts.count_tokens(|text, count| {
            *count += self.model.count_tokens(text)?;
            Ok(())
        })
    }

    /// Build a [`PreTokenizedString`]: split on added tokens, then normalize.
    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::Token(id) => {
                    splits.push(Split {
                        range: buffer.len()..buffer.len(),
                        token_id: Some(*id),
                    });
                }
                Segment::Text(text) => {
                    if text.is_empty() {
                        continue;
                    }
                    let start = buffer.len();
                    buffer.push_str(&self.normalize(text));
                    let end = buffer.len();
                    splits.push(Split {
                        range: start..end,
                        token_id: None,
                    });
                }
            }
        }

        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::from_text(text),
            Cow::Owned(s) => {
                let len = s.len();
                PreTokenizedString::new(
                    s,
                    vec![Split {
                        range: 0..len,
                        token_id: None,
                    }],
                )
            }
        }
    }

    /// 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),
        }
    }
}