macrocc/MiniCPM5-1B-rust
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

chore(代码变更): 1. 清理无用代码 2. 更新依赖库 3. 优化项目结构 4. 修复小型bug 5. 增加注释以提高可读性

Browse Source
This commit is contained in:
王宏建
2026-07-08 15:07:21 +08:00
parent d88afb5fe7
commit 3f952c5fe9
29 changed files with 654718 additions and 2587 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/minicpm-inference/Cargo.toml
+1 -2
View File
@@ -6,10 +6,9 @@ publish = ["gitea"]
[dependencies]
minicpm-core = { path = "../minicpm-core", version = "0.1.2", registry = "gitea" }
burn = { version = "0.21", default-features = false, features = ["std"] }
burn = { version = "0.21", default-features = false, features = ["std", "autotune"] }
tokenizers = "0.20"
rand = "0.8"
anyhow = "1.0"
memmap2 = "0.9"
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
crates/minicpm-inference/src/generator.rs
+7 -14
View File
@@ -5,18 +5,8 @@ use minicpm_core::config::EosTokenId;
use minicpm_core::model::{LlamaForCausalLM, LlamaKVCache};
use crate::config::GenerationConfig;
use crate::sampling::{sample, sample_last};
use crate::sampling::{CachedSampler, 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,
@@ -28,6 +18,7 @@ pub struct TokenStream<'a, B: Backend> {
finished: bool,
first_step: bool,
input_ids: Vec<u32>,
sampler: CachedSampler,
}
impl<'a, B: Backend> TokenStream<'a, B> {
@@ -49,6 +40,7 @@ impl<'a, B: Backend> TokenStream<'a, B> {
finished: false,
first_step: true,
input_ids: input_ids.to_vec(),
sampler: CachedSampler::new(),
}
}
}
@@ -122,11 +114,12 @@ impl<'a, B: Backend> TokenStream<'a, B> {
.slice([0..1, 0..1, 0..vocab_size])
.reshape([vocab_size]);
Some(sample(&next_token_logits, self.config.temperature, self.config.top_p))
let data = next_token_logits.to_data();
let values: Vec<f32> = data.as_slice().unwrap().to_vec();
Some(self.sampler.sample(&values, self.config.temperature, self.config.top_p))
}
}
/// 非流式生成 —— 返回所有生成的 token id
pub fn generate_tokens<B: Backend>(
model: &LlamaForCausalLM<B>,
input_ids: &[u32],
@@ -136,4 +129,4 @@ pub fn generate_tokens<B: Backend>(
) -> Vec<u32> {
let stream = TokenStream::new(model, input_ids, config, eos_token_id, device);
stream.collect()
}
}
crates/minicpm-inference/src/lib.rs
+3 -30
View File
@@ -5,14 +5,14 @@ pub mod sampling;
pub mod tokenizer;
pub use config::GenerationConfig;
pub use generator::TokenStream;
pub use generator::{TokenStream};
pub use minicpm_core::config::{EosTokenId, LlamaConfig};
pub use minicpm_core::model::{LlamaForCausalLM, LlamaKVCache};
pub use tokenizer::TokenizerWrapper;
use burn::tensor::backend::Backend;
use loader::{load_model, load_model_q8_from_dir};
use loader::load_model;
pub struct MiniCPM<B: Backend> {
model: LlamaForCausalLM<B>,
@@ -40,24 +40,6 @@ impl<B: Backend> MiniCPM<B> {
})
}
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 {
model,
config,
tokenizer,
device: device.clone(),
})
}
pub fn generate(
&self,
prompt: &str,
@@ -111,15 +93,6 @@ impl<B: Backend> MiniCPM<B> {
}
}
/// 文本流式迭代器 —— 解码 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,
@@ -147,4 +120,4 @@ impl<'a, B: Backend> TextStream<'a, B> {
pub fn all_tokens(&self) -> &[u32] {
&self.buffer
}
}
}
crates/minicpm-inference/src/loader/mod.rs
-12
View File
@@ -1,5 +1,3 @@
pub mod q8;
use burn::module::Module;
use burn::record::{FullPrecisionSettings, NamedMpkFileRecorder};
use burn::tensor::backend::Backend;
@@ -7,8 +5,6 @@ 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,
@@ -19,11 +15,3 @@ pub fn load_model<B: Backend>(
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)
}
crates/minicpm-inference/src/loader/q8.rs
-201
View File
@@ -1,201 +0,0 @@
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)
}
crates/minicpm-inference/src/sampling.rs
+90 -7
View File
@@ -1,6 +1,91 @@
use burn::tensor::backend::Backend;
use burn::tensor::{Float, Tensor};
use rand::Rng;
use rand::{rngs::ThreadRng, Rng};
struct Sampler {
rng: ThreadRng,
}
impl Sampler {
fn new() -> Self {
Self {
rng: rand::thread_rng(),
}
}
fn sample(&mut self, logits: &[f32], temperature: f32, top_p: f32) -> u32 {
if temperature <= 0.001 {
let mut max_idx = 0;
let mut max_val = f32::NEG_INFINITY;
for (i, &val) in logits.iter().enumerate() {
if val > max_val {
max_val = val;
max_idx = i;
}
}
return max_idx as u32;
}
let max_val = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
let exp_sum: f32 = logits.iter().map(|&v| (v - max_val).exp()).sum();
let mut probs: Vec<f32> = logits.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 r: f32 = self.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
}
}
pub struct CachedSampler {
sampler: std::cell::RefCell<Sampler>,
}
impl CachedSampler {
pub fn new() -> Self {
Self {
sampler: std::cell::RefCell::new(Sampler::new()),
}
}
pub fn sample(&self, logits: &[f32], temperature: f32, top_p: f32) -> u32 {
self.sampler.borrow_mut().sample(logits, temperature, top_p)
}
}
pub fn sample_last<B: Backend>(logits: &Tensor<B, 3, Float>, temperature: f32, top_p: f32) -> u32 {
let shape = logits.shape();
@@ -32,11 +117,9 @@ pub fn sample<B: Backend>(logits: &Tensor<B, 1, Float>, temperature: f32, top_p:
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();
let max_val = values.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
let exp_sum: f32 = values.iter().map(|&v| (v - max_val).exp()).sum();
let mut probs: Vec<f32> = values.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();
@@ -77,4 +160,4 @@ pub fn sample<B: Backend>(logits: &Tensor<B, 1, Float>, temperature: f32, top_p:
}
}
probs.len() as u32 - 1
}
}