chore(代码变更): 1. 清理无用代码 2. 更新依赖库 3. 优化项目结构 4. 修复小型bug 5. 增加注释以提高可读性
This commit is contained in:
@@ -9,4 +9,7 @@ minicpm-core = { path = "../minicpm-core", version = "0.1.2", registry = "gitea"
|
||||
burn = { version = "0.21", default-features = false, features = ["std"] }
|
||||
tokenizers = "0.20"
|
||||
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::model::{LlamaForCausalLM, LlamaKVCache};
|
||||
pub use tokenizer::TokenizerWrapper;
|
||||
|
||||
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
|
||||
}
|
||||
|
||||
/// 流式生成：每生成一个新 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> {
|
||||
model: LlamaForCausalLM<B>,
|
||||
@@ -317,10 +30,25 @@ impl<B: Backend> MiniCPM<B> {
|
||||
) -> anyhow::Result<Self> {
|
||||
let config = LlamaConfig::from_json(config_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);
|
||||
let recorder = NamedMpkFileRecorder::<FullPrecisionSettings>::new();
|
||||
let model = model.load_file(model_path, &recorder, device)?;
|
||||
Ok(Self {
|
||||
model,
|
||||
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 {
|
||||
model,
|
||||
@@ -337,7 +65,7 @@ impl<B: Backend> MiniCPM<B> {
|
||||
config: &GenerationConfig,
|
||||
) -> anyhow::Result<String> {
|
||||
let input_ids = self.tokenizer.apply_chat_template(prompt, think)?;
|
||||
let output_ids = generate_with_cache(
|
||||
let output_ids = generator::generate_tokens(
|
||||
&self.model,
|
||||
&input_ids,
|
||||
config,
|
||||
@@ -348,27 +76,26 @@ impl<B: Backend> MiniCPM<B> {
|
||||
self.tokenizer.decode(new_ids, true)
|
||||
}
|
||||
|
||||
/// 流式生成：每生成一个 token 立即调用 `on_token` 回调，
|
||||
/// 参数为该 token 解码后的文本。
|
||||
pub fn generate_stream(
|
||||
&self,
|
||||
pub fn generate_stream<'a>(
|
||||
&'a self,
|
||||
prompt: &str,
|
||||
think: bool,
|
||||
config: &GenerationConfig,
|
||||
on_token: impl FnMut(&str),
|
||||
) -> anyhow::Result<String> {
|
||||
config: &'a GenerationConfig,
|
||||
) -> anyhow::Result<TextStream<'a, B>> {
|
||||
let input_ids = self.tokenizer.apply_chat_template(prompt, think)?;
|
||||
let output_ids = generate_stream(
|
||||
let token_stream = TokenStream::new(
|
||||
&self.model,
|
||||
&self.tokenizer,
|
||||
&input_ids,
|
||||
config,
|
||||
&self.config.eos_token_id,
|
||||
&self.device,
|
||||
on_token,
|
||||
);
|
||||
let new_ids = &output_ids[input_ids.len()..];
|
||||
self.tokenizer.decode(new_ids, true)
|
||||
|
||||
Ok(TextStream {
|
||||
token_stream,
|
||||
tokenizer: &self.tokenizer,
|
||||
buffer: Vec::new(),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn config(&self) -> &LlamaConfig {
|
||||
@@ -383,3 +110,41 @@ impl<B: Backend> MiniCPM<B> {
|
||||
&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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user