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