1use crate::model_config;
9use super::traits::*;
10use anyhow::Result;
11use serde::{Serialize, Deserialize};
12
13model_config!(OPTConfig {
14 vocab_size: usize = 50272,
15 hidden_size: usize = 768,
16 intermediate_size: usize = 3072,
17 num_hidden_layers: usize = 12,
18 num_attention_heads: usize = 12,
19 num_key_value_heads: usize = 12,
20 hidden_act: String = "relu".to_string(),
21 max_position_embeddings: usize = 2048,
22 initializer_range: f32 = 0.02,
23 layer_norm_eps: f32 = 1e-5,
24 use_cache: bool = true,
25 pad_token_id: i64 = 1,
26 bos_token_id: i64 = 2,
27 eos_token_id: i64 = 2,
28 tie_word_embeddings: bool = false,
29 word_embed_proj_dim: usize = 768,
30 do_layer_norm_before: bool = true,
31});
32
33impl OPTConfig {
34 pub fn from_gguf_config(gguf: &crate::weight_loader_core::GGUFModelConfig) -> Self {
35 Self {
36 vocab_size: gguf.vocab_size,
37 hidden_size: gguf.hidden_size,
38 intermediate_size: gguf.intermediate_size,
39 num_hidden_layers: gguf.num_hidden_layers,
40 num_attention_heads: gguf.num_attention_heads,
41 num_key_value_heads: gguf.num_key_value_heads,
42 max_position_embeddings: gguf.max_position_embeddings,
43 ..Default::default()
44 }
45 }
46}
47
48pub struct OPTModelV2 {
49 config: OPTConfig,
50 device: Device,
51 embed_tokens: Tensor,
52 embed_positions: Tensor,
53 project_in: Option<Tensor>,
54 project_out: Option<Tensor>,
55 layers: Vec<OPTLayer>,
56 final_layer_norm: Tensor,
57 lm_head: Tensor,
58}
59
60pub struct OPTLayer {
61 self_attn: OPTAttention,
62 fc1: Tensor,
63 fc2: Tensor,
64 self_attn_layer_norm: Tensor,
65 final_layer_norm: Tensor,
66 do_layer_norm_before: bool,
67}
68
69pub struct OPTAttention {
70 q_proj: Tensor,
71 k_proj: Tensor,
72 v_proj: Tensor,
73 out_proj: Tensor,
74 num_heads: usize,
75 head_dim: usize,
76 scale: f32,
77}
78
79impl Model for OPTModelV2 {
80 type Config = OPTConfig;
81
82 fn new(config: OPTConfig) -> Result<Self> {
83 let device = Device::CPU;
84 let embed_tokens = ops_fn::zeros(&[config.vocab_size, config.word_embed_proj_dim], DataType::Float32, &device)?;
85 let embed_positions = ops_fn::zeros(&[config.max_position_embeddings + 2, config.hidden_size], DataType::Float32, &device)?;
86
87 let project_in = if config.word_embed_proj_dim != config.hidden_size {
88 Some(ops_fn::zeros(&[config.word_embed_proj_dim, config.hidden_size], DataType::Float32, &device)?)
89 } else { None };
90
91 let project_out = if config.word_embed_proj_dim != config.hidden_size {
92 Some(ops_fn::zeros(&[config.hidden_size, config.word_embed_proj_dim], DataType::Float32, &device)?)
93 } else { None };
94
95 let final_layer_norm = ops_fn::zeros(&[config.hidden_size], DataType::Float32, &device)?;
96 let lm_head = ops_fn::zeros(&[config.word_embed_proj_dim, config.vocab_size], DataType::Float32, &device)?;
97
98 let mut layers = Vec::with_capacity(config.num_hidden_layers);
99 for _ in 0..config.num_hidden_layers {
100 layers.push(OPTLayer::new(&config, &device)?);
101 }
102
103 Ok(Self { config, device, embed_tokens, embed_positions, project_in, project_out, layers, final_layer_norm, lm_head })
104 }
105
106 fn from_weights(config: OPTConfig, weights: ModelWeights) -> Result<Self> {
107 let mut model = Self::new(config)?;
108 if let Some(w) = weights.get("model.decoder.embed_tokens.weight") { model.embed_tokens = w.clone(); }
109 if let Some(w) = weights.get("model.decoder.embed_positions.weight") { model.embed_positions = w.clone(); }
110 if let Some(w) = weights.get("model.decoder.project_in.weight") { model.project_in = Some(ops_fn::transpose(w)?); }
111 if let Some(w) = weights.get("model.decoder.project_out.weight") { model.project_out = Some(ops_fn::transpose(w)?); }
112 if let Some(w) = weights.get("model.decoder.final_layer_norm.weight") { model.final_layer_norm = w.clone(); }
113 if let Some(w) = weights.get("lm_head.weight") { model.lm_head = ops_fn::transpose(w)?; }
114 for (i, layer) in model.layers.iter_mut().enumerate() { layer.load_weights(&weights, i)?; }
115 Ok(model)
116 }
117
118 fn forward(&self, inputs: &ModelInputs) -> Result<ModelOutputs> {
119 match inputs {
120 ModelInputs::Text { input_ids, .. } => {
121 let shape = input_ids.shape();
122 let seq_len = shape[1];
123
124 let mut hidden = ops_fn::embedding(input_ids, &self.embed_tokens)?;
125 if let Some(ref proj) = self.project_in {
126 hidden = ops_fn::matmul(&hidden, proj)?;
127 }
128
129 let positions: Vec<i64> = (2..(seq_len + 2) as i64).collect();
130 let pos_tensor = Tensor::from_i64_slice(&positions, &[1, seq_len], &self.device)?;
131 let pos_embeds = ops_fn::embedding(&pos_tensor, &self.embed_positions)?;
132 hidden = ops_fn::add(&hidden, &pos_embeds)?;
133
134 for layer in &self.layers {
135 hidden = layer.forward(&hidden)?;
136 }
137
138 if self.config.do_layer_norm_before {
139 hidden = ops_fn::layer_norm(&hidden, &self.final_layer_norm, None, self.config.layer_norm_eps)?;
140 }
141
142 if let Some(ref proj) = self.project_out {
143 hidden = ops_fn::matmul(&hidden, proj)?;
144 }
145
146 let logits = ops_fn::matmul(&hidden, &self.lm_head)?;
147 Ok(ModelOutputs::Logits { logits, hidden_states: None })
148 }
149 _ => Err(anyhow::anyhow!("OPT only supports text inputs")),
150 }
151 }
152
153 fn generate(&self, prompt: &str, config: &GenerationConfig) -> Result<String> {
154 use crate::tokenizer::Tokenizer;
155 use rand::Rng;
156 let tokenizer = Tokenizer::new();
157 let mut tokens: Vec<u32> = tokenizer.encode(prompt);
158 for _ in 0..config.max_new_tokens {
159 let tokens_i64: Vec<i64> = tokens.iter().map(|&t| t as i64).collect();
160 let input = Tensor::from_i64_slice(&tokens_i64, &[1, tokens.len()], &self.device)?;
161 let outputs = self.forward(&ModelInputs::text(input))?;
162 let logits = match outputs { ModelOutputs::Logits { logits, .. } => logits, _ => return Err(anyhow::anyhow!("Expected logits")) };
163 let logits_candle = logits.to_candle()?;
164 let last = logits_candle.narrow(1, logits_candle.dims()[1] - 1, 1)?.squeeze(1)?.squeeze(0)?;
165 let logits_vec: Vec<f32> = last.to_vec1()?;
166 let next = if config.do_sample && config.temperature > 0.0 {
167 let scaled: Vec<f32> = logits_vec.iter().map(|&x| x / config.temperature).collect();
168 let max_v = scaled.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
169 let exp_sum: f32 = scaled.iter().map(|&x| (x - max_v).exp()).sum();
170 let probs: Vec<f32> = scaled.iter().map(|&x| (x - max_v).exp() / exp_sum).collect();
171 let mut rng = rand::thread_rng();
172 let r: f32 = rng.gen();
173 let mut cum = 0.0;
174 let mut s = 0u32;
175 for (i, &p) in probs.iter().enumerate() { cum += p; if r <= cum { s = i as u32; break; } }
176 s
177 } else {
178 logits_vec.iter().enumerate().max_by(|a, b| a.1.partial_cmp(b.1).unwrap()).map(|(i, _)| i as u32).unwrap_or(0)
179 };
180 if next == config.eos_token_id { break; }
181 tokens.push(next);
182 }
183 Ok(tokenizer.decode(&tokens))
184 }
185
186 fn config(&self) -> &Self::Config { &self.config }
187 fn memory_requirements(&self) -> MemoryRequirements {
188 let p = self.config.vocab_size * self.config.hidden_size + self.config.num_hidden_layers * 8 * self.config.hidden_size.pow(2);
189 MemoryRequirements { gpu_memory: p * 4, cpu_memory: p, kv_cache_memory: 2 * self.config.num_hidden_layers * self.config.max_position_embeddings * self.config.hidden_size * 4, peak_memory: p * 5 }
190 }
191 fn to_device(&mut self, device: &Device) -> Result<()> {
192 self.embed_tokens = self.embed_tokens.to_device(device)?;
193 self.embed_positions = self.embed_positions.to_device(device)?;
194 if let Some(ref mut p) = self.project_in { *p = p.to_device(device)?; }
195 if let Some(ref mut p) = self.project_out { *p = p.to_device(device)?; }
196 self.final_layer_norm = self.final_layer_norm.to_device(device)?;
197 self.lm_head = self.lm_head.to_device(device)?;
198 for l in &mut self.layers { l.to_device(device)?; }
199 self.device = device.clone();
200 Ok(())
201 }
202}
203
204impl OPTLayer {
205 fn new(config: &OPTConfig, device: &Device) -> Result<Self> {
206 Ok(Self {
207 self_attn: OPTAttention::new(config, device)?,
208 fc1: ops_fn::zeros(&[config.hidden_size, config.intermediate_size], DataType::Float32, device)?,
209 fc2: ops_fn::zeros(&[config.intermediate_size, config.hidden_size], DataType::Float32, device)?,
210 self_attn_layer_norm: ops_fn::zeros(&[config.hidden_size], DataType::Float32, device)?,
211 final_layer_norm: ops_fn::zeros(&[config.hidden_size], DataType::Float32, device)?,
212 do_layer_norm_before: config.do_layer_norm_before,
213 })
214 }
215
216 fn forward(&self, hidden_states: &Tensor) -> Result<Tensor> {
217 let residual = hidden_states.clone();
218 let h = if self.do_layer_norm_before {
219 ops_fn::layer_norm(hidden_states, &self.self_attn_layer_norm, None, 1e-5)?
220 } else { hidden_states.clone() };
221 let attn_out = self.self_attn.forward(&h)?;
222 let h = ops_fn::add(&residual, &attn_out)?;
223 let h = if !self.do_layer_norm_before {
224 ops_fn::layer_norm(&h, &self.self_attn_layer_norm, None, 1e-5)?
225 } else { h };
226
227 let residual = h.clone();
228 let h = if self.do_layer_norm_before {
229 ops_fn::layer_norm(&h, &self.final_layer_norm, None, 1e-5)?
230 } else { h };
231 let fc1_out = ops_fn::matmul(&h, &self.fc1)?;
232 let activated = {
233 let x = fc1_out.to_candle()?;
234 Tensor::from_candle(x.relu()?)
235 };
236 let fc2_out = ops_fn::matmul(&activated, &self.fc2)?;
237 let h = ops_fn::add(&residual, &fc2_out)?;
238 if !self.do_layer_norm_before {
239 ops_fn::layer_norm(&h, &self.final_layer_norm, None, 1e-5)
240 } else { Ok(h) }
241 }
242
243 fn load_weights(&mut self, weights: &ModelWeights, idx: usize) -> Result<()> {
244 let p = format!("model.decoder.layers.{}", idx);
245 if let Some(w) = weights.get(&format!("{}.self_attn.q_proj.weight", p)) { self.self_attn.q_proj = ops_fn::transpose(w)?; }
246 if let Some(w) = weights.get(&format!("{}.self_attn.k_proj.weight", p)) { self.self_attn.k_proj = ops_fn::transpose(w)?; }
247 if let Some(w) = weights.get(&format!("{}.self_attn.v_proj.weight", p)) { self.self_attn.v_proj = ops_fn::transpose(w)?; }
248 if let Some(w) = weights.get(&format!("{}.self_attn.out_proj.weight", p)) { self.self_attn.out_proj = ops_fn::transpose(w)?; }
249 if let Some(w) = weights.get(&format!("{}.fc1.weight", p)) { self.fc1 = ops_fn::transpose(w)?; }
250 if let Some(w) = weights.get(&format!("{}.fc2.weight", p)) { self.fc2 = ops_fn::transpose(w)?; }
251 if let Some(w) = weights.get(&format!("{}.self_attn_layer_norm.weight", p)) { self.self_attn_layer_norm = w.clone(); }
252 if let Some(w) = weights.get(&format!("{}.final_layer_norm.weight", p)) { self.final_layer_norm = w.clone(); }
253 Ok(())
254 }
255
256 fn to_device(&mut self, device: &Device) -> Result<()> {
257 self.self_attn.to_device(device)?;
258 self.fc1 = self.fc1.to_device(device)?;
259 self.fc2 = self.fc2.to_device(device)?;
260 self.self_attn_layer_norm = self.self_attn_layer_norm.to_device(device)?;
261 self.final_layer_norm = self.final_layer_norm.to_device(device)?;
262 Ok(())
263 }
264}
265
266impl OPTAttention {
267 fn new(config: &OPTConfig, device: &Device) -> Result<Self> {
268 let head_dim = config.hidden_size / config.num_attention_heads;
269 Ok(Self {
270 q_proj: ops_fn::zeros(&[config.hidden_size, config.hidden_size], DataType::Float32, device)?,
271 k_proj: ops_fn::zeros(&[config.hidden_size, config.hidden_size], DataType::Float32, device)?,
272 v_proj: ops_fn::zeros(&[config.hidden_size, config.hidden_size], DataType::Float32, device)?,
273 out_proj: ops_fn::zeros(&[config.hidden_size, config.hidden_size], DataType::Float32, device)?,
274 num_heads: config.num_attention_heads,
275 head_dim,
276 scale: 1.0 / (head_dim as f32).sqrt(),
277 })
278 }
279
280 fn forward(&self, hidden_states: &Tensor) -> Result<Tensor> {
281 let shape = hidden_states.shape();
282 let (batch, seq_len, _) = (shape[0], shape[1], shape[2]);
283
284 let q = ops_fn::matmul(hidden_states, &self.q_proj)?.to_candle()?.reshape(&[batch, seq_len, self.num_heads, self.head_dim])?.transpose(1, 2)?;
285 let k = ops_fn::matmul(hidden_states, &self.k_proj)?.to_candle()?.reshape(&[batch, seq_len, self.num_heads, self.head_dim])?.transpose(1, 2)?;
286 let v = ops_fn::matmul(hidden_states, &self.v_proj)?.to_candle()?.reshape(&[batch, seq_len, self.num_heads, self.head_dim])?.transpose(1, 2)?;
287
288 let q = q.contiguous()?;
289 let k_t = k.transpose(2, 3)?.contiguous()?;
290 let scores = (q.matmul(&k_t)? * (self.scale as f64))?;
291 let device = scores.device();
292 let mask = {
293 let mut m = vec![0.0f32; seq_len * seq_len];
294 for i in 0..seq_len { for j in (i+1)..seq_len { m[i*seq_len+j] = f32::NEG_INFINITY; } }
295 candle_core::Tensor::from_vec(m, &[1, 1, seq_len, seq_len], device)?
296 };
297 let v = v.contiguous()?;
298 let attn = candle_nn::ops::softmax_last_dim(&scores.broadcast_add(&mask)?)?.matmul(&v)?;
299 let out = attn.transpose(1, 2)?.reshape(&[batch, seq_len, self.num_heads * self.head_dim])?;
300 ops_fn::matmul(&Tensor::from_candle(out), &self.out_proj)
301 }
302
303 fn to_device(&mut self, device: &Device) -> Result<()> {
304 self.q_proj = self.q_proj.to_device(device)?;
305 self.k_proj = self.k_proj.to_device(device)?;
306 self.v_proj = self.v_proj.to_device(device)?;
307 self.out_proj = self.out_proj.to_device(device)?;
308 Ok(())
309 }
310}
311
312#[cfg(test)]
313mod tests {
314 use super::*;
315 #[test]
316 fn test_opt_creation() {
317 let config = OPTConfig { vocab_size: 1000, hidden_size: 128, intermediate_size: 512, num_hidden_layers: 2, num_attention_heads: 4, num_key_value_heads: 4, ..Default::default() };
318 let model = OPTModelV2::new(config).unwrap();
319 assert_eq!(model.config().vocab_size(), 1000);
320 }
321 #[test]
322 fn test_opt_forward() {
323 let config = OPTConfig { vocab_size: 100, hidden_size: 64, intermediate_size: 256, num_hidden_layers: 1, num_attention_heads: 4, num_key_value_heads: 4, word_embed_proj_dim: 64, ..Default::default() };
324 let model = OPTModelV2::new(config).unwrap();
325 let inputs = ModelInputs::text(ops_fn::zeros(&[2, 8], DataType::Int64, &Device::CPU).unwrap());
326 match model.forward(&inputs).unwrap() { ModelOutputs::Logits { logits, .. } => assert_eq!(logits.shape(), &[2, 8, 100]), _ => panic!() }
327 }
328}