cargo_context_core/
tokenize.rs

1//! Tokenizers.
2//!
3//! Dispatch matrix:
4//!
5//! | Variant               | Backend              | Fidelity                              |
6//! | --------------------- | -------------------- | ------------------------------------- |
7//! | `HfLlama3 { path }`   | HuggingFace `tokenizers` (lazy-loaded from a local `tokenizer.json`) | Exact, matches Llama3 billing |
8//! | `TiktokenCl100k`      | `tiktoken-rs`        | Exact (GPT-4, GPT-4o)                 |
9//! | `TiktokenO200k`       | `tiktoken-rs`        | Exact (GPT-5 family)                  |
10//! | `Claude`              | `tiktoken-rs`        | Approximation via `cl100k_base`       |
11//! | `Llama3`              | calibrated           | ~3.5 chars/token (mixed text)         |
12//! | `Llama2`              | calibrated           | ~3.3 chars/token                      |
13//! | `CharsDiv4`           | built-in             | Offline fallback                      |
14//!
15//! `HfLlama3` lets users point at any local `tokenizer.json` (Meta-Llama-3,
16//! Mistral, Qwen, etc. — anything the HuggingFace `tokenizers` crate can
17//! load). The file is loaded once per path and cached in a process-global
18//! `HashMap`, so repeated `count()` calls pay tokenization cost only.
19//!
20//! A missing or unparseable vocab path silently falls back to the calibrated
21//! heuristic; the scrubber/budget never aborts a pack build because the user
22//! pointed at a bad path.
23
24use std::collections::HashMap;
25use std::path::{Path, PathBuf};
26use std::sync::{Arc, Mutex, OnceLock};
27
28use tiktoken_rs::CoreBPE;
29
30#[derive(Debug, Default, Clone, PartialEq, Eq)]
31pub enum Tokenizer {
32    #[default]
33    Llama3,
34    Llama2,
35    TiktokenCl100k,
36    TiktokenO200k,
37    Claude,
38    CharsDiv4,
39    /// Exact counting via a locally-available `tokenizer.json` (HuggingFace
40    /// format). Pointed at Meta-Llama-3's vocab it produces bill-accurate
41    /// counts for Llama-family models; works equally well with any other
42    /// tokenizer.json.
43    HfLlama3 {
44        vocab_path: PathBuf,
45    },
46}
47
48impl Tokenizer {
49    pub fn label(&self) -> &'static str {
50        match self {
51            Tokenizer::Llama3 => "llama3",
52            Tokenizer::Llama2 => "llama2",
53            Tokenizer::TiktokenCl100k => "tiktoken-cl100k",
54            Tokenizer::TiktokenO200k => "tiktoken-o200k",
55            Tokenizer::Claude => "claude",
56            Tokenizer::CharsDiv4 => "chars-div-4",
57            Tokenizer::HfLlama3 { .. } => "hf-llama3",
58        }
59    }
60
61    /// Estimate tokens for `text`. Never panics; on backend init failure,
62    /// falls back to the calibrated or `CharsDiv4` heuristic.
63    pub fn count(&self, text: &str) -> usize {
64        if text.is_empty() {
65            return 0;
66        }
67        match self {
68            Tokenizer::TiktokenCl100k | Tokenizer::Claude => cl100k()
69                .map(|bpe| bpe.encode_with_special_tokens(text).len())
70                .unwrap_or_else(|| chars_div_4(text)),
71            Tokenizer::TiktokenO200k => o200k()
72                .map(|bpe| bpe.encode_with_special_tokens(text).len())
73                .unwrap_or_else(|| chars_div_4(text)),
74            Tokenizer::Llama3 => calibrated(text, 3.5),
75            Tokenizer::Llama2 => calibrated(text, 3.3),
76            Tokenizer::CharsDiv4 => chars_div_4(text),
77            Tokenizer::HfLlama3 { vocab_path } => count_hf(vocab_path, text),
78        }
79    }
80}
81
82/// `ceil(chars / 4)`.
83fn chars_div_4(text: &str) -> usize {
84    text.chars().count().div_ceil(4)
85}
86
87/// `ceil(chars / chars_per_token)`.
88fn calibrated(text: &str, chars_per_token: f64) -> usize {
89    let chars = text.chars().count() as f64;
90    (chars / chars_per_token).ceil() as usize
91}
92
93fn cl100k() -> Option<&'static CoreBPE> {
94    static CL100K: OnceLock<Option<CoreBPE>> = OnceLock::new();
95    CL100K
96        .get_or_init(|| tiktoken_rs::cl100k_base().ok())
97        .as_ref()
98}
99
100fn o200k() -> Option<&'static CoreBPE> {
101    static O200K: OnceLock<Option<CoreBPE>> = OnceLock::new();
102    O200K
103        .get_or_init(|| tiktoken_rs::o200k_base().ok())
104        .as_ref()
105}
106
107/// Process-global cache of loaded HuggingFace tokenizers keyed by absolute
108/// vocab path. Loading is expensive (parses a multi-MB JSON and compiles
109/// regex patterns); this makes repeated `count()` calls cheap.
110fn hf_cache() -> &'static Mutex<HashMap<PathBuf, Arc<tokenizers::Tokenizer>>> {
111    static CACHE: OnceLock<Mutex<HashMap<PathBuf, Arc<tokenizers::Tokenizer>>>> = OnceLock::new();
112    CACHE.get_or_init(|| Mutex::new(HashMap::new()))
113}
114
115fn load_hf(path: &Path) -> Option<Arc<tokenizers::Tokenizer>> {
116    let key = path.to_path_buf();
117    let mut cache = hf_cache().lock().ok()?;
118    if let Some(t) = cache.get(&key) {
119        return Some(Arc::clone(t));
120    }
121    let t = tokenizers::Tokenizer::from_file(path).ok()?;
122    let arc = Arc::new(t);
123    cache.insert(key, Arc::clone(&arc));
124    Some(arc)
125}
126
127fn count_hf(vocab_path: &Path, text: &str) -> usize {
128    match load_hf(vocab_path) {
129        Some(t) => match t.encode(text, false) {
130            Ok(enc) => enc.len(),
131            Err(_) => calibrated(text, 3.5),
132        },
133        None => calibrated(text, 3.5),
134    }
135}
136
137#[cfg(test)]
138mod tests {
139    use super::*;
140
141    #[test]
142    fn empty_string_counts_zero() {
143        for t in [
144            Tokenizer::Llama3,
145            Tokenizer::Llama2,
146            Tokenizer::TiktokenCl100k,
147            Tokenizer::TiktokenO200k,
148            Tokenizer::Claude,
149            Tokenizer::CharsDiv4,
150            Tokenizer::HfLlama3 {
151                vocab_path: PathBuf::from("/nonexistent"),
152            },
153        ] {
154            assert_eq!(t.count(""), 0, "tokenizer {t:?}");
155        }
156    }
157
158    #[test]
159    fn nonempty_counts_at_least_one() {
160        assert!(Tokenizer::Llama3.count("hello world") >= 1);
161        assert!(Tokenizer::TiktokenCl100k.count("hello world") >= 1);
162    }
163
164    #[test]
165    fn cl100k_matches_known_tokenization() {
166        let n = Tokenizer::TiktokenCl100k.count("hello world");
167        assert_eq!(n, 2, "cl100k count for 'hello world' should be 2, got {n}");
168    }
169
170    #[test]
171    fn o200k_matches_known_tokenization() {
172        let n = Tokenizer::TiktokenO200k.count("hello world");
173        assert_eq!(n, 2, "o200k count for 'hello world' should be 2, got {n}");
174    }
175
176    #[test]
177    fn claude_aliases_cl100k() {
178        let s = "The quick brown fox jumps over the lazy dog.";
179        assert_eq!(
180            Tokenizer::Claude.count(s),
181            Tokenizer::TiktokenCl100k.count(s)
182        );
183    }
184
185    #[test]
186    fn llama_calibration_is_reasonable() {
187        let text = "The quick brown fox jumps over the lazy dog. ".repeat(10);
188        let n = Tokenizer::Llama3.count(&text);
189        let chars = text.chars().count();
190        let expected_low = (chars as f64 / 4.0) as usize;
191        let expected_high = (chars as f64 / 3.0) as usize;
192        assert!(
193            n >= expected_low && n <= expected_high,
194            "llama3 count {n} outside sane range [{expected_low}, {expected_high}] for {chars} chars"
195        );
196    }
197
198    #[test]
199    fn label_is_stable() {
200        // Stable labels are part of the pack's schema; don't change without
201        // a schema version bump.
202        assert_eq!(Tokenizer::Llama3.label(), "llama3");
203        assert_eq!(Tokenizer::TiktokenCl100k.label(), "tiktoken-cl100k");
204        assert_eq!(Tokenizer::TiktokenO200k.label(), "tiktoken-o200k");
205        assert_eq!(Tokenizer::Claude.label(), "claude");
206        assert_eq!(
207            Tokenizer::HfLlama3 {
208                vocab_path: PathBuf::from("/nowhere")
209            }
210            .label(),
211            "hf-llama3"
212        );
213    }
214
215    #[test]
216    fn hf_llama3_missing_path_falls_back_gracefully() {
217        let t = Tokenizer::HfLlama3 {
218            vocab_path: PathBuf::from("/definitely/not/a/real/path/tokenizer.json"),
219        };
220        let n = t.count("hello world");
221        // Falls back to calibrated ~3.5 chars/token → "hello world" = 11 chars → 4 tokens.
222        assert!(n >= 1, "fallback should still produce a positive count");
223    }
224
225    #[test]
226    fn hf_llama3_loads_real_tokenizer_without_falling_back() {
227        // Build a minimal BPE tokenizer, persist it, load via HfLlama3.
228        // The default BPE has an empty vocab so it counts 0 tokens — what we
229        // verify is that the loader did not fall back to the calibrated
230        // heuristic (which would give a positive count for non-empty input).
231        use tokenizers::Tokenizer as HfTokenizer;
232        use tokenizers::models::bpe::BPE;
233
234        let tmp = tempfile::tempdir().unwrap();
235        let path = tmp.path().join("tokenizer.json");
236        let hf = HfTokenizer::new(BPE::default());
237        hf.save(&path, false).expect("save tokenizer fixture");
238
239        let our = Tokenizer::HfLlama3 {
240            vocab_path: path.clone(),
241        };
242        let fallback = Tokenizer::HfLlama3 {
243            vocab_path: PathBuf::from("/does/not/exist"),
244        };
245
246        // Real loader: 0 (empty vocab). Fallback: calibrated chars/3.5 > 0.
247        assert_eq!(
248            our.count("hello"),
249            0,
250            "real loader should short-circuit on empty vocab"
251        );
252        assert!(
253            fallback.count("hello") >= 1,
254            "missing path must fall back to calibrated"
255        );
256    }
257}
cargo_context_core/tokenize.rs

cargo_context_core/
tokenize.rs
