1use crate::model_config;
4use super::traits::*;
5use anyhow::Result;
6use serde::{Serialize, Deserialize};
7
8model_config!(FlorenceConfig {
9 vocab_size: usize = 51289,
10 hidden_size: usize = 768,
11 intermediate_size: usize = 3072,
12 num_hidden_layers: usize = 12,
13 num_attention_heads: usize = 12,
14 vision_hidden_size: usize = 768,
15 vision_num_hidden_layers: usize = 12,
16 vision_num_attention_heads: usize = 12,
17 vision_patch_size: usize = 16,
18 vision_image_size: usize = 384,
19 layer_norm_eps: f32 = 1e-6,
20 pad_token_id: i64 = 1,
21 bos_token_id: i64 = 0,
22 eos_token_id: i64 = 2,
23});
24
25impl FlorenceConfig {
26 pub fn from_gguf_config(gguf: &crate::weight_loader_core::GGUFModelConfig) -> Self {
27 Self {
28 vocab_size: gguf.vocab_size,
29 hidden_size: gguf.hidden_size,
30 num_hidden_layers: gguf.num_hidden_layers,
31 num_attention_heads: gguf.num_attention_heads,
32 ..Default::default()
33 }
34 }
35}
36
37pub struct FlorenceModelV2 {
38 config: FlorenceConfig,
39 device: Device,
40 vision_tower: FlorenceVisionTower,
41 language_model: FlorenceLanguageModel,
42 projector: Tensor,
43}
44
45pub struct FlorenceVisionTower {
46 patch_embed: Tensor,
47 pos_embed: Tensor,
48 blocks: Vec<FlorenceVisionBlock>,
49 norm: Tensor,
50 hidden_size: usize,
51}
52
53pub struct FlorenceVisionBlock {
54 norm1: Tensor,
55 attn_qkv: Tensor,
56 attn_proj: Tensor,
57 norm2: Tensor,
58 mlp_fc1: Tensor,
59 mlp_fc2: Tensor,
60 num_heads: usize,
61}
62
63pub struct FlorenceLanguageModel {
64 embed_tokens: Tensor,
65 layers: Vec<FlorenceDecoderLayer>,
66 norm: Tensor,
67 lm_head: Tensor,
68}
69
70pub struct FlorenceDecoderLayer {
71 self_attn_q: Tensor,
72 self_attn_k: Tensor,
73 self_attn_v: Tensor,
74 self_attn_o: Tensor,
75 mlp_fc1: Tensor,
76 mlp_fc2: Tensor,
77 norm1: Tensor,
78 norm2: Tensor,
79 num_heads: usize,
80 head_dim: usize,
81}
82
83impl Model for FlorenceModelV2 {
84 type Config = FlorenceConfig;
85
86 fn new(config: FlorenceConfig) -> Result<Self> {
87 let device = Device::CPU;
88 let vision_tower = FlorenceVisionTower::new(&config, &device)?;
89 let language_model = FlorenceLanguageModel::new(&config, &device)?;
90 let projector = ops_fn::zeros(&[config.vision_hidden_size, config.hidden_size], DataType::Float32, &device)?;
91
92 Ok(Self { config, device, vision_tower, language_model, projector })
93 }
94
95 fn from_weights(config: FlorenceConfig, weights: ModelWeights) -> Result<Self> {
96 let mut model = Self::new(config)?;
97 model.language_model.load_weights(&weights)?;
98 Ok(model)
99 }
100
101 fn forward(&self, inputs: &ModelInputs) -> Result<ModelOutputs> {
102 match inputs {
103 ModelInputs::Text { input_ids, .. } => {
104 let hidden = ops_fn::embedding(input_ids, &self.language_model.embed_tokens)?;
105 let logits = self.language_model.forward(&hidden)?;
106 Ok(ModelOutputs::Logits { logits, hidden_states: None })
107 }
108 _ => Err(anyhow::anyhow!("Florence requires text input")),
109 }
110 }
111
112 fn generate(&self, prompt: &str, config: &GenerationConfig) -> Result<String> {
113 use crate::tokenizer::Tokenizer;
114 let tokenizer = Tokenizer::new();
115 let mut tokens: Vec<u32> = tokenizer.encode(prompt);
116
117 for _ in 0..config.max_new_tokens {
118 let input_ids = Tensor::from_i64_slice(&tokens.iter().map(|&t| t as i64).collect::<Vec<_>>(), &[1, tokens.len()], &self.device)?;
119 let outputs = self.forward(&ModelInputs::text(input_ids))?;
120 let logits = match outputs { ModelOutputs::Logits { logits, .. } => logits, _ => return Err(anyhow::anyhow!("Expected logits")) };
121
122 let logits_vec: Vec<f32> = logits.to_candle()?.flatten_all()?.to_vec1()?;
123 let start = (tokens.len() - 1) * self.config.vocab_size;
124 let next_token = logits_vec[start..start + self.config.vocab_size].iter()
125 .enumerate().max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap()).map(|(i, _)| i as u32).unwrap_or(0);
126
127 if next_token == config.eos_token_id { break; }
128 tokens.push(next_token);
129 }
130
131 Ok(tokenizer.decode(&tokens))
132 }
133
134 fn config(&self) -> &Self::Config { &self.config }
135 fn memory_requirements(&self) -> MemoryRequirements {
136 let p = (self.config.vocab_size * self.config.hidden_size) * 4;
137 MemoryRequirements { gpu_memory: p, cpu_memory: p / 4, kv_cache_memory: p / 8, peak_memory: p * 2 }
138 }
139 fn to_device(&mut self, device: &Device) -> Result<()> {
140 self.device = device.clone();
141 self.projector = self.projector.to_device(device)?;
142 Ok(())
143 }
144}
145
146impl FlorenceVisionTower {
147 fn new(config: &FlorenceConfig, device: &Device) -> Result<Self> {
148 let patch_dim = 3 * config.vision_patch_size * config.vision_patch_size;
149 let num_patches = (config.vision_image_size / config.vision_patch_size).pow(2);
150
151 let patch_embed = ops_fn::zeros(&[patch_dim, config.vision_hidden_size], DataType::Float32, device)?;
152 let pos_embed = ops_fn::zeros(&[1, num_patches + 1, config.vision_hidden_size], DataType::Float32, device)?;
153 let norm = ops_fn::zeros(&[config.vision_hidden_size], DataType::Float32, device)?;
154
155 let mut blocks = Vec::with_capacity(config.vision_num_hidden_layers);
156 for _ in 0..config.vision_num_hidden_layers {
157 blocks.push(FlorenceVisionBlock::new(config, device)?);
158 }
159
160 Ok(Self { patch_embed, pos_embed, blocks, norm, hidden_size: config.vision_hidden_size })
161 }
162}
163
164impl FlorenceVisionBlock {
165 fn new(config: &FlorenceConfig, device: &Device) -> Result<Self> {
166 Ok(Self {
167 norm1: ops_fn::zeros(&[config.vision_hidden_size], DataType::Float32, device)?,
168 attn_qkv: ops_fn::zeros(&[config.vision_hidden_size, config.vision_hidden_size * 3], DataType::Float32, device)?,
169 attn_proj: ops_fn::zeros(&[config.vision_hidden_size, config.vision_hidden_size], DataType::Float32, device)?,
170 norm2: ops_fn::zeros(&[config.vision_hidden_size], DataType::Float32, device)?,
171 mlp_fc1: ops_fn::zeros(&[config.vision_hidden_size, config.vision_hidden_size * 4], DataType::Float32, device)?,
172 mlp_fc2: ops_fn::zeros(&[config.vision_hidden_size * 4, config.vision_hidden_size], DataType::Float32, device)?,
173 num_heads: config.vision_num_attention_heads,
174 })
175 }
176}
177
178impl FlorenceLanguageModel {
179 fn new(config: &FlorenceConfig, device: &Device) -> Result<Self> {
180 let embed_tokens = ops_fn::zeros(&[config.vocab_size, config.hidden_size], DataType::Float32, device)?;
181 let norm = ops_fn::zeros(&[config.hidden_size], DataType::Float32, device)?;
182 let lm_head = ops_fn::zeros(&[config.hidden_size, config.vocab_size], DataType::Float32, device)?;
183
184 let mut layers = Vec::with_capacity(config.num_hidden_layers);
185 for _ in 0..config.num_hidden_layers {
186 layers.push(FlorenceDecoderLayer::new(config, device)?);
187 }
188
189 Ok(Self { embed_tokens, layers, norm, lm_head })
190 }
191
192 fn forward(&self, hidden_states: &Tensor) -> Result<Tensor> {
193 let mut hidden = hidden_states.clone();
194 for layer in &self.layers { hidden = layer.forward(&hidden)?; }
195 hidden = ops_fn::layer_norm(&hidden, &self.norm, None, 1e-6)?;
196 ops_fn::matmul(&hidden, &self.lm_head)
197 }
198
199 fn load_weights(&mut self, weights: &ModelWeights) -> Result<()> {
200 if let Some(w) = weights.get("model.embed_tokens.weight") { self.embed_tokens = w.clone(); }
201 if let Some(w) = weights.get("model.norm.weight") { self.norm = w.clone(); }
202 if let Some(w) = weights.get("lm_head.weight") { self.lm_head = ops_fn::transpose(w)?; }
203 Ok(())
204 }
205}
206
207impl FlorenceDecoderLayer {
208 fn new(config: &FlorenceConfig, device: &Device) -> Result<Self> {
209 let head_dim = config.hidden_size / config.num_attention_heads;
210 Ok(Self {
211 self_attn_q: ops_fn::zeros(&[config.hidden_size, config.hidden_size], DataType::Float32, device)?,
212 self_attn_k: ops_fn::zeros(&[config.hidden_size, config.hidden_size], DataType::Float32, device)?,
213 self_attn_v: ops_fn::zeros(&[config.hidden_size, config.hidden_size], DataType::Float32, device)?,
214 self_attn_o: ops_fn::zeros(&[config.hidden_size, config.hidden_size], DataType::Float32, device)?,
215 mlp_fc1: ops_fn::zeros(&[config.hidden_size, config.intermediate_size], DataType::Float32, device)?,
216 mlp_fc2: ops_fn::zeros(&[config.intermediate_size, config.hidden_size], DataType::Float32, device)?,
217 norm1: ops_fn::zeros(&[config.hidden_size], DataType::Float32, device)?,
218 norm2: ops_fn::zeros(&[config.hidden_size], DataType::Float32, device)?,
219 num_heads: config.num_attention_heads,
220 head_dim,
221 })
222 }
223
224 fn forward(&self, hidden_states: &Tensor) -> Result<Tensor> {
225 let shape = hidden_states.shape();
226 let (batch_size, seq_len, _) = (shape[0], shape[1], shape[2]);
227
228 let residual = hidden_states.clone();
229 let hidden = ops_fn::layer_norm(hidden_states, &self.norm1, None, 1e-6)?;
230
231 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)?;
232 let k = ops_fn::matmul(&hidden, &self.self_attn_k)?.to_candle()?.reshape(&[batch_size, seq_len, self.num_heads, self.head_dim])?.transpose(1, 2)?;
233 let v = ops_fn::matmul(&hidden, &self.self_attn_v)?.to_candle()?.reshape(&[batch_size, seq_len, self.num_heads, self.head_dim])?.transpose(1, 2)?;
234
235 let scale = (self.head_dim as f32).powf(-0.5);
236 let scores = (q.contiguous()?.matmul(&k.transpose(2, 3)?.contiguous()?)? * (scale as f64))?;
237 let device = scores.device();
238 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)? };
239 let attn = candle_nn::ops::softmax_last_dim(&scores.broadcast_add(&mask)?)?.matmul(&v.contiguous()?)?
240 .transpose(1, 2)?.reshape(&[batch_size, seq_len, self.num_heads * self.head_dim])?;
241 let hidden = ops_fn::add(&residual, &ops_fn::matmul(&Tensor::from_candle(attn), &self.self_attn_o)?)?;
242
243 let residual = hidden.clone();
244 let hidden = ops_fn::layer_norm(&hidden, &self.norm2, None, 1e-6)?;
245 let hidden = ops_fn::gelu(&ops_fn::matmul(&hidden, &self.mlp_fc1)?)?;
246 ops_fn::add(&residual, &ops_fn::matmul(&hidden, &self.mlp_fc2)?)
247 }
248}
249
250#[cfg(test)]
251mod tests {
252 use super::*;
253
254 #[test]
255 fn test_florence_forward() {
256 let config = FlorenceConfig { vocab_size: 100, hidden_size: 32, intermediate_size: 128, num_hidden_layers: 1, num_attention_heads: 2, ..Default::default() };
257 let model = FlorenceModelV2::new(config).unwrap();
258 let outputs = model.forward(&ModelInputs::text(ops_fn::zeros(&[1, 4], DataType::Int64, &Device::CPU).unwrap())).unwrap();
259 match outputs { ModelOutputs::Logits { logits, .. } => assert_eq!(logits.shape(), &[1, 4, 100]), _ => panic!() }
260 }
261}