toklen 0.2.0

A single-threaded, lightweight, and fast token counter.
Documentation
#[cfg(feature = "nfc")]
mod nfc;

#[cfg(feature = "nfd")]
mod nfd;

#[cfg(feature = "nfkc")]
mod nfkc;

#[cfg(feature = "nfkd")]
mod nfkd;

#[cfg(feature = "lowercase")]
mod lowercase;

#[cfg(feature = "strip")]
mod strip;

#[cfg(feature = "replace")]
mod replace;

#[cfg(feature = "prepend")]
mod prepend;

#[cfg(feature = "strip_accents")]
mod strip_accents;

#[cfg(feature = "byte_level")]
mod byte_level;

// ── NormalizerConfig ───────────────────────────────────────────────────────

use serde::Deserialize;

/// Normalizer pipeline step from `tokenizer.json`.
#[derive(Clone, Debug, Deserialize)]
#[serde(tag = "type")]
#[non_exhaustive]
pub enum NormalizerConfig {
    Sequence {
        #[serde(default)]
        normalizers: Vec<NormalizerConfig>,
    },
    #[cfg(feature = "nfc")]
    NFC,
    #[cfg(feature = "nfd")]
    NFD,
    #[cfg(feature = "nfkc")]
    NFKC,
    #[cfg(feature = "nfkd")]
    NFKD,
    #[cfg(feature = "lowercase")]
    Lowercase,
    #[cfg(feature = "strip_accents")]
    StripAccents,
    #[cfg(feature = "strip")]
    Strip {
        #[serde(default)]
        strip_left: bool,
        #[serde(default)]
        strip_right: bool,
    },
    #[cfg(feature = "replace")]
    Replace {
        pattern: replace::ReplacePattern,
        #[serde(default)]
        content: String,
    },
    #[cfg(feature = "prepend")]
    Prepend {
        #[serde(default)]
        prepend: String,
    },
    #[cfg(feature = "byte_level")]
    #[serde(rename = "ByteLevel")]
    ByteLevel,

    // ------------------------------------------------------------------------
    // Unsupported settings.
    #[serde(untagged)]
    Other(serde_json::Value),
    // BertNormalizer,
}

// ── Normalizer enum ───────────────────────────────────────────────────────

use std::borrow::Cow;

/// Errors from constructing a normalizer.
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum Error {
    #[error("unsupported normalizer type: {0}")]
    Unsupported(String),
    #[error("invalid normalizer config: {0}")]
    InvalidConfig(String),
}

/// A compiled normalizer ready for use.
#[derive(Debug)]
#[non_exhaustive]
pub enum Normalizer {
    #[cfg(feature = "nfc")]
    NFC,
    #[cfg(feature = "nfd")]
    NFD,
    #[cfg(feature = "nfkc")]
    NFKC,
    #[cfg(feature = "nfkd")]
    NFKD,
    #[cfg(feature = "lowercase")]
    Lowercase,
    #[cfg(feature = "byte_level")]
    ByteLevel,
    #[cfg(feature = "strip_accents")]
    StripAccents,
    #[cfg(feature = "strip")]
    Strip(strip::Strip),
    #[cfg(feature = "replace")]
    Replace(replace::Replace),
    #[cfg(feature = "prepend")]
    Prepend(prepend::Prepend),
    Sequence(Vec<Normalizer>),
}

impl Normalizer {
    /// Build a normalizer from its JSON configuration.
    pub fn from_config(config: NormalizerConfig) -> Result<Self, Error> {
        match config {
            #[cfg(feature = "nfc")]
            NormalizerConfig::NFC => Ok(Self::NFC),
            #[cfg(feature = "nfd")]
            NormalizerConfig::NFD => Ok(Self::NFD),
            #[cfg(feature = "nfkc")]
            NormalizerConfig::NFKC => Ok(Self::NFKC),
            #[cfg(feature = "nfkd")]
            NormalizerConfig::NFKD => Ok(Self::NFKD),
            #[cfg(feature = "lowercase")]
            NormalizerConfig::Lowercase => Ok(Self::Lowercase),
            #[cfg(feature = "byte_level")]
            NormalizerConfig::ByteLevel => Ok(Self::ByteLevel),
            #[cfg(feature = "strip_accents")]
            NormalizerConfig::StripAccents => Ok(Self::StripAccents),
            #[cfg(feature = "strip")]
            NormalizerConfig::Strip {
                strip_left,
                strip_right,
            } => Ok(Self::Strip(strip::Strip::new(strip_left, strip_right))),
            #[cfg(feature = "replace")]
            NormalizerConfig::Replace { pattern, content } => {
                Ok(Self::Replace(replace::Replace::new(pattern, &content)?))
            }
            #[cfg(feature = "prepend")]
            NormalizerConfig::Prepend { prepend } => {
                Ok(Self::Prepend(prepend::Prepend::new(prepend)))
            }
            NormalizerConfig::Sequence { normalizers } => {
                let mut steps: Vec<Self> = Vec::with_capacity(normalizers.len());
                for normalizer in normalizers {
                    steps.push(Self::from_config(normalizer)?);
                }
                Ok(Self::Sequence(steps))
            }
            NormalizerConfig::Other(v) => Err(Error::Unsupported(v.to_string())),
        }
    }

    /// Normalize `input`, returning `Cow::Borrowed` when unchanged.
    pub fn normalize<'a>(&self, input: &'a str) -> Cow<'a, str> {
        match self {
            #[cfg(feature = "nfc")]
            Self::NFC => nfc::normalize(input),
            #[cfg(feature = "nfd")]
            Self::NFD => nfd::normalize(input),
            #[cfg(feature = "nfkc")]
            Self::NFKC => nfkc::normalize(input),
            #[cfg(feature = "nfkd")]
            Self::NFKD => nfkd::normalize(input),
            #[cfg(feature = "lowercase")]
            Self::Lowercase => lowercase::normalize(input),
            #[cfg(feature = "byte_level")]
            Self::ByteLevel => byte_level::normalize(input),
            #[cfg(feature = "strip_accents")]
            Self::StripAccents => strip_accents::normalize(input),
            #[cfg(feature = "strip")]
            Self::Strip(n) => n.normalize(input),
            #[cfg(feature = "replace")]
            Self::Replace(r) => r.normalize(input),
            #[cfg(feature = "prepend")]
            Self::Prepend(p) => p.normalize(input),
            Self::Sequence(steps) => {
                let mut current: Cow<'_, str> = Cow::Borrowed(input);
                for step in steps {
                    current = match current {
                        Cow::Borrowed(s) => step.normalize(s),
                        Cow::Owned(s) => step.normalize(&s).into_owned().into(),
                    };
                }
                current
            }
        }
    }
}