[Rust][Benchmark] Prevent invalid token IDs in random benchmarks (#50058)
Signed-off-by: reidliu41 <reid201711@gmail.com>
Reid committed
Aug 1, 2026 at 11:20 UTC
62195e9784ebec1ece42b88a861734e0702cc2d5
1 file changed
+80
-18
rust/src/bench/src/tiktoken.rs
+80
-18
@@ -19,9 +19,9 @@ const NUM_RESERVED_SPECIAL_TOKENS: u32 = 256;
19
pub struct TiktokenTokenizer {
20
bpe: tiktoken_rs::CoreBPE,
21
vocab_size: u32,
22
- #[allow(dead_code)]
22
num_base_tokens: u32,
24
- /// IDs of all special tokens (for filtering from allowed_tokens)
23
+ /// Registered special token IDs. File-based tokenizers exclude them from
24
+ /// allowed tokens; built-in tokenizers retain the previously allowed IDs.
25
special_token_ids: Vec<u32>,
26
/// Reverse mapping: token_id -> byte sequence, for lossy UTF-8 decoding.
27
/// Empty for built-in encodings (use bpe.decode instead).
@@ -113,12 +113,17 @@ impl TiktokenTokenizer {
113
}
114
115
/// Create from a built-in tiktoken encoding (o200k_base, cl100k_base, etc.)
116
- pub fn from_builtin_bpe(bpe: tiktoken_rs::CoreBPE, vocab_size: u32) -> Self {
116
+ pub fn from_builtin_bpe(
117
+ bpe: tiktoken_rs::CoreBPE,
118
+ vocab_size: u32,
119
+ num_base_tokens: u32,
120
+ special_token_ids: Vec<u32>,
121
+ ) -> Self {
122
Self {
123
bpe,
124
vocab_size,
120
- num_base_tokens: vocab_size,
121
- special_token_ids: Vec::new(),
125
+ num_base_tokens,
126
+ special_token_ids,
127
decoder: Vec::new(),
128
is_builtin: true,
129
}
@@ -152,14 +157,15 @@ impl TiktokenTokenizer {
157
self.vocab_size
158
}
159
155
- /// Get non-special token IDs whose byte representation is valid UTF-8.
156
- /// Excludes ALL special tokens (like Python's `set(all_tokens) - set(prohibited_tokens)`).
160
+ /// Get token IDs that can be decoded safely.
161
+ ///
162
+ /// File-based tokenizers exclude special tokens and tokens whose byte
163
+ /// representation is not valid UTF-8.
164
pub fn get_allowed_tokens(&self) -> Vec<u32> {
165
if self.is_builtin {
159
- // For built-in encodings, return full token range.
160
- // Note: when used with random dataset, these IDs are sent to vLLM as-is;
161
- // ensure the model's tokenizer is compatible (e.g. GPT-4o for o200k_base).
162
- return (0..self.vocab_size).collect();
166
+ return (0..self.num_base_tokens)
167
+ .chain(self.special_token_ids.iter().copied())
168
+ .collect();
169
}
170
let special_set: std::collections::HashSet<u32> =
171
self.special_token_ids.iter().copied().collect();
@@ -185,12 +191,22 @@ impl TiktokenTokenizer {
191
/// These encodings are bundled with tiktoken-rs — no network download required.
192
/// Useful for consistent cross-model token counting (e.g. Artificial Analysis methodology).
193
pub fn load_builtin_tiktoken(encoding: &str) -> Result<TiktokenTokenizer> {
188
- let (bpe, vocab_size) = match encoding {
189
- "o200k_base" => (tiktoken_rs::o200k_base(), 200_275u32),
190
- "cl100k_base" => (tiktoken_rs::cl100k_base(), 100_277u32),
191
- "p50k_base" => (tiktoken_rs::p50k_base(), 50_281u32),
192
- "p50k_edit" => (tiktoken_rs::p50k_edit(), 50_281u32),
193
- "r50k_base" | "gpt2" => (tiktoken_rs::r50k_base(), 50_257u32),
194
+ let (bpe, vocab_size, num_base_tokens, special_token_ids) = match encoding {
195
+ "o200k_base" => (
196
+ tiktoken_rs::o200k_base(),
197
+ 200_275u32,
198
+ 199_998u32,
199
+ vec![199_999, 200_018],
200
+ ),
201
+ "cl100k_base" => (
202
+ tiktoken_rs::cl100k_base(),
203
+ 100_277u32,
204
+ 100_256u32,
205
+ vec![100_257, 100_258, 100_259, 100_260, 100_276],
206
+ ),
207
+ "p50k_base" => (tiktoken_rs::p50k_base(), 50_281u32, 50_281u32, vec![]),
208
+ "p50k_edit" => (tiktoken_rs::p50k_edit(), 50_281u32, 50_281u32, vec![]),
209
+ "r50k_base" | "gpt2" => (tiktoken_rs::r50k_base(), 50_257u32, 50_256u32, vec![50_256]),
210
_ => {
211
return Err(BenchError::Tokenizer(format!(
212
"Unknown built-in tiktoken encoding: '{encoding}'. \
@@ -205,7 +221,12 @@ pub fn load_builtin_tiktoken(encoding: &str) -> Result<TiktokenTokenizer> {
221
vocab_size,
222
"loaded tokenizer"
223
);
208
- Ok(TiktokenTokenizer::from_builtin_bpe(bpe, vocab_size))
224
+ Ok(TiktokenTokenizer::from_builtin_bpe(
225
+ bpe,
226
+ vocab_size,
227
+ num_base_tokens,
228
+ special_token_ids,
229
+ ))
230
}
231
232
/// Try to load a tiktoken tokenizer from a local directory or HuggingFace model repo.
@@ -552,3 +573,44 @@ fn find_unescaped(s: &str, ch: char) -> Option<usize> {
573
}
574
None
575
}
576
+
577
+#[cfg(test)]
578
+mod tests {
579
+ use super::*;
580
+
581
+ #[test]
582
+ fn builtin_allowed_tokens_are_decodable() {
583
+ let cases = [
584
+ ("o200k_base", 200_000),
585
+ ("cl100k_base", 100_261),
586
+ ("p50k_base", 50_281),
587
+ ("p50k_edit", 50_281),
588
+ ("r50k_base", 50_257),
589
+ ("gpt2", 50_257),
590
+ ];
591
+
592
+ for (encoding, expected_count) in cases {
593
+ let tokenizer = load_builtin_tiktoken(encoding).unwrap();
594
+ let allowed_tokens = tokenizer.get_allowed_tokens();
595
+
596
+ assert_eq!(allowed_tokens.len(), expected_count, "{encoding}");
597
+ for token_id in allowed_tokens {
598
+ tokenizer
599
+ .decode(&[token_id])
600
+ .unwrap_or_else(|error| panic!("{encoding} token {token_id}: {error}"));
601
+ }
602
+ }
603
+ }
604
+
605
+ #[test]
606
+ fn sparse_builtin_allowed_tokens_skip_unassigned_ids() {
607
+ let o200k = load_builtin_tiktoken("o200k_base").unwrap().get_allowed_tokens();
608
+ assert_eq!(&o200k[199_997..], &[199_997, 199_999, 200_018]);
609
+
610
+ let cl100k = load_builtin_tiktoken("cl100k_base").unwrap().get_allowed_tokens();
611
+ assert_eq!(
612
+ &cl100k[100_255..],
613
+ &[100_255, 100_257, 100_258, 100_259, 100_260, 100_276]
614
+ );
615
+ }
616
+}