Skip to main content

maskforge_core/vocab/
vocabulary.rs

1//! Creates `Vocabulary` manually or from pretrained large language model.
2
3use rustc_hash::{FxHashMap as HashMap, FxHashSet as HashSet};
4#[cfg(feature = "tokenizer-processing")]
5use tokenizers::normalizers::Sequence;
6#[cfg(feature = "tokenizer-processing")]
7use tokenizers::{NormalizerWrapper, Tokenizer};
8
9use super::error::{Error, Result};
10#[cfg(feature = "huggingface-hub")]
11use super::locator::{HFLocator, Locator};
12#[cfg(feature = "tokenizer-processing")]
13use super::processor::TokenProcessor;
14#[cfg(feature = "huggingface-hub")]
15use super::FromPretrainedParameters;
16use super::{Token, TokenId};
17
18/// Largest tokenizer JSON accepted before parsing.
19#[cfg(feature = "tokenizer-processing")]
20pub const MAX_TOKENIZER_JSON_BYTES: usize = 64 * 1024 * 1024;
21
22/// Vocabulary for a large language model.
23#[cfg_attr(
24    feature = "huggingface-hub",
25    doc = r##"
26### Create a vocabulary from a pretrained model.
27```no_run
28use maskforge_core::Vocabulary;
29
30let vocabulary = Vocabulary::from_pretrained("openai-community/gpt2", None);
31```
32
33### Create a vocabulary from a pretrained model with some additional parameters.
34```no_run
35use maskforge_core::{FromPretrainedParameters, Vocabulary};
36
37let params = FromPretrainedParameters {
38    revision: "607a30d783dfa663caf39e06633721c8d4cfcd7e".to_string(),
39    ..Default::default()
40};
41let vocabulary = Vocabulary::from_pretrained("openai-community/gpt2", Some(params));
42```
43
44### Create an empty vocabulary and manually insert some tokens.
45```
46use maskforge_core::Vocabulary;
47
48let eos_token_id = 1;
49let mut vocabulary = Vocabulary::new(eos_token_id);
50
51vocabulary.try_insert("token", 0).expect("New token inserted");
52assert_eq!(vocabulary.token_ids("token"), Some(&[0][..]));
53assert_eq!(vocabulary.tokens().len(), 1);
54assert_eq!(vocabulary.eos_token_id(), eos_token_id);
55
56vocabulary.remove("token");
57assert_eq!(vocabulary.token_ids("token"), None);
58```
59"##
60)]
61#[derive(Clone, Debug, Default, PartialEq)]
62pub struct Vocabulary {
63    eos_token_id: TokenId,
64    tokens: HashMap<Token, Vec<TokenId>>,
65}
66
67impl Vocabulary {
68    /// Creates an empty vocabulary.
69    pub fn new(eos_token_id: TokenId) -> Self {
70        Self {
71            eos_token_id,
72            tokens: HashMap::default(),
73        }
74    }
75
76    /// Creates the vocabulary of pre-trained model from Hugging Face Hub.
77    #[cfg(feature = "huggingface-hub")]
78    pub fn from_pretrained(
79        model: &str,
80        parameters: Option<FromPretrainedParameters>,
81    ) -> Result<Self> {
82        Self::from_pretrained_with_locator::<HFLocator>(model, parameters)
83    }
84
85    #[doc(hidden)]
86    #[inline(always)]
87    #[cfg(feature = "huggingface-hub")]
88    fn from_pretrained_with_locator<L: Locator>(
89        model: &str,
90        parameters: Option<FromPretrainedParameters>,
91    ) -> Result<Self> {
92        let tokenizer = Tokenizer::from_pretrained(model, parameters.clone())?;
93
94        // Locate eos_token_id in defined locations.
95        let eos_token_id = match L::locate_eos_token_id(model, &tokenizer, &parameters) {
96            Ok(id) => id,
97            Err(tried) => {
98                let reason = if tried.is_empty() {
99                    "EOS token id".to_string()
100                } else {
101                    format!("EOS token id not found; tried: {}", tried.join("; "))
102                };
103                return Err(Error::UnsupportedTokenizer {
104                    model: model.to_string(),
105                    reason,
106                });
107            }
108        };
109
110        Self::from_tokenizer(tokenizer, eos_token_id, model)
111    }
112
113    /// Builds from fast-tokenizer JSON with a caller-resolved EOS id.
114    /// Oversized input is rejected before parsing.
115    #[cfg(feature = "tokenizer-processing")]
116    pub fn from_tokenizer_json(json: &str, eos_token_id: TokenId) -> Result<Self> {
117        if json.len() > MAX_TOKENIZER_JSON_BYTES {
118            return Err(Error::TokenizerJsonTooLarge {
119                actual: json.len(),
120                limit: MAX_TOKENIZER_JSON_BYTES,
121            });
122        }
123        let tokenizer = Tokenizer::from_bytes(json.as_bytes())?;
124        Self::from_tokenizer(tokenizer, eos_token_id, "<inline tokenizer JSON>")
125    }
126
127    /// The one token-processing pipeline both tokenizer constructors share.
128    #[cfg(feature = "tokenizer-processing")]
129    fn from_tokenizer(
130        mut tokenizer: Tokenizer,
131        eos_token_id: TokenId,
132        source: &str,
133    ) -> Result<Self> {
134        Self::filter_prepend_normalizers(&mut tokenizer);
135
136        // Start building the vocabulary from eos_token_id and added tokens.
137        let mut vocabulary = Vocabulary::new(eos_token_id);
138        for (id, added_token) in tokenizer.get_added_tokens_decoder().iter() {
139            if !added_token.special && id != &eos_token_id {
140                vocabulary.try_insert(added_token.content.clone(), *id)?
141            }
142        }
143
144        // Process each vocabulary token according to the tokenizer's level.
145        let Ok(processor) = TokenProcessor::new(&tokenizer) else {
146            return Err(Error::UnsupportedTokenizer {
147                model: source.to_string(),
148                reason: "Token processor".to_string(),
149            });
150        };
151        for (token, token_id) in tokenizer.get_vocab(false) {
152            if token_id != eos_token_id {
153                let processed_token = processor.process(&token)?;
154                vocabulary.try_insert(processed_token, token_id)?;
155            }
156        }
157
158        Ok(vocabulary)
159    }
160
161    /// Returns all tokens with their token ids in vocabulary.
162    pub fn tokens(&self) -> &HashMap<Token, Vec<TokenId>> {
163        &self.tokens
164    }
165
166    /// Returns all token ids per provided token if available in the vocabulary.
167    pub fn token_ids(&self, token: impl AsRef<[u8]>) -> Option<&[TokenId]> {
168        self.tokens.get(token.as_ref()).map(Vec::as_slice)
169    }
170
171    /// Gets the identifier of the special end of the sentence token.
172    pub fn eos_token_id(&self) -> TokenId {
173        self.eos_token_id
174    }
175
176    /// Inserts a token with the given id. Idempotent: repeating `(token, id)` is a no-op.
177    /// A different id under the same bytes is still kept.
178    pub fn try_insert(&mut self, token: impl Into<Token>, id: TokenId) -> Result<(), Error> {
179        if id == self.eos_token_id {
180            return Err(Error::EOSTokenDisallowed);
181        }
182        let token = token.into();
183        if token.is_empty() {
184            return Err(Error::EmptyToken);
185        }
186        let ids = self.tokens.entry(token).or_default();
187        if !ids.contains(&id) {
188            ids.push(id);
189        }
190        Ok(())
191    }
192
193    /// Removes a given token from the vocabulary.
194    pub fn remove(&mut self, token: impl Into<Token>) {
195        let token = token.into();
196        self.tokens.remove(&token);
197    }
198
199    /// Returns ordinary token ids plus EOS, not the sparse mask width.
200    /// Overflowing counts saturate instead of wrapping.
201    pub fn len(&self) -> usize {
202        self.tokens
203            .values()
204            .try_fold(0usize, |acc, ids| acc.checked_add(ids.len()))
205            .and_then(|n| n.checked_add(1))
206            .unwrap_or(usize::MAX)
207    }
208
209    /// True when no ordinary tokens are present (EOS is not counted).
210    pub fn is_empty(&self) -> bool {
211        self.tokens.is_empty()
212    }
213
214    /// Filters out `Prepend` kind of tokenizer's normalizers.
215    #[cfg(feature = "tokenizer-processing")]
216    fn filter_prepend_normalizers(tokenizer: &mut Tokenizer) {
217        // Remove prepend normalizers so their marker characters do not enter token bytes.
218        if let Some(normalizer) = tokenizer.get_normalizer() {
219            match normalizer {
220                NormalizerWrapper::Sequence(normalization_sequence) => {
221                    let new_sequence = Sequence::new(
222                        normalization_sequence
223                            .as_ref()
224                            .iter()
225                            .filter_map(|normalizer| match normalizer {
226                                NormalizerWrapper::Prepend(_) => None,
227                                _ => Some(normalizer.clone()),
228                            })
229                            .collect(),
230                    );
231                    tokenizer.with_normalizer(new_sequence.into());
232                }
233                NormalizerWrapper::Prepend(_) => {
234                    tokenizer.with_normalizer(None::<NormalizerWrapper>);
235                }
236                _ => {}
237            }
238        }
239    }
240}
241
242/// Removes duplicate ids in place, keeping first-occurrence order (matches `try_insert`). Small
243/// lists (the norm) dedup without a heap allocation; only large lists build a `HashSet`.
244fn dedup_ids(ids: &mut Vec<TokenId>) {
245    if ids.len() <= 1 {
246        return;
247    }
248    // A short id list (the norm even for merged tokens) deduplicates in place, preserving
249    // first-occurrence order, with no allocation - `contains` over <=8 kept ids is cheaper than a set.
250    if ids.len() <= 8 {
251        let mut write = 1;
252        for read in 1..ids.len() {
253            if !ids[..write].contains(&ids[read]) {
254                ids[write] = ids[read];
255                write += 1;
256            }
257        }
258        ids.truncate(write);
259        return;
260    }
261    let mut seen = HashSet::default();
262    ids.retain(|id| seen.insert(*id));
263}
264
265impl std::fmt::Display for Vocabulary {
266    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
267        writeln!(
268            f,
269            "Vocabulary object with eos_token_id={:?} and the following tokens to token_ids:",
270            self.eos_token_id
271        )?;
272        for (token, token_ids) in self.tokens.iter() {
273            writeln!(
274                f,
275                "{:?} -> {:?}",
276                token
277                    .iter()
278                    .map(|b| format!("0x{:02X}", b))
279                    .collect::<Vec<_>>(),
280                token_ids
281            )?;
282        }
283        Ok(())
284    }
285}
286
287impl Vocabulary {
288    /// Raw map constructor, not part of the public API; `build::build_vocabulary` validates first.
289    /// Duplicate ids within one token are collapsed (see `try_insert`).
290    pub(crate) fn from_raw_bytes_map(
291        eos_token_id: TokenId,
292        mut tokens: HashMap<Token, Vec<TokenId>>,
293    ) -> Result<Self, Error> {
294        if tokens.keys().any(|t| t.is_empty()) {
295            return Err(Error::EmptyToken);
296        }
297        if tokens.values().any(Vec::is_empty) {
298            return Err(Error::EmptyTokenIds);
299        }
300        if tokens.iter().any(|(_, ids)| ids.contains(&eos_token_id)) {
301            return Err(Error::EOSTokenDisallowed);
302        }
303        for ids in tokens.values_mut() {
304            dedup_ids(ids);
305        }
306        Ok(Vocabulary {
307            eos_token_id,
308            tokens,
309        })
310    }
311
312    /// Raw map constructor over `String` keys. Same visibility rule as `from_raw_bytes_map`.
313    pub(crate) fn from_raw_string_map(
314        eos_token_id: TokenId,
315        tokens: HashMap<String, Vec<TokenId>>,
316    ) -> Result<Self, Error> {
317        Ok(Vocabulary {
318            eos_token_id,
319            tokens: tokens
320                .into_iter()
321                .map(|(k, mut v)| {
322                    if k.is_empty() {
323                        Err(Error::EmptyToken)
324                    } else if v.is_empty() {
325                        Err(Error::EmptyTokenIds)
326                    } else if v.contains(&eos_token_id) {
327                        Err(Error::EOSTokenDisallowed)
328                    } else {
329                        dedup_ids(&mut v);
330                        Ok((k.into_bytes(), v))
331                    }
332                })
333                .collect::<Result<HashMap<Token, Vec<TokenId>>, _>>()?,
334        })
335    }
336}
337
338#[cfg(test)]
339mod tests {
340    use super::*;
341
342    #[test]
343    fn repeated_identical_insertion_is_idempotent() {
344        let mut v = Vocabulary::new(0);
345        let before_len = v.len();
346        v.try_insert("a", 5).unwrap();
347        let after_first = v.len();
348        v.try_insert("a", 5).unwrap();
349        v.try_insert("a", 5).unwrap();
350        assert_eq!(
351            v.len(),
352            after_first,
353            "repeated identical insert must not grow len()"
354        );
355        assert_eq!(v.token_ids("a"), Some(&[5][..]));
356        assert_eq!(after_first, before_len + 1);
357    }
358
359    #[test]
360    fn dedup_ids_small_and_large_paths_agree_across_the_boundary() {
361        // A reference dedup (order-preserving, always the HashSet path) the fast paths must match on
362        // every length, so the len<=1 / len<=8 shortcuts never drift from the general case (M053).
363        fn reference(ids: &[TokenId]) -> Vec<TokenId> {
364            let mut seen = HashSet::default();
365            ids.iter().copied().filter(|id| seen.insert(*id)).collect()
366        }
367        // Deterministic pseudo-random ids with heavy duplication, spanning the <=1, <=8, and >8 arms.
368        let mut state = 0x1234_5678u32;
369        for len in 0..40usize {
370            let raw: Vec<TokenId> = (0..len)
371                .map(|_| {
372                    state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
373                    state % 6 // small id space forces frequent duplicates
374                })
375                .collect();
376            let mut got = raw.clone();
377            dedup_ids(&mut got);
378            assert_eq!(got, reference(&raw), "len={len} raw={raw:?}");
379        }
380    }
381
382    #[test]
383    fn same_bytes_two_different_ids_are_both_kept() {
384        let mut v = Vocabulary::new(0);
385        v.try_insert("a", 5).unwrap();
386        v.try_insert("a", 6).unwrap();
387        let ids = v.token_ids("a").unwrap();
388        assert_eq!(ids.len(), 2);
389        assert!(ids.contains(&5) && ids.contains(&6));
390    }
391
392    #[test]
393    fn token_id_u32_max_is_accepted_as_an_ordinary_id() {
394        let mut v = Vocabulary::new(0);
395        v.try_insert("a", u32::MAX).unwrap();
396        assert_eq!(v.token_ids("a"), Some(&[u32::MAX][..]));
397    }
398
399    #[test]
400    fn sparse_ids_do_not_inflate_len_beyond_stored_entries() {
401        let mut v = Vocabulary::new(0);
402        v.try_insert("low", 1).unwrap();
403        v.try_insert("high", 1000).unwrap();
404        assert_eq!(v.len(), 3); // 2 tokens + 1 for EOS, not 1001
405    }
406
407    #[test]
408    fn from_raw_bytes_map_dedups_duplicate_ids_within_one_token() {
409        let mut m: HashMap<Token, Vec<TokenId>> = HashMap::default();
410        m.insert(b"a".to_vec(), vec![5, 5, 6, 5]);
411        let v = Vocabulary::from_raw_bytes_map(0, m).unwrap();
412        let ids = v.token_ids("a").unwrap();
413        assert_eq!(ids.len(), 2);
414        assert!(ids.contains(&5) && ids.contains(&6));
415    }
416
417    #[test]
418    fn from_raw_string_map_dedups_duplicate_ids_within_one_token() {
419        let mut m: HashMap<String, Vec<TokenId>> = HashMap::default();
420        m.insert("a".to_string(), vec![5, 5, 6, 5]);
421        let v = Vocabulary::from_raw_string_map(0, m).unwrap();
422        let ids = v.token_ids("a").unwrap();
423        assert_eq!(ids.len(), 2);
424        assert!(ids.contains(&5) && ids.contains(&6));
425    }
426
427    #[test]
428    fn duplicate_removal_preserves_first_occurrence_order_not_sorted() {
429        // Deliberately out of numeric order: 9, 3, 7. A sort would reorder to 3, 7, 9.
430        let mut v = Vocabulary::new(0);
431        for id in [9, 3, 7, 9, 3] {
432            v.try_insert("a", id).unwrap();
433        }
434        assert_eq!(v.token_ids("a"), Some(&[9, 3, 7][..]));
435
436        let mut m: HashMap<Token, Vec<TokenId>> = HashMap::default();
437        m.insert(b"a".to_vec(), vec![9, 3, 7, 9, 3]);
438        let raw = Vocabulary::from_raw_bytes_map(0, m).unwrap();
439        assert_eq!(raw.token_ids("a"), Some(&[9, 3, 7][..]));
440    }
441
442    #[test]
443    fn raw_bytes_and_string_map_construction_agree_on_id_order() {
444        let ids = vec![9, 3, 7, 9, 3];
445        let mut bytes_map: HashMap<Token, Vec<TokenId>> = HashMap::default();
446        bytes_map.insert(b"a".to_vec(), ids.clone());
447        let by_bytes = Vocabulary::from_raw_bytes_map(0, bytes_map).unwrap();
448
449        let mut str_map: HashMap<String, Vec<TokenId>> = HashMap::default();
450        str_map.insert("a".to_string(), ids);
451        let by_str = Vocabulary::from_raw_string_map(0, str_map).unwrap();
452
453        assert_eq!(by_bytes.token_ids("a"), by_str.token_ids("a"));
454        assert_eq!(by_bytes.token_ids("a"), Some(&[9, 3, 7][..]));
455    }
456
457    #[test]
458    fn basic_interface() {
459        let eos_token_id = 3;
460        let mut vocabulary = Vocabulary::new(eos_token_id);
461
462        match vocabulary.try_insert("eos-token", eos_token_id) {
463            Err(Error::EOSTokenDisallowed) => {}
464            _ => unreachable!(),
465        }
466
467        // New empty vocabulary.
468        assert_eq!(vocabulary.eos_token_id, eos_token_id);
469        assert!(vocabulary.tokens.is_empty());
470
471        for (token, id) in [("zero", 0), ("one", 1), ("two", 2)] {
472            vocabulary.try_insert(token, id).expect("Insert failed");
473            assert_eq!(vocabulary.token_ids(token), Some(&[id][..]));
474        }
475        assert_eq!(vocabulary.tokens.len(), 3);
476        assert_eq!(vocabulary.tokens().len(), 3);
477
478        // Confirm different types.
479        vocabulary.try_insert(b"four", 4).expect("Insert failed");
480        assert_eq!(vocabulary.token_ids("four"), Some(&[4][..]));
481
482        vocabulary
483            .try_insert(b"five".to_vec(), 5)
484            .expect("Insert failed");
485        assert_eq!(vocabulary.token_ids("five"), Some(&[5][..]));
486
487        vocabulary
488            .try_insert("six".to_string(), 6)
489            .expect("Insert failed");
490        assert_eq!(vocabulary.token_ids("six"), Some(&[6][..]));
491
492        vocabulary.remove(b"four");
493        assert_eq!(vocabulary.token_ids("four"), None);
494
495        vocabulary.remove(b"five".to_vec());
496        assert_eq!(vocabulary.token_ids("five"), None);
497
498        vocabulary.remove("six".to_string());
499        assert_eq!(vocabulary.token_ids("six"), None);
500    }
501
502    #[test]
503    fn new_empty_vocabulary_from_hashmap() {
504        let map: HashMap<Token, Vec<TokenId>> = HashMap::default();
505        let vocabulary = Vocabulary::from_raw_bytes_map(1_u32, map).expect("Vocabulary failed");
506        assert_eq!(vocabulary.eos_token_id, 1);
507        assert!(vocabulary.tokens.is_empty());
508    }
509
510    #[cfg(feature = "huggingface-hub")]
511    #[test]
512    fn supported_pretrained_models() {
513        // Support is expected for these:
514        for model in [
515            // GPT 2
516            "openai-community/gpt2",
517            // Llama 2
518            "hf-internal-testing/Llama-2-7B-GPTQ",
519            // Llama 3
520            // OpenCoder: shares llama tokenizers
521            "hf-internal-testing/llama-3-8b-internal",
522            // Qwen
523            "Qwen/Qwen2-7B-Instruct",
524            // Salamandra
525            "BSC-LT/salamandra-2b",
526        ] {
527            let vocabulary = Vocabulary::from_pretrained(model, None);
528            match vocabulary {
529                Ok(v) => {
530                    assert_eq!(v.eos_token_id, v.eos_token_id());
531                    assert!(!v.tokens.is_empty());
532                }
533                Err(_) => unreachable!(),
534            }
535        }
536    }
537
538    #[cfg(feature = "huggingface-hub")]
539    #[test]
540    fn pretrained_from_gpt2() {
541        let model = "openai-community/gpt2";
542        let tokenizer = Tokenizer::from_pretrained(model, None).expect("Tokenizer failed");
543        let vocabulary = Vocabulary::from_pretrained(model, None).expect("Vocabulary failed");
544
545        let v_eos = vocabulary.eos_token_id;
546        assert_eq!(v_eos, vocabulary.eos_token_id());
547        assert_eq!(v_eos, 50256);
548        assert_eq!(
549            tokenizer.id_to_token(v_eos).expect("Token not found"),
550            "<|endoftext|>"
551        );
552
553        let token = "Ġal";
554        let btoken = token.as_bytes().to_vec();
555        assert!(vocabulary.token_ids(&btoken).is_none());
556        assert!(tokenizer.token_to_id(token).is_some());
557
558        for (v_token, t_token_expected) in [("abc", "abc"), (" O", "ĠO")] {
559            let v_ids = vocabulary.token_ids(v_token.as_bytes());
560            assert!(v_ids.is_some());
561            for v_id in v_ids.unwrap() {
562                let t_token = tokenizer
563                    .id_to_token(*v_id)
564                    .expect("Token id not found in tokenizer");
565                assert_eq!(&t_token, t_token_expected);
566            }
567        }
568    }
569
570    #[cfg(feature = "huggingface-hub")]
571    #[test]
572    fn from_tokenizer_json_matches_from_pretrained_exactly() {
573        let model = "openai-community/gpt2";
574        let tokenizer = Tokenizer::from_pretrained(model, None).expect("Tokenizer failed");
575        let json = tokenizer.to_string(false).expect("serialize tokenizer");
576        let via_pretrained = Vocabulary::from_pretrained(model, None).expect("from_pretrained");
577        let via_json = Vocabulary::from_tokenizer_json(&json, via_pretrained.eos_token_id())
578            .expect("from_tokenizer_json");
579        assert_eq!(via_json, via_pretrained);
580    }
581
582    #[cfg(feature = "huggingface-hub")]
583    #[test]
584    fn pretrained_from_llama() {
585        use rustc_hash::FxHashSet as HashSet;
586
587        let model = "hf-internal-testing/llama-tokenizer";
588        let tokenizer = Tokenizer::from_pretrained(model, None).expect("Tokenizer failed");
589        let vocabulary = Vocabulary::from_pretrained(model, None).expect("Vocabulary failed");
590
591        let v_eos = vocabulary.eos_token_id;
592        assert_eq!(v_eos, vocabulary.eos_token_id());
593        assert_eq!(v_eos, 2);
594        assert_eq!(
595            tokenizer.id_to_token(v_eos).expect("Token not found"),
596            "</s>"
597        );
598
599        let tests: &[(Vec<u8>, &[&str])] = &[
600            ("abc".as_bytes().to_vec(), &["abc"]),
601            (" al".as_bytes().to_vec(), &["▁al"]),
602            (" O".as_bytes().to_vec(), &["▁O"]),
603            ("   ".as_bytes().to_vec(), &["▁▁▁"]),
604            (" ".as_bytes().to_vec(), &["▁", "<0x20>"]),
605            ("a".as_bytes().to_vec(), &["a", "<0x61>"]),
606            (vec![0xFF], &["<0xFF>"]),
607            (vec![0x20], &["▁", "<0x20>"]),
608        ];
609        for (v_token, t_tokens_expected) in tests {
610            let v_ids = vocabulary.token_ids(v_token);
611            assert!(v_ids.is_some());
612
613            let t_tokens = v_ids
614                .unwrap()
615                .iter()
616                .map(|v_id| {
617                    tokenizer
618                        .id_to_token(*v_id)
619                        .expect("Token id not found in tokenizer")
620                })
621                .collect::<HashSet<String>>();
622            let expected = HashSet::from_iter(t_tokens_expected.iter().map(|s| s.to_string()));
623            assert_eq!(t_tokens, expected)
624        }
625    }
626
627    #[cfg(feature = "huggingface-hub")]
628    #[test]
629    fn token_processor_error() {
630        let model = "hf-internal-testing/tiny-random-XLMRobertaXLForCausalLM";
631        let vocabulary = Vocabulary::from_pretrained(model, None);
632
633        match vocabulary {
634            Err(Error::UnsupportedTokenizer { model, reason }) => {
635                assert_eq!(model, model.to_string());
636                assert_eq!(&reason, "Token processor");
637            }
638            _ => unreachable!(),
639        }
640    }
641
642    #[cfg(feature = "huggingface-hub")]
643    #[test]
644    fn tokenizer_error() {
645        let model = "hf-internal-testing/some-non-existent-model";
646        let vocabulary = Vocabulary::from_pretrained(model, None);
647
648        match vocabulary {
649            Err(Error::TokenizersError(e)) => assert!(!e.to_string().is_empty()),
650            _ => unreachable!(),
651        }
652    }
653
654    #[cfg(feature = "huggingface-hub")]
655    struct NoneLocator;
656    #[cfg(feature = "huggingface-hub")]
657    impl Locator for NoneLocator {
658        fn locate_eos_token_id(
659            _model: &str,
660            _tokenizer: &Tokenizer,
661            _parameters: &Option<FromPretrainedParameters>,
662        ) -> Result<TokenId, Vec<String>> {
663            Err(Vec::new())
664        }
665    }
666
667    #[cfg(feature = "huggingface-hub")]
668    #[test]
669    fn unable_to_locate_eos_token_id_error() {
670        let model = "hf-internal-testing/tiny-random-XLMRobertaXLForCausalLM";
671        let vocabulary = Vocabulary::from_pretrained_with_locator::<NoneLocator>(model, None);
672
673        match vocabulary {
674            Err(Error::UnsupportedTokenizer { model, reason }) => {
675                assert_eq!(model, model.to_string());
676                assert_eq!(&reason, "EOS token id");
677            }
678            _ => unreachable!(),
679        }
680    }
681
682    #[test]
683    #[cfg(feature = "huggingface-hub")]
684    fn prepend_normalizers_filtered_out() {
685        use tokenizers::normalizers::{Prepend, Sequence};
686
687        let prepend = Prepend::new("_".to_string());
688        let prepend_normalizer = NormalizerWrapper::Prepend(prepend);
689        let sequence = Sequence::new(vec![prepend_normalizer.clone()]);
690        let sequence_normalizer = NormalizerWrapper::Sequence(sequence);
691
692        let model = "hf-internal-testing/llama-tokenizer";
693        let tokenizer = Tokenizer::from_pretrained(model, None).expect("Tokenizer failed");
694
695        for normalizer in [prepend_normalizer, sequence_normalizer] {
696            let mut normalized_t = tokenizer.clone();
697            normalized_t.with_normalizer(Some(normalizer));
698            Vocabulary::filter_prepend_normalizers(&mut normalized_t);
699            if let Some(n) = normalized_t.get_normalizer() {
700                match n {
701                    NormalizerWrapper::Sequence(seq) => {
702                        for n in seq.as_ref() {
703                            if let NormalizerWrapper::Prepend(_) = n {
704                                unreachable!()
705                            }
706                        }
707                    }
708                    NormalizerWrapper::Prepend(_) => unreachable!(),
709                    _ => {}
710                }
711            }
712        }
713    }
714
715    #[test]
716    #[cfg(feature = "huggingface-hub")]
717    fn other_normalizers_being_kept() {
718        use tokenizers::normalizers::BertNormalizer;
719
720        let model = "hf-internal-testing/llama-tokenizer";
721        let normalizer = NormalizerWrapper::BertNormalizer(BertNormalizer::default());
722        let mut tokenizer = Tokenizer::from_pretrained(model, None).expect("Tokenizer failed");
723        tokenizer.with_normalizer(Some(normalizer));
724
725        Vocabulary::filter_prepend_normalizers(&mut tokenizer);
726
727        assert!(tokenizer.get_normalizer().is_some());
728    }
729}