chore(代码变更): 1. 清理无用代码 2. 更新依赖库 3. 优化项目结构 4. 修复小型bug 5. 增加注释以提高可读性
This commit is contained in:
@@ -19,3 +19,17 @@
|
|||||||
|
|
||||||
- `minicpm-inference/src/lib.rs`: 新增 `generate_stream` 函数和 `MiniCPM::generate_stream` 方法,每 token 通过回调 `impl FnMut(&str)` 输出解码文本
|
- `minicpm-inference/src/lib.rs`: 新增 `generate_stream` 函数和 `MiniCPM::generate_stream` 方法,每 token 通过回调 `impl FnMut(&str)` 输出解码文本
|
||||||
- `examples/minimal-inference/src/main.rs`: 改为流式调用,`print!` + `flush` 实时输出
|
- `examples/minimal-inference/src/main.rs`: 改为流式调用,`print!` + `flush` 实时输出
|
||||||
|
|
||||||
|
## 升级 0.1.2 重新发布(含流式功能)
|
||||||
|
|
||||||
|
- minicpm-core, minicpm-convert, minicpm-inference 全部升级到 0.1.2
|
||||||
|
- 发布全部成功
|
||||||
|
|
||||||
|
## Q8 量化导出和加载
|
||||||
|
|
||||||
|
- `minicpm-convert`: 新增 `export_model_q8`,2D 权重 per-tensor INT8 量化(scale = max(|f|)/127),1D norm 权重存 F16;输出 `model.q8.json` + `model.q8.bin`
|
||||||
|
- `minicpm-inference`: 新增 `MiniCPM::load_q8`,加载 Q8 模型并在内存中反量化为 f32 推理;新增依赖 `memmap2`、`serde`、`serde_json`
|
||||||
|
- `examples/convert`: 支持 Q8 导出
|
||||||
|
- `examples/minimal-inference`: 支持 `--full|--half|--q8` 三种模式
|
||||||
|
- 量化方式:`q = clamp(round(f / scale), -127, 127)`,`f' = q * scale`
|
||||||
|
|
||||||
|
|||||||
Generated
+23
-9
@@ -999,6 +999,7 @@ dependencies = [
|
|||||||
"aligned-vec",
|
"aligned-vec",
|
||||||
"bon",
|
"bon",
|
||||||
"burn-cubecl",
|
"burn-cubecl",
|
||||||
|
"burn-flex",
|
||||||
"burn-ir",
|
"burn-ir",
|
||||||
"burn-tensor",
|
"burn-tensor",
|
||||||
"bytemuck",
|
"bytemuck",
|
||||||
@@ -2573,6 +2574,15 @@ dependencies = [
|
|||||||
"miniz_oxide",
|
"miniz_oxide",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "flex-backend"
|
||||||
|
version = "0.1.0"
|
||||||
|
dependencies = [
|
||||||
|
"anyhow",
|
||||||
|
"burn",
|
||||||
|
"minicpm-inference",
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "float-ord"
|
name = "float-ord"
|
||||||
version = "0.3.2"
|
version = "0.3.2"
|
||||||
@@ -3966,6 +3976,7 @@ dependencies = [
|
|||||||
"burn",
|
"burn",
|
||||||
"memmap2",
|
"memmap2",
|
||||||
"minicpm-core",
|
"minicpm-core",
|
||||||
|
"serde",
|
||||||
"serde_json",
|
"serde_json",
|
||||||
]
|
]
|
||||||
|
|
||||||
@@ -3985,20 +3996,14 @@ version = "0.1.2"
|
|||||||
dependencies = [
|
dependencies = [
|
||||||
"anyhow",
|
"anyhow",
|
||||||
"burn",
|
"burn",
|
||||||
|
"memmap2",
|
||||||
"minicpm-core",
|
"minicpm-core",
|
||||||
"rand 0.8.6",
|
"rand 0.8.6",
|
||||||
|
"serde",
|
||||||
|
"serde_json",
|
||||||
"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"
|
||||||
@@ -7222,6 +7227,15 @@ dependencies = [
|
|||||||
"wgpu-types",
|
"wgpu-types",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "wgpu-backend"
|
||||||
|
version = "0.1.0"
|
||||||
|
dependencies = [
|
||||||
|
"anyhow",
|
||||||
|
"burn",
|
||||||
|
"minicpm-inference",
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "wgpu-core"
|
name = "wgpu-core"
|
||||||
version = "29.0.3"
|
version = "29.0.3"
|
||||||
|
|||||||
+2
-1
@@ -3,7 +3,8 @@ members = [
|
|||||||
"crates/minicpm-core",
|
"crates/minicpm-core",
|
||||||
"crates/minicpm-convert",
|
"crates/minicpm-convert",
|
||||||
"crates/minicpm-inference",
|
"crates/minicpm-inference",
|
||||||
"examples/minimal-inference",
|
|
||||||
"examples/convert",
|
"examples/convert",
|
||||||
|
"examples/flex-backend",
|
||||||
|
"examples/wgpu-backend",
|
||||||
]
|
]
|
||||||
resolver = "2"
|
resolver = "2"
|
||||||
|
|||||||
@@ -10,3 +10,4 @@ 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"
|
||||||
|
serde = { version = "1.0", features = ["derive"] }
|
||||||
@@ -2,7 +2,7 @@ use minicpm_core::config::LlamaConfig;
|
|||||||
use minicpm_core::model::LlamaForCausalLM;
|
use minicpm_core::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::record::{FullPrecisionSettings, HalfPrecisionSettings, 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;
|
||||||
@@ -40,6 +40,39 @@ pub fn export_model<B: Backend>(
|
|||||||
println!("模型已成功导出到: {:?}", output_dir);
|
println!("模型已成功导出到: {:?}", output_dir);
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// 导出模型为半精度(f16),输出文件约为全精度的一半大小。
|
||||||
|
pub fn export_model_half<B: Backend>(
|
||||||
|
safetensors_path: &Path,
|
||||||
|
config_path: &Path,
|
||||||
|
tokenizer_path: &Path,
|
||||||
|
output_dir: &Path,
|
||||||
|
device: &B::Device,
|
||||||
|
) -> anyhow::Result<()> {
|
||||||
|
println!("开始转换 MiniCPM 模型为 Burn 格式(半精度 f16)...");
|
||||||
|
|
||||||
|
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::<HalfPrecisionSettings>::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,
|
||||||
@@ -256,3 +289,223 @@ fn f16_to_f32(f16: u16) -> f32 {
|
|||||||
let f32_bits = ((sign as u32) << 31) | (f32_exp << 23) | ((mant as u32) << 13);
|
let f32_bits = ((sign as u32) << 31) | (f32_exp << 23) | ((mant as u32) << 13);
|
||||||
f32::from_bits(f32_bits)
|
f32::from_bits(f32_bits)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ==================== Q8 量化导出 ====================
|
||||||
|
|
||||||
|
/// Q8 量化单个 f32 切片,返回 (scale, i8_data)
|
||||||
|
fn quantize_q8(f32_data: &[f32]) -> (f32, Vec<i8>) {
|
||||||
|
let max_abs = f32_data
|
||||||
|
.iter()
|
||||||
|
.map(|&x| x.abs())
|
||||||
|
.fold(0.0f32, |a, b| a.max(b));
|
||||||
|
let scale = if max_abs > 0.0 { max_abs / 127.0 } else { 1.0 };
|
||||||
|
let q8: Vec<i8> = f32_data
|
||||||
|
.iter()
|
||||||
|
.map(|&x| (x / scale).round().clamp(-127.0, 127.0) as i8)
|
||||||
|
.collect();
|
||||||
|
(scale, q8)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Q8 导出元数据中的 tensor 条目
|
||||||
|
#[derive(serde::Serialize, serde::Deserialize)]
|
||||||
|
struct Q8TensorMeta {
|
||||||
|
shape: Vec<usize>,
|
||||||
|
scale: f32,
|
||||||
|
offset: u64,
|
||||||
|
numel: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Q8 导出元数据
|
||||||
|
#[derive(serde::Serialize, serde::Deserialize)]
|
||||||
|
struct Q8Metadata {
|
||||||
|
tensors: std::collections::HashMap<String, Q8TensorMeta>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 导出模型为 Q8 量化格式(INT8),输出 model.q8.json + model.q8.bin。
|
||||||
|
/// 2D 权重(Linear/Embedding)量化为 INT8;1D norm 权重存为 F16。
|
||||||
|
pub fn export_model_q8(
|
||||||
|
safetensors_path: &Path,
|
||||||
|
config_path: &Path,
|
||||||
|
tokenizer_path: &Path,
|
||||||
|
output_dir: &Path,
|
||||||
|
) -> anyhow::Result<()> {
|
||||||
|
println!("开始转换 MiniCPM 模型为 Q8 量化格式...");
|
||||||
|
|
||||||
|
println!("解析 safetensors 文件: {:?}", safetensors_path);
|
||||||
|
let file = std::fs::File::open(safetensors_path)?;
|
||||||
|
let mmap = unsafe { Mmap::map(&file)? };
|
||||||
|
|
||||||
|
let header_len = u64::from_le_bytes(mmap[..8].try_into().unwrap()) as usize;
|
||||||
|
let header_bytes = &mmap[8..8 + header_len];
|
||||||
|
let header: serde_json::Value = serde_json::from_slice(header_bytes)?;
|
||||||
|
|
||||||
|
let data_offset = 8 + header_len;
|
||||||
|
|
||||||
|
// 收集 tensor 信息
|
||||||
|
let mut tensor_infos: Vec<(String, Vec<usize>, String, usize, usize)> = Vec::new();
|
||||||
|
if let Some(obj) = header.as_object() {
|
||||||
|
for (name, info) in obj {
|
||||||
|
if name == "__metadata__" {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
if let Some(info_obj) = info.as_object() {
|
||||||
|
let mut shape = Vec::new();
|
||||||
|
if let Some(s) = info_obj.get("shape") {
|
||||||
|
if let Some(arr) = s.as_array() {
|
||||||
|
shape = arr.iter().map(|v| v.as_u64().unwrap() as usize).collect();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let mut dtype = String::new();
|
||||||
|
if let Some(d) = info_obj.get("dtype") {
|
||||||
|
dtype = d.as_str().unwrap_or("").to_string();
|
||||||
|
}
|
||||||
|
let mut start = 0;
|
||||||
|
let mut end = 0;
|
||||||
|
if let Some(offsets) = info_obj.get("data_offsets") {
|
||||||
|
if let Some(arr) = offsets.as_array() {
|
||||||
|
start = arr[0].as_u64().unwrap() as usize;
|
||||||
|
end = arr[1].as_u64().unwrap() as usize;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
tensor_infos.push((name.clone(), shape, dtype, start, end));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 定义需要量化的 2D 权重名称(与 load_safetensors 中的列表对应)
|
||||||
|
let mut q8_tensors: Vec<String> = Vec::new();
|
||||||
|
for i in 0..32 {
|
||||||
|
let p = format!("model.layers.{i}");
|
||||||
|
q8_tensors.push(format!("{p}.self_attn.q_proj.weight"));
|
||||||
|
q8_tensors.push(format!("{p}.self_attn.k_proj.weight"));
|
||||||
|
q8_tensors.push(format!("{p}.self_attn.v_proj.weight"));
|
||||||
|
q8_tensors.push(format!("{p}.self_attn.o_proj.weight"));
|
||||||
|
q8_tensors.push(format!("{p}.mlp.gate_proj.weight"));
|
||||||
|
q8_tensors.push(format!("{p}.mlp.up_proj.weight"));
|
||||||
|
q8_tensors.push(format!("{p}.mlp.down_proj.weight"));
|
||||||
|
}
|
||||||
|
q8_tensors.push("model.embed_tokens.weight".to_string());
|
||||||
|
q8_tensors.push("lm_head.weight".to_string());
|
||||||
|
|
||||||
|
let q8_set: std::collections::HashSet<String> = q8_tensors.into_iter().collect();
|
||||||
|
|
||||||
|
std::fs::create_dir_all(output_dir)?;
|
||||||
|
let json_path = output_dir.join("model.q8.json");
|
||||||
|
let bin_path = output_dir.join("model.q8.bin");
|
||||||
|
|
||||||
|
let mut metadata: std::collections::HashMap<String, Q8TensorMeta> = std::collections::HashMap::new();
|
||||||
|
let mut bin_file = std::fs::File::create(&bin_path)?;
|
||||||
|
let mut current_offset: u64 = 0;
|
||||||
|
|
||||||
|
// 先处理 2D 权重:量化为 INT8
|
||||||
|
for (name, shape, dtype, start, end) in &tensor_infos {
|
||||||
|
if shape.len() != 2 {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
if !q8_set.contains(name.as_str()) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
println!("量化 2D 权重: {name} shape={shape:?} dtype={dtype}");
|
||||||
|
let raw = &mmap[data_offset + *start..data_offset + *end];
|
||||||
|
let f32_data = convert_to_f32(raw, dtype);
|
||||||
|
let numel = f32_data.len();
|
||||||
|
let (scale, q8_data) = quantize_q8(&f32_data);
|
||||||
|
|
||||||
|
// 写入 INT8 数据
|
||||||
|
let bytes: &[u8] = unsafe {
|
||||||
|
std::slice::from_raw_parts(q8_data.as_ptr() as *const u8, q8_data.len())
|
||||||
|
};
|
||||||
|
use std::io::Write;
|
||||||
|
bin_file.write_all(bytes)?;
|
||||||
|
|
||||||
|
metadata.insert(
|
||||||
|
name.clone(),
|
||||||
|
Q8TensorMeta {
|
||||||
|
shape: shape.clone(),
|
||||||
|
scale,
|
||||||
|
offset: current_offset,
|
||||||
|
numel,
|
||||||
|
},
|
||||||
|
);
|
||||||
|
current_offset += q8_data.len() as u64;
|
||||||
|
}
|
||||||
|
|
||||||
|
// 再处理 1D norm 权重:存为 F16
|
||||||
|
for (name, shape, dtype, start, end) in &tensor_infos {
|
||||||
|
if shape.len() != 1 {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let raw = &mmap[data_offset + *start..data_offset + *end];
|
||||||
|
let f32_data = convert_to_f32(raw, dtype);
|
||||||
|
let numel = f32_data.len();
|
||||||
|
|
||||||
|
// 转 F16 存储
|
||||||
|
let mut f16_bytes = Vec::with_capacity(numel * 2);
|
||||||
|
for &v in &f32_data {
|
||||||
|
let f16_bits = f32_to_f16_bits(v);
|
||||||
|
f16_bytes.extend_from_slice(&f16_bits.to_le_bytes());
|
||||||
|
}
|
||||||
|
|
||||||
|
use std::io::Write;
|
||||||
|
bin_file.write_all(&f16_bytes)?;
|
||||||
|
|
||||||
|
// 用负数 scale 表示 F16 类型(加载时 scale < 0 表示 F16)
|
||||||
|
metadata.insert(
|
||||||
|
name.clone(),
|
||||||
|
Q8TensorMeta {
|
||||||
|
shape: shape.clone(),
|
||||||
|
scale: -1.0, // 标记为 F16
|
||||||
|
offset: current_offset,
|
||||||
|
numel,
|
||||||
|
},
|
||||||
|
);
|
||||||
|
current_offset += f16_bytes.len() as u64;
|
||||||
|
}
|
||||||
|
|
||||||
|
// 写 JSON 元数据
|
||||||
|
let meta = Q8Metadata { tensors: metadata };
|
||||||
|
let json_str = serde_json::to_string_pretty(&meta)?;
|
||||||
|
std::fs::write(&json_path, json_str)?;
|
||||||
|
|
||||||
|
println!("拷贝配置文件和 tokenizer...");
|
||||||
|
std::fs::copy(config_path, output_dir.join("config.json"))?;
|
||||||
|
std::fs::copy(tokenizer_path, output_dir.join("tokenizer.json"))?;
|
||||||
|
|
||||||
|
println!("Q8 量化模型已导出到: {:?}", output_dir);
|
||||||
|
println!(" - model.q8.json (元数据)",);
|
||||||
|
println!(" - model.q8.bin (INT8 + F16 权重)",);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// f32 → f16 位表示(返回 u16)
|
||||||
|
fn f32_to_f16_bits(v: f32) -> u16 {
|
||||||
|
let bits = v.to_bits();
|
||||||
|
let sign = bits >> 31;
|
||||||
|
let exp = (bits >> 23) & 0xFF;
|
||||||
|
let mant = bits & 0x7FFFFF;
|
||||||
|
|
||||||
|
let res: u32 = if exp == 0xFF {
|
||||||
|
// NaN / Inf
|
||||||
|
let f16_exp = 0x1F;
|
||||||
|
let f16_mant = if mant == 0 { 0 } else { 0x3FF };
|
||||||
|
(sign << 15) | (f16_exp << 10) | f16_mant
|
||||||
|
} else if exp < 103 {
|
||||||
|
// 下溢,返回 0
|
||||||
|
0
|
||||||
|
} else {
|
||||||
|
let new_exp = (exp as i32) - 127 + 15;
|
||||||
|
if new_exp >= 31 {
|
||||||
|
// 上溢
|
||||||
|
(sign << 15) | (0x1F << 10)
|
||||||
|
} else if new_exp <= 0 {
|
||||||
|
// 次正规数
|
||||||
|
let mant_new = (mant | 0x800000) >> (113 - new_exp);
|
||||||
|
(sign << 15) | (mant_new & 0x3FF)
|
||||||
|
} else {
|
||||||
|
let f16_mant = mant >> 13;
|
||||||
|
(sign << 15) | ((new_exp as u32) << 10) | f16_mant
|
||||||
|
}
|
||||||
|
};
|
||||||
|
res as u16
|
||||||
|
}
|
||||||
|
|||||||
@@ -10,3 +10,6 @@ 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"
|
||||||
|
memmap2 = "0.9"
|
||||||
|
serde = { version = "1.0", features = ["derive"] }
|
||||||
|
serde_json = "1.0"
|
||||||
@@ -0,0 +1,16 @@
|
|||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct GenerationConfig {
|
||||||
|
pub max_new_tokens: Option<usize>,
|
||||||
|
pub temperature: f32,
|
||||||
|
pub top_p: f32,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for GenerationConfig {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
max_new_tokens: None,
|
||||||
|
temperature: 1.0,
|
||||||
|
top_p: 1.0,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,139 @@
|
|||||||
|
use burn::tensor::backend::Backend;
|
||||||
|
use burn::tensor::{Int, Tensor};
|
||||||
|
|
||||||
|
use minicpm_core::config::EosTokenId;
|
||||||
|
use minicpm_core::model::{LlamaForCausalLM, LlamaKVCache};
|
||||||
|
|
||||||
|
use crate::config::GenerationConfig;
|
||||||
|
use crate::sampling::{sample, sample_last};
|
||||||
|
|
||||||
|
/// 流式生成迭代器 —— 用户可以自行遍历处理每个 token
|
||||||
|
///
|
||||||
|
/// 用法示例:
|
||||||
|
/// ```ignore
|
||||||
|
/// let stream = TokenStream::new(model, input_ids, config, eos_token_id, device);
|
||||||
|
/// for token in stream {
|
||||||
|
/// // 自行处理每个 token
|
||||||
|
/// println!("{}", token);
|
||||||
|
/// }
|
||||||
|
/// ```
|
||||||
|
pub struct TokenStream<'a, B: Backend> {
|
||||||
|
model: &'a LlamaForCausalLM<B>,
|
||||||
|
config: &'a GenerationConfig,
|
||||||
|
eos_token_id: &'a EosTokenId,
|
||||||
|
device: &'a B::Device,
|
||||||
|
cache: Option<LlamaKVCache<B>>,
|
||||||
|
last_token: Option<u32>,
|
||||||
|
count: usize,
|
||||||
|
finished: bool,
|
||||||
|
first_step: bool,
|
||||||
|
input_ids: Vec<u32>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<'a, B: Backend> TokenStream<'a, B> {
|
||||||
|
pub fn new(
|
||||||
|
model: &'a LlamaForCausalLM<B>,
|
||||||
|
input_ids: &[u32],
|
||||||
|
config: &'a GenerationConfig,
|
||||||
|
eos_token_id: &'a EosTokenId,
|
||||||
|
device: &'a B::Device,
|
||||||
|
) -> Self {
|
||||||
|
Self {
|
||||||
|
model,
|
||||||
|
config,
|
||||||
|
eos_token_id,
|
||||||
|
device,
|
||||||
|
cache: None,
|
||||||
|
last_token: None,
|
||||||
|
count: 0,
|
||||||
|
finished: false,
|
||||||
|
first_step: true,
|
||||||
|
input_ids: input_ids.to_vec(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<'a, B: Backend> Iterator for TokenStream<'a, B> {
|
||||||
|
type Item = u32;
|
||||||
|
|
||||||
|
fn next(&mut self) -> Option<Self::Item> {
|
||||||
|
if self.finished {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(max) = self.config.max_new_tokens {
|
||||||
|
if self.count >= max {
|
||||||
|
self.finished = true;
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let next_token = if self.first_step {
|
||||||
|
self.first_step = false;
|
||||||
|
self.step_first()
|
||||||
|
} else {
|
||||||
|
self.step_next()
|
||||||
|
};
|
||||||
|
|
||||||
|
match next_token {
|
||||||
|
Some(token) => {
|
||||||
|
self.count += 1;
|
||||||
|
self.last_token = Some(token);
|
||||||
|
|
||||||
|
if self.eos_token_id.contains(token) {
|
||||||
|
self.finished = true;
|
||||||
|
}
|
||||||
|
|
||||||
|
Some(token)
|
||||||
|
}
|
||||||
|
None => {
|
||||||
|
self.finished = true;
|
||||||
|
None
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<'a, B: Backend> TokenStream<'a, B> {
|
||||||
|
fn step_first(&mut self) -> Option<u32> {
|
||||||
|
let input_ints: Vec<i64> = self.input_ids.iter().map(|&x| x as i64).collect();
|
||||||
|
let input_tensor =
|
||||||
|
Tensor::<B, 1, Int>::from_ints(input_ints.as_slice(), self.device)
|
||||||
|
.unsqueeze::<2>();
|
||||||
|
|
||||||
|
let (logits, new_cache) = self.model.forward_with_cache(input_tensor, self.cache.as_ref());
|
||||||
|
self.cache = Some(new_cache);
|
||||||
|
|
||||||
|
Some(sample_last(&logits, self.config.temperature, self.config.top_p))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn step_next(&mut self) -> Option<u32> {
|
||||||
|
let last_token = self.last_token?;
|
||||||
|
let input_ints: Vec<i64> = vec![last_token as i64];
|
||||||
|
let input_tensor =
|
||||||
|
Tensor::<B, 1, Int>::from_ints(input_ints.as_slice(), self.device)
|
||||||
|
.unsqueeze::<2>();
|
||||||
|
|
||||||
|
let (logits, new_cache) = self.model.forward_with_cache(input_tensor, self.cache.as_ref());
|
||||||
|
self.cache = Some(new_cache);
|
||||||
|
|
||||||
|
let vocab_size = self.model.config.vocab_size;
|
||||||
|
let next_token_logits = logits
|
||||||
|
.slice([0..1, 0..1, 0..vocab_size])
|
||||||
|
.reshape([vocab_size]);
|
||||||
|
|
||||||
|
Some(sample(&next_token_logits, self.config.temperature, self.config.top_p))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 非流式生成 —— 返回所有生成的 token id
|
||||||
|
pub fn generate_tokens<B: Backend>(
|
||||||
|
model: &LlamaForCausalLM<B>,
|
||||||
|
input_ids: &[u32],
|
||||||
|
config: &GenerationConfig,
|
||||||
|
eos_token_id: &EosTokenId,
|
||||||
|
device: &B::Device,
|
||||||
|
) -> Vec<u32> {
|
||||||
|
let stream = TokenStream::new(model, input_ids, config, eos_token_id, device);
|
||||||
|
stream.collect()
|
||||||
|
}
|
||||||
@@ -1,305 +1,18 @@
|
|||||||
|
pub mod config;
|
||||||
|
pub mod generator;
|
||||||
|
pub mod loader;
|
||||||
|
pub mod sampling;
|
||||||
|
pub mod tokenizer;
|
||||||
|
|
||||||
|
pub use config::GenerationConfig;
|
||||||
|
pub use generator::TokenStream;
|
||||||
pub use minicpm_core::config::{EosTokenId, LlamaConfig};
|
pub use minicpm_core::config::{EosTokenId, LlamaConfig};
|
||||||
pub use minicpm_core::model::{LlamaForCausalLM, LlamaKVCache};
|
pub use minicpm_core::model::{LlamaForCausalLM, LlamaKVCache};
|
||||||
|
pub use tokenizer::TokenizerWrapper;
|
||||||
|
|
||||||
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 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)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
pub struct GenerationConfig {
|
|
||||||
pub max_new_tokens: Option<usize>,
|
|
||||||
pub temperature: f32,
|
|
||||||
pub top_p: f32,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl Default for GenerationConfig {
|
|
||||||
fn default() -> Self {
|
|
||||||
Self {
|
|
||||||
max_new_tokens: None,
|
|
||||||
temperature: 1.0,
|
|
||||||
top_p: 1.0,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn generate_with_cache<B: Backend>(
|
|
||||||
model: &LlamaForCausalLM<B>,
|
|
||||||
input_ids: &[u32],
|
|
||||||
config: &GenerationConfig,
|
|
||||||
eos_token_id: &EosTokenId,
|
|
||||||
device: &B::Device,
|
|
||||||
) -> Vec<u32> {
|
|
||||||
let mut output_ids = input_ids.to_vec();
|
|
||||||
let mut cache: Option<LlamaKVCache<B>> = None;
|
|
||||||
|
|
||||||
// 第一步:输入完整 prompt,建立初始 cache
|
|
||||||
{
|
|
||||||
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);
|
|
||||||
|
|
||||||
if eos_token_id.contains(next_token) {
|
|
||||||
return output_ids;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// 后续步骤:每次只输入 1 个 token,使用 cache
|
|
||||||
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);
|
|
||||||
|
|
||||||
if eos_token_id.contains(next_token) {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
|
|
||||||
count += 1;
|
|
||||||
}
|
|
||||||
|
|
||||||
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 {
|
|
||||||
let shape = logits.shape();
|
|
||||||
let dims: [usize; 3] = shape.dims();
|
|
||||||
let seq_len = dims[1];
|
|
||||||
let vocab_size = dims[2];
|
|
||||||
|
|
||||||
let last_logits = logits
|
|
||||||
.clone()
|
|
||||||
.slice([0..1, seq_len - 1..seq_len, 0..vocab_size])
|
|
||||||
.reshape([vocab_size]);
|
|
||||||
|
|
||||||
sample(&last_logits, temperature, top_p)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn sample<B: Backend>(logits: &Tensor<B, 1, Float>, temperature: f32, top_p: f32) -> u32 {
|
|
||||||
let data = logits.clone().to_data();
|
|
||||||
let values: Vec<f32> = data.as_slice().unwrap().to_vec();
|
|
||||||
|
|
||||||
// temperature = 0 时退化为 greedy
|
|
||||||
if temperature <= 0.001 {
|
|
||||||
let mut max_idx = 0;
|
|
||||||
let mut max_val = f32::NEG_INFINITY;
|
|
||||||
for (i, &val) in values.iter().enumerate() {
|
|
||||||
if val > max_val {
|
|
||||||
max_val = val;
|
|
||||||
max_idx = i;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
return max_idx as u32;
|
|
||||||
}
|
|
||||||
|
|
||||||
// 应用 temperature
|
|
||||||
let scaled: Vec<f32> = values.iter().map(|&v| v / temperature).collect();
|
|
||||||
|
|
||||||
// softmax
|
|
||||||
let max_val = scaled.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
|
|
||||||
let exp_sum: f32 = scaled.iter().map(|&v| (v - max_val).exp()).sum();
|
|
||||||
let mut probs: Vec<f32> = scaled.iter().map(|&v| (v - max_val).exp() / exp_sum).collect();
|
|
||||||
|
|
||||||
// top_p 过滤
|
|
||||||
if top_p < 1.0 {
|
|
||||||
let mut indexed: Vec<(usize, f32)> = probs.iter().enumerate().map(|(i, &p)| (i, p)).collect();
|
|
||||||
indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
|
|
||||||
|
|
||||||
let mut cumsum = 0.0;
|
|
||||||
let mut cutoff = indexed.len();
|
|
||||||
for (i, &(_, p)) in indexed.iter().enumerate() {
|
|
||||||
cumsum += p;
|
|
||||||
if cumsum > top_p {
|
|
||||||
cutoff = i + 1;
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
let keep: std::collections::HashSet<usize> = indexed[..cutoff].iter().map(|(i, _)| *i).collect();
|
|
||||||
for (i, p) in probs.iter_mut().enumerate() {
|
|
||||||
if !keep.contains(&i) {
|
|
||||||
*p = 0.0;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
let sum: f32 = probs.iter().sum();
|
|
||||||
if sum > 0.0 {
|
|
||||||
for p in probs.iter_mut() {
|
|
||||||
*p /= sum;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// 多项式采样
|
|
||||||
let mut rng = rand::thread_rng();
|
|
||||||
let r: f32 = rng.gen();
|
|
||||||
let mut cumsum = 0.0;
|
|
||||||
for (i, &p) in probs.iter().enumerate() {
|
|
||||||
cumsum += p;
|
|
||||||
if r < cumsum {
|
|
||||||
return i as u32;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
probs.len() as u32 - 1
|
|
||||||
}
|
|
||||||
|
|
||||||
|
use loader::{load_model, load_model_q8_from_dir};
|
||||||
|
|
||||||
pub struct MiniCPM<B: Backend> {
|
pub struct MiniCPM<B: Backend> {
|
||||||
model: LlamaForCausalLM<B>,
|
model: LlamaForCausalLM<B>,
|
||||||
@@ -317,10 +30,25 @@ impl<B: Backend> MiniCPM<B> {
|
|||||||
) -> anyhow::Result<Self> {
|
) -> anyhow::Result<Self> {
|
||||||
let config = LlamaConfig::from_json(config_path)?;
|
let config = LlamaConfig::from_json(config_path)?;
|
||||||
let tokenizer = TokenizerWrapper::from_file(tokenizer_path)?;
|
let tokenizer = TokenizerWrapper::from_file(tokenizer_path)?;
|
||||||
|
let model = load_model::<B>(model_path, &config, device)?;
|
||||||
|
|
||||||
let model = LlamaForCausalLM::<B>::new(config.clone(), device);
|
Ok(Self {
|
||||||
let recorder = NamedMpkFileRecorder::<FullPrecisionSettings>::new();
|
model,
|
||||||
let model = model.load_file(model_path, &recorder, device)?;
|
config,
|
||||||
|
tokenizer,
|
||||||
|
device: device.clone(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn load_q8(
|
||||||
|
model_dir: &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 = load_model_q8_from_dir::<B>(model_dir, &config, device)?;
|
||||||
|
|
||||||
Ok(Self {
|
Ok(Self {
|
||||||
model,
|
model,
|
||||||
@@ -337,7 +65,7 @@ impl<B: Backend> MiniCPM<B> {
|
|||||||
config: &GenerationConfig,
|
config: &GenerationConfig,
|
||||||
) -> anyhow::Result<String> {
|
) -> anyhow::Result<String> {
|
||||||
let input_ids = self.tokenizer.apply_chat_template(prompt, think)?;
|
let input_ids = self.tokenizer.apply_chat_template(prompt, think)?;
|
||||||
let output_ids = generate_with_cache(
|
let output_ids = generator::generate_tokens(
|
||||||
&self.model,
|
&self.model,
|
||||||
&input_ids,
|
&input_ids,
|
||||||
config,
|
config,
|
||||||
@@ -348,27 +76,26 @@ impl<B: Backend> MiniCPM<B> {
|
|||||||
self.tokenizer.decode(new_ids, true)
|
self.tokenizer.decode(new_ids, true)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// 流式生成:每生成一个 token 立即调用 `on_token` 回调,
|
pub fn generate_stream<'a>(
|
||||||
/// 参数为该 token 解码后的文本。
|
&'a self,
|
||||||
pub fn generate_stream(
|
|
||||||
&self,
|
|
||||||
prompt: &str,
|
prompt: &str,
|
||||||
think: bool,
|
think: bool,
|
||||||
config: &GenerationConfig,
|
config: &'a GenerationConfig,
|
||||||
on_token: impl FnMut(&str),
|
) -> anyhow::Result<TextStream<'a, B>> {
|
||||||
) -> anyhow::Result<String> {
|
|
||||||
let input_ids = self.tokenizer.apply_chat_template(prompt, think)?;
|
let input_ids = self.tokenizer.apply_chat_template(prompt, think)?;
|
||||||
let output_ids = generate_stream(
|
let token_stream = TokenStream::new(
|
||||||
&self.model,
|
&self.model,
|
||||||
&self.tokenizer,
|
|
||||||
&input_ids,
|
&input_ids,
|
||||||
config,
|
config,
|
||||||
&self.config.eos_token_id,
|
&self.config.eos_token_id,
|
||||||
&self.device,
|
&self.device,
|
||||||
on_token,
|
|
||||||
);
|
);
|
||||||
let new_ids = &output_ids[input_ids.len()..];
|
|
||||||
self.tokenizer.decode(new_ids, true)
|
Ok(TextStream {
|
||||||
|
token_stream,
|
||||||
|
tokenizer: &self.tokenizer,
|
||||||
|
buffer: Vec::new(),
|
||||||
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn config(&self) -> &LlamaConfig {
|
pub fn config(&self) -> &LlamaConfig {
|
||||||
@@ -383,3 +110,41 @@ impl<B: Backend> MiniCPM<B> {
|
|||||||
&self.model
|
&self.model
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// 文本流式迭代器 —— 解码 token 并返回文本片段
|
||||||
|
///
|
||||||
|
/// 用法示例:
|
||||||
|
/// ```ignore
|
||||||
|
/// let mut stream = model.generate_stream("你好", false, &config)?;
|
||||||
|
/// for text in stream {
|
||||||
|
/// print!("{}", text);
|
||||||
|
/// }
|
||||||
|
/// ```
|
||||||
|
pub struct TextStream<'a, B: Backend> {
|
||||||
|
token_stream: TokenStream<'a, B>,
|
||||||
|
tokenizer: &'a TokenizerWrapper,
|
||||||
|
buffer: Vec<u32>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<'a, B: Backend> Iterator for TextStream<'a, B> {
|
||||||
|
type Item = String;
|
||||||
|
|
||||||
|
fn next(&mut self) -> Option<Self::Item> {
|
||||||
|
match self.token_stream.next() {
|
||||||
|
Some(token) => {
|
||||||
|
self.buffer.push(token);
|
||||||
|
match self.tokenizer.decode(&[token], true) {
|
||||||
|
Ok(text) => Some(text),
|
||||||
|
Err(_) => Some(String::new()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
None => None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<'a, B: Backend> TextStream<'a, B> {
|
||||||
|
pub fn all_tokens(&self) -> &[u32] {
|
||||||
|
&self.buffer
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,29 @@
|
|||||||
|
pub mod q8;
|
||||||
|
|
||||||
|
use burn::module::Module;
|
||||||
|
use burn::record::{FullPrecisionSettings, NamedMpkFileRecorder};
|
||||||
|
use burn::tensor::backend::Backend;
|
||||||
|
|
||||||
|
use minicpm_core::config::LlamaConfig;
|
||||||
|
use minicpm_core::model::LlamaForCausalLM;
|
||||||
|
|
||||||
|
use self::q8::load_model_q8;
|
||||||
|
|
||||||
|
pub fn load_model<B: Backend>(
|
||||||
|
model_path: &str,
|
||||||
|
config: &LlamaConfig,
|
||||||
|
device: &B::Device,
|
||||||
|
) -> anyhow::Result<LlamaForCausalLM<B>> {
|
||||||
|
let model = LlamaForCausalLM::<B>::new(config.clone(), device);
|
||||||
|
let recorder = NamedMpkFileRecorder::<FullPrecisionSettings>::new();
|
||||||
|
let model = model.load_file(model_path, &recorder, device)?;
|
||||||
|
Ok(model)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn load_model_q8_from_dir<B: Backend>(
|
||||||
|
model_dir: &str,
|
||||||
|
config: &LlamaConfig,
|
||||||
|
device: &B::Device,
|
||||||
|
) -> anyhow::Result<LlamaForCausalLM<B>> {
|
||||||
|
load_model_q8(model_dir, config, device)
|
||||||
|
}
|
||||||
@@ -0,0 +1,201 @@
|
|||||||
|
use burn::module::{Module, Param};
|
||||||
|
use burn::nn::{EmbeddingRecord, LinearRecord};
|
||||||
|
use burn::tensor::backend::Backend;
|
||||||
|
use burn::tensor::{Float, Shape, Tensor, TensorData};
|
||||||
|
use memmap2::Mmap;
|
||||||
|
use std::collections::HashMap;
|
||||||
|
|
||||||
|
use minicpm_core::config::LlamaConfig;
|
||||||
|
use minicpm_core::model::LlamaForCausalLM;
|
||||||
|
|
||||||
|
#[derive(serde::Deserialize)]
|
||||||
|
struct Q8TensorMeta {
|
||||||
|
shape: Vec<usize>,
|
||||||
|
scale: f32,
|
||||||
|
offset: u64,
|
||||||
|
numel: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(serde::Deserialize)]
|
||||||
|
struct Q8Metadata {
|
||||||
|
tensors: HashMap<String, Q8TensorMeta>,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn load_model_q8<B: Backend>(
|
||||||
|
model_dir: &str,
|
||||||
|
config: &LlamaConfig,
|
||||||
|
device: &B::Device,
|
||||||
|
) -> anyhow::Result<LlamaForCausalLM<B>> {
|
||||||
|
let json_path = std::path::Path::new(model_dir).join("model.q8.json");
|
||||||
|
let bin_path = std::path::Path::new(model_dir).join("model.q8.bin");
|
||||||
|
|
||||||
|
let json_str = std::fs::read_to_string(&json_path)?;
|
||||||
|
let metadata: Q8Metadata = serde_json::from_str(&json_str)?;
|
||||||
|
|
||||||
|
let file = std::fs::File::open(&bin_path)?;
|
||||||
|
let mmap = unsafe { Mmap::map(&file)? };
|
||||||
|
|
||||||
|
let mut model = LlamaForCausalLM::<B>::new(config.clone(), device);
|
||||||
|
|
||||||
|
if let Some(meta) = metadata.tensors.get("model.embed_tokens.weight") {
|
||||||
|
let w = if meta.scale >= 0.0 {
|
||||||
|
load_q8_tensor_2d::<B>(&mmap, meta, device)
|
||||||
|
} else {
|
||||||
|
load_f16_tensor_2d::<B>(&mmap, meta, device)
|
||||||
|
};
|
||||||
|
let record = EmbeddingRecord { weight: Param::from_tensor(w) };
|
||||||
|
model.model.embed_tokens = model.model.embed_tokens.clone().load_record(record);
|
||||||
|
}
|
||||||
|
|
||||||
|
for i in 0..config.num_hidden_layers {
|
||||||
|
let p = format!("model.layers.{i}");
|
||||||
|
|
||||||
|
if let Some(meta) = metadata.tensors.get(&format!("{p}.self_attn.q_proj.weight")) {
|
||||||
|
let w = load_q8_tensor_2d::<B>(&mmap, meta, device).transpose();
|
||||||
|
let record = LinearRecord { weight: Param::from_tensor(w), bias: None };
|
||||||
|
model.model.layers[i].self_attn.q_proj = model.model.layers[i].self_attn.q_proj.clone().load_record(record);
|
||||||
|
}
|
||||||
|
if let Some(meta) = metadata.tensors.get(&format!("{p}.self_attn.k_proj.weight")) {
|
||||||
|
let w = load_q8_tensor_2d::<B>(&mmap, meta, device).transpose();
|
||||||
|
let record = LinearRecord { weight: Param::from_tensor(w), bias: None };
|
||||||
|
model.model.layers[i].self_attn.k_proj = model.model.layers[i].self_attn.k_proj.clone().load_record(record);
|
||||||
|
}
|
||||||
|
if let Some(meta) = metadata.tensors.get(&format!("{p}.self_attn.v_proj.weight")) {
|
||||||
|
let w = load_q8_tensor_2d::<B>(&mmap, meta, device).transpose();
|
||||||
|
let record = LinearRecord { weight: Param::from_tensor(w), bias: None };
|
||||||
|
model.model.layers[i].self_attn.v_proj = model.model.layers[i].self_attn.v_proj.clone().load_record(record);
|
||||||
|
}
|
||||||
|
if let Some(meta) = metadata.tensors.get(&format!("{p}.self_attn.o_proj.weight")) {
|
||||||
|
let w = load_q8_tensor_2d::<B>(&mmap, meta, device).transpose();
|
||||||
|
let record = LinearRecord { weight: Param::from_tensor(w), bias: None };
|
||||||
|
model.model.layers[i].self_attn.o_proj = model.model.layers[i].self_attn.o_proj.clone().load_record(record);
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(meta) = metadata.tensors.get(&format!("{p}.mlp.gate_proj.weight")) {
|
||||||
|
let w = load_q8_tensor_2d::<B>(&mmap, meta, device).transpose();
|
||||||
|
let record = LinearRecord { weight: Param::from_tensor(w), bias: None };
|
||||||
|
model.model.layers[i].mlp.gate_proj = model.model.layers[i].mlp.gate_proj.clone().load_record(record);
|
||||||
|
}
|
||||||
|
if let Some(meta) = metadata.tensors.get(&format!("{p}.mlp.up_proj.weight")) {
|
||||||
|
let w = load_q8_tensor_2d::<B>(&mmap, meta, device).transpose();
|
||||||
|
let record = LinearRecord { weight: Param::from_tensor(w), bias: None };
|
||||||
|
model.model.layers[i].mlp.up_proj = model.model.layers[i].mlp.up_proj.clone().load_record(record);
|
||||||
|
}
|
||||||
|
if let Some(meta) = metadata.tensors.get(&format!("{p}.mlp.down_proj.weight")) {
|
||||||
|
let w = load_q8_tensor_2d::<B>(&mmap, meta, device).transpose();
|
||||||
|
let record = LinearRecord { weight: Param::from_tensor(w), bias: None };
|
||||||
|
model.model.layers[i].mlp.down_proj = model.model.layers[i].mlp.down_proj.clone().load_record(record);
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(meta) = metadata.tensors.get(&format!("{p}.input_layernorm.weight")) {
|
||||||
|
let w = load_f16_tensor_1d::<B>(&mmap, meta, device);
|
||||||
|
model.model.layers[i].input_layernorm.weight = Param::from_tensor(w);
|
||||||
|
}
|
||||||
|
if let Some(meta) = metadata.tensors.get(&format!("{p}.post_attention_layernorm.weight")) {
|
||||||
|
let w = load_f16_tensor_1d::<B>(&mmap, meta, device);
|
||||||
|
model.model.layers[i].post_attention_layernorm.weight = Param::from_tensor(w);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(meta) = metadata.tensors.get("model.norm.weight") {
|
||||||
|
let w = load_f16_tensor_1d::<B>(&mmap, meta, device);
|
||||||
|
model.model.norm.weight = Param::from_tensor(w);
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(meta) = metadata.tensors.get("lm_head.weight") {
|
||||||
|
let w = if meta.scale >= 0.0 {
|
||||||
|
load_q8_tensor_2d::<B>(&mmap, meta, device)
|
||||||
|
} else {
|
||||||
|
load_f16_tensor_2d::<B>(&mmap, meta, device)
|
||||||
|
};
|
||||||
|
let w_t = w.transpose();
|
||||||
|
let record = LinearRecord { weight: Param::from_tensor(w_t), bias: None };
|
||||||
|
model.lm_head = model.lm_head.clone().load_record(record);
|
||||||
|
} else if config.tie_word_embeddings {
|
||||||
|
let embed_weight = model.model.embed_tokens.weight.val();
|
||||||
|
let w_t = embed_weight.transpose();
|
||||||
|
let record = LinearRecord { weight: Param::from_tensor(w_t), bias: None };
|
||||||
|
model.lm_head = model.lm_head.clone().load_record(record);
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(model)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn load_q8_tensor_2d<B: Backend>(
|
||||||
|
mmap: &Mmap,
|
||||||
|
meta: &Q8TensorMeta,
|
||||||
|
device: &B::Device,
|
||||||
|
) -> Tensor<B, 2, Float> {
|
||||||
|
let bytes = &mmap[meta.offset as usize..meta.offset as usize + meta.numel];
|
||||||
|
let q8: &[i8] = unsafe { std::slice::from_raw_parts(bytes.as_ptr() as *const i8, meta.numel) };
|
||||||
|
let f32_data: Vec<f32> = q8.iter().map(|&q| (q as f32) * meta.scale).collect();
|
||||||
|
let shape: [usize; 2] = [meta.shape[0], meta.shape[1]];
|
||||||
|
Tensor::from_data(TensorData::new(f32_data, Shape::new(shape)), device)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn load_f16_tensor_1d<B: Backend>(
|
||||||
|
mmap: &Mmap,
|
||||||
|
meta: &Q8TensorMeta,
|
||||||
|
device: &B::Device,
|
||||||
|
) -> Tensor<B, 1, Float> {
|
||||||
|
let start = meta.offset as usize;
|
||||||
|
let bytes = &mmap[start..start + meta.numel * 2];
|
||||||
|
let f32_data: Vec<f32> = bytes
|
||||||
|
.chunks_exact(2)
|
||||||
|
.map(|c| {
|
||||||
|
let bits = u16::from_le_bytes([c[0], c[1]]);
|
||||||
|
f16_to_f32(bits)
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
let shape: [usize; 1] = [meta.shape[0]];
|
||||||
|
Tensor::from_data(TensorData::new(f32_data, Shape::new(shape)), device)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn load_f16_tensor_2d<B: Backend>(
|
||||||
|
mmap: &Mmap,
|
||||||
|
meta: &Q8TensorMeta,
|
||||||
|
device: &B::Device,
|
||||||
|
) -> Tensor<B, 2, Float> {
|
||||||
|
let start = meta.offset as usize;
|
||||||
|
let bytes = &mmap[start..start + meta.numel * 2];
|
||||||
|
let f32_data: Vec<f32> = bytes
|
||||||
|
.chunks_exact(2)
|
||||||
|
.map(|c| {
|
||||||
|
let bits = u16::from_le_bytes([c[0], c[1]]);
|
||||||
|
f16_to_f32(bits)
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
let shape: [usize; 2] = [meta.shape[0], meta.shape[1]];
|
||||||
|
Tensor::from_data(TensorData::new(f32_data, Shape::new(shape)), device)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn f16_to_f32(f16: u16) -> f32 {
|
||||||
|
let sign = (f16 >> 15) & 1;
|
||||||
|
let exp = (f16 >> 10) & 0x1F;
|
||||||
|
let mant = f16 & 0x3FF;
|
||||||
|
|
||||||
|
if exp == 0 {
|
||||||
|
if mant == 0 {
|
||||||
|
return if sign == 0 { 0.0 } else { -0.0 };
|
||||||
|
}
|
||||||
|
let mut m = mant as u32;
|
||||||
|
let mut e = 0u32;
|
||||||
|
while (m & 0x400) == 0 {
|
||||||
|
m <<= 1;
|
||||||
|
e += 1;
|
||||||
|
}
|
||||||
|
m &= 0x3FF;
|
||||||
|
let f32_exp = 127 - 14 - e;
|
||||||
|
let f32_bits = ((sign as u32) << 31) | (f32_exp << 23) | ((m as u32) << 13);
|
||||||
|
return f32::from_bits(f32_bits);
|
||||||
|
}
|
||||||
|
|
||||||
|
if exp == 0x1F {
|
||||||
|
let f32_bits = ((sign as u32) << 31) | (0xFFu32 << 23) | ((mant as u32) << 13);
|
||||||
|
return f32::from_bits(f32_bits);
|
||||||
|
}
|
||||||
|
|
||||||
|
let f32_exp = (exp as u32) + 127 - 15;
|
||||||
|
let f32_bits = ((sign as u32) << 31) | (f32_exp << 23) | ((mant as u32) << 13);
|
||||||
|
f32::from_bits(f32_bits)
|
||||||
|
}
|
||||||
@@ -0,0 +1,80 @@
|
|||||||
|
use burn::tensor::backend::Backend;
|
||||||
|
use burn::tensor::{Float, Tensor};
|
||||||
|
use rand::Rng;
|
||||||
|
|
||||||
|
pub fn sample_last<B: Backend>(logits: &Tensor<B, 3, Float>, temperature: f32, top_p: f32) -> u32 {
|
||||||
|
let shape = logits.shape();
|
||||||
|
let dims: [usize; 3] = shape.dims();
|
||||||
|
let seq_len = dims[1];
|
||||||
|
let vocab_size = dims[2];
|
||||||
|
|
||||||
|
let last_logits = logits
|
||||||
|
.clone()
|
||||||
|
.slice([0..1, seq_len - 1..seq_len, 0..vocab_size])
|
||||||
|
.reshape([vocab_size]);
|
||||||
|
|
||||||
|
sample(&last_logits, temperature, top_p)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn sample<B: Backend>(logits: &Tensor<B, 1, Float>, temperature: f32, top_p: f32) -> u32 {
|
||||||
|
let data = logits.clone().to_data();
|
||||||
|
let values: Vec<f32> = data.as_slice().unwrap().to_vec();
|
||||||
|
|
||||||
|
if temperature <= 0.001 {
|
||||||
|
let mut max_idx = 0;
|
||||||
|
let mut max_val = f32::NEG_INFINITY;
|
||||||
|
for (i, &val) in values.iter().enumerate() {
|
||||||
|
if val > max_val {
|
||||||
|
max_val = val;
|
||||||
|
max_idx = i;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return max_idx as u32;
|
||||||
|
}
|
||||||
|
|
||||||
|
let scaled: Vec<f32> = values.iter().map(|&v| v / temperature).collect();
|
||||||
|
|
||||||
|
let max_val = scaled.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
|
||||||
|
let exp_sum: f32 = scaled.iter().map(|&v| (v - max_val).exp()).sum();
|
||||||
|
let mut probs: Vec<f32> = scaled.iter().map(|&v| (v - max_val).exp() / exp_sum).collect();
|
||||||
|
|
||||||
|
if top_p < 1.0 {
|
||||||
|
let mut indexed: Vec<(usize, f32)> = probs.iter().enumerate().map(|(i, &p)| (i, p)).collect();
|
||||||
|
indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
|
||||||
|
|
||||||
|
let mut cumsum = 0.0;
|
||||||
|
let mut cutoff = indexed.len();
|
||||||
|
for (i, &(_, p)) in indexed.iter().enumerate() {
|
||||||
|
cumsum += p;
|
||||||
|
if cumsum > top_p {
|
||||||
|
cutoff = i + 1;
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let keep: std::collections::HashSet<usize> = indexed[..cutoff].iter().map(|(i, _)| *i).collect();
|
||||||
|
for (i, p) in probs.iter_mut().enumerate() {
|
||||||
|
if !keep.contains(&i) {
|
||||||
|
*p = 0.0;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let sum: f32 = probs.iter().sum();
|
||||||
|
if sum > 0.0 {
|
||||||
|
for p in probs.iter_mut() {
|
||||||
|
*p /= sum;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut rng = rand::thread_rng();
|
||||||
|
let r: f32 = rng.gen();
|
||||||
|
let mut cumsum = 0.0;
|
||||||
|
for (i, &p) in probs.iter().enumerate() {
|
||||||
|
cumsum += p;
|
||||||
|
if r < cumsum {
|
||||||
|
return i as u32;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
probs.len() as u32 - 1
|
||||||
|
}
|
||||||
@@ -0,0 +1,48 @@
|
|||||||
|
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!("加载 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!("编码失败: {}", 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!("解码失败: {}", e))?;
|
||||||
|
Ok(text)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn vocab_size(&self) -> usize {
|
||||||
|
self.tokenizer.get_vocab_size(true)
|
||||||
|
}
|
||||||
|
|
||||||
|
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,24 +1,14 @@
|
|||||||
use burn::backend::Wgpu;
|
use minicpm_convert::export_model_q8;
|
||||||
use minicpm_convert::export_model;
|
|
||||||
use std::path::Path;
|
use std::path::Path;
|
||||||
|
|
||||||
fn main() -> anyhow::Result<()> {
|
fn main() -> anyhow::Result<()> {
|
||||||
let device = Default::default();
|
|
||||||
|
|
||||||
let safetensors_path = Path::new("MiniCPM5-1B/model-00000-of-00001.safetensors");
|
let safetensors_path = Path::new("MiniCPM5-1B/model-00000-of-00001.safetensors");
|
||||||
let config_path = Path::new("MiniCPM5-1B/config.json");
|
let config_path = Path::new("MiniCPM5-1B/config.json");
|
||||||
let tokenizer_path = Path::new("MiniCPM5-1B/tokenizer.json");
|
let tokenizer_path = Path::new("MiniCPM5-1B/tokenizer.json");
|
||||||
let output_dir = Path::new("model");
|
|
||||||
|
|
||||||
export_model::<Wgpu>(
|
// Q8 量化导出(INT8,文件约 1/4 大小)
|
||||||
safetensors_path,
|
export_model_q8(safetensors_path, config_path, tokenizer_path, Path::new("model_q8"))?;
|
||||||
config_path,
|
|
||||||
tokenizer_path,
|
|
||||||
output_dir,
|
|
||||||
&device,
|
|
||||||
)?;
|
|
||||||
|
|
||||||
println!("模型转换完成!");
|
println!("模型转换完成!");
|
||||||
|
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,9 +1,9 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "minimal-inference"
|
name = "flex-backend"
|
||||||
version = "0.1.0"
|
version = "0.1.0"
|
||||||
edition = "2021"
|
edition = "2021"
|
||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
minicpm-inference = { path = "../../crates/minicpm-inference" }
|
minicpm-inference = { path = "../../crates/minicpm-inference" }
|
||||||
burn = { version = "0.21", features = ["std", "wgpu"] }
|
burn = { version = "0.21", default-features = false, features = ["std", "flex"] }
|
||||||
anyhow = "1.0"
|
anyhow = "1.0"
|
||||||
@@ -0,0 +1,82 @@
|
|||||||
|
// Flex 后端推理示例 —— 适用于 Tauri / WebAssembly 环境
|
||||||
|
//
|
||||||
|
// burn-flex 是纯 Rust CPU 后端,无 GPU 依赖,完美适配 Tauri 应用。
|
||||||
|
// 支持 SIMD 加速(AVX2/NEON)和多线程(rayon)。
|
||||||
|
//
|
||||||
|
// 运行方式:cargo run --release -p flex-backend [--full|--half|--q8]
|
||||||
|
|
||||||
|
use burn::backend::{flex::FlexDevice, Flex};
|
||||||
|
use minicpm_inference::{GenerationConfig, MiniCPM};
|
||||||
|
use std::io::Write;
|
||||||
|
use std::time::Instant;
|
||||||
|
|
||||||
|
fn main() -> anyhow::Result<()> {
|
||||||
|
let args: Vec<String> = std::env::args().collect();
|
||||||
|
let mode = args.get(1).map(|s| s.as_str()).unwrap_or("--full");
|
||||||
|
|
||||||
|
let (model_dir, config_path, tokenizer_path, is_q8) = match mode {
|
||||||
|
"--half" => (
|
||||||
|
"model_half",
|
||||||
|
"MiniCPM5-1B/config.json",
|
||||||
|
"MiniCPM5-1B/tokenizer.json",
|
||||||
|
false,
|
||||||
|
),
|
||||||
|
"--q8" => (
|
||||||
|
"model_q8",
|
||||||
|
"MiniCPM5-1B/config.json",
|
||||||
|
"MiniCPM5-1B/tokenizer.json",
|
||||||
|
true,
|
||||||
|
),
|
||||||
|
_ => (
|
||||||
|
"model",
|
||||||
|
"MiniCPM5-1B/config.json",
|
||||||
|
"MiniCPM5-1B/tokenizer.json",
|
||||||
|
false,
|
||||||
|
),
|
||||||
|
};
|
||||||
|
|
||||||
|
println!("正在加载模型(后端:Flex,模式:{mode})...");
|
||||||
|
println!("Flex 后端特性:纯 Rust、SIMD 加速、多线程");
|
||||||
|
let start = Instant::now();
|
||||||
|
|
||||||
|
let device = FlexDevice::default();
|
||||||
|
let model: MiniCPM<Flex> = if is_q8 {
|
||||||
|
MiniCPM::<Flex>::load_q8(model_dir, config_path, tokenizer_path, &device)?
|
||||||
|
} else {
|
||||||
|
let model_path = format!("{model_dir}/model");
|
||||||
|
MiniCPM::<Flex>::load(&model_path, config_path, tokenizer_path, &device)?
|
||||||
|
};
|
||||||
|
|
||||||
|
println!("模型加载完成,耗时: {:.2?}", start.elapsed());
|
||||||
|
|
||||||
|
let prompt = "需要处理生成一个河南的旅游计划,要求如下:
|
||||||
|
1. 旅游计划为三天两夜的行程安排。
|
||||||
|
2. 每天的行程安排包括景点、餐饮和住宿建议。
|
||||||
|
3. 景点安排要考虑交通便利性和时间合理性。
|
||||||
|
请生成详细的旅游计划,包含每天的行程安排、景点介绍、餐饮推荐和住宿建议。
|
||||||
|
";
|
||||||
|
println!("\n用户: {prompt}");
|
||||||
|
print!("Assistant: ");
|
||||||
|
std::io::stdout().flush()?;
|
||||||
|
|
||||||
|
let config = GenerationConfig {
|
||||||
|
max_new_tokens: Some(1200),
|
||||||
|
temperature: 0.7,
|
||||||
|
top_p: 0.95,
|
||||||
|
};
|
||||||
|
|
||||||
|
let start = Instant::now();
|
||||||
|
let stream = model.generate_stream(prompt, false, &config)?;
|
||||||
|
|
||||||
|
for text in stream {
|
||||||
|
print!("{text}");
|
||||||
|
std::io::stdout().flush().ok();
|
||||||
|
}
|
||||||
|
|
||||||
|
let elapsed = start.elapsed();
|
||||||
|
|
||||||
|
println!();
|
||||||
|
println!("\n生成耗时: {:.2?}", elapsed);
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
@@ -1,54 +0,0 @@
|
|||||||
// 运行方式:
|
|
||||||
// 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_stream(
|
|
||||||
prompt,
|
|
||||||
false,
|
|
||||||
&config,
|
|
||||||
|token| {
|
|
||||||
print!("{}", token);
|
|
||||||
std::io::Write::flush(&mut std::io::stdout()).ok();
|
|
||||||
},
|
|
||||||
)?;
|
|
||||||
let elapsed = start.elapsed();
|
|
||||||
|
|
||||||
println!();
|
|
||||||
println!("\n生成耗时: {:.2?}", elapsed);
|
|
||||||
|
|
||||||
Ok(())
|
|
||||||
}
|
|
||||||
@@ -0,0 +1,9 @@
|
|||||||
|
[package]
|
||||||
|
name = "wgpu-backend"
|
||||||
|
version = "0.1.0"
|
||||||
|
edition = "2021"
|
||||||
|
|
||||||
|
[dependencies]
|
||||||
|
minicpm-inference = { path = "../../crates/minicpm-inference" }
|
||||||
|
burn = { version = "0.21", default-features = false, features = ["std", "wgpu"] }
|
||||||
|
anyhow = "1.0"
|
||||||
@@ -0,0 +1,82 @@
|
|||||||
|
// Wgpu 后端推理示例 —— GPU 加速推理
|
||||||
|
//
|
||||||
|
// 依赖 WebGPU API,支持 CUDA/Metal/Vulkan 等平台。
|
||||||
|
// 注意:与 Tauri 应用可能存在兼容性冲突,建议使用 Flex 后端。
|
||||||
|
//
|
||||||
|
// 运行方式:cargo run --release -p wgpu-backend [--full|--half|--q8]
|
||||||
|
|
||||||
|
use burn::backend::{wgpu::WgpuDevice, Wgpu};
|
||||||
|
use minicpm_inference::{GenerationConfig, MiniCPM};
|
||||||
|
use std::io::Write;
|
||||||
|
use std::time::Instant;
|
||||||
|
|
||||||
|
fn main() -> anyhow::Result<()> {
|
||||||
|
let args: Vec<String> = std::env::args().collect();
|
||||||
|
let mode = args.get(1).map(|s| s.as_str()).unwrap_or("--full");
|
||||||
|
|
||||||
|
let (model_dir, config_path, tokenizer_path, is_q8) = match mode {
|
||||||
|
"--half" => (
|
||||||
|
"model_half",
|
||||||
|
"MiniCPM5-1B/config.json",
|
||||||
|
"MiniCPM5-1B/tokenizer.json",
|
||||||
|
false,
|
||||||
|
),
|
||||||
|
"--q8" => (
|
||||||
|
"model_q8",
|
||||||
|
"MiniCPM5-1B/config.json",
|
||||||
|
"MiniCPM5-1B/tokenizer.json",
|
||||||
|
true,
|
||||||
|
),
|
||||||
|
_ => (
|
||||||
|
"model",
|
||||||
|
"MiniCPM5-1B/config.json",
|
||||||
|
"MiniCPM5-1B/tokenizer.json",
|
||||||
|
false,
|
||||||
|
),
|
||||||
|
};
|
||||||
|
|
||||||
|
println!("正在加载模型(后端:Wgpu,模式:{mode})...");
|
||||||
|
println!("Wgpu 后端特性:GPU 加速(CUDA/Metal/Vulkan)");
|
||||||
|
let start = Instant::now();
|
||||||
|
|
||||||
|
let device = WgpuDevice::default();
|
||||||
|
let model: MiniCPM<Wgpu> = if is_q8 {
|
||||||
|
MiniCPM::<Wgpu>::load_q8(model_dir, config_path, tokenizer_path, &device)?
|
||||||
|
} else {
|
||||||
|
let model_path = format!("{model_dir}/model");
|
||||||
|
MiniCPM::<Wgpu>::load(&model_path, config_path, tokenizer_path, &device)?
|
||||||
|
};
|
||||||
|
|
||||||
|
println!("模型加载完成,耗时: {:.2?}", start.elapsed());
|
||||||
|
|
||||||
|
let prompt = "需要处理生成一个河南的旅游计划,要求如下:
|
||||||
|
1. 旅游计划为三天两夜的行程安排。
|
||||||
|
2. 每天的行程安排包括景点、餐饮和住宿建议。
|
||||||
|
3. 景点安排要考虑交通便利性和时间合理性。
|
||||||
|
请生成详细的旅游计划,包含每天的行程安排、景点介绍、餐饮推荐和住宿建议。
|
||||||
|
";
|
||||||
|
println!("\n用户: {prompt}");
|
||||||
|
print!("Assistant: ");
|
||||||
|
std::io::stdout().flush()?;
|
||||||
|
|
||||||
|
let config = GenerationConfig {
|
||||||
|
max_new_tokens: Some(1200),
|
||||||
|
temperature: 0.7,
|
||||||
|
top_p: 0.95,
|
||||||
|
};
|
||||||
|
|
||||||
|
let start = Instant::now();
|
||||||
|
let stream = model.generate_stream(prompt, false, &config)?;
|
||||||
|
|
||||||
|
for text in stream {
|
||||||
|
print!("{text}");
|
||||||
|
std::io::stdout().flush().ok();
|
||||||
|
}
|
||||||
|
|
||||||
|
let elapsed = start.elapsed();
|
||||||
|
|
||||||
|
println!();
|
||||||
|
println!("\n生成耗时: {:.2?}", elapsed);
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
@@ -1,30 +0,0 @@
|
|||||||
{
|
|
||||||
"_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
Reference in New Issue
Block a user