feat(inference): 添加流式输出功能,支持实时解码文本
This commit is contained in:
@@ -14,3 +14,8 @@
|
|||||||
- 三个 crate 版本升级到 0.1.1,convert/inference 的 minicpm-core 依赖也更新到 0.1.1
|
- 三个 crate 版本升级到 0.1.1,convert/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` 实时输出
|
||||||
|
|||||||
@@ -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
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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(())
|
||||||
|
|||||||
Reference in New Issue
Block a user