toklen 0.2.0

A single-threaded, lightweight, and fast token counter.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87

#![allow(clippy::large_enum_variant)]

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

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

// ── ModelConfig ─────────────────────────────────────────────────────────────

use serde::Deserialize;

/// Tokenization model from `tokenizer.json`.
#[derive(Clone, Debug, Deserialize)]
#[serde(tag = "type")]
pub enum ModelConfig {
    #[cfg(feature = "bpe")]
    BPE(Box<bpe::Bpe>),

    #[cfg(feature = "word_level")]
    WordLevel(word_level::WordLevel),

    // ------------------------------------------------------------------------
    // Unsupported settings.
    #[serde(untagged)]
    Other(serde_json::Value),
    // WordPiece {
    //     #[serde(default)]
    //     vocab: HashMap<String, u32>,
    //     #[serde(default)]
    //     unk_token: String,
    //     #[serde(default)]
    //     continuing_subword_prefix: String,
    //     max_input_chars_per_word: Option<u64>,
    // },
    // Unigram {
    //     #[serde(default)]
    //     vocab: Vec<(String, f64)>,
    //     unk_id: Option<u32>,
    //     #[serde(default)]
    //     byte_fallback: bool,
    // },
}

// ── Model ───────────────────────────────────────────────────────────────────

/// A constructed tokenization model ready for encoding.
#[derive(Debug)]
#[non_exhaustive]
pub enum Model {
    #[cfg(feature = "bpe")]
    BPE(bpe::Bpe),
    #[cfg(feature = "word_level")]
    WordLevel(word_level::WordLevel),
    #[allow(unused)]
    #[cfg(not(any(feature = "bpe", feature = "word_level")))]
    Placeholder,
}

impl Model {
    /// Build a model from its JSON configuration.
    #[inline]
    pub fn from_config(config: ModelConfig) -> Result<Self, String> {
        match config {
            #[cfg(feature = "bpe")]
            ModelConfig::BPE(m) => Ok(Self::BPE(*m)),
            #[cfg(feature = "word_level")]
            ModelConfig::WordLevel(m) => Ok(Self::WordLevel(m)),
            ModelConfig::Other(value) => {
                let info = value.to_string();
                Err(format!("unsupported model type: `{info}`"))
            }
        }
    }

    #[inline]
    pub fn count_tokens(&self, input: &str) -> Result<usize, String> {
        match self {
            #[cfg(feature = "bpe")]
            Self::BPE(m) => m.count_tokens(input),
            #[cfg(feature = "word_level")]
            Self::WordLevel(m) => m.count_tokens(input),
            #[cfg(not(any(feature = "bpe", feature = "word_level")))]
            Self::Placeholder => unreachable!(),
        }
    }
}