Compare commits

..

2 Commits

7 changed files with 136 additions and 14 deletions
+12
View File
@@ -7,3 +7,15 @@
- `minicpm-core@0.1.0` 已存在于 registry,跳过发布 - `minicpm-core@0.1.0` 已存在于 registry,跳过发布
- `minicpm-convert@0.1.0` 发布成功 - `minicpm-convert@0.1.0` 发布成功
- `minicpm-inference@0.1.0` 发布成功 - `minicpm-inference@0.1.0` 发布成功
## burn 最小依赖 + 重发布 (0.1.1)
- burn 改为 `default-features = false, features = ["std"]`,移除 `wgpu`wgpu 交给下游 binary 选择)
- 三个 crate 版本升级到 0.1.1convert/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` 实时输出
Generated
+3 -3
View File
@@ -3960,7 +3960,7 @@ checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a"
[[package]] [[package]]
name = "minicpm-convert" name = "minicpm-convert"
version = "0.1.0" version = "0.1.1"
dependencies = [ dependencies = [
"anyhow", "anyhow",
"burn", "burn",
@@ -3971,7 +3971,7 @@ dependencies = [
[[package]] [[package]]
name = "minicpm-core" name = "minicpm-core"
version = "0.1.0" version = "0.1.1"
dependencies = [ dependencies = [
"anyhow", "anyhow",
"burn", "burn",
@@ -3981,7 +3981,7 @@ dependencies = [
[[package]] [[package]]
name = "minicpm-inference" name = "minicpm-inference"
version = "0.1.0" version = "0.1.1"
dependencies = [ dependencies = [
"anyhow", "anyhow",
"burn", "burn",
+3 -3
View File
@@ -1,12 +1,12 @@
[package] [package]
name = "minicpm-convert" name = "minicpm-convert"
version = "0.1.0" version = "0.1.1"
edition = "2021" edition = "2021"
publish = ["gitea"] publish = ["gitea"]
[dependencies] [dependencies]
minicpm-core = { path = "../minicpm-core", version = "0.1.0", registry = "gitea" } minicpm-core = { path = "../minicpm-core", version = "0.1.1", registry = "gitea" }
burn = { version = "0.21", features = ["std", "wgpu"] } burn = { version = "0.21", default-features = false, features = ["std"] }
memmap2 = "0.9" memmap2 = "0.9"
anyhow = "1.0" anyhow = "1.0"
serde_json = "1.0" serde_json = "1.0"
+2 -2
View File
@@ -1,11 +1,11 @@
[package] [package]
name = "minicpm-core" name = "minicpm-core"
version = "0.1.0" version = "0.1.1"
edition = "2021" edition = "2021"
publish = ["gitea"] publish = ["gitea"]
[dependencies] [dependencies]
burn = { version = "0.21", features = ["std", "wgpu"] } burn = { version = "0.21", default-features = false, features = ["std"] }
serde = { version = "1.0", features = ["derive"] } serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0" serde_json = "1.0"
anyhow = "1.0" anyhow = "1.0"
+3 -3
View File
@@ -1,12 +1,12 @@
[package] [package]
name = "minicpm-inference" name = "minicpm-inference"
version = "0.1.0" version = "0.1.1"
edition = "2021" edition = "2021"
publish = ["gitea"] publish = ["gitea"]
[dependencies] [dependencies]
minicpm-core = { path = "../minicpm-core", version = "0.1.0", registry = "gitea" } minicpm-core = { path = "../minicpm-core", version = "0.1.1", registry = "gitea" }
burn = { version = "0.21", features = ["std", "wgpu"] } burn = { version = "0.21", default-features = false, features = ["std"] }
tokenizers = "0.20" tokenizers = "0.20"
rand = "0.8" rand = "0.8"
anyhow = "1.0" anyhow = "1.0"
+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(())