refactor(项目结构) 重构为 workspace 多成员项目
- 将项目拆分为三个 crate：minicpm-core（核心模型）、minicpm-convert（转换功能）、minicpm-inference（推理功能） - 添加两个示例：minimal-inference（最小推理）和 convert（模型转换） - 转换后自动拷贝 config.json 和 tokenizer.json 到 model 目录 - 更新 README 说明 workspace 结构和使用方式
This commit is contained in:
@@ -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"
|
||||
@@ -0,0 +1,283 @@
|
||||
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::{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
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
|
||||
|
||||
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
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user