diff --git a/crates/inference/src/model/qwen35/detokenize.rs b/crates/inference/src/model/qwen35/detokenize.rs index a05e9e513b..1fa149a328 100644 --- a/crates/inference/src/model/qwen35/detokenize.rs +++ b/crates/inference/src/model/qwen35/detokenize.rs @@ -1,27 +1,140 @@ use crate::tokenizer::bpe::BpeTokenizer; use std::collections::HashMap; -/// Decode token IDs back to text using the BPE tokenizer's reverse lookup. -pub(crate) fn decode_tokens(tokenizer: &BpeTokenizer, ids: &[u32]) -> String { +/// Build the GPT-2 byte-level reverse map (unicode placeholder char → original byte). +pub(crate) fn build_byte_decoder() -> HashMap { let byte_encoder = bytes_to_unicode(); let mut byte_decoder: HashMap = HashMap::new(); for (byte_val, &ch) in byte_encoder.iter().enumerate() { byte_decoder.insert(ch, byte_val as u8); } + byte_decoder +} +/// Append the raw decoded bytes for a single token id onto `out`. +/// +/// Byte-level BPE decoding is context-free per token, so concatenating the +/// per-token byte strings equals decoding the whole id sequence at once. +fn append_token_bytes( + tokenizer: &BpeTokenizer, + id: u32, + byte_decoder: &HashMap, + out: &mut Vec, +) { + if let Some(token_str) = tokenizer.token_for_id(id) { + for ch in token_str.chars() { + if let Some(&b) = byte_decoder.get(&ch) { + out.push(b); + } + // Skip characters not in byte_decoder (e.g., special tokens) + } + } +} + +/// Decode token IDs back to text using the BPE tokenizer's reverse lookup. +pub(crate) fn decode_tokens(tokenizer: &BpeTokenizer, ids: &[u32]) -> String { + let byte_decoder = build_byte_decoder(); let mut bytes = Vec::new(); for &id in ids { - if let Some(token_str) = tokenizer.token_for_id(id) { - for ch in token_str.chars() { - if let Some(&b) = byte_decoder.get(&ch) { - bytes.push(b); + append_token_bytes(tokenizer, id, &byte_decoder, &mut bytes); + } + String::from_utf8_lossy(&bytes).to_string() +} + +/// Incremental, UTF-8-boundary-safe detokenizer for streaming generation. +/// +/// Byte-level BPE frequently splits a single multibyte codepoint (CJK, emoji, +/// accented Latin) across two or more tokens. Decoding the running token list +/// with `from_utf8_lossy` after every token would substitute `U+FFFD` for the +/// incomplete trailing bytes and stream that replacement char to the caller, +/// who cannot retract it. This type instead buffers raw bytes and emits only +/// the longest *complete* UTF-8 prefix after each token, holding incomplete +/// trailing bytes until a later token completes the codepoint. The concatenation +/// of every `push` delta plus the final `finish` equals `decode_tokens` over the +/// same ids. +pub(crate) struct IncrementalDetokenizer { + byte_decoder: HashMap, + bytes: Vec, + flushed: usize, +} + +impl IncrementalDetokenizer { + pub(crate) fn new() -> Self { + Self { + byte_decoder: build_byte_decoder(), + bytes: Vec::new(), + flushed: 0, + } + } + + /// Feed one token id; return the complete-UTF-8 text it reveals (which may be + /// empty when the token only contributed continuation bytes of a codepoint + /// still being assembled). + pub(crate) fn push(&mut self, tokenizer: &BpeTokenizer, id: u32) -> String { + append_token_bytes(tokenizer, id, &self.byte_decoder, &mut self.bytes); + self.flush_complete() + } + + fn flush_complete(&mut self) -> String { + let mut out = String::new(); + loop { + match std::str::from_utf8(&self.bytes[self.flushed..]) { + // Entire remaining buffer is valid UTF-8: emit it all. + Ok(s) => { + out.push_str(s); + self.flushed = self.bytes.len(); + return out; + } + Err(e) => { + let valid = e.valid_up_to(); + if valid > 0 { + // [flushed .. flushed + valid] is guaranteed valid by + // valid_up_to, so from_utf8_lossy substitutes nothing. + out.push_str( + String::from_utf8_lossy( + &self.bytes[self.flushed..self.flushed + valid], + ) + .as_ref(), + ); + self.flushed += valid; + } + match e.error_len() { + // None: the error is an incomplete trailing codepoint, + // not a malformed one. A later token may complete it, so + // keep the tail buffered and emit no replacement char. + None => return out, + // Some(len): a genuinely invalid sequence of `len` bytes + // that no future token can repair. Emit one U+FFFD (matching + // from_utf8_lossy's maximal-subpart substitution), skip past + // it, and continue — without this the stream would stall on + // the bad byte until finish(). + Some(len) => { + out.push('\u{FFFD}'); + self.flushed += len; + } + } } - // Skip characters not in byte_decoder (e.g., special tokens) } } } - String::from_utf8_lossy(&bytes).to_string() + /// Flush any trailing incomplete bytes lossily. Call once after the final + /// token so the concatenated stream matches `decode_tokens` exactly even when + /// a generation is truncated mid-codepoint. + pub(crate) fn finish(&mut self) -> String { + if self.flushed < self.bytes.len() { + let tail = String::from_utf8_lossy(&self.bytes[self.flushed..]).into_owned(); + self.flushed = self.bytes.len(); + tail + } else { + String::new() + } + } + + /// The full decoded text so far (equivalent to `decode_tokens` over every fed id). + pub(crate) fn text(&self) -> String { + String::from_utf8_lossy(&self.bytes).to_string() + } } /// GPT-2 byte-to-unicode mapping (same as in bpe.rs). @@ -48,3 +161,81 @@ pub fn bytes_to_unicode() -> Vec { } table } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn incremental_flush_reconstructs_split_codepoints() { + // 好 = E5 A5 BD, 😀 = F0 9F 98 80 — split across simulated token boundaries + // so several intermediate buffers end mid-codepoint. + let chunks: &[&[u8]] = &[&[0xE5, 0xA5], &[0xBD], &[0xF0, 0x9F], &[0x98, 0x80], b"ok"]; + let mut d = IncrementalDetokenizer::new(); + let mut out = String::new(); + for c in chunks { + d.bytes.extend_from_slice(c); + out.push_str(&d.flush_complete()); + } + out.push_str(&d.finish()); + assert_eq!(out, "好😀ok"); + assert_eq!(out, d.text()); + } + + #[test] + fn incremental_flush_truncated_mid_codepoint_matches_lossy() { + // Generation truncated after only the first 2 of 好's 3 bytes: finish() + // flushes the incomplete tail lossily, matching decode_tokens' behavior. + let mut d = IncrementalDetokenizer::new(); + let mut out = String::new(); + d.bytes.extend_from_slice(&[b'h', b'i', 0xE5, 0xA5]); + out.push_str(&d.flush_complete()); + out.push_str(&d.finish()); + assert_eq!(out, String::from_utf8_lossy(&[b'h', b'i', 0xE5, 0xA5])); + assert_eq!(out, d.text()); + } + + #[test] + fn incremental_flush_invalid_byte_does_not_stall_stream() { + // A lone continuation byte (0x80) is a malformed sequence no later token + // can complete. flush_complete must emit U+FFFD for it immediately and + // keep going, not buffer it until finish() (which would stall the stream). + let mut d = IncrementalDetokenizer::new(); + d.bytes.extend_from_slice(&[0x80]); + let first = d.flush_complete(); + assert_eq!(first, "\u{FFFD}", "invalid byte must flush immediately"); + d.bytes.extend_from_slice(b"A"); + let second = d.flush_complete(); + assert_eq!(second, "A", "valid byte after invalid one must flush"); + // Concatenated stream equals from_utf8_lossy over the whole byte buffer. + assert_eq!( + format!("{first}{second}"), + String::from_utf8_lossy(&[0x80, b'A']) + ); + } + + #[test] + fn incremental_flush_invalid_between_valid_matches_lossy() { + // Valid · invalid · valid in one buffer: one U+FFFD for the bad run, both + // valid runs verbatim — byte-for-byte equal to from_utf8_lossy. + let raw = [b'h', b'i', 0xFF, b'y', b'o']; + let mut d = IncrementalDetokenizer::new(); + d.bytes.extend_from_slice(&raw); + let out = d.flush_complete(); + assert_eq!(out, String::from_utf8_lossy(&raw)); + assert_eq!(out, d.text()); + } + + #[test] + fn incremental_flush_ascii_is_exact_per_chunk() { + // Pure-ASCII path: every chunk is whole codepoints, so each flush emits it all. + let mut d = IncrementalDetokenizer::new(); + let mut out = String::new(); + for c in [b"He".as_slice(), b"llo", b"!"] { + d.bytes.extend_from_slice(c); + out.push_str(&d.flush_complete()); + } + out.push_str(&d.finish()); + assert_eq!(out, "Hello!"); + } +} diff --git a/crates/inference/src/model/qwen35/generation.rs b/crates/inference/src/model/qwen35/generation.rs index 77be4a89c8..20e1b443dd 100644 --- a/crates/inference/src/model/qwen35/generation.rs +++ b/crates/inference/src/model/qwen35/generation.rs @@ -1,5 +1,5 @@ use super::cache::{ForwardScratch, KvCache}; -use super::detokenize::decode_tokens; +use super::detokenize::{IncrementalDetokenizer, decode_tokens}; use super::model::Qwen35Model; use super::sampling::sample_token; use crate::attention::gdn::GatedDeltaNetState; @@ -85,6 +85,143 @@ impl Qwen35Model { generated_tokens: generated_ids.len(), }) } + + /// Streaming variant of [`generate`] — identical token sequence, but invokes + /// `on_token` with incremental text deltas after each generated token. + /// + /// # Parity safety + /// + /// The body below is a deliberate copy of `generate` rather than a refactor of + /// the shared path. This ensures that no change here can silently alter the + /// non-streaming `generate` path, which is pinned by the e2e-parity CI gate + /// (greedy token match vs HF transformers). `on_token` is the only addition. + pub fn generate_streaming( + &self, + prompt: &str, + gen_cfg: &GenerateConfig, + mut on_token: impl FnMut(&str), + ) -> Result { + let cfg = &self.config; + + let mut rng_state = initial_rng_state(gen_cfg.seed); + + let input = self.tokenizer.tokenize(prompt); + let prompt_ids: Vec = input.input_ids[..input.real_length].to_vec(); + let prompt_len = prompt_ids.len(); + + if prompt_len == 0 { + return Err(InferenceError::Inference("empty prompt".into())); + } + + // max_new_tokens == 0 means "generate nothing": return before sampling so + // we never emit a token the caller did not ask for. + if gen_cfg.max_new_tokens == 0 { + return Ok(GenerateOutput { + text: String::new(), + token_ids: vec![], + prompt_tokens: prompt_len, + generated_tokens: 0, + }); + } + + let num_linear = cfg.num_linear_attention_layers(); + let num_full = cfg.num_full_attention_layers(); + let mut gdn_states: Vec = (0..num_linear) + .map(|_| GatedDeltaNetState::new(cfg)) + .collect(); + let mut kv_cache = KvCache::new(num_full); + let mut scratch = ForwardScratch::new(); + + let mut generated_ids: Vec = Vec::with_capacity(gen_cfg.max_new_tokens); + let mut all_ids = prompt_ids.clone(); + + prefill_tokens( + self, + &prompt_ids, + &mut gdn_states, + &mut kv_cache, + &mut scratch, + ); + kv_cache.seq_len = prompt_len; + + let next_id = sample_token( + &scratch.logits[..cfg.vocab_size], + gen_cfg, + &all_ids, + &mut rng_state, + ); + + if should_stop_token(cfg, gen_cfg, next_id) { + return Ok(GenerateOutput { + text: String::new(), + token_ids: vec![], + prompt_tokens: prompt_len, + generated_tokens: 0, + }); + } + + generated_ids.push(next_id); + all_ids.push(next_id); + + // Incremental detokenization: emit only complete-UTF-8 text deltas. A + // byte-level BPE codepoint can span several tokens, so we buffer raw bytes + // and never stream a partial codepoint (see IncrementalDetokenizer). + let mut detok = IncrementalDetokenizer::new(); + let delta = detok.push(&self.tokenizer, next_id); + if !delta.is_empty() { + on_token(&delta); + } + + // Decode loop (mirrors decode_loop free function exactly). + for _ in 1..gen_cfg.max_new_tokens { + let pos = kv_cache.seq_len; + let Some(&last_token) = all_ids.last() else { + return Err(InferenceError::Inference("empty generation state".into())); + }; + + self.forward_step( + last_token, + pos, + &mut gdn_states, + &mut kv_cache, + &mut scratch, + ); + kv_cache.seq_len += 1; + + let next_id = sample_token( + &scratch.logits[..cfg.vocab_size], + gen_cfg, + &all_ids, + &mut rng_state, + ); + + if should_stop_token(cfg, gen_cfg, next_id) { + break; + } + + generated_ids.push(next_id); + all_ids.push(next_id); + + let delta = detok.push(&self.tokenizer, next_id); + if !delta.is_empty() { + on_token(&delta); + } + } + + // Flush any trailing incomplete bytes (generation truncated mid-codepoint) + // so the streamed deltas concatenate to exactly the returned text. + let tail = detok.finish(); + if !tail.is_empty() { + on_token(&tail); + } + + Ok(GenerateOutput { + text: detok.text(), + token_ids: generated_ids.clone(), + prompt_tokens: prompt_len, + generated_tokens: generated_ids.len(), + }) + } } fn initial_rng_state(seed: Option) -> u64 {