add minicpm with a bug

2025-09-25 12:09:25 +08:00
parent ec989f397f
commit 4fdab3e7cd
15 changed files with 900 additions and 49 deletions
@@ -0,0 +1,32 @@
+use candle_nn::Activation;
+
+#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
+pub struct RopeScalingConfig {
+    pub rope_type: String,
+    pub long_factor: Vec<f32>,
+    pub short_factor: Vec<f32>,
+    pub original_max_position_embeddings: usize,
+}
+
+#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
+pub struct MiniCPM4Config {
+    pub bos_token_id: u32,
+    pub eos_token_id: Vec<u32>,
+    pub hidden_act: Activation,
+    pub hidden_size: usize,
+    pub intermediate_size: usize,
+    pub max_position_embeddings: usize,
+    pub num_attention_heads: usize,
+    pub num_hidden_layers: usize,
+    pub num_key_value_heads: usize,
+    pub rms_norm_eps: f64,
+    pub rope_scaling: RopeScalingConfig,
+    pub torch_dtype: String,
+    pub vocab_size: usize,
+    // pub use_mup: bool,
+    pub scale_emb:f32,
+    pub dim_model_base: usize,
+    pub scale_depth: f32,
+    // pub rope_theta: f32,
+    // pub kv_channels: i32,
+}
@@ -0,0 +1,136 @@
+use crate::models::minicpm4::config::MiniCPM4Config;
+use crate::models::minicpm4::model::MiniCPMModel;
+// use crate::models::GenerateStream;
+use crate::utils::utils::{
+    build_completion_chunk_response, build_completion_response, find_safetensors_files, get_device,
+    get_dtype, get_logit_processor,
+};
+use crate::{
+    chat_template::chat_template::ChatTemplate, models::GenerateModel,
+    tokenizer::tokenizer::TokenizerModel,
+};
+use anyhow::{Result, anyhow};
+use candle_core::{D, DType, Device, IndexOp, Tensor};
+use candle_nn::VarBuilder;
+use openai_dive::v1::resources::chat::{
+    ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse,
+};
+use rocket::async_stream::stream;
+use rocket::futures::Stream;
+
+pub struct MiniCPMGenerateModel<'a> {
+    chat_template: ChatTemplate<'a>,
+    tokenizer: TokenizerModel,
+    minicpm: MiniCPMModel,
+    device: Device,
+    endoftext_id: u32,
+    im_end_id: u32,
+}
+
+impl<'a> GenerateModel for MiniCPMGenerateModel<'a> {
+    fn init(path: &str, device: Option<&Device>, dtype: Option<DType>) -> Result<Self> {
+        let chat_template = ChatTemplate::init(path)?;
+        let tokenizer = TokenizerModel::init(path)?;
+        let config_path = path.to_string() + "/config.json";
+        let cfg: MiniCPM4Config = serde_json::from_slice(&std::fs::read(config_path)?)?;
+        let device = &get_device(device);
+        let cfg_dtype = cfg.torch_dtype.as_str();
+        let dtype = get_dtype(dtype, cfg_dtype);
+        let endoftext_id = cfg.eos_token_id[0];
+        let im_end_id = cfg.eos_token_id[1];
+        let model_list = find_safetensors_files(&path)?;
+        let vb = unsafe { VarBuilder::from_mmaped_safetensors(&model_list, dtype, device)? };
+        let minicpm = MiniCPMModel::new(vb, cfg)?;
+
+        Ok(MiniCPMGenerateModel {
+            chat_template,
+            tokenizer,
+            minicpm,
+            device: device.clone(),
+            endoftext_id,
+            im_end_id,
+        })
+    }
+
+    fn generate(&mut self, mes: ChatCompletionParameters) -> Result<ChatCompletionResponse> {
+        let mut logit_processor = get_logit_processor(mes.temperature, mes.top_p);
+        let mes_render = self.chat_template.apply_chat_template(&mes)?;
+        let mut input_ids = self.tokenizer.text_encode(mes_render, &self.device)?;
+        let mut seq_len = input_ids.dim(1)?;
+        let mut seqlen_offset = 0;
+        let mut generate = Vec::new();
+        let sample_len = match mes.max_tokens {
+            Some(max) => max,
+            None => 512,
+        };
+        for _ in 0..sample_len {
+            let logits = self.minicpm.forward_step(&input_ids, seqlen_offset)?;
+            let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?;
+            let next_token = logit_processor.sample(&logits)?;
+            generate.push(next_token);
+            if next_token == self.endoftext_id || next_token == self.im_end_id {
+                break;
+            }
+            seqlen_offset += seq_len;
+            seq_len = 1;
+            input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?;
+        }
+        let res = self.tokenizer.token_decode(generate)?;
+        self.minicpm.clear_kv_cache();
+        let response = build_completion_response(res, "minicpm");
+        Ok(response)
+    }
+    fn generate_stream(
+        &mut self,
+        mes: ChatCompletionParameters,
+    ) -> Result<impl Stream<Item = Result<ChatCompletionChunkResponse, anyhow::Error>>> {
+        let mut logit_processor = get_logit_processor(mes.temperature, mes.top_p);
+        let mes_render = self.chat_template.apply_chat_template(&mes)?;
+        let mut input_ids = self.tokenizer.text_encode(mes_render, &self.device)?;
+        let mut seq_len = input_ids.dim(1)?;
+        let mut seqlen_offset = 0;
+        let sample_len = match mes.max_tokens {
+            Some(max) => max,
+            None => 512,
+        };
+        let stream = stream! {
+            let mut error_tokens = Vec::new();
+            for _ in 0..sample_len {
+                let logits = self.minicpm.forward_step(
+                    &input_ids,
+                    seqlen_offset,
+                )?;
+                let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?;
+                let next_token = logit_processor.sample(&logits)?;
+                let mut decode_ids = Vec::new();
+                if error_tokens.len() > 0 {
+                    decode_ids.extend_from_slice(&error_tokens);
+                }
+                decode_ids.push(next_token);
+                let decoded_token = self.tokenizer.token_decode(decode_ids).map_err(|e| anyhow!(format!("stream decode error{}", e)))?;
+                if decoded_token.contains("�") {
+                    error_tokens.push(next_token);
+                    if error_tokens.len() > 3 {
+                        error_tokens.clear();
+                    }
+                    seqlen_offset += seq_len;
+                    seq_len = 1;
+                    input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?;
+                    continue;
+                }
+                error_tokens.clear();
+                let chunk = build_completion_chunk_response(decoded_token, "minicpm", None, None);
+                yield Ok(chunk);
+                if next_token == self.endoftext_id || next_token == self.im_end_id {
+                    break;
+                }
+                seqlen_offset += seq_len;
+                seq_len = 1;
+                input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?;
+
+            }
+            self.minicpm.clear_kv_cache();
+        };
+        Ok(stream)
+    }
+}
@@ -0,0 +1,3 @@
+pub mod config;
+pub mod model;
+pub mod generate;
@@ -0,0 +1,257 @@
+use crate::{
+    models::{
+        base_modules::{AttentionNobias, MLPNoBias},
+        minicpm4::config::MiniCPM4Config,
+    },
+    position_embed::rope::compute_default_rope_parameters,
+    utils::tensor_utils::prepare_causal_attention_mask,
+};
+use anyhow::{Ok, Result};
+use candle_core::{D, DType, Device, Tensor, Var};
+use candle_nn::{embedding, rms_norm, Embedding, Linear, Module, RmsNorm, VarBuilder};
+
+pub struct MiniCPMLongRoPE {
+    head_dim: usize,
+    rope_theta: f32,
+    max_position_embeddings: usize,
+    short_factor: Vec<f32>,
+    long_factor: Vec<f32>,
+    original_max_position_embeddings: usize,
+    inv_freq: Tensor,
+    cos_cached: Tensor,
+    sin_cached: Tensor,
+}
+impl MiniCPMLongRoPE {
+    pub fn new(cfg: &MiniCPM4Config, device: &Device) -> Result<Self> {
+        let head_dim = cfg.hidden_size / cfg.num_attention_heads;
+        let rope_theta = 10000.0;
+        let max_position_embeddings = cfg.max_position_embeddings;
+        let short_factor = cfg.rope_scaling.short_factor.clone();
+        let long_factor = cfg.rope_scaling.short_factor.clone();
+        let original_max_position_embeddings = cfg.rope_scaling.original_max_position_embeddings;
+        let scale = max_position_embeddings / original_max_position_embeddings;
+        let scaling_factor =
+            (1.0 + (scale as f64).ln() + (original_max_position_embeddings as f64).ln()).sqrt();
+        let inv_freq = compute_default_rope_parameters(head_dim, rope_theta);
+        let inv_freq = Tensor::from_slice(&inv_freq, (1, inv_freq.len()), device)?;
+        let t = Tensor::arange(0.0_f32, max_position_embeddings as f32, device)?
+            .reshape((max_position_embeddings, 1))?;
+        // short_factor.len() = 32
+        // head_dim = 1024 / 16 = 64, inv_freq.len() = 32
+        let ext_factors = Tensor::from_slice(&short_factor, (1, short_factor.len()), device)?;
+        let ext_factors = Tensor::ones_like(&ext_factors)?.div(&ext_factors)?;
+        // (seq_len, 1) matmul (1, 32) -> (seq_len, 32) * (1, 32)-> (seq_len, 32)
+        let freqs = t.matmul(&ext_factors)?.broadcast_mul(&inv_freq)?;
+
+        let emb = Tensor::cat(&[&freqs, &freqs], D::Minus1)?;
+        let cos_cached = emb.cos()?.affine(scaling_factor, 0.0)?;
+        let sin_cached = emb.sin()?.affine(scaling_factor, 0.0)?;
+        Ok(Self {
+            head_dim,
+            rope_theta,
+            max_position_embeddings,
+            short_factor,
+            long_factor,
+            original_max_position_embeddings,
+            inv_freq,
+            cos_cached,
+            sin_cached,
+        })
+    }
+    pub fn update_cos_sin_cache(&mut self, seqlen: usize, device: &Device) -> Result<()> {
+        let t = Tensor::arange(0.0_f32, seqlen as f32, device)?.reshape((seqlen, 1))?;
+        let mut ext_factors =
+            Tensor::from_slice(&self.short_factor, (1, self.short_factor.len()), device)?;
+        if seqlen > self.original_max_position_embeddings {
+            ext_factors =
+                Tensor::from_slice(&self.long_factor, (1, self.long_factor.len()), device)?;
+        }
+        let ext_factors = Tensor::ones_like(&ext_factors)?.div(&ext_factors)?;
+        let freqs = t.matmul(&ext_factors)?.broadcast_mul(&self.inv_freq)?;
+        let emb = Tensor::cat(&[&freqs, &freqs], D::Minus1)?;
+        let scale = seqlen / self.original_max_position_embeddings;
+        let scaling_factor =
+            (1.0 + (scale as f64).ln() + (self.original_max_position_embeddings as f64).ln())
+                .sqrt();
+        let cos_cached = emb.cos()?.affine(scaling_factor, 0.0)?;
+        let sin_cached = emb.sin()?.affine(scaling_factor, 0.0)?;
+        self.cos_cached = cos_cached;
+        self.sin_cached = sin_cached;
+        Ok(())
+    }
+    pub fn forward(&self, pos_offset: usize, seqlen: usize) -> Result<(Tensor, Tensor)> {
+        let cos = self.cos_cached.narrow(0, pos_offset, seqlen)?;
+        let sin = self.sin_cached.narrow(0, pos_offset, seqlen)?;
+        Ok((cos, sin))
+    }
+}
+
+pub struct MiniCPMDecoderLayer {
+    self_attn: AttentionNobias,
+    mlp: MLPNoBias,
+    input_layernorm: RmsNorm,
+    post_attention_layernorm: RmsNorm,
+    scale_depth: f32,
+    num_hidden_layers: usize,
+}
+
+impl MiniCPMDecoderLayer {
+    pub fn new(vb: VarBuilder, cfg: &MiniCPM4Config) -> Result<Self> {
+        let self_attn = AttentionNobias::new(
+            vb.pp("self_attn"),
+            cfg.hidden_size,
+            cfg.num_attention_heads,
+            cfg.num_key_value_heads,
+        )?;
+        let mlp = MLPNoBias::new(
+            vb.pp("mlp"),
+            cfg.hidden_size,
+            cfg.intermediate_size,
+            cfg.hidden_act,
+        )?;
+        let input_layernorm =
+            rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?;
+        let post_attention_layernorm = rms_norm(
+            cfg.hidden_size,
+            cfg.rms_norm_eps,
+            vb.pp("post_attention_layernorm"),
+        )?;
+        Ok(Self {
+            self_attn,
+            mlp,
+            input_layernorm,
+            post_attention_layernorm,
+            scale_depth: cfg.scale_depth,
+            num_hidden_layers: cfg.num_hidden_layers,
+        })
+    }
+
+    pub fn forward(
+        &self,
+        xs: &Tensor,
+        cos: &Tensor,
+        sin: &Tensor,
+        attention_mask: Option<&Tensor>,
+    ) -> Result<Tensor> {
+        let residual = xs;
+        let xs = self.input_layernorm.forward(xs)?;
+        let xs = self.self_attn.forward(&xs, cos, sin, attention_mask)?;
+        let xs = (xs + residual)?;
+        let residual = &xs;
+        let xs = xs.apply(&self.post_attention_layernorm)?.apply(&self.mlp)?;
+        let xs = (residual + xs)?;
+        Ok(xs)
+    }
+
+    pub fn forward_step(
+        &mut self,
+        xs: &Tensor,
+        cos: &Tensor,
+        sin: &Tensor,
+        attention_mask: Option<&Tensor>,
+    ) -> Result<Tensor> {
+        let residual = xs;
+        let xs = self.input_layernorm.forward(xs)?;
+        let xs = self.self_attn.forward_step(&xs, cos, sin, attention_mask)?;
+        let xs = (xs + residual)?;
+        let residual = &xs;
+        let xs = xs.apply(&self.post_attention_layernorm)?.apply(&self.mlp)?;
+        let xs = (residual + xs)?;
+        Ok(xs)
+    }
+    pub fn clear_kv_cache(&mut self) {
+        self.self_attn.clear_kv_cache();
+    }
+}
+
+pub struct MiniCPMModel {
+    cfg: MiniCPM4Config,
+    embed_tokens: Embedding,
+    layers: Vec<MiniCPMDecoderLayer>,
+    norm: RmsNorm,
+    rope_emb: MiniCPMLongRoPE,
+    lm_head: Linear,
+}
+
+impl MiniCPMModel {
+    pub fn new(vb: VarBuilder, cfg: MiniCPM4Config) -> Result<Self> {
+        let embed_tokens = embedding(cfg.vocab_size, cfg.hidden_size, vb.pp("embed_tokens"))?;
+        let mut layers = Vec::with_capacity(cfg.num_hidden_layers);
+        let vb_layers = vb.pp("layers");
+        for i in 0..cfg.num_hidden_layers {
+            let layer = MiniCPMDecoderLayer::new(vb_layers.pp(i), &cfg)?;
+            layers.push(layer);
+        }
+        let norm = rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("norm"))?;
+        let rope_emb = MiniCPMLongRoPE::new(&cfg, vb.device())?;
+        let lm_head = Linear::new(embed_tokens.embeddings().clone(), None);
+        Ok(Self {
+            cfg,
+            embed_tokens,
+            layers,
+            norm,
+            rope_emb,
+            lm_head
+        })
+    }
+
+    pub fn forward(&self, input_ids: &Tensor, position_id: usize) -> Result<Tensor> {
+        let (bs, seq_len) = input_ids.dims2()?;
+        let input_embeds = self.embed_tokens.forward(&input_ids)?;
+        let attention_mask: Option<&Tensor> = {
+            if seq_len <= 1 {
+                None
+            } else {
+                Some(&prepare_causal_attention_mask(
+                    bs,
+                    seq_len,
+                    position_id,
+                    input_ids.device(),
+                )?)
+            }
+        };
+
+        let (cos, sin) = self.rope_emb.forward(position_id, seq_len)?;
+        let mut hidden_states = input_embeds;
+        for decode_layer in &self.layers {
+            hidden_states = decode_layer.forward(&hidden_states, &cos, &sin, attention_mask)?;
+        }
+        hidden_states = self.norm.forward(&hidden_states)?;
+        let hidden_state = hidden_states.narrow(1, seq_len - 1, 1)?;
+        let logits = self.lm_head.forward(&hidden_state)?;
+        Ok(logits)
+    }
+
+    pub fn forward_step(&mut self, input_ids: &Tensor, position_id: usize) -> Result<Tensor> {
+        let (bs, seq_len) = input_ids.dims2()?;
+        let input_embeds = self.embed_tokens.forward(&input_ids)?;
+        let attention_mask: Option<&Tensor> = {
+            if seq_len <= 1 {
+                None
+            } else {
+                Some(&prepare_causal_attention_mask(
+                    bs,
+                    seq_len,
+                    position_id,
+                    input_ids.device(),
+                )?)
+            }
+        };
+
+        let (cos, sin) = self.rope_emb.forward(position_id, seq_len)?;
+        let mut hidden_states = input_embeds;
+        for decode_layer in &mut self.layers {
+            hidden_states = decode_layer.forward_step(&hidden_states, &cos, &sin, attention_mask)?;
+        }
+        hidden_states = self.norm.forward(&hidden_states)?;
+        let hidden_state = hidden_states.narrow(1, seq_len - 1, 1)?;
+        let logits = self.lm_head.forward(&hidden_state)?;
+        Ok(logits)
+    }
+
+    pub fn clear_kv_cache(&mut self) {
+        for layer in self.layers.iter_mut() {
+            layer.clear_kv_cache()
+        }
+    }
+}