lean_ctx/core/embeddings/
tokenizer.rs

1//! Minimal WordPiece tokenizer for BERT-style embedding models.
2//!
3//! Implements the standard BERT tokenization pipeline:
4//! 1. Lowercase + accent stripping
5//! 2. Whitespace + punctuation splitting
6//! 3. WordPiece subword tokenization
7//! 4. Special token insertion (`[CLS]`, `[SEP]`)
8//!
9//! Optimized for code search: handles camelCase, snake_case, and common
10//! programming punctuation correctly.
11
12use std::collections::HashMap;
13use std::path::Path;
14
15pub struct WordPieceTokenizer {
16    vocab: HashMap<String, i32>,
17    cls_id: i32,
18    sep_id: i32,
19    pad_id: i32,
20    unk_id: i32,
21    max_word_chars: usize,
22}
23
24#[derive(Debug, Clone)]
25pub struct TokenizedInput {
26    pub input_ids: Vec<i32>,
27    pub attention_mask: Vec<i32>,
28    pub token_type_ids: Vec<i32>,
29}
30
31impl TokenizedInput {
32    /// Pad the input to a fixed length.
33    pub fn pad_to(&mut self, target_len: usize, pad_id: i32) {
34        while self.input_ids.len() < target_len {
35            self.input_ids.push(pad_id);
36            self.attention_mask.push(0);
37            self.token_type_ids.push(0);
38        }
39    }
40}
41
42impl WordPieceTokenizer {
43    /// Load vocabulary from a standard vocab.txt file (one token per line).
44    pub fn from_file(path: &Path) -> anyhow::Result<Self> {
45        let content = std::fs::read_to_string(path)
46            .map_err(|e| anyhow::anyhow!("Failed to read vocab file {}: {}", path.display(), e))?;
47        Self::from_vocab_str(&content)
48    }
49
50    /// Build tokenizer from vocabulary string (one token per line).
51    pub fn from_vocab_str(vocab_str: &str) -> anyhow::Result<Self> {
52        Self::build(vocab_str, true)
53    }
54
55    /// `require_markers`: transformer inputs are framed by `[CLS]`/`[SEP]`, so
56    /// those must exist. Static embedding models (model2vec) average plain
57    /// token vectors and may ship no sentence markers at all; they are encoded
58    /// with [`Self::encode_plain`] only, which never reads the markers.
59    fn build(vocab_str: &str, require_markers: bool) -> anyhow::Result<Self> {
60        let vocab: HashMap<String, i32> = vocab_str
61            .lines()
62            .enumerate()
63            .map(|(i, line)| (line.to_string(), i as i32))
64            .collect();
65
66        let marker = |name: &str| -> anyhow::Result<i32> {
67            match vocab.get(name) {
68                Some(id) => Ok(*id),
69                None if require_markers => anyhow::bail!("Vocabulary missing {name} token"),
70                None => Ok(-1),
71            }
72        };
73        let cls_id = marker("[CLS]")?;
74        let sep_id = marker("[SEP]")?;
75        let pad_id = *vocab
76            .get("[PAD]")
77            .ok_or_else(|| anyhow::anyhow!("Vocabulary missing [PAD] token"))?;
78        let unk_id = *vocab
79            .get("[UNK]")
80            .ok_or_else(|| anyhow::anyhow!("Vocabulary missing [UNK] token"))?;
81
82        Ok(Self {
83            vocab,
84            cls_id,
85            sep_id,
86            pad_id,
87            unk_id,
88            max_word_chars: 200,
89        })
90    }
91
92    /// Encode text into token IDs with `[CLS]` prefix and `[SEP]` suffix.
93    pub fn encode(&self, text: &str, max_len: usize) -> TokenizedInput {
94        let words = self.pre_tokenize(text);
95        let mut ids = vec![self.cls_id];
96
97        for word in &words {
98            if ids.len() >= max_len - 1 {
99                break;
100            }
101            let subword_ids = self.wordpiece_encode(word);
102            for id in subword_ids {
103                if ids.len() >= max_len - 1 {
104                    break;
105                }
106                ids.push(id);
107            }
108        }
109
110        ids.push(self.sep_id);
111
112        let len = ids.len();
113        TokenizedInput {
114            input_ids: ids,
115            attention_mask: vec![1; len],
116            token_type_ids: vec![0; len],
117        }
118    }
119
120    /// Token ids of `text` without sentence markers and without `[UNK]` —
121    /// the input of a static (model2vec) embedding, which averages token
122    /// vectors: an unknown piece carries no meaning to average in.
123    pub fn encode_plain(&self, text: &str, max_tokens: usize) -> Vec<i32> {
124        let mut ids = Vec::new();
125        for word in self.pre_tokenize(text) {
126            for id in self.wordpiece_encode(&word) {
127                if ids.len() >= max_tokens {
128                    return ids;
129                }
130                if id != self.unk_id {
131                    ids.push(id);
132                }
133            }
134        }
135        ids
136    }
137
138    pub fn pad_id(&self) -> i32 {
139        self.pad_id
140    }
141
142    pub fn vocab_size(&self) -> usize {
143        self.vocab.len()
144    }
145
146    /// Split text into word-level tokens.
147    /// Handles: whitespace splitting, punctuation splitting,
148    /// camelCase splitting, underscore-separated identifiers, then lowercase.
149    fn pre_tokenize(&self, text: &str) -> Vec<String> {
150        let mut words = Vec::new();
151        let mut current = String::new();
152
153        for ch in text.chars() {
154            if ch.is_whitespace() {
155                if !current.is_empty() {
156                    words.extend(self.split_identifier(&current));
157                    current.clear();
158                }
159            } else if is_bert_punctuation(ch) {
160                if !current.is_empty() {
161                    words.extend(self.split_identifier(&current));
162                    current.clear();
163                }
164                words.push(ch.to_string());
165            } else {
166                current.push(ch);
167            }
168        }
169        if !current.is_empty() {
170            words.extend(self.split_identifier(&current));
171        }
172
173        words.iter().map(|w| w.to_lowercase()).collect()
174    }
175
176    /// Split programming identifiers (camelCase, snake_case) into subwords.
177    /// Called BEFORE lowercasing so case boundaries are still visible.
178    fn split_identifier(&self, word: &str) -> Vec<String> {
179        let lower = word.to_lowercase();
180        if self.vocab.contains_key(&lower) {
181            return vec![word.to_string()];
182        }
183
184        let mut parts = Vec::new();
185        let mut current = String::new();
186        let chars: Vec<char> = word.chars().collect();
187
188        for (i, &ch) in chars.iter().enumerate() {
189            if ch == '_' || ch == '-' {
190                if !current.is_empty() {
191                    parts.push(current.clone());
192                    current.clear();
193                }
194            } else if i > 0 && ch.is_ascii_uppercase() && chars[i - 1].is_ascii_lowercase() {
195                if !current.is_empty() {
196                    parts.push(current.clone());
197                    current.clear();
198                }
199                current.push(ch);
200            } else {
201                current.push(ch);
202            }
203        }
204        if !current.is_empty() {
205            parts.push(current);
206        }
207
208        if parts.is_empty() {
209            vec![word.to_string()]
210        } else {
211            parts
212        }
213    }
214
215    /// Apply WordPiece algorithm to a single word.
216    fn wordpiece_encode(&self, word: &str) -> Vec<i32> {
217        if word.chars().count() > self.max_word_chars {
218            return vec![self.unk_id];
219        }
220
221        let chars: Vec<char> = word.chars().collect();
222        let mut tokens = Vec::new();
223        let mut start = 0;
224
225        while start < chars.len() {
226            let mut end = chars.len();
227            let mut matched = false;
228
229            while start < end {
230                let substr: String = chars[start..end].iter().collect();
231                let candidate = if start > 0 {
232                    format!("##{substr}")
233                } else {
234                    substr
235                };
236
237                if let Some(&id) = self.vocab.get(&candidate) {
238                    tokens.push(id);
239                    matched = true;
240                    start = end;
241                    break;
242                }
243                end -= 1;
244            }
245
246            if !matched {
247                tokens.push(self.unk_id);
248                start += 1;
249            }
250        }
251
252        tokens
253    }
254}
255
256/// BERT-style punctuation detection.
257fn is_bert_punctuation(ch: char) -> bool {
258    if ch.is_ascii() {
259        matches!(
260            ch,
261            '!' | '"'
262                | '#'
263                | '$'
264                | '%'
265                | '&'
266                | '\''
267                | '('
268                | ')'
269                | '*'
270                | '+'
271                | ','
272                | '-'
273                | '.'
274                | '/'
275                | ':'
276                | ';'
277                | '<'
278                | '='
279                | '>'
280                | '?'
281                | '@'
282                | '['
283                | '\\'
284                | ']'
285                | '^'
286                | '_'
287                | '`'
288                | '{'
289                | '|'
290                | '}'
291                | '~'
292        )
293    } else {
294        ch.is_ascii_punctuation()
295    }
296}
297
298/// Wrapper for HuggingFace `tokenizer.json` files.
299///
300/// Parses the JSON to extract the vocabulary map, then delegates to
301/// the existing `WordPieceTokenizer` for actual tokenization.
302/// Supports both WordPiece and BPE vocab formats (both map tokens → IDs).
303pub struct HfTokenizerWrapper {
304    inner: WordPieceTokenizer,
305}
306
307impl HfTokenizerWrapper {
308    /// Load from a HuggingFace `tokenizer.json` file.
309    pub fn from_file(path: &Path) -> anyhow::Result<Self> {
310        let content = std::fs::read_to_string(path).map_err(|e| {
311            anyhow::anyhow!("Failed to read tokenizer.json {}: {}", path.display(), e)
312        })?;
313        Self::from_json(&content, true)
314    }
315
316    /// Load the tokenizer of a static embedding model, which may have no
317    /// `[CLS]`/`[SEP]`; use [`Self::encode_plain`] with it.
318    pub fn from_file_static(path: &Path) -> anyhow::Result<Self> {
319        let content = std::fs::read_to_string(path).map_err(|e| {
320            anyhow::anyhow!("Failed to read tokenizer.json {}: {}", path.display(), e)
321        })?;
322        Self::from_json(&content, false)
323    }
324
325    fn from_json(json_str: &str, require_markers: bool) -> anyhow::Result<Self> {
326        let parsed: serde_json::Value = serde_json::from_str(json_str)
327            .map_err(|e| anyhow::anyhow!("Invalid tokenizer.json: {e}"))?;
328
329        let vocab_obj = parsed
330            .get("model")
331            .and_then(|m| m.get("vocab"))
332            .and_then(|v| v.as_object())
333            .ok_or_else(|| anyhow::anyhow!("tokenizer.json missing model.vocab object"))?;
334
335        let mut vocab_lines: Vec<(String, i32)> = vocab_obj
336            .iter()
337            .filter_map(|(token, id)| id.as_i64().map(|id| (token.clone(), id as i32)))
338            .collect();
339        vocab_lines.sort_by_key(|(_, id)| *id);
340
341        // GL #397 / #498: BPE/RoBERTa-style tokenizer.json files ship
342        // different special token names than BERT's `[CLS]`, `[SEP]`, etc.
343        // Remap them here so the downstream `WordPieceTokenizer` can find
344        // its expected sentinel tokens — the sort-order preserves the
345        // correct original IDs.
346        for (token, _) in &mut vocab_lines {
347            let mapped: &str = match token.as_str() {
348                "<s>" => "[CLS]",
349                "</s>" => "[SEP]",
350                "<pad>" => "[PAD]",
351                "<unk>" => "[UNK]",
352                // <mask> is unused by WordPieceTokenizer; map it for completeness.
353                "<mask>" => "[MASK]",
354                _ => continue,
355            };
356            *token = mapped.to_string();
357        }
358
359        let vocab_str: String = vocab_lines
360            .into_iter()
361            .map(|(token, _)| token)
362            .collect::<Vec<_>>()
363            .join("\n");
364
365        let inner = WordPieceTokenizer::build(&vocab_str, require_markers)?;
366        Ok(Self { inner })
367    }
368
369    pub fn encode_plain(&self, text: &str, max_tokens: usize) -> Vec<i32> {
370        self.inner.encode_plain(text, max_tokens)
371    }
372
373    pub fn vocab_size(&self) -> usize {
374        self.inner.vocab_size()
375    }
376
377    pub fn encode(&self, text: &str, max_len: usize) -> TokenizedInput {
378        self.inner.encode(text, max_len)
379    }
380}
381
382#[cfg(test)]
383mod tests {
384    use super::*;
385
386    fn test_vocab() -> WordPieceTokenizer {
387        let vocab = "[PAD]\n[UNK]\n[CLS]\n[SEP]\nhello\nworld\nfn\nvalidate\ntoken\n##s\n##ing\nauth\n##enticate\nuser\nhandle\nrequest\n##er\nprocess\ndata\n.\n,\n(\n)\n{";
388        WordPieceTokenizer::from_vocab_str(vocab).unwrap()
389    }
390
391    #[test]
392    fn encode_basic() {
393        let tok = test_vocab();
394        let input = tok.encode("hello world", 512);
395        assert_eq!(input.input_ids[0], tok.cls_id);
396        assert_eq!(*input.input_ids.last().unwrap(), tok.sep_id);
397        assert!(input.input_ids.len() >= 4); // [CLS] hello world [SEP]
398    }
399
400    #[test]
401    fn encode_attention_mask() {
402        let tok = test_vocab();
403        let input = tok.encode("hello", 512);
404        assert!(input.attention_mask.iter().all(|&m| m == 1));
405        assert_eq!(input.attention_mask.len(), input.input_ids.len());
406    }
407
408    #[test]
409    fn encode_token_type_ids_are_zero() {
410        let tok = test_vocab();
411        let input = tok.encode("hello", 512);
412        assert!(input.token_type_ids.iter().all(|&t| t == 0));
413    }
414
415    #[test]
416    fn encode_respects_max_len() {
417        let tok = test_vocab();
418        let input = tok.encode("hello world hello world hello world", 6);
419        assert!(input.input_ids.len() <= 6);
420        assert_eq!(input.input_ids[0], tok.cls_id);
421        assert_eq!(*input.input_ids.last().unwrap(), tok.sep_id);
422    }
423
424    #[test]
425    fn wordpiece_subwords() {
426        let tok = test_vocab();
427        // "tokens" should split into "token" + "##s"
428        let ids = tok.wordpiece_encode("tokens");
429        assert_eq!(ids.len(), 2);
430        assert_eq!(ids[0], *tok.vocab.get("token").unwrap());
431        assert_eq!(ids[1], *tok.vocab.get("##s").unwrap());
432    }
433
434    #[test]
435    fn wordpiece_unknown() {
436        let tok = test_vocab();
437        let ids = tok.wordpiece_encode("xyzzyplugh");
438        assert!(ids.contains(&tok.unk_id));
439    }
440
441    #[test]
442    fn pre_tokenize_camel_case() {
443        let tok = test_vocab();
444        let words = tok.pre_tokenize("handleRequest");
445        assert!(words.contains(&"handle".to_string()));
446        assert!(words.contains(&"request".to_string()));
447    }
448
449    #[test]
450    fn pre_tokenize_snake_case() {
451        let tok = test_vocab();
452        let words = tok.pre_tokenize("validate_token");
453        assert!(words.contains(&"validate".to_string()));
454        assert!(words.contains(&"token".to_string()));
455    }
456
457    #[test]
458    fn pre_tokenize_punctuation() {
459        let tok = test_vocab();
460        let words = tok.pre_tokenize("fn(x)");
461        assert!(words.contains(&"fn".to_string()));
462        assert!(words.contains(&"(".to_string()));
463        assert!(words.contains(&")".to_string()));
464    }
465
466    #[test]
467    fn pad_to_extends() {
468        let tok = test_vocab();
469        let mut input = tok.encode("hello", 512);
470        let original_len = input.input_ids.len();
471        input.pad_to(10, tok.pad_id);
472        assert_eq!(input.input_ids.len(), 10);
473        assert_eq!(input.attention_mask[original_len], 0);
474    }
475
476    #[test]
477    fn vocab_size() {
478        let tok = test_vocab();
479        assert_eq!(tok.vocab_size(), 24);
480    }
481
482    #[test]
483    fn empty_input() {
484        let tok = test_vocab();
485        let input = tok.encode("", 512);
486        assert_eq!(input.input_ids.len(), 2); // [CLS] [SEP]
487    }
488
489    #[test]
490    fn bert_punctuation_detection() {
491        assert!(is_bert_punctuation('.'));
492        assert!(is_bert_punctuation('('));
493        assert!(is_bert_punctuation('{'));
494        assert!(!is_bert_punctuation('a'));
495        assert!(!is_bert_punctuation('0'));
496    }
497
498    #[test]
499    fn hf_tokenizer_remaps_bpe_special_tokens() {
500        // BPE/RoBERTa-style tokenizer.json with <s>/</s>/<pad>/<unk>
501        let json = r#"{
502            "version": "1.0",
503            "model": {
504                "type": "BPE",
505                "vocab": {
506                    "<s>": 0, "<pad>": 1, "</s>": 2, "<unk>": 3,
507                    "hello": 4, "world": 5, "fn": 6
508                }
509            }
510        }"#;
511        let tok = HfTokenizerWrapper::from_json(json, true).unwrap();
512
513        // After remapping: <s>→[CLS], <pad>→[PAD], </s>→[SEP], <unk>→[UNK]
514        // IDs come from sort order (0,1,2,3,4,5,6) which matches original IDs.
515        let input = tok.encode("hello world", 512);
516        assert_eq!(
517            input.input_ids[0], 0,
518            "first token should be [CLS] (remapped from <s>)"
519        );
520        assert_eq!(
521            *input.input_ids.last().unwrap(),
522            2,
523            "last token should be [SEP] (remapped from </s>)"
524        );
525        assert_eq!(input.input_ids.len(), 4); // [CLS] hello world [SEP]
526    }
527
528    #[test]
529    fn hf_tokenizer_from_json() {
530        let json = r#"{
531            "version": "1.0",
532            "model": {
533                "type": "WordPiece",
534                "vocab": {
535                    "[PAD]": 0, "[UNK]": 1, "[CLS]": 2, "[SEP]": 3,
536                    "hello": 4, "world": 5, "fn": 6
537                }
538            }
539        }"#;
540        let tok = HfTokenizerWrapper::from_json(json, true).unwrap();
541        let input = tok.encode("hello world", 512);
542        assert_eq!(input.input_ids[0], 2); // [CLS]
543        assert_eq!(*input.input_ids.last().unwrap(), 3); // [SEP]
544        assert!(input.input_ids.len() >= 4);
545    }
546
547    #[test]
548    fn hf_tokenizer_invalid_json() {
549        assert!(HfTokenizerWrapper::from_json("not json", true).is_err());
550    }
551
552    #[test]
553    fn hf_tokenizer_missing_vocab() {
554        let json = r#"{"model": {"type": "WordPiece"}}"#;
555        assert!(HfTokenizerWrapper::from_json(json, true).is_err());
556    }
557
558    /// model2vec vocabularies (potion-*) carry `[PAD]`/`[UNK]` but no sentence
559    /// markers: they load for plain encoding, which drops unknown pieces, and
560    /// stay rejected for transformer framing.
561    #[test]
562    fn static_vocab_without_markers_encodes_plain_tokens() {
563        let json = r###"{"model": {"type": "WordPiece", "vocab": {
564            "[PAD]": 0, "[UNK]": 1, "read": 2, "file": 3, "##s": 4}}}"###;
565        assert!(HfTokenizerWrapper::from_json(json, true).is_err());
566        let tok = HfTokenizerWrapper::from_json(json, false).unwrap();
567        assert_eq!(tok.encode_plain("read files", 64), vec![2, 3, 4]);
568        assert_eq!(tok.encode_plain("read zzz", 64), vec![2], "[UNK] dropped");
569        assert_eq!(tok.encode_plain("read files", 2), vec![2, 3], "token cap");
570        {
571            let found = &tok.encode_plain("", 64);
572            assert!(found.is_empty(), "{found:?}");
573        }
574    }
575}
lean_ctx/core/embeddings/tokenizer.rs

lean_ctx/core/embeddings/
tokenizer.rs
