deepstrike_tokenizer/
lib.rs

1//! # DeepStrike Tokenizer
2//!
3//! Multi-model token counting engine.
4
5use tiktoken_rs::{CoreBPE, cl100k_base, o200k_base};
6
7/// Supported tokenizer backends.
8pub enum TokenizerBackend {
9    /// cl100k_base — GPT-4, GPT-3.5
10    Cl100k,
11    /// o200k_base — GPT-4o
12    O200k,
13}
14
15/// Token counter with cached BPE instance.
16pub struct Tokenizer {
17    bpe: CoreBPE,
18}
19
20impl Tokenizer {
21    pub fn new(backend: TokenizerBackend) -> Self {
22        let bpe = match backend {
23            TokenizerBackend::Cl100k => cl100k_base().expect("failed to load cl100k_base"),
24            TokenizerBackend::O200k => o200k_base().expect("failed to load o200k_base"),
25        };
26        Self { bpe }
27    }
28
29    /// Count tokens in text.
30    pub fn count(&self, text: &str) -> u32 {
31        self.bpe.encode_ordinary(text).len() as u32
32    }
33
34    /// Count tokens for multiple texts.
35    pub fn count_batch(&self, texts: &[&str]) -> Vec<u32> {
36        texts.iter().map(|t| self.count(t)).collect()
37    }
38
39    /// Truncate text to fit within a token budget.
40    /// Returns the longest prefix that fits.
41    pub fn truncate<'a>(&self, text: &'a str, max_tokens: u32) -> &'a str {
42        let tokens = self.bpe.encode_ordinary(text);
43        if tokens.len() as u32 <= max_tokens {
44            return text;
45        }
46
47        // Binary search for the right byte boundary
48        let target = &tokens[..max_tokens as usize];
49        let decoded = self.bpe.decode(target.to_vec()).unwrap_or_default();
50        let byte_len = decoded.len().min(text.len());
51
52        // Ensure we land on a valid UTF-8 boundary
53        let mut end = byte_len;
54        while end > 0 && !text.is_char_boundary(end) {
55            end -= 1;
56        }
57        &text[..end]
58    }
59}
60
61/// Quick helper: count tokens with cl100k_base.
62pub fn count_tokens(text: &str) -> u32 {
63    // Note: in production, cache the Tokenizer instance
64    let t = Tokenizer::new(TokenizerBackend::Cl100k);
65    t.count(text)
66}
67
68#[cfg(test)]
69mod tests {
70    use super::*;
71
72    #[test]
73    fn counts_tokens() {
74        let t = Tokenizer::new(TokenizerBackend::Cl100k);
75        let count = t.count("Hello, world!");
76        assert!(count > 0);
77        assert!(count < 10);
78    }
79
80    #[test]
81    fn truncate_respects_budget() {
82        let t = Tokenizer::new(TokenizerBackend::Cl100k);
83        let text = "The quick brown fox jumps over the lazy dog. ".repeat(100);
84        let truncated = t.truncate(&text, 10);
85        assert!(t.count(truncated) <= 10);
86        assert!(!truncated.is_empty());
87    }
88
89    #[test]
90    fn batch_counting() {
91        let t = Tokenizer::new(TokenizerBackend::Cl100k);
92        let counts = t.count_batch(&["hello", "world", "foo bar baz"]);
93        assert_eq!(counts.len(), 3);
94        assert!(counts.iter().all(|&c| c > 0));
95    }
96}
deepstrike_tokenizer/lib.rs

deepstrike_tokenizer/
lib.rs
