[Bugfix][Rust Frontend] Tolerate out-of-vocab prompt ids in detokenizer (#44682)
Co-authored-by: Bugen Zhao <i@bugenzhao.com> Signed-off-by: Ting Sun <suntcrick@gmail.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>
Ting SUN committed
Jul 3, 2026 at 10:41 UTC
1aeabec355ebf2ea9efc5f60f68730d2316ce9f0
1 file changed
+103
-10
rust/src/tokenizer/src/incremental.rs
+103
-10
@@ -32,6 +32,7 @@ pub(crate) struct DecodeStream<'a, T: Tokenizer + ?Sized> {
32
ids: Vec<u32>,
33
prefix: String,
34
prefix_index: usize,
35
+ prefix_seeded: bool,
36
cumulative_output: String,
37
output_index: usize,
38
}
@@ -50,6 +51,7 @@ impl<'a, T: Tokenizer + ?Sized> DecodeStream<'a, T> {
51
ids: prompt_token_ids.to_vec(),
52
prefix: String::new(),
53
prefix_index: 0,
54
+ prefix_seeded: prompt_token_ids.is_empty(),
55
cumulative_output: String::new(),
56
output_index: 0,
57
}
@@ -63,27 +65,54 @@ const SAFE_SUFFIX_MIN: usize = 4;
65
const SAFE_SUFFIX_MAX: usize = 6;
66
67
impl<T: Tokenizer + ?Sized> DecodeStream<'_, T> {
66
- /// Seed `self.prefix` from the shortest trailing suffix whose decoded text
67
- /// has no U+FFFD — a clean decode means the suffix starts and ends at
68
- /// valid UTF-8/token boundaries, so priming from it is equivalent to
69
- /// priming from the full prompt.
68
+ /// Decode prompt-only context for prefix seeding.
69
+ ///
70
+ /// Prompt ids may come from the model vocabulary rather than the local
71
+ /// tokenizer vocabulary. For this seeding path, ids that cannot be mapped
72
+ /// back to raw token text are dropped before retrying strict decode. This
73
+ /// tolerance is intentionally limited to prompt context; generated ids are
74
+ /// decoded later through the normal strict path.
75
+ fn decode_prompt_context(&self, ids: &[u32]) -> Result<(String, Vec<u32>)> {
76
+ match self.tokenizer.decode(ids, self.skip_special_tokens) {
77
+ Ok(decoded) => Ok((decoded, ids.to_vec())),
78
+ Err(error) => {
79
+ let filtered = ids
80
+ .iter()
81
+ .copied()
82
+ .filter(|&id| self.tokenizer.id_to_token(id).is_some())
83
+ .collect::<Vec<_>>();
84
+ if filtered.len() == ids.len() {
85
+ return Err(error);
86
+ }
87
+ self.tokenizer
88
+ .decode(&filtered, self.skip_special_tokens)
89
+ .map(|decoded| (decoded, filtered))
90
+ }
91
+ }
92
+ }
93
+
94
+ /// Seed `self.prefix` from the shortest trailing prompt suffix whose
95
+ /// filtered context is still long enough and whose decoded text has no
96
+ /// U+FFFD. A clean decode means the suffix starts and ends at valid
97
+ /// UTF-8/token boundaries, so priming from it is equivalent to priming from
98
+ /// the full prompt.
99
fn seed_prefix(&mut self) -> Result<()> {
100
let prompt_len = self.ids.len();
101
if prompt_len > SAFE_SUFFIX_MIN {
102
let max_try = SAFE_SUFFIX_MAX.min(prompt_len - 1);
103
for suffix_len in SAFE_SUFFIX_MIN..=max_try {
104
let start = prompt_len - suffix_len;
76
- let decoded =
77
- self.tokenizer.decode(&self.ids[start..], self.skip_special_tokens)?;
78
- if !decoded.contains('\u{FFFD}') {
105
+ let (decoded, context_ids) = self.decode_prompt_context(&self.ids[start..])?;
106
+ if !decoded.contains('\u{FFFD}') && context_ids.len() >= SAFE_SUFFIX_MIN {
107
self.prefix = decoded;
80
- self.ids.drain(..start);
108
+ self.ids = context_ids;
109
self.prefix_index = self.ids.len();
110
return Ok(());
111
}
112
}
113
}
86
- let decoded = self.tokenizer.decode(&self.ids, self.skip_special_tokens)?;
114
+ let (decoded, context_ids) = self.decode_prompt_context(&self.ids)?;
115
+ self.ids = context_ids;
116
if !decoded.ends_with('\u{FFFD}') {
117
self.prefix = decoded;
118
self.prefix_index = self.ids.len();
@@ -94,8 +123,9 @@ impl<T: Tokenizer + ?Sized> DecodeStream<'_, T> {
123
124
impl<T: Tokenizer + ?Sized> IncrementalDecoder for DecodeStream<'_, T> {
125
fn push_token(&mut self, token_id: u32) -> Result<usize> {
97
- if self.prefix.is_empty() && !self.ids.is_empty() {
126
+ if !self.prefix_seeded && !self.ids.is_empty() {
127
self.seed_prefix()?;
128
+ self.prefix_seeded = true;
129
}
130
131
self.ids.push(token_id);
@@ -131,6 +161,7 @@ impl<T: Tokenizer + ?Sized> IncrementalDecoder for DecodeStream<'_, T> {
161
self.ids.clear();
162
self.prefix.clear();
163
self.prefix_index = 0;
164
+ self.prefix_seeded = true;
165
// Ensure we split at a utf-8 char boundary.
166
self.cumulative_output
167
.push_str(&string[string.floor_char_boundary(prefix_len)..]);
@@ -152,6 +183,7 @@ impl<T: Tokenizer + ?Sized> IncrementalDecoder for DecodeStream<'_, T> {
183
#[cfg(test)]
184
mod tests {
185
use super::*;
186
+ use crate::test_utils::TestTokenizer;
187
188
/// Backend that treats each token ID as a raw byte, producing lossy UTF-8.
189
#[derive(Debug)]
@@ -281,6 +313,67 @@ mod tests {
313
assert_eq!(decoder.output(), "!");
314
}
315
316
+ #[test]
317
+ fn prompt_context_filters_unknown_ids() {
318
+ let tokenizer = TestTokenizer::new();
319
+ assert_eq!(tokenizer.id_to_token(10_000), None);
320
+
321
+ let cases: &[(&str, &[u32], u32, &str)] = &[
322
+ (
323
+ "suffix seed",
324
+ &[
325
+ b'a' as u32,
326
+ b'b' as u32,
327
+ b'c' as u32,
328
+ 10_000,
329
+ b'H' as u32,
330
+ b'i' as u32,
331
+ ],
332
+ b'!' as u32,
333
+ "!",
334
+ ),
335
+ ("all unknown", &[10_000], b'!' as u32, "!"),
336
+ (
337
+ "unknown before incomplete utf-8",
338
+ &[10_000, 0xe4, 0xbd],
339
+ 0xa0,
340
+ "你",
341
+ ),
342
+ (
343
+ "incomplete utf-8 before filtered suffix",
344
+ &[0xe4, 0xbd, 10_000, 10_001, 10_002, 10_003, 10_004, 10_005],
345
+ 0xa0,
346
+ "你",
347
+ ),
348
+ ];
349
+
350
+ for &(name, prompt, token_id, output) in cases {
351
+ let mut decoder = tokenizer.create_decode_stream(prompt, false, 0);
352
+ assert_eq!(
353
+ decoder.push_token(token_id).unwrap(),
354
+ output.len(),
355
+ "{name}"
356
+ );
357
+ assert_eq!(decoder.output(), output, "{name}");
358
+ }
359
+ }
360
+
361
+ #[test]
362
+ fn generated_unknown_ids_still_return_decode_error() {
363
+ let tokenizer = TestTokenizer::new();
364
+ assert_eq!(tokenizer.id_to_token(10_000), None);
365
+
366
+ let prompt = &[10_000, b'H' as u32, b'i' as u32];
367
+ let mut decoder = tokenizer.create_decode_stream(prompt, false, 0);
368
+
369
+ let error = decoder.push_token(10_000).unwrap_err();
370
+ assert!(
371
+ error
372
+ .to_string()
373
+ .contains("test tokenizer cannot decode unknown token id 10000")
374
+ );
375
+ }
376
+
377
#[test]
378
fn chunks_concatenate_to_full_text() {
379
let backend = Utf8Backend;