combs_models/
whisper.rs

1//! Whisper encoder–decoder for speech-to-text.
2//!
3//! Encoder: two GELU conv1d stems (stride 1 then 2), checkpoint-stored
4//! sinusoidal positions, pre-LN MHA stack, final LayerNorm. Decoder:
5//! learned positions, causal self-attention, cross-attention against the
6//! encoder states, tied output head. Attention mirrors the HF layout: q/v
7//! and the output projection carry biases, k_proj does not.
8//!
9//! v1 simplification, stated plainly: `decode_step` recomputes the decoder
10//! over the whole token prefix each step (no KV cache). Whisper contexts
11//! are ≤ 448 tokens and transcription prefixes are usually far shorter, so
12//! the quadratic cost is negligible at base/small scale; a decoder KV
13//! cache slots in behind the same trait without API changes.
14
15use burn::tensor::activation::{gelu, softmax};
16use burn::tensor::module::conv1d;
17use burn::tensor::ops::ConvOptions;
18use burn::tensor::{Device, Int, Tensor, TensorData, backend::Backend};
19use combs_formats::{ModelMetadata, ModelSource};
20
21use crate::llama::{linear, load_tensor};
22use crate::matmul::safe_matmul;
23use crate::norm::layer_norm;
24use crate::precision::{to_f32, to_float};
25use crate::traits::SpeechToTextModel;
26use crate::{ModelError, Result};
27
28/// LayerNorm epsilon Whisper checkpoints are trained with.
29const LN_EPS: f64 = 1e-5;
30
31/// One attention block's projections. `k_b` is `None` for Whisper
32/// (k_proj is bias-free in every export).
33struct Attn<B: Backend> {
34    q_w: Tensor<B, 2>,
35    q_b: Tensor<B, 1>,
36    k_w: Tensor<B, 2>,
37    v_w: Tensor<B, 2>,
38    v_b: Tensor<B, 1>,
39    o_w: Tensor<B, 2>,
40    o_b: Tensor<B, 1>,
41}
42
43struct EncoderLayer<B: Backend> {
44    ln1_w: Tensor<B, 1>,
45    ln1_b: Tensor<B, 1>,
46    attn: Attn<B>,
47    ln2_w: Tensor<B, 1>,
48    ln2_b: Tensor<B, 1>,
49    fc1_w: Tensor<B, 2>,
50    fc1_b: Tensor<B, 1>,
51    fc2_w: Tensor<B, 2>,
52    fc2_b: Tensor<B, 1>,
53}
54
55struct DecoderLayer<B: Backend> {
56    ln1_w: Tensor<B, 1>,
57    ln1_b: Tensor<B, 1>,
58    self_attn: Attn<B>,
59    ln_x_w: Tensor<B, 1>,
60    ln_x_b: Tensor<B, 1>,
61    cross_attn: Attn<B>,
62    ln2_w: Tensor<B, 1>,
63    ln2_b: Tensor<B, 1>,
64    fc1_w: Tensor<B, 2>,
65    fc1_b: Tensor<B, 1>,
66    fc2_w: Tensor<B, 2>,
67    fc2_b: Tensor<B, 1>,
68}
69
70pub struct WhisperModel<B: Backend> {
71    meta: ModelMetadata,
72    device: Device<B>,
73    heads: usize,
74    head_dim: usize,
75    scale: f64,
76    /// Mel bins expected by conv1 (from its weight shape).
77    n_mels: usize,
78    /// Max encoder frames after the stride-2 conv (from embed_positions).
79    n_audio_ctx: usize,
80    conv1_w: Tensor<B, 3>,
81    conv1_b: Tensor<B, 1>,
82    conv2_w: Tensor<B, 3>,
83    conv2_b: Tensor<B, 1>,
84    enc_pos: Tensor<B, 2>,
85    enc_layers: Vec<EncoderLayer<B>>,
86    enc_ln_w: Tensor<B, 1>,
87    enc_ln_b: Tensor<B, 1>,
88    embed_tokens: Tensor<B, 2>,
89    dec_pos: Tensor<B, 2>,
90    dec_layers: Vec<DecoderLayer<B>>,
91    dec_ln_w: Tensor<B, 1>,
92    dec_ln_b: Tensor<B, 1>,
93}
94
95impl<B: Backend> WhisperModel<B> {
96    pub fn load(source: &dyn ModelSource, device: &Device<B>) -> Result<Self> {
97        let meta = source.metadata().clone();
98        let t = |name: &str| -> Result<Tensor<B, 2>> { load_tensor(source, device, name) };
99        let v = |name: &str| -> Result<Tensor<B, 1>> { load_tensor(source, device, name) };
100
101        let attn = |p: &str| -> Result<Attn<B>> {
102            Ok(Attn {
103                q_w: t(&format!("{p}.q_proj.weight"))?,
104                q_b: v(&format!("{p}.q_proj.bias"))?,
105                k_w: t(&format!("{p}.k_proj.weight"))?,
106                v_w: t(&format!("{p}.v_proj.weight"))?,
107                v_b: v(&format!("{p}.v_proj.bias"))?,
108                o_w: t(&format!("{p}.out_proj.weight"))?,
109                o_b: v(&format!("{p}.out_proj.bias"))?,
110            })
111        };
112
113        let conv1_w: Tensor<B, 3> = load_tensor(source, device, "model.encoder.conv1.weight")?;
114        let conv2_w: Tensor<B, 3> = load_tensor(source, device, "model.encoder.conv2.weight")?;
115        let enc_pos: Tensor<B, 2> =
116            load_tensor(source, device, "model.encoder.embed_positions.weight")?;
117        let [_, n_mels, _] = conv1_w.dims();
118        let [n_audio_ctx, _] = enc_pos.dims();
119
120        let n_layers = meta.num_hidden_layers;
121        let mut enc_layers = Vec::with_capacity(n_layers);
122        for i in 0..n_layers {
123            let p = format!("model.encoder.layers.{i}");
124            enc_layers.push(EncoderLayer {
125                ln1_w: v(&format!("{p}.self_attn_layer_norm.weight"))?,
126                ln1_b: v(&format!("{p}.self_attn_layer_norm.bias"))?,
127                attn: attn(&format!("{p}.self_attn"))?,
128                ln2_w: v(&format!("{p}.final_layer_norm.weight"))?,
129                ln2_b: v(&format!("{p}.final_layer_norm.bias"))?,
130                fc1_w: t(&format!("{p}.fc1.weight"))?,
131                fc1_b: v(&format!("{p}.fc1.bias"))?,
132                fc2_w: t(&format!("{p}.fc2.weight"))?,
133                fc2_b: v(&format!("{p}.fc2.bias"))?,
134            });
135        }
136
137        let mut dec_layers = Vec::with_capacity(n_layers);
138        for i in 0..n_layers {
139            let p = format!("model.decoder.layers.{i}");
140            dec_layers.push(DecoderLayer {
141                ln1_w: v(&format!("{p}.self_attn_layer_norm.weight"))?,
142                ln1_b: v(&format!("{p}.self_attn_layer_norm.bias"))?,
143                self_attn: attn(&format!("{p}.self_attn"))?,
144                ln_x_w: v(&format!("{p}.encoder_attn_layer_norm.weight"))?,
145                ln_x_b: v(&format!("{p}.encoder_attn_layer_norm.bias"))?,
146                cross_attn: attn(&format!("{p}.encoder_attn"))?,
147                ln2_w: v(&format!("{p}.final_layer_norm.weight"))?,
148                ln2_b: v(&format!("{p}.final_layer_norm.bias"))?,
149                fc1_w: t(&format!("{p}.fc1.weight"))?,
150                fc1_b: v(&format!("{p}.fc1.bias"))?,
151                fc2_w: t(&format!("{p}.fc2.weight"))?,
152                fc2_b: v(&format!("{p}.fc2.bias"))?,
153            });
154        }
155
156        let heads = meta.num_attention_heads;
157        let head_dim = meta.head_dim;
158        Ok(WhisperModel {
159            device: device.clone(),
160            heads,
161            head_dim,
162            scale: 1.0 / (head_dim as f64).sqrt(),
163            n_mels,
164            n_audio_ctx,
165            conv1_w,
166            conv1_b: v("model.encoder.conv1.bias")?,
167            conv2_w,
168            conv2_b: v("model.encoder.conv2.bias")?,
169            enc_pos,
170            enc_layers,
171            enc_ln_w: v("model.encoder.layer_norm.weight")?,
172            enc_ln_b: v("model.encoder.layer_norm.bias")?,
173            embed_tokens: t("model.decoder.embed_tokens.weight")?,
174            dec_pos: t("model.decoder.embed_positions.weight")?,
175            dec_layers,
176            dec_ln_w: v("model.decoder.layer_norm.weight")?,
177            dec_ln_b: v("model.decoder.layer_norm.bias")?,
178            meta,
179        })
180    }
181
182    /// Multi-head attention over `x` (queries) and `kv` (keys/values), with
183    /// an optional additive mask on the score matrix. Scores and softmax run
184    /// in f32 (no-op in f32 builds).
185    fn attention(
186        &self,
187        attn: &Attn<B>,
188        x: Tensor<B, 3>,
189        kv: Tensor<B, 3>,
190        mask: Option<&Tensor<B, 2>>,
191    ) -> Tensor<B, 3> {
192        let [batch, q_len, _] = x.dims();
193        let [_, kv_len, _] = kv.dims();
194        let (h, d) = (self.heads, self.head_dim);
195
196        let q = linear(x, &attn.q_w, Some(&attn.q_b))
197            .reshape([batch, q_len, h, d])
198            .swap_dims(1, 2);
199        let k = linear(kv.clone(), &attn.k_w, None)
200            .reshape([batch, kv_len, h, d])
201            .swap_dims(1, 2);
202        let v = linear(kv, &attn.v_w, Some(&attn.v_b))
203            .reshape([batch, kv_len, h, d])
204            .swap_dims(1, 2);
205
206        let out_dtype = q.dtype();
207        let (q, k, v) = (to_f32(q), to_f32(k), to_f32(v));
208        let mut scores = safe_matmul(q, k.transpose()).mul_scalar(self.scale);
209        if let Some(m) = mask {
210            scores = scores + m.clone().reshape([1, 1, q_len, kv_len]);
211        }
212        let ctx = to_float(safe_matmul(softmax(scores, 3), v), out_dtype);
213        let ctx = ctx.swap_dims(1, 2).reshape([batch, q_len, h * d]);
214        linear(ctx, &attn.o_w, Some(&attn.o_b))
215    }
216
217    fn mlp(
218        &self,
219        x: Tensor<B, 3>,
220        fc1_w: &Tensor<B, 2>,
221        fc1_b: &Tensor<B, 1>,
222        fc2_w: &Tensor<B, 2>,
223        fc2_b: &Tensor<B, 1>,
224    ) -> Tensor<B, 3> {
225        linear(gelu(linear(x, fc1_w, Some(fc1_b))), fc2_w, Some(fc2_b))
226    }
227
228    /// Additive causal mask `[n, n]`: 0 on and below the diagonal, a large
229    /// negative above. Built host-side — n is ≤ the 448 decoder context.
230    fn causal_mask(&self, n: usize) -> Tensor<B, 2> {
231        let mut data = vec![0.0f32; n * n];
232        for r in 0..n {
233            for c in (r + 1)..n {
234                data[r * n + c] = f32::MIN / 2.0;
235            }
236        }
237        Tensor::from_data(TensorData::new(data, [n, n]), &self.device)
238    }
239}
240
241impl<B: Backend> SpeechToTextModel<B> for WhisperModel<B> {
242    fn metadata(&self) -> &ModelMetadata {
243        &self.meta
244    }
245
246    fn n_mels(&self) -> usize {
247        self.n_mels
248    }
249
250    fn encode_audio(&self, mel: Tensor<B, 3>) -> Result<Tensor<B, 3>> {
251        let [_, mels, _] = mel.dims();
252        if mels != self.n_mels {
253            return Err(ModelError::BadShape {
254                tensor: "mel spectrogram".into(),
255                expected: vec![1, self.n_mels],
256                got: vec![1, mels],
257            });
258        }
259        let x = gelu(conv1d(
260            mel,
261            self.conv1_w.clone(),
262            Some(self.conv1_b.clone()),
263            ConvOptions::new([1], [1], [1], 1),
264        ));
265        let x = gelu(conv1d(
266            x,
267            self.conv2_w.clone(),
268            Some(self.conv2_b.clone()),
269            ConvOptions::new([2], [1], [1], 1),
270        ));
271        // [1, d, frames/2] -> [1, frames/2, d] + sinusoidal positions.
272        let mut x = x.swap_dims(1, 2);
273        let [_, t, d] = x.dims();
274        if t > self.n_audio_ctx {
275            return Err(ModelError::BadShape {
276                tensor: "encoder frames".into(),
277                expected: vec![self.n_audio_ctx],
278                got: vec![t],
279            });
280        }
281        x = x + self
282            .enc_pos
283            .clone()
284            .slice([0..t, 0..d])
285            .reshape([1, t, d]);
286
287        for layer in &self.enc_layers {
288            let normed = layer_norm(x.clone(), layer.ln1_w.clone(), layer.ln1_b.clone(), LN_EPS);
289            x = x + self.attention(&layer.attn, normed.clone(), normed, None);
290            let normed = layer_norm(x.clone(), layer.ln2_w.clone(), layer.ln2_b.clone(), LN_EPS);
291            x = x + self.mlp(normed, &layer.fc1_w, &layer.fc1_b, &layer.fc2_w, &layer.fc2_b);
292        }
293        Ok(layer_norm(x, self.enc_ln_w.clone(), self.enc_ln_b.clone(), LN_EPS))
294    }
295
296    fn decode_step(&self, tokens: &[u32], encoded: &Tensor<B, 3>) -> Result<Tensor<B, 1>> {
297        let n = tokens.len();
298        let [_, _, d] = encoded.dims();
299        if n == 0 {
300            return Err(ModelError::BadShape {
301                tensor: "decoder tokens".into(),
302                expected: vec![1],
303                got: vec![0],
304            });
305        }
306        let [max_ctx, _] = self.dec_pos.dims();
307        if n > max_ctx {
308            return Err(ModelError::BadShape {
309                tensor: "decoder context".into(),
310                expected: vec![max_ctx],
311                got: vec![n],
312            });
313        }
314
315        let ids: Vec<i32> = tokens.iter().map(|&t| t as i32).collect();
316        let ids: Tensor<B, 2, Int> =
317            Tensor::from_data(TensorData::new(ids, [1, n]), &self.device);
318        let mut x = self
319            .embed_tokens
320            .clone()
321            .select(0, ids.reshape([n]))
322            .reshape([1, n, d])
323            + self.dec_pos.clone().slice([0..n, 0..d]).reshape([1, n, d]);
324
325        let mask = self.causal_mask(n);
326        for layer in &self.dec_layers {
327            let normed = layer_norm(x.clone(), layer.ln1_w.clone(), layer.ln1_b.clone(), LN_EPS);
328            x = x + self.attention(&layer.self_attn, normed.clone(), normed, Some(&mask));
329            let normed =
330                layer_norm(x.clone(), layer.ln_x_w.clone(), layer.ln_x_b.clone(), LN_EPS);
331            x = x + self.attention(&layer.cross_attn, normed, encoded.clone(), None);
332            let normed = layer_norm(x.clone(), layer.ln2_w.clone(), layer.ln2_b.clone(), LN_EPS);
333            x = x + self.mlp(normed, &layer.fc1_w, &layer.fc1_b, &layer.fc2_w, &layer.fc2_b);
334        }
335        let x = layer_norm(x, self.dec_ln_w.clone(), self.dec_ln_b.clone(), LN_EPS);
336
337        // Tied head over the final position only.
338        let last = x.slice([0..1, (n - 1)..n, 0..d]).reshape([1, d]);
339        let logits = safe_matmul(to_f32(last), to_f32(self.embed_tokens.clone().transpose()));
340        let vocab = self.meta.vocab_size;
341        Ok(logits.reshape([vocab]))
342    }
343}
344
345/// Loads a speech model by `metadata.architecture` — the ASR counterpart of
346/// the text registry's loader map.
347pub fn load_speech_model<B: Backend>(
348    source: &dyn ModelSource,
349    device: &Device<B>,
350) -> Result<Box<dyn SpeechToTextModel<B>>> {
351    match source.metadata().architecture.as_str() {
352        "whisper" => Ok(Box::new(WhisperModel::<B>::load(source, device)?)),
353        other => Err(ModelError::UnsupportedArchitecture(other.to_string())),
354    }
355}
356
357#[cfg(test)]
358mod tests {
359    use super::*;
360    use combs_core::init_device;
361
362    type TB = combs_core::CombsBackend;
363
364    #[test]
365    fn causal_mask_shape_and_triangle() {
366        if crate::skip_no_gpu() {
367            return;
368        }
369        let device = init_device();
370        let model_mask = |n: usize| {
371            // Build via the same routine the decoder uses.
372            let mut data = vec![0.0f32; n * n];
373            for r in 0..n {
374                for c in (r + 1)..n {
375                    data[r * n + c] = f32::MIN / 2.0;
376                }
377            }
378            Tensor::<TB, 2>::from_data(TensorData::new(data, [n, n]), &device)
379        };
380        let m = model_mask(4).into_data().to_vec::<f32>().unwrap();
381        assert_eq!(m[0 * 4 + 0], 0.0);
382        assert!(m[0 * 4 + 1] < -1e30);
383        assert_eq!(m[3 * 4 + 3], 0.0);
384        assert_eq!(m[3 * 4 + 0], 0.0);
385    }
386}
combs_models/whisper.rs

combs_models/
whisper.rs
