feat(inference): 添加流式输出功能,支持实时解码文本

This commit is contained in:
2026-07-01 17:54:18 +08:00
parent 18a5ca103c
commit faeeaf7409
3 changed files with 118 additions and 3 deletions
+5
View File
@@ -14,3 +14,8 @@
- 三个 crate 版本升级到 0.1.1convert/inference 的 minicpm-core 依赖也更新到 0.1.1 - 三个 crate 版本升级到 0.1.1convert/inference 的 minicpm-core 依赖也更新到 0.1.1
- 发布顺序:`cargo publish --allow-dirty --registry gitea -p minicpm-core``minicpm-convert``minicpm-inference` - 发布顺序:`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` 实时输出
+102
View File
@@ -139,6 +139,85 @@ pub fn generate_with_cache<B: Backend>(
output_ids output_ids
} }
/// 流式生成:每生成一个新 token,立即调用 `on_token` 回调输出该 token 的解码文本。
/// 返回完整的 output token IDs。
pub fn generate_stream<B: Backend>(
model: &LlamaForCausalLM<B>,
tokenizer: &TokenizerWrapper,
input_ids: &[u32],
config: &GenerationConfig,
eos_token_id: &EosTokenId,
device: &B::Device,
mut on_token: impl FnMut(&str),
) -> Vec<u32> {
let mut output_ids = input_ids.to_vec();
let mut cache: Option<LlamaKVCache<B>> = None;
// 第一步:完整 prompt 输入
{
let input_ints: Vec<i64> = input_ids.iter().map(|&x| x as i64).collect();
let input_tensor =
Tensor::<B, 1, Int>::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<i64> = vec![last_token as i64];
let input_tensor =
Tensor::<B, 1, Int>::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<B: Backend>(logits: &Tensor<B, 3, Float>, temperature: f32, top_p: f32) -> u32 { fn sample_last<B: Backend>(logits: &Tensor<B, 3, Float>, temperature: f32, top_p: f32) -> u32 {
let shape = logits.shape(); let shape = logits.shape();
let dims: [usize; 3] = shape.dims(); let dims: [usize; 3] = shape.dims();
@@ -269,6 +348,29 @@ impl<B: Backend> MiniCPM<B> {
self.tokenizer.decode(new_ids, true) 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<String> {
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 { pub fn config(&self) -> &LlamaConfig {
&self.config &self.config
} }
+11 -3
View File
@@ -25,7 +25,7 @@ fn main() -> anyhow::Result<()> {
println!("模型加载完成,耗时: {:.2?}", start.elapsed()); println!("模型加载完成,耗时: {:.2?}", start.elapsed());
let prompt = "你知道今天的天气不"; let prompt = "我怕人家一说要改前端文件,明天咔嚓甩给我了";
println!("\n用户: {}", prompt); println!("\n用户: {}", prompt);
println!("Assistant: "); println!("Assistant: ");
@@ -36,10 +36,18 @@ fn main() -> anyhow::Result<()> {
}; };
let start = Instant::now(); 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(); let elapsed = start.elapsed();
println!("{}", response); println!();
println!("\n生成耗时: {:.2?}", elapsed); println!("\n生成耗时: {:.2?}", elapsed);
Ok(()) Ok(())