From faeeaf740984908643d673d547e484cd93e6932d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=8E=8B=E5=AE=8F=E5=BB=BA?= Date: Wed, 1 Jul 2026 17:54:18 +0800 Subject: [PATCH] =?UTF-8?q?feat(inference):=20=E6=B7=BB=E5=8A=A0=E6=B5=81?= =?UTF-8?q?=E5=BC=8F=E8=BE=93=E5=87=BA=E5=8A=9F=E8=83=BD=EF=BC=8C=E6=94=AF?= =?UTF-8?q?=E6=8C=81=E5=AE=9E=E6=97=B6=E8=A7=A3=E7=A0=81=E6=96=87=E6=9C=AC?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .workbuddy/memory/2026-07-01.md | 5 ++ crates/minicpm-inference/src/lib.rs | 102 +++++++++++++++++++++++++ examples/minimal-inference/src/main.rs | 14 +++- 3 files changed, 118 insertions(+), 3 deletions(-) diff --git a/.workbuddy/memory/2026-07-01.md b/.workbuddy/memory/2026-07-01.md index b55842a..eaacc47 100644 --- a/.workbuddy/memory/2026-07-01.md +++ b/.workbuddy/memory/2026-07-01.md @@ -14,3 +14,8 @@ - 三个 crate 版本升级到 0.1.1，convert/inference 的 minicpm-core 依赖也更新到 0.1.1 - 发布顺序：`cargo publish --allow-dirty --registry gitea -p minicpm-core` → `minicpm-convert` → `minicpm-inference` - 全部发布成功 + +## 添加流式输出 + +- `minicpm-inference/src/lib.rs`: 新增 `generate_stream` 函数和 `MiniCPM::generate_stream` 方法，每 token 通过回调 `impl FnMut(&str)` 输出解码文本 +- `examples/minimal-inference/src/main.rs`: 改为流式调用，`print!` + `flush` 实时输出 diff --git a/crates/minicpm-inference/src/lib.rs b/crates/minicpm-inference/src/lib.rs index 1ed0737..1bbcfad 100644 --- a/crates/minicpm-inference/src/lib.rs +++ b/crates/minicpm-inference/src/lib.rs @@ -139,6 +139,85 @@ pub fn generate_with_cache( output_ids } +/// 流式生成：每生成一个新 token，立即调用 `on_token` 回调输出该 token 的解码文本。 +/// 返回完整的 output token IDs。 +pub fn generate_stream( + model: &LlamaForCausalLM, + tokenizer: &TokenizerWrapper, + input_ids: &[u32], + config: &GenerationConfig, + eos_token_id: &EosTokenId, + device: &B::Device, + mut on_token: impl FnMut(&str), +) -> Vec { + let mut output_ids = input_ids.to_vec(); + let mut cache: Option> = None; + + // 第一步：完整 prompt 输入 + { + let input_ints: Vec = input_ids.iter().map(|&x| x as i64).collect(); + let input_tensor = + Tensor::::from_ints(input_ints.as_slice(), device) + .unsqueeze::<2>(); + + let (logits, new_cache) = model.forward_with_cache(input_tensor, cache.as_ref()); + cache = Some(new_cache); + + let next_token = sample_last(&logits, config.temperature, config.top_p); + output_ids.push(next_token); + + // 流式输出第一个 token + if let Ok(text) = tokenizer.decode(&[next_token], true) { + on_token(&text); + } + + if eos_token_id.contains(next_token) { + return output_ids; + } + } + + // 后续步骤：每步生成一个 token + let mut count = 1; + loop { + if let Some(max) = config.max_new_tokens { + if count >= max { + break; + } + } + + let last_token = output_ids[output_ids.len() - 1]; + let input_ints: Vec = vec![last_token as i64]; + let input_tensor = + Tensor::::from_ints(input_ints.as_slice(), device) + .unsqueeze::<2>(); + + let (logits, new_cache) = model.forward_with_cache(input_tensor, cache.as_ref()); + cache = Some(new_cache); + + let vocab_size = model.config.vocab_size; + let next_token_logits = logits + .clone() + .slice([0..1, 0..1, 0..vocab_size]) + .reshape([vocab_size]); + + let next_token = sample(&next_token_logits, config.temperature, config.top_p); + output_ids.push(next_token); + + // 流式输出新 token + if let Ok(text) = tokenizer.decode(&[next_token], true) { + on_token(&text); + } + + if eos_token_id.contains(next_token) { + break; + } + + count += 1; + } + + output_ids +} + fn sample_last(logits: &Tensor, temperature: f32, top_p: f32) -> u32 { let shape = logits.shape(); let dims: [usize; 3] = shape.dims(); @@ -269,6 +348,29 @@ impl MiniCPM { self.tokenizer.decode(new_ids, true) } + /// 流式生成：每生成一个 token 立即调用 `on_token` 回调， + /// 参数为该 token 解码后的文本。 + pub fn generate_stream( + &self, + prompt: &str, + think: bool, + config: &GenerationConfig, + on_token: impl FnMut(&str), + ) -> anyhow::Result { + let input_ids = self.tokenizer.apply_chat_template(prompt, think)?; + let output_ids = generate_stream( + &self.model, + &self.tokenizer, + &input_ids, + config, + &self.config.eos_token_id, + &self.device, + on_token, + ); + let new_ids = &output_ids[input_ids.len()..]; + self.tokenizer.decode(new_ids, true) + } + pub fn config(&self) -> &LlamaConfig { &self.config } diff --git a/examples/minimal-inference/src/main.rs b/examples/minimal-inference/src/main.rs index d265c5c..1bf3469 100644 --- a/examples/minimal-inference/src/main.rs +++ b/examples/minimal-inference/src/main.rs @@ -25,7 +25,7 @@ fn main() -> anyhow::Result<()> { println!("模型加载完成，耗时: {:.2?}", start.elapsed()); - let prompt = "你知道今天的天气不"; + let prompt = "我怕人家一说要改前端文件，明天咔嚓甩给我了"; println!("\n用户: {}", prompt); println!("Assistant: "); @@ -36,10 +36,18 @@ fn main() -> anyhow::Result<()> { }; let start = Instant::now(); - let response = model.generate(prompt, false, &config)?; + let _response = model.generate_stream( + prompt, + false, + &config, + |token| { + print!("{}", token); + std::io::Write::flush(&mut std::io::stdout()).ok(); + }, + )?; let elapsed = start.elapsed(); - println!("{}", response); + println!(); println!("\n生成耗时: {:.2?}", elapsed); Ok(())