From 21aeff0c66687576dac8998bddebd8846402e3c3 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Mon, 6 Apr 2026 03:45:39 -0700 Subject: [PATCH 1/4] refactor(tokenizer): introduce DecodeResult enum to Decoder trait Replace hardcoded U+FFFD replacement character string checks with a typed DecodeResult enum (Complete | Partial) in the Decoder trait. This makes partial-decode state explicit in the type system, eliminating brittle `ends_with("U+FFFD")` checks in DecodeStream::step() and Sequence::append_token_id(). Addresses: https://github.com/ai-dynamo/dynamo/pull/6996#discussion_r2894253655 --- lib/llm/src/backend.rs | 2 +- lib/llm/src/tokenizers.rs | 90 ++++++++++++++++++++++++----- lib/llm/src/tokenizers/fastokens.rs | 13 +++-- lib/llm/src/tokenizers/hf.rs | 6 +- lib/llm/src/tokenizers/tiktoken.rs | 57 +++++++++--------- lib/llm/tests/test_stop_behavior.rs | 7 ++- lib/llm/tests/tokenizers.rs | 15 +++-- 7 files changed, 131 insertions(+), 59 deletions(-) diff --git a/lib/llm/src/backend.rs b/lib/llm/src/backend.rs index 4b639f117a06..5f682976537e 100644 --- a/lib/llm/src/backend.rs +++ b/lib/llm/src/backend.rs @@ -638,7 +638,7 @@ mod tests { &self, _token_ids: &[TokenIdType], _skip_special_tokens: bool, - ) -> anyhow::Result { + ) -> anyhow::Result { Err(anyhow::anyhow!( "Unable to decode into a valid UTF-8 string: incomplete utf-8 byte sequence from index 6" )) diff --git a/lib/llm/src/tokenizers.rs b/lib/llm/src/tokenizers.rs index 21b411cc9de1..ff191ba43418 100644 --- a/lib/llm/src/tokenizers.rs +++ b/lib/llm/src/tokenizers.rs @@ -19,6 +19,7 @@ pub use anyhow::{Error, Result}; pub use fastokens::FastTokenizer; pub use hf::HuggingFaceTokenizer; pub use tiktoken::TikTokenTokenizer; +pub use traits::DecodeResult; /// Represents the type of tokenizer being used #[derive(Debug)] @@ -62,12 +63,60 @@ pub mod traits { fn encode_batch(&self, inputs: &[&str]) -> Result>; } + /// Result of decoding token IDs to text. + /// + /// Distinguishes between fully valid UTF-8 output and output that contains + /// trailing incomplete multi-byte sequences (represented as U+FFFD). + /// This lets callers like `DecodeStream::step()` decide whether to emit or + /// buffer without resorting to hardcoded replacement-character string checks. + #[derive(Debug, Clone, PartialEq, Eq)] + pub enum DecodeResult { + /// All bytes decoded to valid UTF-8 -- no trailing incomplete sequences. + Complete(String), + /// The decoded string has incomplete trailing multi-byte bytes (ends with U+FFFD). + Partial(String), + } + + impl DecodeResult { + /// Unwrap the inner string regardless of variant. + pub fn into_string(self) -> String { + match self { + DecodeResult::Complete(s) | DecodeResult::Partial(s) => s, + } + } + + /// Returns `true` if the decode produced incomplete trailing bytes. + pub fn is_partial(&self) -> bool { + matches!(self, DecodeResult::Partial(_)) + } + + /// Returns a reference to the inner string. + pub fn as_str(&self) -> &str { + match self { + DecodeResult::Complete(s) | DecodeResult::Partial(s) => s, + } + } + + /// Classify a decoded string: `Partial` if it ends with U+FFFD, else `Complete`. + pub fn classify(text: String) -> Self { + if text.ends_with('\u{FFFD}') { + DecodeResult::Partial(text) + } else { + DecodeResult::Complete(text) + } + } + } + /// Implementations **must** use lossy UTF-8 conversion (e.g. `String::from_utf8_lossy`) - /// so that partial multi-byte sequences produce U+FFFD (`�`) rather than returning `Err`. - /// `DecodeStream::step()` relies on the replacement character to detect incomplete + /// so that partial multi-byte sequences produce U+FFFD (`\u{FFFD}`) rather than returning `Err`. + /// `DecodeStream::step()` relies on `DecodeResult::Partial` to detect incomplete /// sequences and buffer tokens until the full character arrives. pub trait Decoder: Send + Sync { - fn decode(&self, token_ids: &[TokenIdType], skip_special_tokens: bool) -> Result; + fn decode( + &self, + token_ids: &[TokenIdType], + skip_special_tokens: bool, + ) -> Result; } pub trait Tokenizer: Encoder + Decoder { @@ -219,23 +268,27 @@ impl DecodeStream { pub fn step(&mut self, id: u32) -> Result> { self.all_token_ids.push(id); - let prefix_text = self.tokenizer.decode( - &self.all_token_ids[self.prefix_offset..self.read_offset], - self.skip_special_tokens, - )?; + let prefix_text = self + .tokenizer + .decode( + &self.all_token_ids[self.prefix_offset..self.read_offset], + self.skip_special_tokens, + )? + .into_string(); - let new_text = self.tokenizer.decode( + let new_result = self.tokenizer.decode( &self.all_token_ids[self.prefix_offset..], self.skip_special_tokens, )?; - if new_text.len() > prefix_text.len() && !new_text.ends_with("�") { - let new_text = new_text[prefix_text.len()..].to_string(); + let new_text = new_result.as_str(); + if new_text.len() > prefix_text.len() && !new_result.is_partial() { + let emitted = new_text[prefix_text.len()..].to_string(); self.prefix_offset = self.read_offset; self.read_offset = self.all_token_ids.len(); - Ok(Some(new_text)) + Ok(Some(emitted)) } else { Ok(None) } @@ -324,12 +377,15 @@ impl Sequence { let prefix_text = self .tokenizer - .decode(&self.token_ids[self.prefix_offset..self.read_offset], false)?; + .decode(&self.token_ids[self.prefix_offset..self.read_offset], false)? + .into_string(); - let new_text = self + let new_result = self .tokenizer .decode(&self.token_ids[self.prefix_offset..], false)?; + let new_text = new_result.as_str(); + // if the end character of the previous returned sequence is a multi-byte character // then we can not split the text on that byte offset, so we roll back to the byte offset // of the start of that character @@ -340,11 +396,13 @@ impl Sequence { let prefix_text_len = prefix_text_len; if new_text.len() > prefix_text.len() { - if new_text.ends_with("�") { + if new_result.is_partial() { return Ok("".to_string()); } else { // shift and update the state - let new_text = new_text[prefix_text_len..].to_string().replace("�", ""); + let new_text = new_text[prefix_text_len..] + .to_string() + .replace('\u{FFFD}', ""); self.prefix_offset = self.read_offset; self.read_offset = self.token_ids.len(); return Ok(new_text); @@ -366,7 +424,7 @@ impl Sequence { // let tokenizer = self.tokenizer.read().map_err(|err| { // Error::msg(format!("Failed to acquire read lock on tokenizer: {}", err)) // })?; - self.tokenizer.decode(&self.token_ids, false) + Ok(self.tokenizer.decode(&self.token_ids, false)?.into_string()) } } diff --git a/lib/llm/src/tokenizers/fastokens.rs b/lib/llm/src/tokenizers/fastokens.rs index 83c5c7e3bade..80bed68b7c2f 100644 --- a/lib/llm/src/tokenizers/fastokens.rs +++ b/lib/llm/src/tokenizers/fastokens.rs @@ -14,7 +14,7 @@ use rayon::prelude::*; use super::{ Encoding, Error, Result, TokenIdType, hf::HuggingFaceTokenizer, - traits::{Decoder, Encoder, Tokenizer}, + traits::{DecodeResult, Decoder, Encoder, Tokenizer}, }; /// Hybrid tokenizer: fast BPE encoding via `fastokens`, decoding via HuggingFace. @@ -52,7 +52,7 @@ impl Encoder for FastTokenizer { } impl Decoder for FastTokenizer { - fn decode(&self, token_ids: &[TokenIdType], skip_special_tokens: bool) -> Result { + fn decode(&self, token_ids: &[TokenIdType], skip_special_tokens: bool) -> Result { self.hf_decoder.decode(token_ids, skip_special_tokens) } } @@ -81,7 +81,10 @@ mod tests { let text = "Hello, world!"; let encoding = tokenizer.encode(text).unwrap(); assert!(!encoding.token_ids().is_empty()); - let decoded = tokenizer.decode(encoding.token_ids(), true).unwrap(); + let decoded = tokenizer + .decode(encoding.token_ids(), true) + .unwrap() + .into_string(); assert!(!decoded.is_empty()); // The decoded text should contain the same non-space characters let enc_chars: String = text.chars().filter(|c| !c.is_whitespace()).collect(); @@ -149,8 +152,8 @@ mod tests { // decode(continuation) which lacks the surrounding context. let mut all_ids = prompt_ids.clone(); all_ids.extend_from_slice(&cont_ids); - let full_text = wrapper.decode(&all_ids, true).unwrap(); - let prompt_text = wrapper.decode(&prompt_ids, true).unwrap(); + let full_text = wrapper.decode(&all_ids, true).unwrap().into_string(); + let prompt_text = wrapper.decode(&prompt_ids, true).unwrap().into_string(); let expected = &full_text[prompt_text.len()..]; assert_eq!( accumulated, expected, diff --git a/lib/llm/src/tokenizers/hf.rs b/lib/llm/src/tokenizers/hf.rs index 0ca8f3fec838..ec1e34b5270c 100644 --- a/lib/llm/src/tokenizers/hf.rs +++ b/lib/llm/src/tokenizers/hf.rs @@ -5,7 +5,7 @@ use tokenizers::tokenizer::Tokenizer as HfTokenizer; use super::{ Encoding, Error, Result, TokenIdType, - traits::{Decoder, Encoder, Tokenizer}, + traits::{DecodeResult, Decoder, Encoder, Tokenizer}, }; pub struct HuggingFaceTokenizer { @@ -52,14 +52,14 @@ impl Encoder for HuggingFaceTokenizer { } impl Decoder for HuggingFaceTokenizer { - fn decode(&self, token_ids: &[TokenIdType], skip_special_tokens: bool) -> Result { + fn decode(&self, token_ids: &[TokenIdType], skip_special_tokens: bool) -> Result { // This calls into the library let text = self .tokenizer .decode(token_ids, skip_special_tokens) .map_err(|err| Error::msg(format!("Error de-tokenizing input: {err}")))?; - Ok(text) + Ok(DecodeResult::classify(text)) } } diff --git a/lib/llm/src/tokenizers/tiktoken.rs b/lib/llm/src/tokenizers/tiktoken.rs index 798af8f1e76e..03b0cedbffe6 100644 --- a/lib/llm/src/tokenizers/tiktoken.rs +++ b/lib/llm/src/tokenizers/tiktoken.rs @@ -11,7 +11,7 @@ use tiktoken_rs::CoreBPE; use super::{ Encoding, Error, Result, TokenIdType, - traits::{Decoder, Encoder, Tokenizer}, + traits::{DecodeResult, Decoder, Encoder, Tokenizer}, }; /// Number of reserved special-token slots to generate when filling gaps in the vocabulary. @@ -89,7 +89,7 @@ impl Encoder for TikTokenTokenizer { } impl Decoder for TikTokenTokenizer { - fn decode(&self, token_ids: &[TokenIdType], skip_special_tokens: bool) -> Result { + fn decode(&self, token_ids: &[TokenIdType], skip_special_tokens: bool) -> Result { let ids: Vec = if skip_special_tokens { token_ids .iter() @@ -100,12 +100,13 @@ impl Decoder for TikTokenTokenizer { token_ids.to_vec() }; - // Use lossy UTF-8 conversion so that partial multi-byte sequences become U+FFFD (�). + // Use lossy UTF-8 conversion so that partial multi-byte sequences become U+FFFD. // This is critical for incremental detokenization: DecodeStream::step() relies on - // the replacement character to detect incomplete sequences and buffer tokens until + // DecodeResult::Partial to detect incomplete sequences and buffer tokens until // a complete character arrives. CoreBPE::decode() would error on invalid UTF-8 instead. let bytes: Vec = self.bpe._decode_native_and_split(ids).flatten().collect(); - Ok(String::from_utf8_lossy(&bytes).into_owned()) + let text = String::from_utf8_lossy(&bytes).into_owned(); + Ok(DecodeResult::classify(text)) } } @@ -351,7 +352,7 @@ mod tests { assert!(!ids.is_empty()); // Test decode roundtrip - let decoded = tokenizer.decode(ids, false).unwrap(); + let decoded = tokenizer.decode(ids, false).unwrap().into_string(); assert_eq!(decoded, "hello world"); } @@ -393,11 +394,11 @@ mod tests { ids.push(22); // [EOS] // Decode with skip_special_tokens=true should strip special tokens - let decoded_skip = tokenizer.decode(&ids, true).unwrap(); + let decoded_skip = tokenizer.decode(&ids, true).unwrap().into_string(); assert_eq!(decoded_skip, "hello"); // Decode with skip_special_tokens=false should include them - let decoded_all = tokenizer.decode(&ids, false).unwrap(); + let decoded_all = tokenizer.decode(&ids, false).unwrap().into_string(); assert!(decoded_all.contains("hello")); } @@ -416,7 +417,7 @@ mod tests { let ids = encoding.token_ids(); assert!(!ids.is_empty()); - let decoded = tokenizer.decode(ids, false).unwrap(); + let decoded = tokenizer.decode(ids, false).unwrap().into_string(); assert_eq!(decoded, "hello world"); } @@ -516,11 +517,11 @@ mod tests { result.is_ok(), "decode() should not error on incomplete UTF-8 bytes" ); - let text = result.unwrap(); + let decode_result = result.unwrap(); assert!( - text.contains('\u{FFFD}'), - "incomplete UTF-8 byte should produce replacement character, got: {:?}", - text + decode_result.is_partial(), + "incomplete UTF-8 byte should produce DecodeResult::Partial, got: {:?}", + decode_result ); } @@ -532,11 +533,11 @@ mod tests { let result = tokenizer.decode(&[100, 101], false); assert!(result.is_ok()); - let text = result.unwrap(); + let decode_result = result.unwrap(); assert!( - text.contains('\u{FFFD}'), - "incomplete 2-of-3 UTF-8 bytes should produce replacement character, got: {:?}", - text + decode_result.is_partial(), + "incomplete 2-of-3 UTF-8 bytes should produce DecodeResult::Partial, got: {:?}", + decode_result ); } @@ -550,7 +551,7 @@ mod tests { let result = tokenizer.decode(&[100, 101, 102], false); assert!(result.is_ok()); - assert_eq!(result.unwrap(), "你"); + assert_eq!(result.unwrap().into_string(), "你"); } /// All 4 emoji bytes together form valid UTF-8, so this passes both before and after @@ -562,7 +563,7 @@ mod tests { let result = tokenizer.decode(&[200, 201, 202, 203], false); assert!(result.is_ok()); - assert_eq!(result.unwrap(), "😀"); + assert_eq!(result.unwrap().into_string(), "😀"); } /// Without the fix, fails with "incomplete utf-8 byte sequence" from CoreBPE::decode(). @@ -573,7 +574,7 @@ mod tests { let result = tokenizer.decode(&[200], false); assert!(result.is_ok()); - assert!(result.unwrap().contains('\u{FFFD}')); + assert!(result.unwrap().is_partial()); } /// Without the fix, fails with "incomplete utf-8 byte sequence" from CoreBPE::decode(). @@ -584,16 +585,17 @@ mod tests { let result = tokenizer.decode(&[5, 100], false); assert!(result.is_ok()); - let text = result.unwrap(); + let decode_result = result.unwrap(); + assert!( + decode_result.is_partial(), + "trailing incomplete byte should produce DecodeResult::Partial" + ); + let text = decode_result.into_string(); assert!( text.starts_with("hello"), "should start with 'hello', got: {:?}", text ); - assert!( - text.contains('\u{FFFD}'), - "trailing incomplete byte should produce U+FFFD" - ); } /// End-to-end incremental detokenization: DecodeStream buffers partial bytes, @@ -656,7 +658,10 @@ mod tests { assert_eq!(encodings.len(), 2); for (encoding, input) in encodings.iter().zip(inputs.iter()) { - let decoded = tokenizer.decode(encoding.token_ids(), false).unwrap(); + let decoded = tokenizer + .decode(encoding.token_ids(), false) + .unwrap() + .into_string(); assert_eq!(decoded, *input); } } diff --git a/lib/llm/tests/test_stop_behavior.rs b/lib/llm/tests/test_stop_behavior.rs index 675b9c49ee3d..7a318fd4c2c9 100644 --- a/lib/llm/tests/test_stop_behavior.rs +++ b/lib/llm/tests/test_stop_behavior.rs @@ -25,8 +25,8 @@ impl tokenizer_traits::Encoder for TestTokenizer { } impl tokenizer_traits::Decoder for TestTokenizer { - fn decode(&self, ids: &[u32], skip_special: bool) -> Result { - Ok(ids + fn decode(&self, ids: &[u32], skip_special: bool) -> Result { + let text: String = ids .iter() .filter_map(|&id| match id { EOS if skip_special => None, @@ -36,7 +36,8 @@ impl tokenizer_traits::Decoder for TestTokenizer { EOS => Some(""), _ => Some("?"), }) - .collect()) + .collect(); + Ok(tokenizer_traits::DecodeResult::classify(text)) } } diff --git a/lib/llm/tests/tokenizers.rs b/lib/llm/tests/tokenizers.rs index 31edc311f025..0e81acc46865 100644 --- a/lib/llm/tests/tokenizers.rs +++ b/lib/llm/tests/tokenizers.rs @@ -115,7 +115,8 @@ fn test_encode_decode_roundtrip(#[case] tokenizer: Arc) { let decoded = tokenizer .decode(encoding.token_ids(), false) - .unwrap_or_else(|e| panic!("Failed to decode '{text}': {e}")); + .unwrap_or_else(|e| panic!("Failed to decode '{text}': {e}")) + .into_string(); assert_eq!(decoded, text, "Roundtrip failed for: '{text}'"); } } @@ -131,7 +132,8 @@ fn test_encode_decode_roundtrip_multibyte(#[case] tokenizer: Arc) let decoded = tokenizer .decode(encoding.token_ids(), false) - .unwrap_or_else(|e| panic!("Failed to decode '{text}': {e}")); + .unwrap_or_else(|e| panic!("Failed to decode '{text}': {e}")) + .into_string(); assert_eq!(decoded, text, "Roundtrip failed for: '{text}'"); } } @@ -149,7 +151,8 @@ fn test_batch_encode_roundtrip(#[case] tokenizer: Arc) { for (encoding, &input) in encodings.iter().zip(inputs.iter()) { let decoded = tokenizer .decode(encoding.token_ids(), false) - .expect("Failed to decode"); + .expect("Failed to decode") + .into_string(); assert_eq!(decoded, input); } } @@ -356,12 +359,14 @@ fn test_decode_with_skip_special_tokens() { // Decode with skip_special_tokens = false (should keep special tokens) let decoded_with_special = tokenizer .decode(&token_ids, false) - .expect("Failed to decode with skip_special_tokens=false"); + .expect("Failed to decode with skip_special_tokens=false") + .into_string(); // Decode with skip_special_tokens = true (should remove special tokens) let decoded_without_special = tokenizer .decode(&token_ids, true) - .expect("Failed to decode with skip_special_tokens=true"); + .expect("Failed to decode with skip_special_tokens=true") + .into_string(); // Validate exact matches on the entire decoded strings assert_eq!(decoded_with_special, " Hello world"); From 28490222607c77572f20c5a4252f44b7cc72f867 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Mon, 6 Apr 2026 03:46:19 -0700 Subject: [PATCH 2/4] docs(tokenizer): clarify DecodeResult and Decoder trait docstrings - Complete variant may still contain interior U+FFFD from mid-stream invalid byte sequences; only trailing status is tracked. - Soften Decoder trait doc to not mandate a specific mechanism (from_utf8_lossy), since HF uses ByteFallback internally. --- lib/llm/src/tokenizers.rs | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/lib/llm/src/tokenizers.rs b/lib/llm/src/tokenizers.rs index ff191ba43418..18e7f1b4b1f3 100644 --- a/lib/llm/src/tokenizers.rs +++ b/lib/llm/src/tokenizers.rs @@ -71,9 +71,12 @@ pub mod traits { /// buffer without resorting to hardcoded replacement-character string checks. #[derive(Debug, Clone, PartialEq, Eq)] pub enum DecodeResult { - /// All bytes decoded to valid UTF-8 -- no trailing incomplete sequences. + /// No trailing incomplete multi-byte sequences (text does not end with U+FFFD). + /// Note: the string may still contain *interior* U+FFFD characters from + /// mid-stream invalid byte sequences; only trailing status is tracked here. Complete(String), - /// The decoded string has incomplete trailing multi-byte bytes (ends with U+FFFD). + /// The decoded string ends with U+FFFD, indicating incomplete trailing + /// multi-byte bytes that may be completed by subsequent tokens. Partial(String), } @@ -107,10 +110,11 @@ pub mod traits { } } - /// Implementations **must** use lossy UTF-8 conversion (e.g. `String::from_utf8_lossy`) - /// so that partial multi-byte sequences produce U+FFFD (`\u{FFFD}`) rather than returning `Err`. - /// `DecodeStream::step()` relies on `DecodeResult::Partial` to detect incomplete - /// sequences and buffer tokens until the full character arrives. + /// Implementations must ensure that partial multi-byte sequences produce U+FFFD + /// (`\u{FFFD}`) in the output rather than returning `Err`. This is commonly achieved + /// via `String::from_utf8_lossy` (tiktoken) or library-internal byte-fallback handling + /// (HuggingFace). `DecodeStream::step()` relies on `DecodeResult::Partial` to detect + /// incomplete sequences and buffer tokens until the full character arrives. pub trait Decoder: Send + Sync { fn decode( &self, From c2e8b519e0513521beb4645ce6d89e2f4890ede6 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Sat, 11 Apr 2026 13:30:38 -0700 Subject: [PATCH 3/4] address PR comments --- lib/llm/src/tokenizers.rs | 42 +++++++++++++++-------------- lib/llm/src/tokenizers/fastokens.rs | 8 +++--- lib/llm/src/tokenizers/hf.rs | 2 +- lib/llm/src/tokenizers/tiktoken.rs | 20 +++++++------- lib/llm/tests/test_stop_behavior.rs | 2 +- lib/llm/tests/tokenizers.rs | 20 +++++++------- 6 files changed, 48 insertions(+), 46 deletions(-) diff --git a/lib/llm/src/tokenizers.rs b/lib/llm/src/tokenizers.rs index 18e7f1b4b1f3..561642965d09 100644 --- a/lib/llm/src/tokenizers.rs +++ b/lib/llm/src/tokenizers.rs @@ -69,7 +69,7 @@ pub mod traits { /// trailing incomplete multi-byte sequences (represented as U+FFFD). /// This lets callers like `DecodeStream::step()` decide whether to emit or /// buffer without resorting to hardcoded replacement-character string checks. - #[derive(Debug, Clone, PartialEq, Eq)] + #[derive(Debug, Clone, PartialEq, Eq, strum::EnumIs)] pub enum DecodeResult { /// No trailing incomplete multi-byte sequences (text does not end with U+FFFD). /// Note: the string may still contain *interior* U+FFFD characters from @@ -81,18 +81,6 @@ pub mod traits { } impl DecodeResult { - /// Unwrap the inner string regardless of variant. - pub fn into_string(self) -> String { - match self { - DecodeResult::Complete(s) | DecodeResult::Partial(s) => s, - } - } - - /// Returns `true` if the decode produced incomplete trailing bytes. - pub fn is_partial(&self) -> bool { - matches!(self, DecodeResult::Partial(_)) - } - /// Returns a reference to the inner string. pub fn as_str(&self) -> &str { match self { @@ -100,8 +88,8 @@ pub mod traits { } } - /// Classify a decoded string: `Partial` if it ends with U+FFFD, else `Complete`. - pub fn classify(text: String) -> Self { + /// Construct from a decoded string: `Partial` if it ends with U+FFFD, else `Complete`. + pub fn from_decoded(text: String) -> Self { if text.ends_with('\u{FFFD}') { DecodeResult::Partial(text) } else { @@ -110,6 +98,20 @@ pub mod traits { } } + impl From for DecodeResult { + fn from(text: String) -> Self { + DecodeResult::from_decoded(text) + } + } + + impl From for String { + fn from(result: DecodeResult) -> Self { + match result { + DecodeResult::Complete(s) | DecodeResult::Partial(s) => s, + } + } + } + /// Implementations must ensure that partial multi-byte sequences produce U+FFFD /// (`\u{FFFD}`) in the output rather than returning `Err`. This is commonly achieved /// via `String::from_utf8_lossy` (tiktoken) or library-internal byte-fallback handling @@ -272,13 +274,13 @@ impl DecodeStream { pub fn step(&mut self, id: u32) -> Result> { self.all_token_ids.push(id); - let prefix_text = self + let prefix_text: String = self .tokenizer .decode( &self.all_token_ids[self.prefix_offset..self.read_offset], self.skip_special_tokens, )? - .into_string(); + .into(); let new_result = self.tokenizer.decode( &self.all_token_ids[self.prefix_offset..], @@ -379,10 +381,10 @@ impl Sequence { self.token_ids.push(token_id); // log::trace!("pushed token_id: {}", token_id); - let prefix_text = self + let prefix_text: String = self .tokenizer .decode(&self.token_ids[self.prefix_offset..self.read_offset], false)? - .into_string(); + .into(); let new_result = self .tokenizer @@ -428,7 +430,7 @@ impl Sequence { // let tokenizer = self.tokenizer.read().map_err(|err| { // Error::msg(format!("Failed to acquire read lock on tokenizer: {}", err)) // })?; - Ok(self.tokenizer.decode(&self.token_ids, false)?.into_string()) + Ok(self.tokenizer.decode(&self.token_ids, false)?.into()) } } diff --git a/lib/llm/src/tokenizers/fastokens.rs b/lib/llm/src/tokenizers/fastokens.rs index 80bed68b7c2f..cc805fccabcf 100644 --- a/lib/llm/src/tokenizers/fastokens.rs +++ b/lib/llm/src/tokenizers/fastokens.rs @@ -81,10 +81,10 @@ mod tests { let text = "Hello, world!"; let encoding = tokenizer.encode(text).unwrap(); assert!(!encoding.token_ids().is_empty()); - let decoded = tokenizer + let decoded: String = tokenizer .decode(encoding.token_ids(), true) .unwrap() - .into_string(); + .into(); assert!(!decoded.is_empty()); // The decoded text should contain the same non-space characters let enc_chars: String = text.chars().filter(|c| !c.is_whitespace()).collect(); @@ -152,8 +152,8 @@ mod tests { // decode(continuation) which lacks the surrounding context. let mut all_ids = prompt_ids.clone(); all_ids.extend_from_slice(&cont_ids); - let full_text = wrapper.decode(&all_ids, true).unwrap().into_string(); - let prompt_text = wrapper.decode(&prompt_ids, true).unwrap().into_string(); + let full_text: String = wrapper.decode(&all_ids, true).unwrap().into(); + let prompt_text: String = wrapper.decode(&prompt_ids, true).unwrap().into(); let expected = &full_text[prompt_text.len()..]; assert_eq!( accumulated, expected, diff --git a/lib/llm/src/tokenizers/hf.rs b/lib/llm/src/tokenizers/hf.rs index ec1e34b5270c..080a775719fe 100644 --- a/lib/llm/src/tokenizers/hf.rs +++ b/lib/llm/src/tokenizers/hf.rs @@ -59,7 +59,7 @@ impl Decoder for HuggingFaceTokenizer { .decode(token_ids, skip_special_tokens) .map_err(|err| Error::msg(format!("Error de-tokenizing input: {err}")))?; - Ok(DecodeResult::classify(text)) + Ok(text.into()) } } diff --git a/lib/llm/src/tokenizers/tiktoken.rs b/lib/llm/src/tokenizers/tiktoken.rs index 03b0cedbffe6..2f376c95caa9 100644 --- a/lib/llm/src/tokenizers/tiktoken.rs +++ b/lib/llm/src/tokenizers/tiktoken.rs @@ -106,7 +106,7 @@ impl Decoder for TikTokenTokenizer { // a complete character arrives. CoreBPE::decode() would error on invalid UTF-8 instead. let bytes: Vec = self.bpe._decode_native_and_split(ids).flatten().collect(); let text = String::from_utf8_lossy(&bytes).into_owned(); - Ok(DecodeResult::classify(text)) + Ok(text.into()) } } @@ -352,7 +352,7 @@ mod tests { assert!(!ids.is_empty()); // Test decode roundtrip - let decoded = tokenizer.decode(ids, false).unwrap().into_string(); + let decoded: String = tokenizer.decode(ids, false).unwrap().into(); assert_eq!(decoded, "hello world"); } @@ -394,11 +394,11 @@ mod tests { ids.push(22); // [EOS] // Decode with skip_special_tokens=true should strip special tokens - let decoded_skip = tokenizer.decode(&ids, true).unwrap().into_string(); + let decoded_skip: String = tokenizer.decode(&ids, true).unwrap().into(); assert_eq!(decoded_skip, "hello"); // Decode with skip_special_tokens=false should include them - let decoded_all = tokenizer.decode(&ids, false).unwrap().into_string(); + let decoded_all: String = tokenizer.decode(&ids, false).unwrap().into(); assert!(decoded_all.contains("hello")); } @@ -417,7 +417,7 @@ mod tests { let ids = encoding.token_ids(); assert!(!ids.is_empty()); - let decoded = tokenizer.decode(ids, false).unwrap().into_string(); + let decoded: String = tokenizer.decode(ids, false).unwrap().into(); assert_eq!(decoded, "hello world"); } @@ -551,7 +551,7 @@ mod tests { let result = tokenizer.decode(&[100, 101, 102], false); assert!(result.is_ok()); - assert_eq!(result.unwrap().into_string(), "你"); + assert_eq!(String::from(result.unwrap()), "你"); } /// All 4 emoji bytes together form valid UTF-8, so this passes both before and after @@ -563,7 +563,7 @@ mod tests { let result = tokenizer.decode(&[200, 201, 202, 203], false); assert!(result.is_ok()); - assert_eq!(result.unwrap().into_string(), "😀"); + assert_eq!(String::from(result.unwrap()), "😀"); } /// Without the fix, fails with "incomplete utf-8 byte sequence" from CoreBPE::decode(). @@ -590,7 +590,7 @@ mod tests { decode_result.is_partial(), "trailing incomplete byte should produce DecodeResult::Partial" ); - let text = decode_result.into_string(); + let text: String = decode_result.into(); assert!( text.starts_with("hello"), "should start with 'hello', got: {:?}", @@ -658,10 +658,10 @@ mod tests { assert_eq!(encodings.len(), 2); for (encoding, input) in encodings.iter().zip(inputs.iter()) { - let decoded = tokenizer + let decoded: String = tokenizer .decode(encoding.token_ids(), false) .unwrap() - .into_string(); + .into(); assert_eq!(decoded, *input); } } diff --git a/lib/llm/tests/test_stop_behavior.rs b/lib/llm/tests/test_stop_behavior.rs index 7a318fd4c2c9..1a45638efd65 100644 --- a/lib/llm/tests/test_stop_behavior.rs +++ b/lib/llm/tests/test_stop_behavior.rs @@ -37,7 +37,7 @@ impl tokenizer_traits::Decoder for TestTokenizer { _ => Some("?"), }) .collect(); - Ok(tokenizer_traits::DecodeResult::classify(text)) + Ok(text.into()) } } diff --git a/lib/llm/tests/tokenizers.rs b/lib/llm/tests/tokenizers.rs index 0e81acc46865..97e612f66d6c 100644 --- a/lib/llm/tests/tokenizers.rs +++ b/lib/llm/tests/tokenizers.rs @@ -113,10 +113,10 @@ fn test_encode_decode_roundtrip(#[case] tokenizer: Arc) { .unwrap_or_else(|e| panic!("Failed to encode '{text}': {e}")); assert!(!encoding.token_ids().is_empty()); - let decoded = tokenizer + let decoded: String = tokenizer .decode(encoding.token_ids(), false) .unwrap_or_else(|e| panic!("Failed to decode '{text}': {e}")) - .into_string(); + .into(); assert_eq!(decoded, text, "Roundtrip failed for: '{text}'"); } } @@ -130,10 +130,10 @@ fn test_encode_decode_roundtrip_multibyte(#[case] tokenizer: Arc) .encode(text) .unwrap_or_else(|e| panic!("Failed to encode '{text}': {e}")); - let decoded = tokenizer + let decoded: String = tokenizer .decode(encoding.token_ids(), false) .unwrap_or_else(|e| panic!("Failed to decode '{text}': {e}")) - .into_string(); + .into(); assert_eq!(decoded, text, "Roundtrip failed for: '{text}'"); } } @@ -149,10 +149,10 @@ fn test_batch_encode_roundtrip(#[case] tokenizer: Arc) { assert_eq!(encodings.len(), inputs.len()); for (encoding, &input) in encodings.iter().zip(inputs.iter()) { - let decoded = tokenizer + let decoded: String = tokenizer .decode(encoding.token_ids(), false) .expect("Failed to decode") - .into_string(); + .into(); assert_eq!(decoded, input); } } @@ -357,16 +357,16 @@ fn test_decode_with_skip_special_tokens() { token_ids.push(2); // // Decode with skip_special_tokens = false (should keep special tokens) - let decoded_with_special = tokenizer + let decoded_with_special: String = tokenizer .decode(&token_ids, false) .expect("Failed to decode with skip_special_tokens=false") - .into_string(); + .into(); // Decode with skip_special_tokens = true (should remove special tokens) - let decoded_without_special = tokenizer + let decoded_without_special: String = tokenizer .decode(&token_ids, true) .expect("Failed to decode with skip_special_tokens=true") - .into_string(); + .into(); // Validate exact matches on the entire decoded strings assert_eq!(decoded_with_special, " Hello world"); From 7f75ece4846e967701a01da2345dd0963be758a7 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Sat, 11 Apr 2026 14:04:26 -0700 Subject: [PATCH 4/4] fix(tokenizer): use from_utf8 fast path to correctly classify legitimate U+FFFD tokens --- lib/llm/src/tokenizers/fastokens.rs | 5 +-- lib/llm/src/tokenizers/tiktoken.rs | 50 +++++++++++++++++++++++++---- 2 files changed, 45 insertions(+), 10 deletions(-) diff --git a/lib/llm/src/tokenizers/fastokens.rs b/lib/llm/src/tokenizers/fastokens.rs index cc805fccabcf..25d8a7d83f11 100644 --- a/lib/llm/src/tokenizers/fastokens.rs +++ b/lib/llm/src/tokenizers/fastokens.rs @@ -81,10 +81,7 @@ mod tests { let text = "Hello, world!"; let encoding = tokenizer.encode(text).unwrap(); assert!(!encoding.token_ids().is_empty()); - let decoded: String = tokenizer - .decode(encoding.token_ids(), true) - .unwrap() - .into(); + let decoded: String = tokenizer.decode(encoding.token_ids(), true).unwrap().into(); assert!(!decoded.is_empty()); // The decoded text should contain the same non-space characters let enc_chars: String = text.chars().filter(|c| !c.is_whitespace()).collect(); diff --git a/lib/llm/src/tokenizers/tiktoken.rs b/lib/llm/src/tokenizers/tiktoken.rs index 2f376c95caa9..9d311c971e35 100644 --- a/lib/llm/src/tokenizers/tiktoken.rs +++ b/lib/llm/src/tokenizers/tiktoken.rs @@ -100,13 +100,22 @@ impl Decoder for TikTokenTokenizer { token_ids.to_vec() }; - // Use lossy UTF-8 conversion so that partial multi-byte sequences become U+FFFD. - // This is critical for incremental detokenization: DecodeStream::step() relies on - // DecodeResult::Partial to detect incomplete sequences and buffer tokens until - // a complete character arrives. CoreBPE::decode() would error on invalid UTF-8 instead. + // Try strict UTF-8 first: valid bytes get `Complete` with zero extra allocation + // (takes ownership of the Vec). This correctly handles vocabulary tokens whose + // raw bytes are EF BF BD (legitimate U+FFFD) -- they are valid UTF-8 and must + // not be confused with incomplete multi-byte sequences. + // + // On failure, fall back to lossy conversion so partial multi-byte sequences + // become U+FFFD, then classify via the trailing-FFFD heuristic. This path is + // only hit during incremental detokenization of byte-fallback tokens. let bytes: Vec = self.bpe._decode_native_and_split(ids).flatten().collect(); - let text = String::from_utf8_lossy(&bytes).into_owned(); - Ok(text.into()) + match String::from_utf8(bytes) { + Ok(text) => Ok(DecodeResult::Complete(text)), + Err(e) => { + let text = String::from_utf8_lossy(e.as_bytes()).into_owned(); + Ok(DecodeResult::from_decoded(text)) + } + } } } @@ -491,6 +500,15 @@ mod tests { content.push_str(&format!("{encoded} {rank}\n")); } + // Legitimate U+FFFD token: valid UTF-8 bytes EF BF BD (replacement character + // as an actual vocabulary entry, not an artifact of lossy conversion) + let fffd_token: Vec<(Vec, u32)> = vec![(vec![0xEF, 0xBF, 0xBD], 300)]; + + for (token, rank) in &fffd_token { + let encoded = engine.encode(token); + content.push_str(&format!("{encoded} {rank}\n")); + } + let file_path = dir.join("tiktoken.model"); let mut file = std::fs::File::create(&file_path).unwrap(); file.write_all(content.as_bytes()).unwrap(); @@ -566,6 +584,26 @@ mod tests { assert_eq!(String::from(result.unwrap()), "😀"); } + /// Regression test: a vocabulary token whose raw bytes are EF BF BD (the valid + /// UTF-8 encoding of U+FFFD) must decode as `Complete`, not `Partial`. Before the + /// from_utf8 fast-path fix, from_utf8_lossy + the trailing-FFFD heuristic would + /// misclassify this as Partial, causing the incremental decoder to suppress it. + #[test] + fn test_decode_legitimate_replacement_char_token_is_complete() { + let dir = tempfile::tempdir().unwrap(); + let tokenizer = create_byte_token_tokenizer(dir.path()); + + let result = tokenizer.decode(&[300], false); + assert!(result.is_ok()); + let decode_result = result.unwrap(); + assert!( + decode_result.is_complete(), + "legitimate U+FFFD vocab token must be Complete, got: {:?}", + decode_result + ); + assert_eq!(decode_result.as_str(), "\u{FFFD}"); + } + /// Without the fix, fails with "incomplete utf-8 byte sequence" from CoreBPE::decode(). #[test] fn test_decode_partial_emoji_does_not_error() {