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

crates/minicpm-inference/src/lib.rs

@@ -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
+    }
+}