runtime/
tokenizer.rs

1//! Tokenizer implementation using HuggingFace tokenizers
2//!
3//! This module provides tokenization using the HuggingFace tokenizers library,
4//! with fallback to a basic tokenizer for simple use cases.
5
6use crate::types::*;
7use std::collections::HashMap;
8use std::path::Path;
9use tokenizers::Tokenizer as HfTokenizer;
10
11/// Special token IDs (defaults, may be overridden by loaded tokenizer)
12pub const BOS_TOKEN_ID: u32 = 1;   // Beginning of sequence
13pub const EOS_TOKEN_ID: u32 = 2;   // End of sequence
14pub const UNK_TOKEN_ID: u32 = 0;   // Unknown token
15pub const PAD_TOKEN_ID: u32 = 3;   // Padding token
16
17/// Special tokens configuration
18#[derive(Debug, Clone)]
19pub struct SpecialTokens {
20    pub bos_token_id: u32,
21    pub eos_token_id: u32,
22    pub unk_token_id: u32,
23    pub pad_token_id: u32,
24}
25
26impl Default for SpecialTokens {
27    fn default() -> Self {
28        Self {
29            bos_token_id: BOS_TOKEN_ID,
30            eos_token_id: EOS_TOKEN_ID,
31            unk_token_id: UNK_TOKEN_ID,
32            pad_token_id: PAD_TOKEN_ID,
33        }
34    }
35}
36
37/// Tokenizer implementation that wraps HuggingFace tokenizers
38pub struct Tokenizer {
39    inner: TokenizerBackend,
40    special_tokens: SpecialTokens,
41}
42
43enum TokenizerBackend {
44    HuggingFace(HfTokenizer),
45    Basic(BasicTokenizer),
46    GGUF(GGUFTokenizerImpl),
47}
48
49impl Tokenizer {
50    /// Create a new tokenizer with basic vocabulary (fallback)
51    pub fn new() -> Self {
52        Self {
53            inner: TokenizerBackend::Basic(BasicTokenizer::new()),
54            special_tokens: SpecialTokens::default(),
55        }
56    }
57
58    /// Load tokenizer from a tokenizer.json file (HuggingFace format)
59    pub fn from_file<P: AsRef<Path>>(path: P) -> ModelResult<Self> {
60        let hf_tokenizer = HfTokenizer::from_file(path.as_ref())
61            .map_err(|e| ModelError::InitializationFailed(format!("Failed to load tokenizer: {}", e)))?;
62
63        // Try to extract special token IDs from the tokenizer
64        let special_tokens = Self::extract_special_tokens(&hf_tokenizer);
65
66        Ok(Self {
67            inner: TokenizerBackend::HuggingFace(hf_tokenizer),
68            special_tokens,
69        })
70    }
71
72    /// Load tokenizer from a model directory (looks for tokenizer.json)
73    pub fn from_model_dir<P: AsRef<Path>>(model_dir: P) -> ModelResult<Self> {
74        let tokenizer_path = model_dir.as_ref().join("tokenizer.json");
75        if tokenizer_path.exists() {
76            Self::from_file(tokenizer_path)
77        } else {
78            // Try tokenizer_config.json for special tokens, use basic tokenizer
79            Ok(Self::new())
80        }
81    }
82
83    /// Create tokenizer from GGUF tokenizer data
84    pub fn from_gguf(gguf_tokenizer: &crate::weight_loader_core::GGUFTokenizer) -> ModelResult<Self> {
85        let impl_tokenizer = GGUFTokenizerImpl::new(gguf_tokenizer);
86        let special_tokens = impl_tokenizer.special_tokens.clone();
87
88        Ok(Self {
89            inner: TokenizerBackend::GGUF(impl_tokenizer),
90            special_tokens,
91        })
92    }
93
94    /// Create tokenizer from ModelWeights (prefers GGUF tokenizer if available)
95    pub fn from_model_weights(weights: &crate::model_core::ModelWeights) -> ModelResult<Self> {
96        // Prefer GGUF tokenizer if available
97        if let Some(ref gguf_tok) = weights.gguf_tokenizer {
98            return Self::from_gguf(gguf_tok);
99        }
100
101        // Fallback to basic tokenizer
102        Ok(Self::new())
103    }
104
105    /// Create tokenizer from pretrained model name (downloads from HuggingFace)
106    /// Note: This requires network access and HuggingFace Hub authentication for gated models
107    pub fn from_pretrained(model_name: &str) -> ModelResult<Self> {
108        // Try to load from local cache first
109        let cache_dir = dirs::cache_dir()
110            .unwrap_or_else(|| std::path::PathBuf::from("."))
111            .join("huggingface")
112            .join("hub");
113
114        // Convert model name to cache path format
115        let model_path = cache_dir.join(format!("models--{}", model_name.replace('/', "--")));
116        let tokenizer_path = model_path.join("snapshots").join("*").join("tokenizer.json");
117
118        // Try to find tokenizer in cache
119        if let Ok(entries) = glob::glob(tokenizer_path.to_str().unwrap_or("")) {
120            for entry in entries.flatten() {
121                if let Ok(tokenizer) = Self::from_file(entry) {
122                    return Ok(tokenizer);
123                }
124            }
125        }
126
127        // Fall back to basic tokenizer
128        Ok(Self::new())
129    }
130
131    /// Extract special token IDs from HuggingFace tokenizer
132    fn extract_special_tokens(hf_tokenizer: &HfTokenizer) -> SpecialTokens {
133        let mut special = SpecialTokens::default();
134
135        // Try to get special token IDs from the tokenizer
136        if let Some(id) = hf_tokenizer.token_to_id("<s>") {
137            special.bos_token_id = id;
138        } else if let Some(id) = hf_tokenizer.token_to_id("<bos>") {
139            special.bos_token_id = id;
140        }
141
142        if let Some(id) = hf_tokenizer.token_to_id("</s>") {
143            special.eos_token_id = id;
144        } else if let Some(id) = hf_tokenizer.token_to_id("<eos>") {
145            special.eos_token_id = id;
146        }
147
148        if let Some(id) = hf_tokenizer.token_to_id("<unk>") {
149            special.unk_token_id = id;
150        }
151
152        if let Some(id) = hf_tokenizer.token_to_id("<pad>") {
153            special.pad_token_id = id;
154        }
155
156        special
157    }
158
159    /// Encode text to token IDs
160    pub fn encode(&self, text: &str) -> Vec<u32> {
161        match &self.inner {
162            TokenizerBackend::HuggingFace(hf) => {
163                match hf.encode(text, false) {
164                    Ok(encoding) => encoding.get_ids().to_vec(),
165                    Err(_) => {
166                        // Fallback to basic encoding
167                        vec![self.special_tokens.bos_token_id, self.special_tokens.eos_token_id]
168                    }
169                }
170            }
171            TokenizerBackend::Basic(basic) => basic.encode(text),
172            TokenizerBackend::GGUF(gguf) => gguf.encode(text),
173        }
174    }
175
176    /// Encode text with special tokens
177    pub fn encode_with_special_tokens(&self, text: &str, add_bos: bool, add_eos: bool) -> Vec<u32> {
178        match &self.inner {
179            TokenizerBackend::HuggingFace(hf) => {
180                match hf.encode(text, add_bos) {
181                    Ok(encoding) => {
182                        let mut ids = encoding.get_ids().to_vec();
183                        if add_eos && !ids.ends_with(&[self.special_tokens.eos_token_id]) {
184                            ids.push(self.special_tokens.eos_token_id);
185                        }
186                        ids
187                    }
188                    Err(_) => {
189                        let mut ids = if add_bos {
190                            vec![self.special_tokens.bos_token_id]
191                        } else {
192                            vec![]
193                        };
194                        if add_eos {
195                            ids.push(self.special_tokens.eos_token_id);
196                        }
197                        ids
198                    }
199                }
200            }
201            TokenizerBackend::Basic(basic) => {
202                let mut ids = basic.encode_raw(text);
203                if add_bos {
204                    ids.insert(0, self.special_tokens.bos_token_id);
205                }
206                if add_eos {
207                    ids.push(self.special_tokens.eos_token_id);
208                }
209                ids
210            }
211            TokenizerBackend::GGUF(gguf) => {
212                let mut ids = gguf.encode_raw(text);
213                if add_bos {
214                    ids.insert(0, self.special_tokens.bos_token_id);
215                }
216                if add_eos {
217                    ids.push(self.special_tokens.eos_token_id);
218                }
219                ids
220            }
221        }
222    }
223
224    /// Decode token IDs back to text
225    pub fn decode(&self, token_ids: &[u32]) -> String {
226        match &self.inner {
227            TokenizerBackend::HuggingFace(hf) => {
228                hf.decode(token_ids, true).unwrap_or_default()
229            }
230            TokenizerBackend::Basic(basic) => basic.decode(token_ids),
231            TokenizerBackend::GGUF(gguf) => gguf.decode(token_ids),
232        }
233    }
234
235    /// Decode token IDs without special token filtering
236    pub fn decode_raw(&self, token_ids: &[u32]) -> String {
237        match &self.inner {
238            TokenizerBackend::HuggingFace(hf) => {
239                hf.decode(token_ids, false).unwrap_or_default()
240            }
241            TokenizerBackend::Basic(basic) => basic.decode(token_ids),
242            TokenizerBackend::GGUF(gguf) => gguf.decode_raw(token_ids),
243        }
244    }
245
246    /// Get vocabulary size
247    pub fn vocab_size(&self) -> usize {
248        match &self.inner {
249            TokenizerBackend::HuggingFace(hf) => hf.get_vocab_size(true),
250            TokenizerBackend::Basic(basic) => basic.vocab_size(),
251            TokenizerBackend::GGUF(gguf) => gguf.vocab_size(),
252        }
253    }
254
255    /// Get special tokens
256    pub fn special_tokens(&self) -> &SpecialTokens {
257        &self.special_tokens
258    }
259
260    /// Get BOS token ID
261    pub fn bos_token_id(&self) -> u32 {
262        self.special_tokens.bos_token_id
263    }
264
265    /// Get EOS token ID
266    pub fn eos_token_id(&self) -> u32 {
267        self.special_tokens.eos_token_id
268    }
269
270    /// Get PAD token ID
271    pub fn pad_token_id(&self) -> u32 {
272        self.special_tokens.pad_token_id
273    }
274
275    /// Get UNK token ID
276    pub fn unk_token_id(&self) -> u32 {
277        self.special_tokens.unk_token_id
278    }
279
280    /// Check if token exists in vocabulary
281    pub fn contains_token(&self, token: &str) -> bool {
282        match &self.inner {
283            TokenizerBackend::HuggingFace(hf) => hf.token_to_id(token).is_some(),
284            TokenizerBackend::Basic(basic) => basic.contains_token(token),
285            TokenizerBackend::GGUF(gguf) => gguf.contains_token(token),
286        }
287    }
288
289    /// Get token ID for a string
290    pub fn token_to_id(&self, token: &str) -> Option<u32> {
291        match &self.inner {
292            TokenizerBackend::HuggingFace(hf) => hf.token_to_id(token),
293            TokenizerBackend::Basic(basic) => basic.token_to_id(token),
294            TokenizerBackend::GGUF(gguf) => gguf.token_to_id(token),
295        }
296    }
297
298    /// Get string for a token ID
299    pub fn id_to_token(&self, id: u32) -> Option<String> {
300        match &self.inner {
301            TokenizerBackend::HuggingFace(hf) => hf.id_to_token(id),
302            TokenizerBackend::Basic(basic) => basic.id_to_token(id).map(|s| s.to_string()),
303            TokenizerBackend::GGUF(gguf) => gguf.id_to_token(id).map(|s| s.to_string()),
304        }
305    }
306}
307
308impl Default for Tokenizer {
309    fn default() -> Self {
310        Self::new()
311    }
312}
313
314/// Basic tokenizer for fallback (same as the original implementation)
315struct BasicTokenizer {
316    vocab: HashMap<String, u32>,
317    id_to_token: HashMap<u32, String>,
318    vocab_size: usize,
319}
320
321impl BasicTokenizer {
322    fn new() -> Self {
323        let mut vocab = HashMap::new();
324        let mut id_to_token = HashMap::new();
325
326        // Add special tokens
327        vocab.insert("<unk>".to_string(), UNK_TOKEN_ID);
328        vocab.insert("<s>".to_string(), BOS_TOKEN_ID);
329        vocab.insert("</s>".to_string(), EOS_TOKEN_ID);
330        vocab.insert("<pad>".to_string(), PAD_TOKEN_ID);
331
332        id_to_token.insert(UNK_TOKEN_ID, "<unk>".to_string());
333        id_to_token.insert(BOS_TOKEN_ID, "<s>".to_string());
334        id_to_token.insert(EOS_TOKEN_ID, "</s>".to_string());
335        id_to_token.insert(PAD_TOKEN_ID, "<pad>".to_string());
336
337        let mut next_id = 4;
338
339        // Add basic vocabulary
340        let basic_vocab = vec![
341            "the", "a", "an", "and", "or", "but", "in", "on", "at", "to", "for", "of", "with", "by",
342            "is", "are", "was", "were", "be", "been", "have", "has", "had", "do", "does", "did",
343            "will", "would", "could", "should", "can", "may", "might", "must",
344            "I", "you", "he", "she", "it", "we", "they", "me", "him", "her", "us", "them",
345            "this", "that", "these", "those", "here", "there", "where", "when", "why", "how",
346            "what", "who", "which", "whose", "all", "some", "any", "no", "not", "yes",
347            "hello", "world", "test", "example", "text", "model", "language", "AI",
348            "quick", "brown", "fox", "dog", "cat", "house", "car", "tree", "book", "water",
349            ".", "!", "?", ",", ";", ":", "'", "\"", "-", "_", "(", ")", "[", "]", "{", "}",
350            "0", "1", "2", "3", "4", "5", "6", "7", "8", "9",
351            "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m",
352            "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z",
353            "A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L", "M",
354            "N", "O", "P", "Q", "R", "S", "T", "U", "V", "W", "X", "Y", "Z",
355            "ing", "ed", "er", "ly", "tion", "ment", "ness", "ity", "ous", "ful",
356            " ", "\n", "\t",
357        ];
358
359        for word in basic_vocab {
360            if !vocab.contains_key(word) {
361                vocab.insert(word.to_string(), next_id);
362                id_to_token.insert(next_id, word.to_string());
363                next_id += 1;
364            }
365        }
366
367        Self {
368            vocab,
369            id_to_token,
370            vocab_size: next_id as usize,
371        }
372    }
373
374    fn encode(&self, text: &str) -> Vec<u32> {
375        let mut tokens = vec![BOS_TOKEN_ID];
376        tokens.extend(self.encode_raw(text));
377        tokens.push(EOS_TOKEN_ID);
378        tokens
379    }
380
381    fn encode_raw(&self, text: &str) -> Vec<u32> {
382        let mut tokens = Vec::new();
383        let words = self.tokenize_text(text);
384
385        for word in words {
386            if let Some(&token_id) = self.vocab.get(&word) {
387                tokens.push(token_id);
388            } else {
389                for ch in word.chars() {
390                    let ch_str = ch.to_string();
391                    if let Some(&token_id) = self.vocab.get(&ch_str) {
392                        tokens.push(token_id);
393                    } else {
394                        tokens.push(UNK_TOKEN_ID);
395                    }
396                }
397            }
398        }
399        tokens
400    }
401
402    fn decode(&self, token_ids: &[u32]) -> String {
403        let mut result = String::new();
404        let mut last_was_char = false;
405
406        for &token_id in token_ids {
407            if token_id == BOS_TOKEN_ID || token_id == EOS_TOKEN_ID || token_id == PAD_TOKEN_ID {
408                continue;
409            }
410
411            if let Some(token) = self.id_to_token.get(&token_id) {
412                let is_single_char = token.len() == 1 && token.chars().next().unwrap().is_alphabetic();
413                let is_whitespace = token.chars().all(|c| c.is_whitespace());
414                let is_punctuation = token.chars().all(|c| c.is_ascii_punctuation());
415
416                if !result.is_empty() && !is_whitespace && !is_punctuation {
417                    if !last_was_char || !is_single_char {
418                        result.push(' ');
419                    }
420                }
421
422                result.push_str(token);
423                last_was_char = is_single_char;
424            } else {
425                if !result.is_empty() {
426                    result.push(' ');
427                }
428                result.push_str("<unk>");
429                last_was_char = false;
430            }
431        }
432
433        result.trim().to_string()
434    }
435
436    fn vocab_size(&self) -> usize {
437        self.vocab_size
438    }
439
440    fn contains_token(&self, token: &str) -> bool {
441        self.vocab.contains_key(token)
442    }
443
444    fn token_to_id(&self, token: &str) -> Option<u32> {
445        self.vocab.get(token).copied()
446    }
447
448    fn id_to_token(&self, id: u32) -> Option<&str> {
449        self.id_to_token.get(&id).map(|s| s.as_str())
450    }
451
452    fn tokenize_text(&self, text: &str) -> Vec<String> {
453        let mut tokens = Vec::new();
454        let mut current_word = String::new();
455
456        for ch in text.chars() {
457            if ch.is_whitespace() {
458                if !current_word.is_empty() {
459                    tokens.push(current_word.clone());
460                    current_word.clear();
461                }
462                tokens.push(ch.to_string());
463            } else if ch.is_ascii_punctuation() {
464                if !current_word.is_empty() {
465                    tokens.push(current_word.clone());
466                    current_word.clear();
467                }
468                tokens.push(ch.to_string());
469            } else {
470                current_word.push(ch);
471            }
472        }
473
474        if !current_word.is_empty() {
475            tokens.push(current_word);
476        }
477
478        tokens
479    }
480}
481
482/// GGUF-based tokenizer implementation
483struct GGUFTokenizerImpl {
484    tokens: Vec<String>,
485    token_to_id: HashMap<String, u32>,
486    id_to_token: HashMap<u32, String>,
487    special_tokens: SpecialTokens,
488    vocab_size: usize,
489}
490
491impl GGUFTokenizerImpl {
492    fn new(gguf_tokenizer: &crate::weight_loader_core::GGUFTokenizer) -> Self {
493        let mut token_to_id = HashMap::new();
494        let mut id_to_token = HashMap::new();
495
496        for (id, token) in gguf_tokenizer.tokens.iter().enumerate() {
497            let id = id as u32;
498            token_to_id.insert(token.clone(), id);
499            id_to_token.insert(id, token.clone());
500        }
501
502        let gguf_special = &gguf_tokenizer.special_tokens;
503        let special_tokens = SpecialTokens {
504            bos_token_id: gguf_special.bos_token_id.unwrap_or(BOS_TOKEN_ID),
505            eos_token_id: gguf_special.eos_token_id.unwrap_or(EOS_TOKEN_ID),
506            unk_token_id: gguf_special.unk_token_id.unwrap_or(UNK_TOKEN_ID),
507            pad_token_id: gguf_special.pad_token_id.unwrap_or(PAD_TOKEN_ID),
508        };
509
510        let vocab_size = gguf_tokenizer.tokens.len();
511
512        Self {
513            tokens: gguf_tokenizer.tokens.clone(),
514            token_to_id,
515            id_to_token,
516            special_tokens,
517            vocab_size,
518        }
519    }
520
521    fn encode(&self, text: &str) -> Vec<u32> {
522        let mut tokens = vec![self.special_tokens.bos_token_id];
523        tokens.extend(self.encode_raw(text));
524        tokens.push(self.special_tokens.eos_token_id);
525        tokens
526    }
527
528    fn encode_raw(&self, text: &str) -> Vec<u32> {
529        // SentencePiece uses ▁ (U+2581) to represent spaces
530        // We need to convert spaces to this character for proper tokenization
531        // The ▁ is prepended to tokens that follow a space (word boundaries)
532        let processed = format!("▁{}", text.replace(' ', "▁"));
533
534        let mut result = Vec::new();
535        let bytes = processed.as_bytes();
536        let mut i = 0;
537
538        while i < bytes.len() {
539            let mut matched = false;
540
541            // Try to match longest token first (greedy)
542            // Start with reasonable max length to avoid O(n^2)
543            let max_len = std::cmp::min(bytes.len() - i, 32);
544
545            for len in (1..=max_len).rev() {
546                if let Ok(substr) = std::str::from_utf8(&bytes[i..i + len]) {
547                    if let Some(&id) = self.token_to_id.get(substr) {
548                        result.push(id);
549                        i += len;
550                        matched = true;
551                        break;
552                    }
553                }
554            }
555
556            if !matched {
557                // Try single byte as fallback
558                let byte = bytes[i];
559
560                // First try the character directly
561                if let Ok(ch) = std::str::from_utf8(&bytes[i..i + 1]) {
562                    if let Some(&id) = self.token_to_id.get(ch) {
563                        result.push(id);
564                        i += 1;
565                        continue;
566                    }
567                }
568
569                // Try byte-level token format: <0xNN>
570                let byte_token = format!("<0x{:02X}>", byte);
571                if let Some(&id) = self.token_to_id.get(&byte_token) {
572                    result.push(id);
573                } else {
574                    // Unknown token
575                    result.push(self.special_tokens.unk_token_id);
576                }
577                i += 1;
578            }
579        }
580
581        result
582    }
583
584    fn decode(&self, token_ids: &[u32]) -> String {
585        let mut result = String::new();
586
587        for &id in token_ids {
588            // Skip special tokens
589            if id == self.special_tokens.bos_token_id
590                || id == self.special_tokens.eos_token_id
591                || id == self.special_tokens.pad_token_id
592            {
593                continue;
594            }
595
596            if let Some(token) = self.id_to_token.get(&id) {
597                result.push_str(token);
598            }
599        }
600
601        // Convert SentencePiece's ▁ back to spaces and trim leading space
602        result.replace('▁', " ").trim_start().to_string()
603    }
604
605    fn decode_raw(&self, token_ids: &[u32]) -> String {
606        let mut result = String::new();
607
608        for &id in token_ids {
609            if let Some(token) = self.id_to_token.get(&id) {
610                result.push_str(token);
611            }
612        }
613
614        // Convert SentencePiece's ▁ back to spaces and trim leading space
615        result.replace('▁', " ").trim_start().to_string()
616    }
617
618    fn vocab_size(&self) -> usize {
619        self.vocab_size
620    }
621
622    fn contains_token(&self, token: &str) -> bool {
623        self.token_to_id.contains_key(token)
624    }
625
626    fn token_to_id(&self, token: &str) -> Option<u32> {
627        self.token_to_id.get(token).copied()
628    }
629
630    fn id_to_token(&self, id: u32) -> Option<&str> {
631        self.id_to_token.get(&id).map(|s| s.as_str())
632    }
633}
634
635/// Batch tokenization for multiple texts
636pub struct BatchTokenizer {
637    tokenizer: Tokenizer,
638}
639
640impl BatchTokenizer {
641    pub fn new(tokenizer: Tokenizer) -> Self {
642        Self { tokenizer }
643    }
644
645    /// Encode multiple texts with padding
646    pub fn encode_batch(&self, texts: &[&str], max_length: Option<usize>) -> (Vec<Vec<u32>>, Vec<Vec<bool>>) {
647        let mut encoded_batch = Vec::new();
648        let mut attention_masks = Vec::new();
649
650        for text in texts {
651            let tokens = self.tokenizer.encode(text);
652            encoded_batch.push(tokens);
653        }
654
655        let max_len = max_length.unwrap_or_else(|| {
656            encoded_batch.iter().map(|tokens| tokens.len()).max().unwrap_or(0)
657        });
658
659        for tokens in &mut encoded_batch {
660            if tokens.len() > max_len {
661                tokens.truncate(max_len);
662            }
663
664            let mut attention_mask = vec![true; tokens.len()];
665
666            while tokens.len() < max_len {
667                tokens.push(self.tokenizer.pad_token_id());
668                attention_mask.push(false);
669            }
670
671            attention_masks.push(attention_mask);
672        }
673
674        (encoded_batch, attention_masks)
675    }
676
677    /// Decode batch of token sequences
678    pub fn decode_batch(&self, token_sequences: &[Vec<u32>]) -> Vec<String> {
679        token_sequences.iter()
680            .map(|tokens| self.tokenizer.decode(tokens))
681            .collect()
682    }
683}
684
685#[cfg(test)]
686mod tests {
687    use super::*;
688
689    #[test]
690    fn test_tokenizer_creation() {
691        let tokenizer = Tokenizer::new();
692        assert!(tokenizer.vocab_size() > 100);
693        assert!(tokenizer.contains_token("the"));
694        assert!(tokenizer.contains_token("<s>"));
695        assert!(tokenizer.contains_token("</s>"));
696    }
697
698    #[test]
699    fn test_special_tokens() {
700        let tokenizer = Tokenizer::new();
701        assert_eq!(tokenizer.token_to_id("<unk>"), Some(UNK_TOKEN_ID));
702        assert_eq!(tokenizer.token_to_id("<s>"), Some(BOS_TOKEN_ID));
703        assert_eq!(tokenizer.token_to_id("</s>"), Some(EOS_TOKEN_ID));
704        assert_eq!(tokenizer.token_to_id("<pad>"), Some(PAD_TOKEN_ID));
705    }
706
707    #[test]
708    fn test_encode_decode_simple() {
709        let tokenizer = Tokenizer::new();
710        let text = "hello world";
711
712        let tokens = tokenizer.encode(text);
713        assert!(tokens.len() > 2);
714
715        let decoded = tokenizer.decode(&tokens);
716        assert!(decoded.contains("hello"));
717        assert!(decoded.contains("world"));
718    }
719
720    #[test]
721    fn test_special_token_ids() {
722        let tokenizer = Tokenizer::new();
723        assert_eq!(tokenizer.bos_token_id(), BOS_TOKEN_ID);
724        assert_eq!(tokenizer.eos_token_id(), EOS_TOKEN_ID);
725        assert_eq!(tokenizer.pad_token_id(), PAD_TOKEN_ID);
726        assert_eq!(tokenizer.unk_token_id(), UNK_TOKEN_ID);
727    }
728
729    #[test]
730    fn test_batch_tokenizer() {
731        let tokenizer = Tokenizer::new();
732        let batch_tokenizer = BatchTokenizer::new(tokenizer);
733
734        let texts = vec!["hello", "hello world", "hello world test"];
735        let (encoded_batch, attention_masks) = batch_tokenizer.encode_batch(&texts, Some(10));
736
737        assert_eq!(encoded_batch.len(), 3);
738        assert_eq!(attention_masks.len(), 3);
739
740        for tokens in &encoded_batch {
741            assert_eq!(tokens.len(), 10);
742        }
743    }
744}
runtime/tokenizer.rs

runtime/
tokenizer.rs
