[Perf][Rust Frontend] Count the tokenizer vocabulary once at construction (#54449)
Co-authored-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Feathbow <feathbow@gmail.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>
Salt Sato committed
Sep 1, 2026 at 02:25 UTC
22df3a34e098738b9e6fbd09f63967d60acc24c9
1 file changed
+21
-4
rust/src/tokenizer/src/hf.rs
+21
-4
@@ -150,6 +150,7 @@ pub struct HuggingFaceTokenizer {
150
backend: Backend,
151
special_token_ids: Arc<[u32]>,
152
added_vocab: Box<[(String, u32)]>,
153
+ vocab_size: usize,
154
}
155
156
impl HuggingFaceTokenizer {
@@ -174,10 +175,13 @@ impl HuggingFaceTokenizer {
175
ids.dedup();
176
Arc::from(ids)
177
};
178
+ // HF materializes the merged vocabulary to count it, so cache the result.
179
+ let vocab_size = tokenizer.get_vocab_size(true);
180
Self {
181
backend: Backend::Hf(Box::new(tokenizer)),
182
special_token_ids,
183
added_vocab,
184
+ vocab_size,
185
}
186
}
187
@@ -205,6 +209,7 @@ impl HuggingFaceTokenizer {
209
Arc::from(ids)
210
};
211
let byte_level = tokenizer.decoder().is_some_and(is_byte_level_only);
212
+ let vocab_size = tokenizer.vocab_size();
213
let backend = if byte_level {
214
Backend::FastokensByteLevel(Box::new(tokenizer))
215
} else {
@@ -214,6 +219,7 @@ impl HuggingFaceTokenizer {
219
backend,
220
special_token_ids,
221
added_vocab,
222
+ vocab_size,
223
}
224
}
225
@@ -299,10 +305,7 @@ impl Tokenizer for HuggingFaceTokenizer {
305
}
306
307
fn vocab_size(&self) -> usize {
302
- match &self.backend {
303
- Backend::Hf(t) => t.get_vocab_size(true),
304
- Backend::Fastokens(t) | Backend::FastokensByteLevel(t) => t.vocab_size(),
305
- }
308
+ self.vocab_size
309
}
310
311
fn id_to_token(&self, id: u32) -> Option<String> {
@@ -551,6 +554,20 @@ mod tests {
554
}
555
}
556
557
+ #[test]
558
+ fn hf_vocab_size_counts_added_tokens() {
559
+ let mut tokenizer = tiny_bpe_tokenizer();
560
+ tokenizer.add_special_tokens(&[AddedToken::from("<|im_end|>", true)]);
561
+ let expected = tokenizer.get_vocab_size(true);
562
+
563
+ let dir = tempdir().expect("create temp dir");
564
+ let path = dir.path().join("tokenizer.json");
565
+ tokenizer.save(&path, false).expect("save tokenizer json");
566
+
567
+ let wrapper = HuggingFaceTokenizer::new_hf(&path).expect("load hf wrapper");
568
+ assert_eq!(wrapper.vocab_size(), expected);
569
+ }
570
+
571
#[test]
572
fn hf_constructor_resolves_added_token_ids() {
573
let mut tokenizer = tiny_bpe_tokenizer();