refactor(项目结构) 重构为 workspace 多成员项目
- 将项目拆分为三个 crate：minicpm-core（核心模型）、minicpm-convert（转换功能）、minicpm-inference（推理功能） - 添加两个示例：minimal-inference（最小推理）和 convert（模型转换） - 转换后自动拷贝 config.json 和 tokenizer.json 到 model 目录 - 更新 README 说明 workspace 结构和使用方式
This commit is contained in:
Generated
+39
-111
@@ -84,56 +84,6 @@ dependencies = [
|
|||||||
"libc",
|
"libc",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "anstream"
|
|
||||||
version = "1.0.0"
|
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
||||||
checksum = "824a212faf96e9acacdbd09febd34438f8f711fb84e09a8916013cd7815ca28d"
|
|
||||||
dependencies = [
|
|
||||||
"anstyle",
|
|
||||||
"anstyle-parse",
|
|
||||||
"anstyle-query",
|
|
||||||
"anstyle-wincon",
|
|
||||||
"colorchoice",
|
|
||||||
"is_terminal_polyfill",
|
|
||||||
"utf8parse",
|
|
||||||
]
|
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "anstyle"
|
|
||||||
version = "1.0.14"
|
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
||||||
checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000"
|
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "anstyle-parse"
|
|
||||||
version = "1.0.0"
|
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
||||||
checksum = "52ce7f38b242319f7cabaa6813055467063ecdc9d355bbb4ce0c68908cd8130e"
|
|
||||||
dependencies = [
|
|
||||||
"utf8parse",
|
|
||||||
]
|
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "anstyle-query"
|
|
||||||
version = "1.1.5"
|
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
||||||
checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc"
|
|
||||||
dependencies = [
|
|
||||||
"windows-sys 0.61.2",
|
|
||||||
]
|
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "anstyle-wincon"
|
|
||||||
version = "3.0.11"
|
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
||||||
checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d"
|
|
||||||
dependencies = [
|
|
||||||
"anstyle",
|
|
||||||
"once_cell_polyfill",
|
|
||||||
"windows-sys 0.61.2",
|
|
||||||
]
|
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "anyhow"
|
name = "anyhow"
|
||||||
version = "1.0.103"
|
version = "1.0.103"
|
||||||
@@ -1282,46 +1232,6 @@ dependencies = [
|
|||||||
"libloading 0.8.9",
|
"libloading 0.8.9",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "clap"
|
|
||||||
version = "4.6.1"
|
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
||||||
checksum = "1ddb117e43bbf7dacf0a4190fef4d345b9bad68dfc649cb349e7d17d28428e51"
|
|
||||||
dependencies = [
|
|
||||||
"clap_builder",
|
|
||||||
"clap_derive",
|
|
||||||
]
|
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "clap_builder"
|
|
||||||
version = "4.6.0"
|
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
||||||
checksum = "714a53001bf66416adb0e2ef5ac857140e7dc3a0c48fb28b2f10762fc4b5069f"
|
|
||||||
dependencies = [
|
|
||||||
"anstream",
|
|
||||||
"anstyle",
|
|
||||||
"clap_lex",
|
|
||||||
"strsim",
|
|
||||||
]
|
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "clap_derive"
|
|
||||||
version = "4.6.1"
|
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
||||||
checksum = "f2ce8604710f6733aa641a2b3731eaa1e8b3d9973d5e3565da11800813f997a9"
|
|
||||||
dependencies = [
|
|
||||||
"heck",
|
|
||||||
"proc-macro2",
|
|
||||||
"quote",
|
|
||||||
"syn 2.0.118",
|
|
||||||
]
|
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "clap_lex"
|
|
||||||
version = "1.1.0"
|
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
||||||
checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9"
|
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "codespan-reporting"
|
name = "codespan-reporting"
|
||||||
version = "0.13.1"
|
version = "0.13.1"
|
||||||
@@ -1339,12 +1249,6 @@ version = "1.1.0"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "3d7b894f5411737b7867f4827955924d7c254fc9f4d91a6aad6b097804b1018b"
|
checksum = "3d7b894f5411737b7867f4827955924d7c254fc9f4d91a6aad6b097804b1018b"
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "colorchoice"
|
|
||||||
version = "1.0.5"
|
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
||||||
checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570"
|
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "colored"
|
name = "colored"
|
||||||
version = "3.1.1"
|
version = "3.1.1"
|
||||||
@@ -1437,6 +1341,15 @@ version = "0.6.1"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "136d3e02915a2cea4d74caa8681e2d44b1c3254bdbf17d11d41d587ff858832c"
|
checksum = "136d3e02915a2cea4d74caa8681e2d44b1c3254bdbf17d11d41d587ff858832c"
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "convert"
|
||||||
|
version = "0.1.0"
|
||||||
|
dependencies = [
|
||||||
|
"anyhow",
|
||||||
|
"burn",
|
||||||
|
"minicpm-convert",
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "convert_case"
|
name = "convert_case"
|
||||||
version = "0.8.0"
|
version = "0.8.0"
|
||||||
@@ -3608,12 +3521,6 @@ version = "2.12.0"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2"
|
checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2"
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "is_terminal_polyfill"
|
|
||||||
version = "1.70.2"
|
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
||||||
checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695"
|
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "itertools"
|
name = "itertools"
|
||||||
version = "0.11.0"
|
version = "0.11.0"
|
||||||
@@ -4052,19 +3959,46 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||||||
checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a"
|
checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a"
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "minicpm-burn"
|
name = "minicpm-convert"
|
||||||
version = "0.1.0"
|
version = "0.1.0"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"anyhow",
|
"anyhow",
|
||||||
"burn",
|
"burn",
|
||||||
"clap",
|
|
||||||
"memmap2",
|
"memmap2",
|
||||||
"rand 0.8.6",
|
"minicpm-core",
|
||||||
|
"serde_json",
|
||||||
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "minicpm-core"
|
||||||
|
version = "0.1.0"
|
||||||
|
dependencies = [
|
||||||
|
"anyhow",
|
||||||
|
"burn",
|
||||||
"serde",
|
"serde",
|
||||||
"serde_json",
|
"serde_json",
|
||||||
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "minicpm-inference"
|
||||||
|
version = "0.1.0"
|
||||||
|
dependencies = [
|
||||||
|
"anyhow",
|
||||||
|
"burn",
|
||||||
|
"minicpm-core",
|
||||||
|
"rand 0.8.6",
|
||||||
"tokenizers 0.20.4",
|
"tokenizers 0.20.4",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "minimal-inference"
|
||||||
|
version = "0.1.0"
|
||||||
|
dependencies = [
|
||||||
|
"anyhow",
|
||||||
|
"burn",
|
||||||
|
"minicpm-inference",
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "minimal-lexical"
|
name = "minimal-lexical"
|
||||||
version = "0.2.1"
|
version = "0.2.1"
|
||||||
@@ -4510,12 +4444,6 @@ version = "1.21.4"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50"
|
checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50"
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "once_cell_polyfill"
|
|
||||||
version = "1.70.2"
|
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
||||||
checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe"
|
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "oneshot"
|
name = "oneshot"
|
||||||
version = "0.2.1"
|
version = "0.2.1"
|
||||||
|
|||||||
+3
-24
@@ -1,24 +1,3 @@
|
|||||||
[package]
|
[workspace]
|
||||||
name = "minicpm-burn"
|
members = ["crates/minicpm-core", "crates/minicpm-convert", "crates/minicpm-inference", "examples/minimal-inference", "examples/convert"]
|
||||||
version = "0.1.0"
|
resolver = "2"
|
||||||
edition = "2021"
|
|
||||||
|
|
||||||
[lib]
|
|
||||||
name = "minicpm_burn_lib"
|
|
||||||
path = "src/lib.rs"
|
|
||||||
|
|
||||||
[[bin]]
|
|
||||||
name = "minicpm-burn"
|
|
||||||
path = "src/main.rs"
|
|
||||||
|
|
||||||
[dependencies]
|
|
||||||
burn = { version = "0.21", features = ["std", "wgpu"] }
|
|
||||||
tokenizers = "0.20"
|
|
||||||
serde = { version = "1.0", features = ["derive"] }
|
|
||||||
serde_json = "1.0"
|
|
||||||
anyhow = "1.0"
|
|
||||||
rand = "0.8"
|
|
||||||
memmap2 = "0.9"
|
|
||||||
clap = { version = "4.5", features = ["derive"] }
|
|
||||||
|
|
||||||
[dev-dependencies]
|
|
||||||
|
|||||||
@@ -0,0 +1,11 @@
|
|||||||
|
[package]
|
||||||
|
name = "minicpm-convert"
|
||||||
|
version = "0.1.0"
|
||||||
|
edition = "2021"
|
||||||
|
|
||||||
|
[dependencies]
|
||||||
|
minicpm-core = { path = "../minicpm-core" }
|
||||||
|
burn = { version = "0.21", features = ["std", "wgpu"] }
|
||||||
|
memmap2 = "0.9"
|
||||||
|
anyhow = "1.0"
|
||||||
|
serde_json = "1.0"
|
||||||
@@ -1,13 +1,45 @@
|
|||||||
use std::collections::HashMap;
|
use minicpm_core::config::LlamaConfig;
|
||||||
|
use minicpm_core::model::LlamaForCausalLM;
|
||||||
use crate::model::LlamaForCausalLM;
|
|
||||||
use burn::module::{Module, Param};
|
use burn::module::{Module, Param};
|
||||||
use burn::nn::{EmbeddingRecord, LinearRecord};
|
use burn::nn::{EmbeddingRecord, LinearRecord};
|
||||||
|
use burn::record::{FullPrecisionSettings, NamedMpkFileRecorder};
|
||||||
use burn::tensor::backend::Backend;
|
use burn::tensor::backend::Backend;
|
||||||
use burn::tensor::{Float, Shape, Tensor, TensorData};
|
use burn::tensor::{Float, Shape, Tensor, TensorData};
|
||||||
use memmap2::Mmap;
|
use memmap2::Mmap;
|
||||||
|
use std::collections::HashMap;
|
||||||
use std::path::Path;
|
use std::path::Path;
|
||||||
|
|
||||||
|
pub fn export_model<B: Backend>(
|
||||||
|
safetensors_path: &Path,
|
||||||
|
config_path: &Path,
|
||||||
|
tokenizer_path: &Path,
|
||||||
|
output_dir: &Path,
|
||||||
|
device: &B::Device,
|
||||||
|
) -> anyhow::Result<()> {
|
||||||
|
println!("开始转换 MiniCPM 模型为 Burn 格式...");
|
||||||
|
|
||||||
|
println!("加载配置文件: {:?}", config_path);
|
||||||
|
let config = LlamaConfig::from_json(config_path.to_str().unwrap())?;
|
||||||
|
|
||||||
|
println!("创建模型结构...");
|
||||||
|
let model = LlamaForCausalLM::<B>::new(config, device);
|
||||||
|
|
||||||
|
println!("加载 safetensors 权重...");
|
||||||
|
let model = load_safetensors(model, safetensors_path, device)?;
|
||||||
|
|
||||||
|
println!("保存为 MPK 格式...");
|
||||||
|
std::fs::create_dir_all(output_dir)?;
|
||||||
|
let output_path = output_dir.join("model");
|
||||||
|
let recorder = NamedMpkFileRecorder::<FullPrecisionSettings>::new();
|
||||||
|
model.save_file(&output_path, &recorder)?;
|
||||||
|
|
||||||
|
println!("拷贝配置文件和 tokenizer...");
|
||||||
|
std::fs::copy(config_path, output_dir.join("config.json"))?;
|
||||||
|
std::fs::copy(tokenizer_path, output_dir.join("tokenizer.json"))?;
|
||||||
|
|
||||||
|
println!("模型已成功导出到: {:?}", output_dir);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
pub fn load_safetensors<B: Backend>(
|
pub fn load_safetensors<B: Backend>(
|
||||||
model: LlamaForCausalLM<B>,
|
model: LlamaForCausalLM<B>,
|
||||||
path: &Path,
|
path: &Path,
|
||||||
@@ -0,0 +1,10 @@
|
|||||||
|
[package]
|
||||||
|
name = "minicpm-core"
|
||||||
|
version = "0.1.0"
|
||||||
|
edition = "2021"
|
||||||
|
|
||||||
|
[dependencies]
|
||||||
|
burn = { version = "0.21", features = ["std", "wgpu"] }
|
||||||
|
serde = { version = "1.0", features = ["derive"] }
|
||||||
|
serde_json = "1.0"
|
||||||
|
anyhow = "1.0"
|
||||||
@@ -0,0 +1,5 @@
|
|||||||
|
pub mod config;
|
||||||
|
pub mod model;
|
||||||
|
|
||||||
|
pub use config::{EosTokenId, LlamaConfig};
|
||||||
|
pub use model::{LlamaForCausalLM, LlamaKVCache};
|
||||||
@@ -5,5 +5,4 @@ pub mod model;
|
|||||||
pub mod norm;
|
pub mod norm;
|
||||||
pub mod rope;
|
pub mod rope;
|
||||||
|
|
||||||
pub use attention::KVCache;
|
pub use model::{LlamaForCausalLM, LlamaKVCache};
|
||||||
pub use model::{LlamaForCausalLM, LlamaKVCache};
|
|
||||||
@@ -0,0 +1,11 @@
|
|||||||
|
[package]
|
||||||
|
name = "minicpm-inference"
|
||||||
|
version = "0.1.0"
|
||||||
|
edition = "2021"
|
||||||
|
|
||||||
|
[dependencies]
|
||||||
|
minicpm-core = { path = "../minicpm-core" }
|
||||||
|
burn = { version = "0.21", features = ["std", "wgpu"] }
|
||||||
|
tokenizers = "0.20"
|
||||||
|
rand = "0.8"
|
||||||
|
anyhow = "1.0"
|
||||||
@@ -1,9 +1,62 @@
|
|||||||
|
pub use minicpm_core::config::{EosTokenId, LlamaConfig};
|
||||||
|
pub use minicpm_core::model::{LlamaForCausalLM, LlamaKVCache};
|
||||||
|
|
||||||
|
use burn::module::Module;
|
||||||
|
use burn::record::{FullPrecisionSettings, NamedMpkFileRecorder};
|
||||||
use burn::tensor::backend::Backend;
|
use burn::tensor::backend::Backend;
|
||||||
use burn::tensor::{Float, Int, Tensor};
|
use burn::tensor::{Float, Int, Tensor};
|
||||||
use rand::Rng;
|
use rand::Rng;
|
||||||
|
use tokenizers::Tokenizer;
|
||||||
|
|
||||||
|
|
||||||
|
pub struct TokenizerWrapper {
|
||||||
|
tokenizer: Tokenizer,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl TokenizerWrapper {
|
||||||
|
pub fn from_file(path: &str) -> anyhow::Result<Self> {
|
||||||
|
let tokenizer = Tokenizer::from_file(path)
|
||||||
|
.map_err(|e| anyhow::anyhow!("Failed to load tokenizer: {}", e))?;
|
||||||
|
Ok(Self { tokenizer })
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn encode(&self, text: &str, add_special_tokens: bool) -> anyhow::Result<Vec<u32>> {
|
||||||
|
let encoding = self
|
||||||
|
.tokenizer
|
||||||
|
.encode(text, add_special_tokens)
|
||||||
|
.map_err(|e| anyhow::anyhow!("Failed to encode: {}", e))?;
|
||||||
|
Ok(encoding.get_ids().to_vec())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn decode(&self, ids: &[u32], skip_special_tokens: bool) -> anyhow::Result<String> {
|
||||||
|
let text = self
|
||||||
|
.tokenizer
|
||||||
|
.decode(ids, skip_special_tokens)
|
||||||
|
.map_err(|e| anyhow::anyhow!("Failed to decode: {}", e))?;
|
||||||
|
Ok(text)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn vocab_size(&self) -> usize {
|
||||||
|
self.tokenizer.get_vocab_size(true)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 应用 MiniCPM5 chat template
|
||||||
|
pub fn apply_chat_template(&self, user_msg: &str, enable_thinking: bool) -> anyhow::Result<Vec<u32>> {
|
||||||
|
let prompt = if enable_thinking {
|
||||||
|
format!(
|
||||||
|
"<s><|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n<think>\n",
|
||||||
|
user_msg
|
||||||
|
)
|
||||||
|
} else {
|
||||||
|
format!(
|
||||||
|
"<s><|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n",
|
||||||
|
user_msg
|
||||||
|
)
|
||||||
|
};
|
||||||
|
self.encode(&prompt, false)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
use crate::config::EosTokenId;
|
|
||||||
use crate::model::{LlamaForCausalLM, LlamaKVCache};
|
|
||||||
|
|
||||||
pub struct GenerationConfig {
|
pub struct GenerationConfig {
|
||||||
pub max_new_tokens: Option<usize>,
|
pub max_new_tokens: Option<usize>,
|
||||||
@@ -167,3 +220,64 @@ fn sample<B: Backend>(logits: &Tensor<B, 1, Float>, temperature: f32, top_p: f32
|
|||||||
}
|
}
|
||||||
probs.len() as u32 - 1
|
probs.len() as u32 - 1
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
pub struct MiniCPM<B: Backend> {
|
||||||
|
model: LlamaForCausalLM<B>,
|
||||||
|
config: LlamaConfig,
|
||||||
|
tokenizer: TokenizerWrapper,
|
||||||
|
device: B::Device,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<B: Backend> MiniCPM<B> {
|
||||||
|
pub fn load(
|
||||||
|
model_path: &str,
|
||||||
|
config_path: &str,
|
||||||
|
tokenizer_path: &str,
|
||||||
|
device: &B::Device,
|
||||||
|
) -> anyhow::Result<Self> {
|
||||||
|
let config = LlamaConfig::from_json(config_path)?;
|
||||||
|
let tokenizer = TokenizerWrapper::from_file(tokenizer_path)?;
|
||||||
|
|
||||||
|
let model = LlamaForCausalLM::<B>::new(config.clone(), device);
|
||||||
|
let recorder = NamedMpkFileRecorder::<FullPrecisionSettings>::new();
|
||||||
|
let model = model.load_file(model_path, &recorder, device)?;
|
||||||
|
|
||||||
|
Ok(Self {
|
||||||
|
model,
|
||||||
|
config,
|
||||||
|
tokenizer,
|
||||||
|
device: device.clone(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn generate(
|
||||||
|
&self,
|
||||||
|
prompt: &str,
|
||||||
|
think: bool,
|
||||||
|
config: &GenerationConfig,
|
||||||
|
) -> anyhow::Result<String> {
|
||||||
|
let input_ids = self.tokenizer.apply_chat_template(prompt, think)?;
|
||||||
|
let output_ids = generate_with_cache(
|
||||||
|
&self.model,
|
||||||
|
&input_ids,
|
||||||
|
config,
|
||||||
|
&self.config.eos_token_id,
|
||||||
|
&self.device,
|
||||||
|
);
|
||||||
|
let new_ids = &output_ids[input_ids.len()..];
|
||||||
|
self.tokenizer.decode(new_ids, true)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn config(&self) -> &LlamaConfig {
|
||||||
|
&self.config
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn tokenizer(&self) -> &TokenizerWrapper {
|
||||||
|
&self.tokenizer
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn inner_model(&self) -> &LlamaForCausalLM<B> {
|
||||||
|
&self.model
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,9 @@
|
|||||||
|
[package]
|
||||||
|
name = "convert"
|
||||||
|
version = "0.1.0"
|
||||||
|
edition = "2021"
|
||||||
|
|
||||||
|
[dependencies]
|
||||||
|
minicpm-convert = { path = "../../crates/minicpm-convert" }
|
||||||
|
burn = { version = "0.21", features = ["std", "wgpu"] }
|
||||||
|
anyhow = "1.0"
|
||||||
@@ -0,0 +1,24 @@
|
|||||||
|
use burn::backend::Wgpu;
|
||||||
|
use minicpm_convert::export_model;
|
||||||
|
use std::path::Path;
|
||||||
|
|
||||||
|
fn main() -> anyhow::Result<()> {
|
||||||
|
let device = Default::default();
|
||||||
|
|
||||||
|
let safetensors_path = Path::new("MiniCPM5-1B/model-00000-of-00001.safetensors");
|
||||||
|
let config_path = Path::new("MiniCPM5-1B/config.json");
|
||||||
|
let tokenizer_path = Path::new("MiniCPM5-1B/tokenizer.json");
|
||||||
|
let output_dir = Path::new("model");
|
||||||
|
|
||||||
|
export_model::<Wgpu>(
|
||||||
|
safetensors_path,
|
||||||
|
config_path,
|
||||||
|
tokenizer_path,
|
||||||
|
output_dir,
|
||||||
|
&device,
|
||||||
|
)?;
|
||||||
|
|
||||||
|
println!("模型转换完成！");
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
@@ -0,0 +1,9 @@
|
|||||||
|
[package]
|
||||||
|
name = "minimal-inference"
|
||||||
|
version = "0.1.0"
|
||||||
|
edition = "2021"
|
||||||
|
|
||||||
|
[dependencies]
|
||||||
|
minicpm-inference = { path = "../../crates/minicpm-inference" }
|
||||||
|
burn = { version = "0.21", features = ["std", "wgpu"] }
|
||||||
|
anyhow = "1.0"
|
||||||
@@ -0,0 +1,46 @@
|
|||||||
|
// 运行方式：
|
||||||
|
// 1. 确保已在项目根目录
|
||||||
|
// 2. 运行：cargo run --release -p minimal-inference
|
||||||
|
// 3. 确保模型文件存在：
|
||||||
|
// - model/model.mpk (模型权重)
|
||||||
|
// - MiniCPM5-1B/config.json (模型配置)
|
||||||
|
// - MiniCPM5-1B/tokenizer.json (分词器)
|
||||||
|
|
||||||
|
use burn::backend::Wgpu;
|
||||||
|
use minicpm_inference::{GenerationConfig, MiniCPM};
|
||||||
|
use std::time::Instant;
|
||||||
|
|
||||||
|
fn main() -> anyhow::Result<()> {
|
||||||
|
let device = Default::default();
|
||||||
|
|
||||||
|
println!("正在加载模型...");
|
||||||
|
let start = Instant::now();
|
||||||
|
|
||||||
|
let model = MiniCPM::<Wgpu>::load(
|
||||||
|
"model/model",
|
||||||
|
"MiniCPM5-1B/config.json",
|
||||||
|
"MiniCPM5-1B/tokenizer.json",
|
||||||
|
&device,
|
||||||
|
)?;
|
||||||
|
|
||||||
|
println!("模型加载完成，耗时: {:.2?}", start.elapsed());
|
||||||
|
|
||||||
|
let prompt = "你知道今天的天气不";
|
||||||
|
println!("\n用户: {}", prompt);
|
||||||
|
println!("Assistant: ");
|
||||||
|
|
||||||
|
let config = GenerationConfig {
|
||||||
|
max_new_tokens: Some(200),
|
||||||
|
temperature: 0.7,
|
||||||
|
top_p: 0.95,
|
||||||
|
};
|
||||||
|
|
||||||
|
let start = Instant::now();
|
||||||
|
let response = model.generate(prompt, false, &config)?;
|
||||||
|
let elapsed = start.elapsed();
|
||||||
|
|
||||||
|
println!("{}", response);
|
||||||
|
println!("\n生成耗时: {:.2?}", elapsed);
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
@@ -0,0 +1,30 @@
|
|||||||
|
{
|
||||||
|
"_name_or_path": "openbmb/MiniCPM5-1B",
|
||||||
|
"architectures": [
|
||||||
|
"LlamaForCausalLM"
|
||||||
|
],
|
||||||
|
"bos_token_id": 0,
|
||||||
|
"eos_token_id": [
|
||||||
|
1,
|
||||||
|
130073
|
||||||
|
],
|
||||||
|
"pad_token_id": 1,
|
||||||
|
"hidden_act": "silu",
|
||||||
|
"hidden_size": 1536,
|
||||||
|
"initializer_range": 0.02,
|
||||||
|
"intermediate_size": 4608,
|
||||||
|
"max_position_embeddings": 131072,
|
||||||
|
"model_type": "llama",
|
||||||
|
"num_attention_heads": 16,
|
||||||
|
"num_hidden_layers": 24,
|
||||||
|
"num_key_value_heads": 2,
|
||||||
|
"head_dim": 128,
|
||||||
|
"rms_norm_eps": 1e-06,
|
||||||
|
"rope_theta": 5000000,
|
||||||
|
"rope_scaling": null,
|
||||||
|
"tie_word_embeddings": false,
|
||||||
|
"torch_dtype": "bfloat16",
|
||||||
|
"transformers_version": "5.6.2",
|
||||||
|
"use_cache": true,
|
||||||
|
"vocab_size": 130560
|
||||||
|
}
|
||||||
Binary file not shown.
+653947
File diff suppressed because it is too large
Load Diff
@@ -1,107 +1,186 @@
|
|||||||
# MiniCPM5-1B Burn
|
# MiniCPM5-1B-rust
|
||||||
|
|
||||||
使用 [Burn](https://github.com/tracel-ai/burn) 框架从零实现的 MiniCPM5-1B 模型。
|
基于 [Burn](https://burn.dev/) 深度学习框架实现的 MiniCPM5-1B 大语言模型推理库，纯 Rust 编写，支持 GPU 加速。
|
||||||
|
|
||||||
## 功能
|
## 模型架构参数
|
||||||
|
|
||||||
- **模型转换**：将 safetensors 格式转换为 Burn 原生 bincode 格式
|
| 参数 | 值 |
|
||||||
- **模型推理**：使用 WGPU 后端运行 MiniCPM5-1B 进行文本生成，支持无上限输出
|
|------|-----|
|
||||||
- **独立输出**：转换后的模型文件放在 `model/` 目录，方便移植到其他项目
|
| 层数 | 24 |
|
||||||
|
| 注意力机制 | GQA (16 Query / 2 KV) |
|
||||||
|
| head_dim | 128 |
|
||||||
|
| hidden_size | 1536 |
|
||||||
|
| intermediate_size | 8960 |
|
||||||
|
| vocab_size | 151936 |
|
||||||
|
| RoPE theta | 5000000 |
|
||||||
|
| 归一化 | RMSNorm |
|
||||||
|
| 激活函数 | SiLU (SwiGLU) |
|
||||||
|
| KV Cache | 支持 |
|
||||||
|
|
||||||
## 架构
|
## Workspace 结构
|
||||||
|
|
||||||
- 24 层 Transformer
|
项目采用 Cargo workspace 多成员结构，包含三个 crate：
|
||||||
- GQA：16 个查询头，2 个 KV 头
|
|
||||||
- RoPE 旋转位置编码（rope_theta = 5000000）
|
|
||||||
- RMSNorm 归一化
|
|
||||||
- SiLU 激活函数（FFN: gate + up + down）
|
|
||||||
- KV Cache 优化
|
|
||||||
|
|
||||||
## 使用
|
### minicpm-core
|
||||||
|
|
||||||
### 准备模型
|
核心模型定义，包含：
|
||||||
|
|
||||||
将 MiniCPM5-1B 模型放到 `MiniCPM5-1B/` 目录：
|
- `LlamaConfig` — 模型配置（从 `config.json` 加载）
|
||||||
```
|
- `LlamaForCausalLM` — 因果语言模型主体
|
||||||
MiniCPM5-1B/
|
- `LlamaKVCache` — KV 缓存结构
|
||||||
├── config.json
|
- `EosTokenId` — EOS token 标识（支持单个或多个）
|
||||||
├── tokenizer.json
|
- 各模块实现：Attention、Decoder、FFN、RMSNorm、RoPE
|
||||||
└── model-00000-of-00001.safetensors
|
|
||||||
```
|
### minicpm-convert
|
||||||
|
|
||||||
|
模型格式转换工具，负责将 HuggingFace safetensors 格式的权重转换为 Burn MPK 格式：
|
||||||
|
|
||||||
|
- `export_model()` — 完整转换流程（加载 safetensors → 构建模型 → 导出 MPK）
|
||||||
|
- 支持 BF16 / F16 / F32 精度输入
|
||||||
|
- 支持 `tie_word_embeddings` 权重共享
|
||||||
|
|
||||||
|
### minicpm-inference
|
||||||
|
|
||||||
|
推理功能封装，提供高层 API：
|
||||||
|
|
||||||
|
- `MiniCPM` — 高层封装，整合模型 + tokenizer + 推理逻辑
|
||||||
|
- `TokenizerWrapper` — tokenizer 封装，支持 MiniCPM5 chat template
|
||||||
|
- `GenerationConfig` — 生成配置（max_new_tokens、temperature、top_p）
|
||||||
|
- `generate_with_cache()` — 带 KV Cache 的自回归生成
|
||||||
|
- 支持 temperature 采样、top-p 采样、greedy 解码
|
||||||
|
|
||||||
|
## 快速开始
|
||||||
|
|
||||||
### 1. 转换模型
|
### 1. 转换模型
|
||||||
|
|
||||||
```bash
|
首先将 HuggingFace 格式的 MiniCPM5-1B 模型转换为 Burn MPK 格式：
|
||||||
# 默认输出到 model/model.bin
|
|
||||||
cargo run --release -- convert
|
|
||||||
|
|
||||||
# 指定输出路径
|
```rust
|
||||||
cargo run --release -- convert --output my_model/model.bin
|
use minicpm_convert::export_model;
|
||||||
|
use burn::backend::Wgpu;
|
||||||
|
use std::path::Path;
|
||||||
|
|
||||||
|
fn main() -> anyhow::Result<()> {
|
||||||
|
let device = Default::default();
|
||||||
|
|
||||||
|
export_model::<Wgpu>(
|
||||||
|
Path::new("MiniCPM5-1B/model.safetensors"),
|
||||||
|
Path::new("MiniCPM5-1B/config.json"),
|
||||||
|
Path::new("MiniCPM5-1B-burn"),
|
||||||
|
&device,
|
||||||
|
)?;
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
```
|
```
|
||||||
|
|
||||||
### 2. 运行推理
|
转换完成后，`MiniCPM5-1B-burn/` 目录下会生成 `model.mp` 文件。
|
||||||
|
|
||||||
```bash
|
### 2. 推理使用
|
||||||
# 无上限生成（直到遇到 eos token）
|
|
||||||
cargo run --release -- run --text "你好"
|
|
||||||
|
|
||||||
# 限制最大生成 token 数
|
使用 `MiniCPM` 高层 API 进行文本生成：
|
||||||
cargo run --release -- run --text "你好" --max-tokens 100
|
|
||||||
|
```rust
|
||||||
|
use minicpm_inference::{MiniCPM, GenerationConfig};
|
||||||
|
use burn::backend::Wgpu;
|
||||||
|
|
||||||
|
fn main() -> anyhow::Result<()> {
|
||||||
|
let device = Default::default();
|
||||||
|
|
||||||
|
// 加载模型
|
||||||
|
let model = MiniCPM::<Wgpu>::load(
|
||||||
|
"MiniCPM5-1B-burn/model",
|
||||||
|
"MiniCPM5-1B/config.json",
|
||||||
|
"MiniCPM5-1B/tokenizer.json",
|
||||||
|
&device,
|
||||||
|
)?;
|
||||||
|
|
||||||
|
// 生成配置
|
||||||
|
let gen_config = GenerationConfig {
|
||||||
|
max_new_tokens: Some(512),
|
||||||
|
temperature: 0.7,
|
||||||
|
top_p: 0.8,
|
||||||
|
};
|
||||||
|
|
||||||
|
// 生成回答（think = true 启用思考模式）
|
||||||
|
let response = model.generate("用 Rust 写一个 Hello World", true, &gen_config)?;
|
||||||
|
|
||||||
|
println!("{}", response);
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
```
|
```
|
||||||
|
|
||||||
## 移植到其他项目
|
## 移植到其他项目
|
||||||
|
|
||||||
转换后的文件：
|
### 依赖配置
|
||||||
```
|
|
||||||
model/
|
在你的 `Cargo.toml` 中添加：
|
||||||
└── model.bin # Burn bincode 格式模型
|
|
||||||
|
```toml
|
||||||
|
[dependencies]
|
||||||
|
minicpm-inference = { path = "../MiniCPM5-1B-rust/crates/minicpm-inference" }
|
||||||
|
burn = { version = "0.21", features = ["std", "wgpu"] }
|
||||||
|
anyhow = "1.0"
|
||||||
```
|
```
|
||||||
|
|
||||||
在其他项目中加载：
|
> 也可以根据需要选择其他 backend（如 `tch-gpu`、`cuda` 等）。
|
||||||
|
|
||||||
|
### 最小示例
|
||||||
|
|
||||||
```rust
|
```rust
|
||||||
use burn::record::{BinFileRecorder, FullPrecisionSettings};
|
use minicpm_inference::{MiniCPM, GenerationConfig};
|
||||||
use minicpm_burn_lib::{LlamaConfig, LlamaForCausalLM};
|
use burn::backend::Wgpu;
|
||||||
|
|
||||||
let config = LlamaConfig::from_json("config.json")?;
|
fn main() -> anyhow::Result<()> {
|
||||||
let model = LlamaForCausalLM::<Backend>::new(config, &device);
|
let device = Default::default();
|
||||||
let model = model.load_file("model/model.bin", &BinFileRecorder::<FullPrecisionSettings>::new())?;
|
|
||||||
|
let model = MiniCPM::<Wgpu>::load(
|
||||||
|
"model",
|
||||||
|
"config.json",
|
||||||
|
"tokenizer.json",
|
||||||
|
&device,
|
||||||
|
)?;
|
||||||
|
|
||||||
|
let config = GenerationConfig {
|
||||||
|
max_new_tokens: Some(256),
|
||||||
|
temperature: 1.0,
|
||||||
|
top_p: 1.0,
|
||||||
|
};
|
||||||
|
|
||||||
|
let output = model.generate("你好", false, &config)?;
|
||||||
|
println!("{}", output);
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
```
|
```
|
||||||
|
|
||||||
## 项目结构
|
## 模型文件准备
|
||||||
|
|
||||||
|
`MiniCPM5-1B/` 目录需要包含以下文件：
|
||||||
|
|
||||||
```
|
```
|
||||||
src/
|
MiniCPM5-1B/
|
||||||
├── lib.rs # 库入口
|
├── config.json # 模型配置文件
|
||||||
├── main.rs # CLI 工具
|
├── model.safetensors # safetensors 格式权重
|
||||||
├── config.rs # 模型配置
|
└── tokenizer.json # tokenizer 配置
|
||||||
├── exporter.rs # 模型转换
|
|
||||||
├── model/ # 模型架构
|
|
||||||
│ ├── mod.rs
|
|
||||||
│ ├── attention.rs # GQA 注意力 + KV Cache
|
|
||||||
│ ├── decoder.rs # 解码器层
|
|
||||||
│ ├── ffn.rs # 前馈网络
|
|
||||||
│ ├── model.rs # LlamaForCausalLM
|
|
||||||
│ ├── norm.rs # RMSNorm
|
|
||||||
│ └── rope.rs # RoPE 位置编码
|
|
||||||
└── inference/ # 推理工具
|
|
||||||
├── mod.rs
|
|
||||||
├── generation.rs # 文本生成（greedy + KV Cache）
|
|
||||||
├── loader.rs # safetensors 加载器
|
|
||||||
└── tokenizer.rs # 分词器封装
|
|
||||||
```
|
```
|
||||||
|
|
||||||
## 依赖
|
转换后 Burn 模型目录结构：
|
||||||
|
|
||||||
- `burn` 0.21 - 深度学习框架（WGPU 后端）
|
```
|
||||||
- `burn-store` - safetensors 支持
|
MiniCPM5-1B-burn/
|
||||||
- `tokenizers` - HuggingFace 分词器
|
└── model.mp # Burn MPK 格式权重
|
||||||
- `clap` - CLI 参数解析
|
```
|
||||||
|
|
||||||
## 性能
|
## 性能说明
|
||||||
|
|
||||||
- WGPU Vulkan 后端：约 10 tokens/s
|
| Backend | 设备 | 速度 | 显存占用 |
|
||||||
- 模型大小：约 4.0 GB（F32）
|
|---------|------|------|----------|
|
||||||
|
| WGPU (Vulkan) | RTX 4060 | ~10 tokens/s | ~4 GB (F32) |
|
||||||
|
|
||||||
|
> 以上数据仅供参考，实际性能因硬件配置、生成长度等因素而异。
|
||||||
|
|
||||||
## 许可证
|
## 许可证
|
||||||
|
|
||||||
MIT
|
MIT
|
||||||
|
|||||||
@@ -1,34 +0,0 @@
|
|||||||
use crate::config::LlamaConfig;
|
|
||||||
use crate::inference::load_safetensors;
|
|
||||||
use crate::model::LlamaForCausalLM;
|
|
||||||
use burn::module::Module;
|
|
||||||
use burn::record::{FullPrecisionSettings, NamedMpkFileRecorder};
|
|
||||||
use burn::tensor::backend::Backend;
|
|
||||||
use std::path::Path;
|
|
||||||
|
|
||||||
pub fn export_model<B: Backend>(
|
|
||||||
safetensors_path: &Path,
|
|
||||||
config_path: &Path,
|
|
||||||
output_dir: &Path,
|
|
||||||
device: &B::Device,
|
|
||||||
) -> anyhow::Result<()> {
|
|
||||||
println!("开始转换 MiniCPM 模型为 Burn 格式...");
|
|
||||||
|
|
||||||
println!("加载配置文件: {:?}", config_path);
|
|
||||||
let config = LlamaConfig::from_json(config_path.to_str().unwrap())?;
|
|
||||||
|
|
||||||
println!("创建模型结构...");
|
|
||||||
let model = LlamaForCausalLM::<B>::new(config, device);
|
|
||||||
|
|
||||||
println!("加载 safetensors 权重...");
|
|
||||||
let model = load_safetensors(model, safetensors_path, device)?;
|
|
||||||
|
|
||||||
println!("保存为 MPK 格式...");
|
|
||||||
std::fs::create_dir_all(output_dir)?;
|
|
||||||
let output_path = output_dir.join("model");
|
|
||||||
let recorder = NamedMpkFileRecorder::<FullPrecisionSettings>::new();
|
|
||||||
model.save_file(&output_path, &recorder)?;
|
|
||||||
|
|
||||||
println!("模型已成功导出到: {:?}", output_dir);
|
|
||||||
Ok(())
|
|
||||||
}
|
|
||||||
@@ -1,7 +0,0 @@
|
|||||||
pub mod generation;
|
|
||||||
pub mod loader;
|
|
||||||
pub mod tokenizer;
|
|
||||||
|
|
||||||
pub use generation::{GenerationConfig, generate_with_cache};
|
|
||||||
pub use loader::load_safetensors;
|
|
||||||
pub use tokenizer::TokenizerWrapper;
|
|
||||||
@@ -1,49 +0,0 @@
|
|||||||
use tokenizers::Tokenizer;
|
|
||||||
|
|
||||||
pub struct TokenizerWrapper {
|
|
||||||
tokenizer: Tokenizer,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl TokenizerWrapper {
|
|
||||||
pub fn from_file(path: &str) -> anyhow::Result<Self> {
|
|
||||||
let tokenizer = Tokenizer::from_file(path)
|
|
||||||
.map_err(|e| anyhow::anyhow!("Failed to load tokenizer: {}", e))?;
|
|
||||||
Ok(Self { tokenizer })
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn encode(&self, text: &str, add_special_tokens: bool) -> anyhow::Result<Vec<u32>> {
|
|
||||||
let encoding = self
|
|
||||||
.tokenizer
|
|
||||||
.encode(text, add_special_tokens)
|
|
||||||
.map_err(|e| anyhow::anyhow!("Failed to encode: {}", e))?;
|
|
||||||
Ok(encoding.get_ids().to_vec())
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn decode(&self, ids: &[u32], skip_special_tokens: bool) -> anyhow::Result<String> {
|
|
||||||
let text = self
|
|
||||||
.tokenizer
|
|
||||||
.decode(ids, skip_special_tokens)
|
|
||||||
.map_err(|e| anyhow::anyhow!("Failed to decode: {}", e))?;
|
|
||||||
Ok(text)
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn vocab_size(&self) -> usize {
|
|
||||||
self.tokenizer.get_vocab_size(true)
|
|
||||||
}
|
|
||||||
|
|
||||||
/// 应用 MiniCPM5 chat template
|
|
||||||
pub fn apply_chat_template(&self, user_msg: &str, enable_thinking: bool) -> anyhow::Result<Vec<u32>> {
|
|
||||||
let prompt = if enable_thinking {
|
|
||||||
format!(
|
|
||||||
"<s><|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n<think>\n",
|
|
||||||
user_msg
|
|
||||||
)
|
|
||||||
} else {
|
|
||||||
format!(
|
|
||||||
"<s><|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n",
|
|
||||||
user_msg
|
|
||||||
)
|
|
||||||
};
|
|
||||||
self.encode(&prompt, false)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,9 +0,0 @@
|
|||||||
pub mod config;
|
|
||||||
pub mod exporter;
|
|
||||||
pub mod model;
|
|
||||||
pub mod inference;
|
|
||||||
|
|
||||||
pub use config::{LlamaConfig, EosTokenId};
|
|
||||||
pub use model::LlamaForCausalLM;
|
|
||||||
pub use inference::{GenerationConfig, generate_with_cache, load_safetensors, TokenizerWrapper};
|
|
||||||
pub use exporter::export_model;
|
|
||||||
-109
@@ -1,109 +0,0 @@
|
|||||||
use anyhow::Result;
|
|
||||||
use burn::backend::Wgpu;
|
|
||||||
use clap::Parser;
|
|
||||||
use minicpm_burn_lib::{
|
|
||||||
config::LlamaConfig,
|
|
||||||
inference::{generation, GenerationConfig, load_safetensors, TokenizerWrapper},
|
|
||||||
model::LlamaForCausalLM,
|
|
||||||
};
|
|
||||||
|
|
||||||
type Backend = Wgpu;
|
|
||||||
|
|
||||||
#[derive(Parser, Debug)]
|
|
||||||
#[command(name = "minicpm-burn")]
|
|
||||||
#[command(version = "0.1.0")]
|
|
||||||
enum Command {
|
|
||||||
/// 模型推理
|
|
||||||
Run {
|
|
||||||
/// 输入文本
|
|
||||||
#[arg(short, long)]
|
|
||||||
text: String,
|
|
||||||
|
|
||||||
/// 最大生成 token 数（不指定则无上限）
|
|
||||||
#[arg(short, long)]
|
|
||||||
max_tokens: Option<usize>,
|
|
||||||
|
|
||||||
/// 启用思考模式（temperature=0.9）
|
|
||||||
#[arg(long, default_value_t = false)]
|
|
||||||
think: bool,
|
|
||||||
},
|
|
||||||
/// 转换模型为 Burn MPK 格式
|
|
||||||
Convert {
|
|
||||||
/// 输出目录
|
|
||||||
#[arg(short, long, default_value = "model")]
|
|
||||||
output: String,
|
|
||||||
},
|
|
||||||
}
|
|
||||||
|
|
||||||
fn main() -> Result<()> {
|
|
||||||
let cmd = Command::parse();
|
|
||||||
|
|
||||||
match cmd {
|
|
||||||
Command::Run { text, max_tokens, think } => run_inference(&text, max_tokens, think),
|
|
||||||
Command::Convert { output } => run_convert(&output),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn run_inference(text: &str, max_tokens: Option<usize>, think: bool) -> Result<()> {
|
|
||||||
use std::time::Instant;
|
|
||||||
|
|
||||||
println!("MiniCPM5-1B 推理 [{}]", if think { "思考模式" } else { "直接回答" });
|
|
||||||
|
|
||||||
let device = burn::backend::wgpu::WgpuDevice::default();
|
|
||||||
burn::backend::wgpu::init_setup::<burn::backend::wgpu::graphics::Vulkan>(
|
|
||||||
&device,
|
|
||||||
Default::default(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let config = LlamaConfig::from_json("MiniCPM5-1B/config.json")?;
|
|
||||||
let tokenizer = TokenizerWrapper::from_file("MiniCPM5-1B/tokenizer.json")?;
|
|
||||||
let model = LlamaForCausalLM::<Backend>::new(config.clone(), &device);
|
|
||||||
let model = load_safetensors(model, std::path::Path::new("MiniCPM5-1B/model-00000-of-00001.safetensors"), &device)?;
|
|
||||||
|
|
||||||
let input_ids = tokenizer.apply_chat_template(text, think)?;
|
|
||||||
let gen_config = GenerationConfig {
|
|
||||||
max_new_tokens: max_tokens,
|
|
||||||
temperature: if think { 0.9 } else { 0.7 },
|
|
||||||
top_p: 0.95,
|
|
||||||
};
|
|
||||||
|
|
||||||
let start = Instant::now();
|
|
||||||
let output_ids = generation::generate_with_cache(
|
|
||||||
&model,
|
|
||||||
&input_ids,
|
|
||||||
&gen_config,
|
|
||||||
&config.eos_token_id,
|
|
||||||
&device,
|
|
||||||
);
|
|
||||||
let elapsed = start.elapsed();
|
|
||||||
|
|
||||||
let new_ids = &output_ids[input_ids.len()..];
|
|
||||||
if !new_ids.is_empty() {
|
|
||||||
println!("\n输出: {}", tokenizer.decode(new_ids, true)?);
|
|
||||||
}
|
|
||||||
|
|
||||||
println!("\n耗时: {:.2}s", elapsed.as_secs_f64());
|
|
||||||
println!("速度: {:.2} tokens/s", (output_ids.len() - input_ids.len()) as f64 / elapsed.as_secs_f64());
|
|
||||||
|
|
||||||
Ok(())
|
|
||||||
}
|
|
||||||
|
|
||||||
fn run_convert(output: &str) -> Result<()> {
|
|
||||||
println!("MiniCPM5-1B 模型转换");
|
|
||||||
|
|
||||||
let device = burn::backend::wgpu::WgpuDevice::default();
|
|
||||||
burn::backend::wgpu::init_setup::<burn::backend::wgpu::graphics::Vulkan>(
|
|
||||||
&device,
|
|
||||||
Default::default(),
|
|
||||||
);
|
|
||||||
|
|
||||||
minicpm_burn_lib::export_model::<Backend>(
|
|
||||||
std::path::Path::new("MiniCPM5-1B/model-00000-of-00001.safetensors"),
|
|
||||||
std::path::Path::new("MiniCPM5-1B/config.json"),
|
|
||||||
std::path::Path::new(output),
|
|
||||||
&device,
|
|
||||||
)?;
|
|
||||||
|
|
||||||
println!("\n转换完成！输出目录: {}", output);
|
|
||||||
Ok(())
|
|
||||||
}
|
|
||||||
Reference in New Issue
Block a user