1use crate::model_config;
4use super::traits::*;
5use anyhow::Result;
6use serde::{Serialize, Deserialize};
7
8model_config!(CogVLMConfig {
9 vocab_size: usize = 32000,
10 hidden_size: usize = 4096,
11 intermediate_size: usize = 11008,
12 num_hidden_layers: usize = 32,
13 num_attention_heads: usize = 32,
14 num_key_value_heads: usize = 32,
15 vision_hidden_size: usize = 1024,
16 vision_num_hidden_layers: usize = 24,
17 vision_num_attention_heads: usize = 16,
18 vision_patch_size: usize = 14,
19 vision_image_size: usize = 224,
20 rms_norm_eps: f32 = 1e-5,
21 rope_theta: f32 = 10000.0,
22 pad_token_id: i64 = 0,
23 bos_token_id: i64 = 1,
24 eos_token_id: i64 = 2,
25});
26
27impl CogVLMConfig {
28 pub fn from_gguf_config(gguf: &crate::weight_loader_core::GGUFModelConfig) -> Self {
29 Self {
30 vocab_size: gguf.vocab_size,
31 hidden_size: gguf.hidden_size,
32 intermediate_size: gguf.intermediate_size,
33 num_hidden_layers: gguf.num_hidden_layers,
34 num_attention_heads: gguf.num_attention_heads,
35 num_key_value_heads: gguf.num_key_value_heads,
36 rms_norm_eps: gguf.rms_norm_eps,
37 rope_theta: gguf.rope_theta,
38 ..Default::default()
39 }
40 }
41}
42
43pub struct CogVLMModelV2 {
44 config: CogVLMConfig,
45 device: Device,
46 embed_tokens: Tensor,
47 layers: Vec<CogVLMLayer>,
48 norm: Tensor,
49 lm_head: Tensor,
50}
51
52pub struct CogVLMLayer {
53 q_proj: Tensor,
54 k_proj: Tensor,
55 v_proj: Tensor,
56 o_proj: Tensor,
57 gate_proj: Tensor,
58 up_proj: Tensor,
59 down_proj: Tensor,
60 input_layernorm: Tensor,
61 post_attention_layernorm: Tensor,
62 num_heads: usize,
63 head_dim: usize,
64}
65
66impl Model for CogVLMModelV2 {
67 type Config = CogVLMConfig;
68
69 fn new(config: CogVLMConfig) -> Result<Self> {
70 let device = Device::CPU;
71 let embed_tokens = ops_fn::zeros(&[config.vocab_size, config.hidden_size], DataType::Float32, &device)?;
72 let norm = ops_fn::zeros(&[config.hidden_size], DataType::Float32, &device)?;
73 let lm_head = ops_fn::zeros(&[config.hidden_size, config.vocab_size], DataType::Float32, &device)?;
74
75 let mut layers = Vec::with_capacity(config.num_hidden_layers);
76 for _ in 0..config.num_hidden_layers {
77 layers.push(CogVLMLayer::new(&config, &device)?);
78 }
79
80 Ok(Self { config, device, embed_tokens, layers, norm, lm_head })
81 }
82
83 fn from_weights(config: CogVLMConfig, weights: ModelWeights) -> Result<Self> {
84 let mut model = Self::new(config)?;
85 if let Some(w) = weights.get("model.embed_tokens.weight") { model.embed_tokens = w.clone(); }
86 if let Some(w) = weights.get("model.norm.weight") { model.norm = w.clone(); }
87 if let Some(w) = weights.get("lm_head.weight") { model.lm_head = ops_fn::transpose(w)?; }
88 Ok(model)
89 }
90
91 fn forward(&self, inputs: &ModelInputs) -> Result<ModelOutputs> {
92 match inputs {
93 ModelInputs::Text { input_ids, .. } => {
94 let mut hidden = ops_fn::embedding(input_ids, &self.embed_tokens)?;
95 for layer in &self.layers { hidden = layer.forward(&hidden)?; }
96 hidden = ops_fn::rms_norm(&hidden, &self.norm, self.config.rms_norm_eps)?;
97 let logits = ops_fn::matmul(&hidden, &self.lm_head)?;
98 Ok(ModelOutputs::Logits { logits, hidden_states: None })
99 }
100 _ => Err(anyhow::anyhow!("CogVLM requires text input")),
101 }
102 }
103
104 fn generate(&self, prompt: &str, config: &GenerationConfig) -> Result<String> {
105 use crate::tokenizer::Tokenizer;
106 let tokenizer = Tokenizer::new();
107 let mut tokens: Vec<u32> = tokenizer.encode(prompt);
108
109 for _ in 0..config.max_new_tokens {
110 let input_ids = Tensor::from_i64_slice(&tokens.iter().map(|&t| t as i64).collect::<Vec<_>>(), &[1, tokens.len()], &self.device)?;
111 let outputs = self.forward(&ModelInputs::text(input_ids))?;
112 let logits = match outputs { ModelOutputs::Logits { logits, .. } => logits, _ => return Err(anyhow::anyhow!("Expected logits")) };
113
114 let logits_vec: Vec<f32> = logits.to_candle()?.flatten_all()?.to_vec1()?;
115 let start = (tokens.len() - 1) * self.config.vocab_size;
116 let next_token = logits_vec[start..start + self.config.vocab_size].iter()
117 .enumerate().max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap()).map(|(i, _)| i as u32).unwrap_or(0);
118
119 if next_token == config.eos_token_id { break; }
120 tokens.push(next_token);
121 }
122
123 Ok(tokenizer.decode(&tokens))
124 }
125
126 fn config(&self) -> &Self::Config { &self.config }
127 fn memory_requirements(&self) -> MemoryRequirements {
128 let p = (self.config.vocab_size * self.config.hidden_size) * 4;
129 MemoryRequirements { gpu_memory: p, cpu_memory: p / 4, kv_cache_memory: p / 8, peak_memory: p * 2 }
130 }
131 fn to_device(&mut self, device: &Device) -> Result<()> {
132 self.device = device.clone();
133 self.embed_tokens = self.embed_tokens.to_device(device)?;
134 self.norm = self.norm.to_device(device)?;
135 self.lm_head = self.lm_head.to_device(device)?;
136 Ok(())
137 }
138}
139
140impl CogVLMLayer {
141 fn new(config: &CogVLMConfig, device: &Device) -> Result<Self> {
142 let head_dim = config.hidden_size / config.num_attention_heads;
143 Ok(Self {
144 q_proj: ops_fn::zeros(&[config.hidden_size, config.hidden_size], DataType::Float32, device)?,
145 k_proj: ops_fn::zeros(&[config.hidden_size, config.hidden_size], DataType::Float32, device)?,
146 v_proj: ops_fn::zeros(&[config.hidden_size, config.hidden_size], DataType::Float32, device)?,
147 o_proj: ops_fn::zeros(&[config.hidden_size, config.hidden_size], DataType::Float32, device)?,
148 gate_proj: ops_fn::zeros(&[config.hidden_size, config.intermediate_size], DataType::Float32, device)?,
149 up_proj: ops_fn::zeros(&[config.hidden_size, config.intermediate_size], DataType::Float32, device)?,
150 down_proj: ops_fn::zeros(&[config.intermediate_size, config.hidden_size], DataType::Float32, device)?,
151 input_layernorm: ops_fn::zeros(&[config.hidden_size], DataType::Float32, device)?,
152 post_attention_layernorm: ops_fn::zeros(&[config.hidden_size], DataType::Float32, device)?,
153 num_heads: config.num_attention_heads,
154 head_dim,
155 })
156 }
157
158 fn forward(&self, hidden_states: &Tensor) -> Result<Tensor> {
159 let shape = hidden_states.shape();
160 let (batch_size, seq_len, _) = (shape[0], shape[1], shape[2]);
161
162 let residual = hidden_states.clone();
163 let hidden = ops_fn::rms_norm(hidden_states, &self.input_layernorm, 1e-5)?;
164
165 let q = ops_fn::matmul(&hidden, &self.q_proj)?.to_candle()?.reshape(&[batch_size, seq_len, self.num_heads, self.head_dim])?.transpose(1, 2)?;
166 let k = ops_fn::matmul(&hidden, &self.k_proj)?.to_candle()?.reshape(&[batch_size, seq_len, self.num_heads, self.head_dim])?.transpose(1, 2)?;
167 let v = ops_fn::matmul(&hidden, &self.v_proj)?.to_candle()?.reshape(&[batch_size, seq_len, self.num_heads, self.head_dim])?.transpose(1, 2)?;
168
169 let scale = (self.head_dim as f32).powf(-0.5);
170 let scores = (q.contiguous()?.matmul(&k.transpose(2, 3)?.contiguous()?)? * (scale as f64))?;
171 let device = scores.device();
172 let mask = { let mut m = vec![0.0f32; seq_len * seq_len]; for i in 0..seq_len { for j in (i+1)..seq_len { m[i*seq_len+j] = f32::NEG_INFINITY; } } candle_core::Tensor::from_vec(m, &[1,1,seq_len,seq_len], device)? };
173 let attn = candle_nn::ops::softmax_last_dim(&scores.broadcast_add(&mask)?)?.matmul(&v.contiguous()?)?
174 .transpose(1, 2)?.reshape(&[batch_size, seq_len, self.num_heads * self.head_dim])?;
175 let hidden = ops_fn::add(&residual, &ops_fn::matmul(&Tensor::from_candle(attn), &self.o_proj)?)?;
176
177 let residual = hidden.clone();
178 let hidden = ops_fn::rms_norm(&hidden, &self.post_attention_layernorm, 1e-5)?;
179 let gate = ops_fn::silu(&ops_fn::matmul(&hidden, &self.gate_proj)?)?;
180 let up = ops_fn::matmul(&hidden, &self.up_proj)?;
181 ops_fn::add(&residual, &ops_fn::matmul(&ops_fn::mul(&gate, &up)?, &self.down_proj)?)
182 }
183}
184
185#[cfg(test)]
186mod tests {
187 use super::*;
188
189 #[test]
190 fn test_cogvlm_forward() {
191 let config = CogVLMConfig { vocab_size: 100, hidden_size: 32, intermediate_size: 128, num_hidden_layers: 1, num_attention_heads: 2, num_key_value_heads: 2, ..Default::default() };
192 let model = CogVLMModelV2::new(config).unwrap();
193 let outputs = model.forward(&ModelInputs::text(ops_fn::zeros(&[1, 4], DataType::Int64, &Device::CPU).unwrap())).unwrap();
194 match outputs { ModelOutputs::Logits { logits, .. } => assert_eq!(logits.shape(), &[1, 4, 100]), _ => panic!() }
195 }
196}