From 7c1ce53e1aaf7b5bfa6e569a088e11f7d14c5f80 Mon Sep 17 00:00:00 2001 From: fandral Date: Fri, 23 Oct 2026 10:54:50 +0800 Subject: [PATCH] fix(stream): align logprobs with retained parser segments --- deepseek-recipe/src/stream/mod.rs | 8 +- deepseek-recipe/src/stream/processor.rs | 73 +++- deepseek-recipe/tests/logprob_alignment.rs | 367 +++++++++++++++++++++ docs/streaming.md | 7 + docs/streaming.zh.md | 2 + 5 files changed, 443 insertions(+), 14 deletions(-) create mode 100644 deepseek-recipe/tests/logprob_alignment.rs diff --git a/deepseek-recipe/src/stream/mod.rs b/deepseek-recipe/src/stream/mod.rs index ce2f60e..5efdef2 100644 --- a/deepseek-recipe/src/stream/mod.rs +++ b/deepseek-recipe/src/stream/mod.rs @@ -59,9 +59,11 @@ impl std::error::Error for StreamError {} /// `ToolCallArgumentsDelta` chunks; output truncated inside a tool call ends /// that sequence early. /// -/// `logprobs` holds the log probabilities of the tokens of the decoded text the -/// content comes from. Content split from one decoded text carries the log -/// probabilities of the whole text. +/// `logprobs` holds original token probabilities only for complete tokens +/// inside the emitted source-text segment. Tokens crossing a parser segment +/// boundary are omitted, as are probabilities for replaced text (such as JSON +/// padding). If vocabulary bytes do not match decoded text, probabilities for +/// that decoded piece are omitted. No token probability is duplicated. #[derive(Debug, Clone, PartialEq)] pub enum OutputChunk { /// Output begins. diff --git a/deepseek-recipe/src/stream/processor.rs b/deepseek-recipe/src/stream/processor.rs index 870a419..a8a80f7 100644 --- a/deepseek-recipe/src/stream/processor.rs +++ b/deepseek-recipe/src/stream/processor.rs @@ -150,7 +150,7 @@ where .map(|decoded| OutputPiece { content: decoded.content, token_count: decoded.token_count, - logprobs: decoded.logprobs.map(Arc::from), + logprobs: decoded.logprobs, kind: if decoded.is_special_token { ChunkKind::SpecialToken } else { @@ -237,15 +237,61 @@ where struct OutputPiece { content: String, token_count: usize, - logprobs: Option>, + logprobs: Option>, kind: ChunkKind, } -/// Output text waiting for its parsing actions, with the log probabilities of -/// the text it was split from. +/// Output text waiting for its parsing actions and the remaining token probabilities. struct BufferedChunk { content: String, - logprobs: Option>, + logprobs: Option, +} + +/// Consume each token probability once, using its original byte range. +struct BufferedLogprobs { + tokens: VecDeque, + /// Bytes already consumed from the front token. A split token is omitted. + offset: usize, +} + +impl BufferedLogprobs { + fn new(content: &str, logprobs: Vec) -> Option { + // A custom decoder can transform vocabulary bytes. Without matching + // bytes there is no reliable alignment; do not invent probabilities. + if logprobs.iter().any(|token| token.bytes.is_empty()) + || !logprobs + .iter() + .flat_map(|token| token.bytes.iter().copied()) + .eq(content.bytes()) + { + return None; + } + Some(Self { + tokens: logprobs.into(), + offset: 0, + }) + } + + fn take(&mut self, mut byte_len: usize) -> Option> { + let mut taken = Vec::new(); + while let Some(token) = self.tokens.front() { + let remaining = token.bytes.len() - self.offset; + if remaining > byte_len { + self.offset += byte_len; + break; + } + let token = self.tokens.pop_front().expect("the front token exists"); + if self.offset == 0 { + taken.push(token); + } + self.offset = 0; + byte_len -= remaining; + if byte_len == 0 { + break; + } + } + (!taken.is_empty()).then(|| taken.into()) + } } /// Output text that the parser has not fully assigned yet, and the state @@ -280,7 +326,8 @@ impl PendingOutput { } } - fn push(&mut self, content: String, logprobs: Option>) { + fn push(&mut self, content: String, logprobs: Option>) { + let logprobs = logprobs.and_then(|logprobs| BufferedLogprobs::new(&content, logprobs)); self.chunks.push_back(BufferedChunk { content, logprobs }); } @@ -290,20 +337,24 @@ impl PendingOutput { /// Take the first `segment.byte_len` bytes of the front buffered chunk and /// convert them by the segment's action. `segment.byte_len` must not exceed - /// the length of the front chunk. Text split off from a chunk keeps the - /// chunk's log probabilities. + /// the length of the front chunk. Only complete tokens inside the taken + /// byte range keep their probabilities; a token crossing the split is omitted. fn take_piece(&mut self, segment: OutputActionSegment) -> Option { let BufferedChunk { content: mut piece, - logprobs, + mut logprobs, } = self.chunks.pop_front()?; + let piece_logprobs = logprobs + .as_mut() + .and_then(|logprobs| logprobs.take(segment.byte_len)); if segment.byte_len < piece.len() { let remaining = piece.split_off(segment.byte_len); self.chunks.push_front(BufferedChunk { content: remaining, - logprobs: logprobs.clone(), + logprobs, }); } + let logprobs = piece_logprobs; let is_first_piece = self.previous_piece_action != segment.action; self.previous_piece_action = segment.action; @@ -321,7 +372,7 @@ impl PendingOutput { } OutputAction::ToSpace => Some(OutputChunk::Content { content: " ".to_string(), - logprobs, + logprobs: None, }), OutputAction::Skip | OutputAction::SkipInvalid { .. } => None, OutputAction::StopSequence => { diff --git a/deepseek-recipe/tests/logprob_alignment.rs b/deepseek-recipe/tests/logprob_alignment.rs new file mode 100644 index 0000000..253c031 --- /dev/null +++ b/deepseek-recipe/tests/logprob_alignment.rs @@ -0,0 +1,367 @@ +use deepseek_recipe::stream::{ + ChunkGenerator, InferenceChunk, InferenceFinishReason, InferenceTokenLogprob, OutputChunk, + ParsingOptions, StreamProcessor, StreamTokenizer, +}; +use tokio_stream::StreamExt; + +struct AsciiTokenizer; +impl StreamTokenizer for AsciiTokenizer { + fn decode_ids(&self, ids: &[u32], _: bool) -> Result { + Ok(ids.iter().map(|&id| char::from_u32(id).unwrap()).collect()) + } + fn id_to_token(&self, id: u32) -> Option { + char::from_u32(id).map(|ch| ch.to_string()) + } +} +struct Passthrough; +impl ChunkGenerator for Passthrough { + type Chunk = OutputChunk; + async fn generate(&mut self, chunk: OutputChunk) -> Vec { + vec![chunk] + } +} +async fn run(text: &str, batch: bool, stop: Vec) -> (String, Vec) { + let ids: Vec = text.chars().map(u32::from).collect(); + let batches = if batch { + vec![ids] + } else { + ids.into_iter().map(|id| vec![id]).collect() + }; + let mut input: Vec = batches + .into_iter() + .map(|token_ids| { + let logprobs = token_ids + .iter() + .map(|_| InferenceTokenLogprob { + logprob: -1.0, + top_logprobs: vec![], + }) + .collect(); + InferenceChunk::Tokens { + token_ids, + logprobs: Some(logprobs), + } + }) + .collect(); + input.push(InferenceChunk::Finish { + finish_reason: InferenceFinishReason::Stop, + }); + let processor = StreamProcessor::new( + Passthrough, + ParsingOptions { + stop_sequences: stop, + ..Default::default() + }, + ) + .with_tokenizer(AsciiTokenizer); + let stream = processor.into_protocol_stream(tokio_stream::iter(input)); + let chunks: Vec<_> = std::pin::pin!(stream).collect().await; + let mut text = String::new(); + let mut tokens = vec![]; + for chunk in chunks { + if let OutputChunk::Content { content, logprobs } = chunk.unwrap() { + text.push_str(&content); + if let Some(entries) = logprobs { + tokens.extend(entries.iter().map(|entry| entry.token.clone())); + } + } + } + (text, tokens) +} +#[tokio::test] +async fn stop_filtered_logprobs_do_not_depend_on_backend_batching() { + let batched = run("ABSTOPZ", true, vec!["STOP".into()]).await; + let single = run("ABSTOPZ", false, vec!["STOP".into()]).await; + assert_eq!(batched.0, "AB"); + assert_eq!(batched, single); + assert_eq!(batched.1, vec!["A", "B"]); +} +#[tokio::test] +async fn stripped_leading_newline_has_no_content_logprob() { + let batched = run("\nAB", true, vec![]).await; + let single = run("\nAB", false, vec![]).await; + assert_eq!(batched.0, "AB"); + assert_eq!(batched, single); +} + +#[tokio::test] +async fn json_padding_and_labels_do_not_duplicate_logprobs() { + let text = "prefix{\"x\":1}\n```suffix"; + let ids: Vec = text.chars().map(u32::from).collect(); + for batch_size in 1..=ids.len() { + let input = ids + .chunks(batch_size) + .map(|ids| InferenceChunk::Tokens { + token_ids: ids.to_vec(), + logprobs: Some( + ids.iter() + .map(|_| InferenceTokenLogprob { + logprob: -1.0, + top_logprobs: vec![], + }) + .collect(), + ), + }) + .chain([InferenceChunk::Finish { + finish_reason: InferenceFinishReason::Stop, + }]); + let processor = StreamProcessor::new( + Passthrough, + ParsingOptions { + parse_json_output: true, + ..Default::default() + }, + ) + .with_tokenizer(AsciiTokenizer); + let stream = processor.into_protocol_stream(tokio_stream::iter(input)); + let output: Vec<_> = std::pin::pin!(stream).collect().await; + let mut tokens = Vec::new(); + for chunk in output { + if let OutputChunk::Content { + logprobs: Some(logprobs), + .. + } = chunk.unwrap() + { + tokens.extend(logprobs.iter().cloned()); + } + } + assert_eq!( + tokens + .iter() + .map(|token| token.token.as_str()) + .collect::(), + "{\"x\":1}", + "batch size {batch_size}" + ); + } +} + +struct WholePieceTokenizer(&'static str); +impl StreamTokenizer for WholePieceTokenizer { + fn decode_ids(&self, _: &[u32], _: bool) -> Result { + Ok("ABSTOPZ".into()) + } + fn id_to_token(&self, _: u32) -> Option { + Some(self.0.into()) + } +} +#[tokio::test] +async fn a_token_crossing_the_stop_boundary_is_not_fabricated_or_repeated() { + let processor = StreamProcessor::new( + Passthrough, + ParsingOptions { + stop_sequences: vec!["STOP".into()], + ..Default::default() + }, + ) + .with_tokenizer(WholePieceTokenizer("ABSTOPZ")); + let input = tokio_stream::iter(vec![ + InferenceChunk::Tokens { + token_ids: vec![1], + logprobs: Some(vec![InferenceTokenLogprob { + logprob: -0.25, + top_logprobs: vec![], + }]), + }, + InferenceChunk::Finish { + finish_reason: InferenceFinishReason::Stop, + }, + ]); + let output: Vec<_> = std::pin::pin!(processor.into_protocol_stream(input)) + .collect() + .await; + let content: Vec<_> = output + .into_iter() + .map(Result::unwrap) + .filter_map(|chunk| { + if let OutputChunk::Content { content, logprobs } = chunk { + Some((content, logprobs)) + } else { + None + } + }) + .collect(); + assert_eq!(content.len(), 1); + assert_eq!(content[0].0, "AB"); + assert!(content[0].1.is_none()); +} + +#[tokio::test] +async fn multibyte_token_logprobs_survive_stop_filtering_at_every_token_split() { + let tokenizer = tokenizers::Tokenizer::from_file(concat!( + env!("CARGO_MANIFEST_DIR"), + "/../static/tokenizers/v41/tokenizer.json" + )) + .unwrap(); + let text = "你好🤭STOPsuffix"; + let ids = tokenizer.encode(text, false).unwrap().get_ids().to_vec(); + for split in 0..=ids.len() { + let input = [&ids[..split], &ids[split..]] + .into_iter() + .map(|ids| InferenceChunk::Tokens { + token_ids: ids.to_vec(), + logprobs: Some( + ids.iter() + .map(|_| InferenceTokenLogprob { + logprob: -1.0, + top_logprobs: vec![], + }) + .collect(), + ), + }) + .chain([InferenceChunk::Finish { + finish_reason: InferenceFinishReason::Stop, + }]); + let processor = StreamProcessor::new( + Passthrough, + ParsingOptions { + stop_sequences: vec!["STOP".into()], + ..Default::default() + }, + ) + .with_tokenizer(tokenizer.clone()); + let output: Vec<_> = + std::pin::pin!(processor.into_protocol_stream(tokio_stream::iter(input))) + .collect() + .await; + let mut content = String::new(); + let mut bytes = Vec::new(); + for chunk in output { + if let OutputChunk::Content { + content: part, + logprobs, + } = chunk.unwrap() + { + content.push_str(&part); + if let Some(logprobs) = logprobs { + bytes.extend( + logprobs + .iter() + .flat_map(|token| token.bytes.iter().copied()), + ); + } + } + } + assert_eq!(content, "你好🤭", "split {split}"); + assert_eq!(bytes, "你好🤭".as_bytes(), "split {split}"); + } +} + +#[tokio::test] +async fn vocabulary_bytes_that_do_not_match_decoded_text_are_not_reported() { + let processor = StreamProcessor::new(Passthrough, ParsingOptions::default()) + .with_tokenizer(WholePieceTokenizer("different bytes")); + let input = tokio_stream::iter(vec![ + InferenceChunk::Tokens { + token_ids: vec![1], + logprobs: Some(vec![InferenceTokenLogprob { + logprob: -0.25, + top_logprobs: vec![], + }]), + }, + InferenceChunk::Finish { + finish_reason: InferenceFinishReason::Stop, + }, + ]); + let output: Vec<_> = std::pin::pin!(processor.into_protocol_stream(input)) + .collect() + .await; + for chunk in output { + if let OutputChunk::Content { content, logprobs } = chunk.unwrap() { + assert_eq!(content, "ABSTOPZ"); + assert!(logprobs.is_none()); + } + } +} + +async fn public_chunks( + generator: G, + options: ParsingOptions, +) -> Vec +where + G::Chunk: serde::Serialize + Send, +{ + let ids: Vec = "ABSTOPZ".chars().map(u32::from).collect(); + let input = tokio_stream::iter(vec![ + InferenceChunk::Tokens { + logprobs: Some( + ids.iter() + .map(|_| InferenceTokenLogprob { + logprob: -1.0, + top_logprobs: vec![], + }) + .collect(), + ), + token_ids: ids, + }, + InferenceChunk::Finish { + finish_reason: InferenceFinishReason::Stop, + }, + ]); + let processor = StreamProcessor::new(generator, options).with_tokenizer(AsciiTokenizer); + std::pin::pin!(processor.into_protocol_stream(input)) + .map(|chunk| serde_json::to_value(chunk.unwrap()).unwrap()) + .collect() + .await +} +#[tokio::test] +async fn chat_completion_wire_logprobs_match_retained_content() { + use deepseek_recipe::openai::ChatCompletionRequest; + use deepseek_recipe::request::{ConversionOptions, ProtocolRequest}; + let request: ChatCompletionRequest = serde_json::from_value(serde_json::json!({ + "model":"DeepSeek-V4.1-Flash", "messages":[{"role":"user","content":"Hello"}], + "thinking":{"type":"disabled"}, "logprobs":true, "stop":"STOP" + })) + .unwrap(); + let converted = request.convert(ConversionOptions::default()).unwrap(); + let generator = ChatCompletionRequest::chunk_generator( + &converted, + "test-id".into(), + "DeepSeek-V4.1-Flash".into(), + ); + let chunks = public_chunks(generator, converted.parsing_options).await; + let text: String = chunks + .iter() + .filter_map(|chunk| chunk["choices"][0]["delta"]["content"].as_str()) + .collect(); + let tokens: String = chunks + .iter() + .filter_map(|chunk| chunk["choices"][0]["logprobs"]["content"].as_array()) + .flatten() + .map(|entry| entry["token"].as_str().unwrap()) + .collect(); + assert_eq!(text, "AB"); + assert_eq!(tokens, "AB"); +} +#[tokio::test] +async fn responses_wire_logprobs_match_retained_content() { + use deepseek_recipe::openai::ResponsesRequest; + use deepseek_recipe::request::{ConversionOptions, ProtocolRequest}; + let request: ResponsesRequest = serde_json::from_value(serde_json::json!({ + "model":"DeepSeek-V4.1-Flash", "input":"Hello", "reasoning":{"effort":"none"}, "top_logprobs":0 + })).unwrap(); + let mut converted = request.convert(ConversionOptions::default()).unwrap(); + // Responses has no request-level stop field; the caller can configure parsing. + converted.parsing_options.stop_sequences = vec!["STOP".into()]; + let generator = ResponsesRequest::chunk_generator( + &converted, + "test-id".into(), + "DeepSeek-V4.1-Flash".into(), + ); + let chunks = public_chunks(generator, converted.parsing_options).await; + let deltas: Vec<_> = chunks + .iter() + .filter(|chunk| chunk["type"] == "response.output_text.delta") + .collect(); + let text: String = deltas + .iter() + .map(|chunk| chunk["delta"].as_str().unwrap()) + .collect(); + let tokens: String = deltas + .iter() + .flat_map(|chunk| chunk["logprobs"].as_array().unwrap()) + .map(|entry| entry["token"].as_str().unwrap()) + .collect(); + assert_eq!(text, "AB"); + assert_eq!(tokens, "AB"); +} diff --git a/docs/streaming.md b/docs/streaming.md index 248f7c3..e9b2901 100644 --- a/docs/streaming.md +++ b/docs/streaming.md @@ -8,6 +8,13 @@ Completions chunks. The application supplies the backend and HTTP transport. In thinking mode, the processor counts `reasoning_tokens` for usage only when a tokenizer is attached; see [Use with tokenizer](tokenizer.md#decoding). +Token log probabilities follow the original bytes retained by the parser. +Only complete tokens within an emitted source-text segment are reported; tokens +crossing a parser segment boundary are omitted. Removed text and replacement +spaces around JSON output carry no probabilities. If a custom tokenizer's +vocabulary bytes do not match its decoded text, probabilities for that decoded +piece are omitted. Completion usage still counts the processed backend tokens. + ## Rust Add the dependencies to your application: diff --git a/docs/streaming.zh.md b/docs/streaming.zh.md index aa0f2ca..70b7c11 100644 --- a/docs/streaming.zh.md +++ b/docs/streaming.zh.md @@ -4,6 +4,8 @@ `StreamProcessor` 把后端输出转换为 Messages、Chat Completions 或 Responses 事件。以下示例使用 mock 推理并打印 Chat Completions chunk。后端与 HTTP 传输由应用提供。思考模式下，只有为 processor 附加 tokenizer 后，它才会统计 usage 中的 `reasoning_tokens`，见[使用 tokenizer](tokenizer.zh.md#解码)。 +Token 的 log probability 按解析器保留的原始字节归属。只报告完整落在输出原文片段内的 token；跨解析片段边界的 token 不报告概率。被删除的文本和 JSON 输出周围生成的替代空格不带概率。若自定义 tokenizer 的词表字节与解码文本不一致，该解码片段不报告概率。completion usage 仍按已处理的后端 token 计数。 + ## Rust 为你的应用添加依赖：