From b1d42231fba72ff508bc02c34c65feab5409ec4c Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Tue, 4 Jul 2023 07:23:02 +0100
Subject: [PATCH 01/36] Start sketching the whisper model.

---
 candle-examples/examples/whisper/main.rs | 281 +++++++++++++++++++++++
 1 file changed, 281 insertions(+)
 create mode 100644 candle-examples/examples/whisper/main.rs

diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs
new file mode 100644
index 00000000..7b2d0c74
--- /dev/null
+++ b/candle-examples/examples/whisper/main.rs
@@ -0,0 +1,281 @@
+#![allow(dead_code)]
+// https://github.com/openai/whisper/blob/main/whisper/model.py
+
+use anyhow::{Error as E, Result};
+use candle::{safetensors::SafeTensors, DType, Device, Shape, Tensor};
+use clap::Parser;
+use std::collections::HashMap;
+
+const DTYPE: DType = DType::F32;
+
+struct VarBuilder<'a> {
+    safetensors: Option<(HashMap<String, usize>, Vec<SafeTensors<'a>>)>,
+    dtype: DType,
+    device: Device,
+}
+
+impl<'a> VarBuilder<'a> {
+    pub fn from_safetensors(
+        safetensors: Vec<SafeTensors<'a>>,
+        dtype: DType,
+        device: Device,
+    ) -> Self {
+        let mut routing = HashMap::new();
+        for (index, sf) in safetensors.iter().enumerate() {
+            for k in sf.names() {
+                routing.insert(k.to_string(), index);
+            }
+        }
+        Self {
+            safetensors: Some((routing, safetensors)),
+            device,
+            dtype,
+        }
+    }
+
+    pub fn zeros(dtype: DType, device: Device) -> Self {
+        Self {
+            safetensors: None,
+            device,
+            dtype,
+        }
+    }
+
+    pub fn get<S: Into<Shape>>(&self, s: S, tensor_name: &str) -> candle::Result<Tensor> {
+        let s: Shape = s.into();
+        match &self.safetensors {
+            None => Tensor::zeros(s, self.dtype, &self.device),
+            Some((routing, safetensors)) => {
+                // Unwrap or 0  just to let the proper error flow.
+                let index = routing.get(tensor_name).unwrap_or(&0);
+                let tensor = safetensors[*index]
+                    .tensor(tensor_name, &self.device)?
+                    .to_dtype(self.dtype)?;
+                if *tensor.shape() != s {
+                    let msg = format!("shape mismatch for {tensor_name}");
+                    Err(candle::Error::UnexpectedShape {
+                        msg,
+                        expected: s,
+                        got: tensor.shape().clone(),
+                    })?
+                }
+                Ok(tensor)
+            }
+        }
+    }
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum HiddenAct {
+    Gelu,
+    Relu,
+}
+
+impl HiddenAct {
+    fn forward(&self, xs: &Tensor) -> candle::Result<Tensor> {
+        match self {
+            Self::Gelu => xs.gelu(),
+            Self::Relu => xs.relu(),
+        }
+    }
+}
+
+#[derive(Debug, Clone, PartialEq)]
+struct Config {
+    n_mels: usize,
+    n_audio_ctx: usize,
+    n_audio_state: usize,
+    n_audio_head: usize,
+    n_audio_layer: usize,
+    n_vocab: usize,
+    n_text_ctx: usize,
+    n_text_state: usize,
+    n_text_head: usize,
+    n_text_layer: usize,
+}
+
+struct Embedding {
+    embeddings: Tensor,
+    hidden_size: usize,
+}
+
+impl Embedding {
+    fn new(embeddings: Tensor, hidden_size: usize) -> Self {
+        Self {
+            embeddings,
+            hidden_size,
+        }
+    }
+
+    fn load(vocab_size: usize, hidden_size: usize, p: &str, vb: &VarBuilder) -> Result<Self> {
+        let embeddings = vb.get((vocab_size, hidden_size), &format!("{p}.weight"))?;
+        Ok(Self::new(embeddings, hidden_size))
+    }
+
+    fn forward(&self, indexes: &Tensor) -> Result<Tensor> {
+        let mut final_dims = indexes.dims().to_vec();
+        final_dims.push(self.hidden_size);
+        let indexes = indexes.flatten_all()?;
+        let values = Tensor::embedding(&indexes, &self.embeddings)?;
+        let values = values.reshape(final_dims)?;
+        Ok(values)
+    }
+}
+
+struct Linear {
+    weight: Tensor,
+    bias: Option<Tensor>,
+}
+
+impl Linear {
+    fn load(size1: usize, size2: usize, p: &str, vb: &VarBuilder) -> Result<Self> {
+        let weight = vb.get((size2, size1), &format!("{p}.weight"))?;
+        let bias = vb.get(size2, &format!("{p}.bias"))?;
+        Ok(Self {
+            weight,
+            bias: Some(bias),
+        })
+    }
+
+    fn load_no_bias(size1: usize, size2: usize, p: &str, vb: &VarBuilder) -> Result<Self> {
+        let weight = vb.get((size2, size1), &format!("{p}.weight"))?;
+        Ok(Self { weight, bias: None })
+    }
+
+    fn forward(&self, x: &Tensor) -> candle::Result<Tensor> {
+        let (bsize, _, _) = x.shape().r3()?;
+        let w = self.weight.broadcast_left(bsize)?.t()?;
+        let x = x.matmul(&w)?;
+        match &self.bias {
+            None => Ok(x),
+            Some(bias) => x.broadcast_add(bias),
+        }
+    }
+}
+
+struct Dropout {
+    pr: f64,
+}
+
+impl Dropout {
+    fn new(pr: f64) -> Self {
+        Self { pr }
+    }
+
+    fn forward(&self, x: &Tensor) -> Result<Tensor> {
+        // TODO
+        Ok(x.clone())
+    }
+}
+
+// This layer norm version handles both weight and bias so removes the mean.
+struct LayerNorm {
+    weight: Tensor,
+    bias: Tensor,
+    eps: f64,
+}
+
+impl LayerNorm {
+    fn new(weight: Tensor, bias: Tensor, eps: f64) -> Self {
+        Self { weight, bias, eps }
+    }
+
+    fn load(size: usize, eps: f64, p: &str, vb: &VarBuilder) -> Result<Self> {
+        let weight = vb.get(size, &format!("{p}.weight"))?;
+        let bias = vb.get(size, &format!("{p}.bias"))?;
+        Ok(Self::new(weight, bias, eps))
+    }
+
+    fn forward(&self, x: &Tensor) -> Result<Tensor> {
+        let (_bsize, _seq_len, hidden_size) = x.shape().r3()?;
+        let mean_x = (x.sum(&[2])? / hidden_size as f64)?;
+        let x = x.broadcast_sub(&mean_x)?;
+        let norm_x = ((&x * &x)?.sum(&[2])? / hidden_size as f64)?;
+        let x_normed = x.broadcast_div(&(norm_x + self.eps)?.sqrt()?)?;
+        let x = x_normed
+            .broadcast_mul(&self.weight)?
+            .broadcast_add(&self.bias)?;
+        Ok(x)
+    }
+}
+
+// https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L62
+struct MultiHeadAttention {
+    query: Linear,
+    key: Linear,
+    value: Linear,
+    out: Linear,
+    n_head: usize,
+}
+
+impl MultiHeadAttention {
+    fn load(n_head: usize, n: usize, p: &str, vb: &VarBuilder) -> Result<Self> {
+        let query = Linear::load(n, n, &format!("{p}.query"), vb)?;
+        let value = Linear::load_no_bias(n, n, &format!("{p}.value"), vb)?;
+        let key = Linear::load(n, n, &format!("{p}.key"), vb)?;
+        let out = Linear::load(n, n, &format!("{p}.out"), vb)?;
+        Ok(Self {
+            query,
+            key,
+            value,
+            out,
+            n_head,
+        })
+    }
+
+    fn forward(&self, x: &Tensor) -> Result<(Tensor, Tensor)> {
+        let q = self.query.forward(x)?;
+        let k = self.key.forward(x)?;
+        let v = self.value.forward(x)?;
+        let (wv, qk) = self.qkv_attention(&q, &k, &v)?;
+        let out = self.out.forward(&wv)?;
+        Ok((out, qk))
+    }
+
+    fn qkv_attention(&self, q: &Tensor, k: &Tensor, v: &Tensor) -> Result<(Tensor, Tensor)> {
+        let (n_batch, n_ctx, n_state) = q.shape().r3()?;
+        let target_dims = &[n_batch, n_ctx, self.n_head, n_state / self.n_head];
+        let scale = ((n_state / self.n_head) as f64).powf(-0.25);
+        let q = (q.reshape(target_dims)?.transpose(1, 2)? * scale)?;
+        let k = (k.reshape(target_dims)?.transpose(1, 2)?.transpose(2, 3)? * scale)?;
+        let v = v.reshape(target_dims)?.transpose(1, 2)?;
+        let qk = q.matmul(&k)?;
+        let w = qk.softmax(qk.rank() - 1)?;
+        let wv = w.matmul(&v)?.transpose(1, 2)?.flatten(Some(2), None)?;
+        let qk = qk.detach()?;
+        Ok((wv, qk))
+    }
+}
+
+#[derive(Parser, Debug)]
+#[command(author, version, about, long_about = None)]
+struct Args {
+    /// Run on CPU rather than on GPU.
+    #[arg(long)]
+    cpu: bool,
+
+    #[arg(long)]
+    tokenizer_config: String,
+
+    #[arg(long)]
+    weights: String,
+}
+
+fn main() -> Result<()> {
+    use tokenizers::Tokenizer;
+
+    let args = Args::parse();
+    let device = if args.cpu {
+        Device::Cpu
+    } else {
+        Device::new_cuda(0)?
+    };
+
+    let mut tokenizer = Tokenizer::from_file(args.tokenizer_config).map_err(E::msg)?;
+    let _tokenizer = tokenizer.with_padding(None).with_truncation(None);
+
+    let weights = unsafe { candle::safetensors::MmapedFile::new(args.weights)? };
+    let weights = weights.deserialize()?;
+    let _vb = VarBuilder::from_safetensors(vec![weights], DTYPE, device);
+    Ok(())
+}

From 0ca2af6940e2812e65a2e2af7f53ce9b8c1980d4 Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Tue, 4 Jul 2023 07:43:36 +0100
Subject: [PATCH 02/36] Add the residual attention block.

---
 candle-examples/examples/whisper/main.rs | 92 +++++++++++++++++++-----
 1 file changed, 75 insertions(+), 17 deletions(-)

diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs
index 7b2d0c74..1e32c972 100644
--- a/candle-examples/examples/whisper/main.rs
+++ b/candle-examples/examples/whisper/main.rs
@@ -176,14 +176,14 @@ struct LayerNorm {
 }
 
 impl LayerNorm {
-    fn new(weight: Tensor, bias: Tensor, eps: f64) -> Self {
-        Self { weight, bias, eps }
-    }
-
-    fn load(size: usize, eps: f64, p: &str, vb: &VarBuilder) -> Result<Self> {
+    fn load(size: usize, p: &str, vb: &VarBuilder) -> Result<Self> {
         let weight = vb.get(size, &format!("{p}.weight"))?;
         let bias = vb.get(size, &format!("{p}.bias"))?;
-        Ok(Self::new(weight, bias, eps))
+        Ok(Self {
+            weight,
+            bias,
+            eps: 1e-5,
+        })
     }
 
     fn forward(&self, x: &Tensor) -> Result<Tensor> {
@@ -209,11 +209,11 @@ struct MultiHeadAttention {
 }
 
 impl MultiHeadAttention {
-    fn load(n_head: usize, n: usize, p: &str, vb: &VarBuilder) -> Result<Self> {
-        let query = Linear::load(n, n, &format!("{p}.query"), vb)?;
-        let value = Linear::load_no_bias(n, n, &format!("{p}.value"), vb)?;
-        let key = Linear::load(n, n, &format!("{p}.key"), vb)?;
-        let out = Linear::load(n, n, &format!("{p}.out"), vb)?;
+    fn load(n_state: usize, n_head: usize, p: &str, vb: &VarBuilder) -> Result<Self> {
+        let query = Linear::load(n_state, n_state, &format!("{p}.query"), vb)?;
+        let value = Linear::load_no_bias(n_state, n_state, &format!("{p}.value"), vb)?;
+        let key = Linear::load(n_state, n_state, &format!("{p}.key"), vb)?;
+        let out = Linear::load(n_state, n_state, &format!("{p}.out"), vb)?;
         Ok(Self {
             query,
             key,
@@ -223,16 +223,16 @@ impl MultiHeadAttention {
         })
     }
 
-    fn forward(&self, x: &Tensor) -> Result<(Tensor, Tensor)> {
+    fn forward(&self, x: &Tensor) -> Result<Tensor> {
         let q = self.query.forward(x)?;
         let k = self.key.forward(x)?;
         let v = self.value.forward(x)?;
-        let (wv, qk) = self.qkv_attention(&q, &k, &v)?;
+        let wv = self.qkv_attention(&q, &k, &v)?;
         let out = self.out.forward(&wv)?;
-        Ok((out, qk))
+        Ok(out)
     }
 
-    fn qkv_attention(&self, q: &Tensor, k: &Tensor, v: &Tensor) -> Result<(Tensor, Tensor)> {
+    fn qkv_attention(&self, q: &Tensor, k: &Tensor, v: &Tensor) -> Result<Tensor> {
         let (n_batch, n_ctx, n_state) = q.shape().r3()?;
         let target_dims = &[n_batch, n_ctx, self.n_head, n_state / self.n_head];
         let scale = ((n_state / self.n_head) as f64).powf(-0.25);
@@ -242,8 +242,66 @@ impl MultiHeadAttention {
         let qk = q.matmul(&k)?;
         let w = qk.softmax(qk.rank() - 1)?;
         let wv = w.matmul(&v)?.transpose(1, 2)?.flatten(Some(2), None)?;
-        let qk = qk.detach()?;
-        Ok((wv, qk))
+        Ok(wv)
+    }
+}
+
+// https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L111
+struct ResidualAttentionBlock {
+    attn: MultiHeadAttention,
+    attn_ln: LayerNorm,
+    cross_attn: Option<MultiHeadAttention>,
+    cross_attn_ln: Option<LayerNorm>,
+    mlp_linear1: Linear,
+    mlp_linear2: Linear,
+    mlp_ln: LayerNorm,
+}
+
+impl ResidualAttentionBlock {
+    fn load(n_state: usize, n_head: usize, ca: bool, p: &str, vb: &VarBuilder) -> Result<Self> {
+        let attn = MultiHeadAttention::load(n_state, n_head, &format!("{p}.attn"), vb)?;
+        let attn_ln = LayerNorm::load(n_state, &format!("{p}.attn_ln"), vb)?;
+        let (cross_attn, cross_attn_ln) = if ca {
+            let cross_attn =
+                MultiHeadAttention::load(n_state, n_head, &format!("{p}.cross_attn"), vb)?;
+            let cross_attn_ln = LayerNorm::load(n_state, &format!("{p}.cross_attn_ln"), vb)?;
+            (Some(cross_attn), Some(cross_attn_ln))
+        } else {
+            (None, None)
+        };
+        let n_mlp = n_state * 4;
+        let mlp_linear1 = Linear::load(n_state, n_mlp, &format!("{p}.mlp.0"), vb)?;
+        let mlp_linear2 = Linear::load(n_mlp, n_state, &format!("{p}.mlp.2"), vb)?;
+        let mlp_ln = LayerNorm::load(n_state, &format!("{p}.mlp_ln"), vb)?;
+        Ok(Self {
+            attn,
+            attn_ln,
+            cross_attn,
+            cross_attn_ln,
+            mlp_linear1,
+            mlp_linear2,
+            mlp_ln,
+        })
+    }
+
+    fn forward(&self, x: &Tensor) -> Result<Tensor> {
+        let attn = self.attn.forward(&self.attn_ln.forward(x)?)?;
+        let mut x = (x + attn)?;
+        // Cross-Attn
+        if let Some(cross_attn) = &self.cross_attn {
+            x = cross_attn.forward(&x)?
+        }
+        if let Some(cross_attn_ln) = &self.cross_attn_ln {
+            x = cross_attn_ln.forward(&x)?
+        }
+        // Mlp
+        let mlp = self.mlp_linear2.forward(
+            &self
+                .mlp_linear1
+                .forward(&self.mlp_ln.forward(&x)?)?
+                .gelu()?,
+        )?;
+        Ok((x + mlp)?)
     }
 }
 

From 512dd9e4d6c278d72f3942da937a40d6cd45f933 Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Tue, 4 Jul 2023 08:01:06 +0100
Subject: [PATCH 03/36] Flesh out the whisper example.

---
 candle-examples/examples/whisper/main.rs | 84 ++++++++++++++++++++++++
 1 file changed, 84 insertions(+)

diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs
index 1e32c972..e0292e17 100644
--- a/candle-examples/examples/whisper/main.rs
+++ b/candle-examples/examples/whisper/main.rs
@@ -305,6 +305,90 @@ impl ResidualAttentionBlock {
     }
 }
 
+// https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L143
+struct AudioEncoder {
+    conv1: Linear, // TODO
+    conv2: Linear, // TODO
+    blocks: Vec<ResidualAttentionBlock>,
+    ln_post: LayerNorm,
+}
+
+impl AudioEncoder {
+    fn load(p: &str, vb: &VarBuilder, cfg: &Config) -> Result<Self> {
+        let n_state = cfg.n_audio_state;
+        let n_head = cfg.n_audio_head;
+        let conv1 = Linear::load(cfg.n_mels, n_state, &format!("{p}.conv1"), vb)?;
+        let conv2 = Linear::load(n_state, n_state, &format!("{p}.conv2"), vb)?;
+        let blocks = (0..cfg.n_audio_layer)
+            .map(|i| {
+                ResidualAttentionBlock::load(n_state, n_head, false, &format!("{p}.blocks.{i}"), vb)
+            })
+            .collect::<Result<Vec<_>>>()?;
+        let ln_post = LayerNorm::load(n_state, &format!("{p}.ln_post"), vb)?;
+        Ok(Self {
+            conv1,
+            conv2,
+            blocks,
+            ln_post,
+        })
+    }
+    fn forward(&self, _x: &Tensor) -> Result<Tensor> {
+        todo!()
+    }
+}
+
+// https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L176
+struct TextDecoder {
+    token_embedding: Embedding,
+    blocks: Vec<ResidualAttentionBlock>,
+    ln: LayerNorm,
+    mask: Tensor,
+}
+
+impl TextDecoder {
+    fn load(p: &str, vb: &VarBuilder, cfg: &Config) -> Result<Self> {
+        let n_state = cfg.n_text_state;
+        let n_head = cfg.n_text_head;
+        let token_embedding =
+            Embedding::load(cfg.n_vocab, n_state, &format!("{p}.token_embedding"), vb)?;
+        let blocks = (0..cfg.n_text_layer)
+            .map(|i| {
+                ResidualAttentionBlock::load(n_state, n_head, false, &format!("{p}.blocks.{i}"), vb)
+            })
+            .collect::<Result<Vec<_>>>()?;
+        let ln = LayerNorm::load(n_state, &format!("{p}.ln"), vb)?;
+        let mask = Tensor::new(&[0u32], &vb.device)?;
+        Ok(Self {
+            token_embedding,
+            blocks,
+            ln,
+            mask,
+        })
+    }
+    fn forward(&self, _tokens: &Tensor, _enc: &Tensor) -> Result<Tensor> {
+        todo!()
+    }
+}
+
+// https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L221
+struct Whisper {
+    encoder: AudioEncoder,
+    decoder: TextDecoder,
+}
+
+impl Whisper {
+    fn load(p: &str, vb: &VarBuilder, cfg: &Config) -> Result<Self> {
+        let encoder = AudioEncoder::load(&format!("{p}.encoder"), vb, cfg)?;
+        let decoder = TextDecoder::load(&format!("{p}.decoder"), vb, cfg)?;
+        Ok(Self { encoder, decoder })
+    }
+    fn forward(&self, mel: &Tensor, tokens: &Tensor) -> Result<Tensor> {
+        let enc = self.encoder.forward(mel)?;
+        let dec = self.decoder.forward(tokens, &enc)?;
+        Ok(dec)
+    }
+}
+
 #[derive(Parser, Debug)]
 #[command(author, version, about, long_about = None)]
 struct Args {

From 6728a856765ca23f6bf50cd9246bb2f8e63ee41a Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Tue, 4 Jul 2023 08:32:03 +0100
Subject: [PATCH 04/36] Add more to the whisper inference.

---
 candle-examples/examples/whisper/main.rs | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs
index e0292e17..3e780a2c 100644
--- a/candle-examples/examples/whisper/main.rs
+++ b/candle-examples/examples/whisper/main.rs
@@ -1,5 +1,7 @@
 #![allow(dead_code)]
 // https://github.com/openai/whisper/blob/main/whisper/model.py
+// TODO:
+// - kv-cache support?
 
 use anyhow::{Error as E, Result};
 use candle::{safetensors::SafeTensors, DType, Device, Shape, Tensor};
@@ -309,6 +311,7 @@ impl ResidualAttentionBlock {
 struct AudioEncoder {
     conv1: Linear, // TODO
     conv2: Linear, // TODO
+    positional_embedding: Tensor,
     blocks: Vec<ResidualAttentionBlock>,
     ln_post: LayerNorm,
 }
@@ -319,6 +322,7 @@ impl AudioEncoder {
         let n_head = cfg.n_audio_head;
         let conv1 = Linear::load(cfg.n_mels, n_state, &format!("{p}.conv1"), vb)?;
         let conv2 = Linear::load(n_state, n_state, &format!("{p}.conv2"), vb)?;
+        let positional_embedding = Tensor::new(&[0u32], &vb.device)?; // TODO
         let blocks = (0..cfg.n_audio_layer)
             .map(|i| {
                 ResidualAttentionBlock::load(n_state, n_head, false, &format!("{p}.blocks.{i}"), vb)
@@ -328,12 +332,21 @@ impl AudioEncoder {
         Ok(Self {
             conv1,
             conv2,
+            positional_embedding,
             blocks,
             ln_post,
         })
     }
-    fn forward(&self, _x: &Tensor) -> Result<Tensor> {
-        todo!()
+    fn forward(&self, x: &Tensor) -> Result<Tensor> {
+        let x = self.conv1.forward(x)?.gelu()?;
+        let x = self.conv2.forward(&x)?.gelu()?;
+        let x = x.transpose(1, 2)?;
+        let mut x = x.broadcast_add(&self.positional_embedding)?;
+        for block in self.blocks.iter() {
+            x = block.forward(&x)?
+        }
+        let x = self.ln_post.forward(&x)?;
+        Ok(x)
     }
 }
 
@@ -357,7 +370,7 @@ impl TextDecoder {
             })
             .collect::<Result<Vec<_>>>()?;
         let ln = LayerNorm::load(n_state, &format!("{p}.ln"), vb)?;
-        let mask = Tensor::new(&[0u32], &vb.device)?;
+        let mask = Tensor::new(&[0u32], &vb.device)?; // TODO
         Ok(Self {
             token_embedding,
             blocks,

From 998cfda9c4baefd0fb3407724c344158356d2d31 Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Tue, 4 Jul 2023 09:03:25 +0100
Subject: [PATCH 05/36] Sinusoid embeddings.

---
 candle-examples/examples/whisper/main.rs | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs
index 3e780a2c..1930a7a9 100644
--- a/candle-examples/examples/whisper/main.rs
+++ b/candle-examples/examples/whisper/main.rs
@@ -307,6 +307,21 @@ impl ResidualAttentionBlock {
     }
 }
 
+fn sinusoids(length: usize, channels: usize) -> Result<Tensor> {
+    let max_timescale = 10000f32;
+    let log_timescale_increment = max_timescale.ln() / (channels / 2 - 1) as f32;
+    let inv_timescales: Vec<_> = (0..channels / 2)
+        .map(|i| (i as f32 * (-log_timescale_increment)).exp())
+        .collect();
+    let arange: Vec<_> = (0..length).map(|c| c as f32).collect();
+    let inv_timescales = Tensor::new(inv_timescales.as_slice(), &Device::Cpu)?.unsqueeze(0)?;
+    let arange = Tensor::new(arange.as_slice(), &Device::Cpu)?.unsqueeze(1)?;
+    let sh = (length, channels / 2);
+    let scaled_time = (arange.broadcast_as(sh)? * inv_timescales.broadcast_as(sh)?)?;
+    let sincos = Tensor::cat(&[scaled_time.sin()?, scaled_time.cos()?], 1)?;
+    Ok(sincos)
+}
+
 // https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L143
 struct AudioEncoder {
     conv1: Linear, // TODO
@@ -322,7 +337,7 @@ impl AudioEncoder {
         let n_head = cfg.n_audio_head;
         let conv1 = Linear::load(cfg.n_mels, n_state, &format!("{p}.conv1"), vb)?;
         let conv2 = Linear::load(n_state, n_state, &format!("{p}.conv2"), vb)?;
-        let positional_embedding = Tensor::new(&[0u32], &vb.device)?; // TODO
+        let positional_embedding = sinusoids(cfg.n_audio_ctx, n_state)?.to_device(&vb.device)?;
         let blocks = (0..cfg.n_audio_layer)
             .map(|i| {
                 ResidualAttentionBlock::load(n_state, n_head, false, &format!("{p}.blocks.{i}"), vb)

From c09aa4b0f4c1fb074aecbf71aea3e96cadebea60 Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Tue, 4 Jul 2023 09:14:28 +0100
Subject: [PATCH 06/36] Add more to the forward pass and fix the
 cross-attention.

---
 candle-examples/examples/whisper/main.rs | 37 +++++++++++++++++-------
 1 file changed, 26 insertions(+), 11 deletions(-)

diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs
index 1930a7a9..129c73b8 100644
--- a/candle-examples/examples/whisper/main.rs
+++ b/candle-examples/examples/whisper/main.rs
@@ -225,10 +225,10 @@ impl MultiHeadAttention {
         })
     }
 
-    fn forward(&self, x: &Tensor) -> Result<Tensor> {
+    fn forward(&self, x: &Tensor, xa: Option<&Tensor>) -> Result<Tensor> {
         let q = self.query.forward(x)?;
-        let k = self.key.forward(x)?;
-        let v = self.value.forward(x)?;
+        let k = self.key.forward(xa.unwrap_or(x))?;
+        let v = self.value.forward(xa.unwrap_or(x))?;
         let wv = self.qkv_attention(&q, &k, &v)?;
         let out = self.out.forward(&wv)?;
         Ok(out)
@@ -286,16 +286,16 @@ impl ResidualAttentionBlock {
         })
     }
 
-    fn forward(&self, x: &Tensor) -> Result<Tensor> {
-        let attn = self.attn.forward(&self.attn_ln.forward(x)?)?;
+    fn forward(&self, x: &Tensor, xa: Option<&Tensor>) -> Result<Tensor> {
+        let attn = self.attn.forward(&self.attn_ln.forward(x)?, None)?;
         let mut x = (x + attn)?;
         // Cross-Attn
-        if let Some(cross_attn) = &self.cross_attn {
-            x = cross_attn.forward(&x)?
-        }
         if let Some(cross_attn_ln) = &self.cross_attn_ln {
             x = cross_attn_ln.forward(&x)?
         }
+        if let Some(cross_attn) = &self.cross_attn {
+            x = cross_attn.forward(&x, xa)?
+        }
         // Mlp
         let mlp = self.mlp_linear2.forward(
             &self
@@ -358,7 +358,7 @@ impl AudioEncoder {
         let x = x.transpose(1, 2)?;
         let mut x = x.broadcast_add(&self.positional_embedding)?;
         for block in self.blocks.iter() {
-            x = block.forward(&x)?
+            x = block.forward(&x, None)?
         }
         let x = self.ln_post.forward(&x)?;
         Ok(x)
@@ -368,6 +368,7 @@ impl AudioEncoder {
 // https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L176
 struct TextDecoder {
     token_embedding: Embedding,
+    positional_embedding: Tensor,
     blocks: Vec<ResidualAttentionBlock>,
     ln: LayerNorm,
     mask: Tensor,
@@ -377,8 +378,11 @@ impl TextDecoder {
     fn load(p: &str, vb: &VarBuilder, cfg: &Config) -> Result<Self> {
         let n_state = cfg.n_text_state;
         let n_head = cfg.n_text_head;
+        let n_ctx = cfg.n_text_ctx;
         let token_embedding =
             Embedding::load(cfg.n_vocab, n_state, &format!("{p}.token_embedding"), vb)?;
+        let positional_embedding =
+            vb.get((n_ctx, n_state), &format!("{p}.positional_embedding"))?;
         let blocks = (0..cfg.n_text_layer)
             .map(|i| {
                 ResidualAttentionBlock::load(n_state, n_head, false, &format!("{p}.blocks.{i}"), vb)
@@ -388,13 +392,24 @@ impl TextDecoder {
         let mask = Tensor::new(&[0u32], &vb.device)?; // TODO
         Ok(Self {
             token_embedding,
+            positional_embedding,
             blocks,
             ln,
             mask,
         })
     }
-    fn forward(&self, _tokens: &Tensor, _enc: &Tensor) -> Result<Tensor> {
-        todo!()
+    fn forward(&self, x: &Tensor, xa: &Tensor) -> Result<Tensor> {
+        let x_dims = x.dims();
+        let last = x_dims[x_dims.len() - 1];
+        let token_embedding = self.token_embedding.forward(x)?;
+        let positional_embedding = self.positional_embedding.narrow(0, 0, last)?;
+        let mut x = (token_embedding + positional_embedding)?;
+        for block in self.blocks.iter() {
+            x = block.forward(&x, Some(xa))?
+        }
+        let x = self.ln.forward(&x)?;
+        let logits = x.matmul(&self.token_embedding.embeddings.t()?)?;
+        Ok(logits)
     }
 }
 

From d71b31144d286df98d2ae12bde66976794f33640 Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Tue, 4 Jul 2023 09:29:19 +0100
Subject: [PATCH 07/36] Add a weight extraction script.

---
 .../examples/whisper/extract_weights.py         | 13 +++++++++++++
 candle-examples/examples/whisper/main.rs        | 17 +++++++++++++++++
 2 files changed, 30 insertions(+)
 create mode 100644 candle-examples/examples/whisper/extract_weights.py

diff --git a/candle-examples/examples/whisper/extract_weights.py b/candle-examples/examples/whisper/extract_weights.py
new file mode 100644
index 00000000..d6ccffc6
--- /dev/null
+++ b/candle-examples/examples/whisper/extract_weights.py
@@ -0,0 +1,13 @@
+# Get the checkpoint from
+# https://openaipublic.azureedge.net/main/whisper/models/d3dd57d32accea0b295c96e26691aa14d8822fac7d9d27d5dc00b4ca2826dd03/tiny.en.pt
+
+import torch
+from safetensors.torch import save_file
+
+data = torch.load("tiny.en.pt")
+weights = {}
+for k, v in data["model_state_dict"].items():
+    weights[k] = v.contiguous()
+    print(k, v.shape)
+save_file(weights, "tiny.en.safetensors")
+print(data["dims"])
diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs
index 129c73b8..dac6df84 100644
--- a/candle-examples/examples/whisper/main.rs
+++ b/candle-examples/examples/whisper/main.rs
@@ -96,6 +96,23 @@ struct Config {
     n_text_layer: usize,
 }
 
+impl Config {
+    fn tiny() -> Self {
+        Self {
+            n_mels: 80,
+            n_vocab: 51864,
+            n_audio_ctx: 1500,
+            n_audio_state: 384,
+            n_audio_head: 6,
+            n_audio_layer: 4,
+            n_text_ctx: 448,
+            n_text_state: 384,
+            n_text_head: 6,
+            n_text_layer: 4,
+        }
+    }
+}
+
 struct Embedding {
     embeddings: Tensor,
     hidden_size: usize,

From e6b01d0c18d27b2363f5aad7a19da38afc51f7d1 Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Tue, 4 Jul 2023 10:01:05 +0100
Subject: [PATCH 08/36] Add the conv1d layer (but not the op).

---
 candle-examples/examples/whisper/main.rs | 89 ++++++++++++++++++++++--
 1 file changed, 85 insertions(+), 4 deletions(-)

diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs
index dac6df84..75ab2189 100644
--- a/candle-examples/examples/whisper/main.rs
+++ b/candle-examples/examples/whisper/main.rs
@@ -172,6 +172,79 @@ impl Linear {
     }
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+struct ConvConfig {
+    padding: usize,
+    stride: usize,
+}
+
+impl Default for ConvConfig {
+    fn default() -> Self {
+        Self {
+            padding: 0,
+            stride: 1,
+        }
+    }
+}
+
+struct Conv1D {
+    weight: Tensor,
+    bias: Option<Tensor>,
+    config: ConvConfig,
+}
+
+impl Conv1D {
+    fn load(
+        in_channels: usize,
+        out_channels: usize,
+        kernel_size: usize,
+        config: ConvConfig,
+        p: &str,
+        vb: &VarBuilder,
+    ) -> Result<Self> {
+        let weight = vb.get(
+            (out_channels, in_channels, kernel_size),
+            &format!("{p}.weight"),
+        )?;
+        let bias = vb.get(out_channels, &format!("{p}.bias"))?;
+        Ok(Self {
+            weight,
+            bias: Some(bias),
+            config,
+        })
+    }
+
+    fn load_no_bias(
+        in_channels: usize,
+        out_channels: usize,
+        kernel_size: usize,
+        config: ConvConfig,
+        p: &str,
+        vb: &VarBuilder,
+    ) -> Result<Self> {
+        let weight = vb.get(
+            (out_channels, in_channels, kernel_size),
+            &format!("{p}.weight"),
+        )?;
+        Ok(Self {
+            weight,
+            bias: None,
+            config,
+        })
+    }
+
+    fn forward(&self, x: &Tensor) -> candle::Result<Tensor> {
+        let (bsize, _, _) = x.shape().r3()?;
+        let w = self.weight.broadcast_left(bsize)?.t()?;
+        // TODO: Add the conv1d operation
+        let x = x.matmul(&w)?;
+        match &self.bias {
+            None => Ok(x),
+            Some(bias) => x.broadcast_add(bias),
+        }
+    }
+}
+
 struct Dropout {
     pr: f64,
 }
@@ -341,8 +414,8 @@ fn sinusoids(length: usize, channels: usize) -> Result<Tensor> {
 
 // https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L143
 struct AudioEncoder {
-    conv1: Linear, // TODO
-    conv2: Linear, // TODO
+    conv1: Conv1D,
+    conv2: Conv1D,
     positional_embedding: Tensor,
     blocks: Vec<ResidualAttentionBlock>,
     ln_post: LayerNorm,
@@ -352,8 +425,16 @@ impl AudioEncoder {
     fn load(p: &str, vb: &VarBuilder, cfg: &Config) -> Result<Self> {
         let n_state = cfg.n_audio_state;
         let n_head = cfg.n_audio_head;
-        let conv1 = Linear::load(cfg.n_mels, n_state, &format!("{p}.conv1"), vb)?;
-        let conv2 = Linear::load(n_state, n_state, &format!("{p}.conv2"), vb)?;
+        let cfg1 = ConvConfig {
+            padding: 1,
+            stride: 1,
+        };
+        let cfg2 = ConvConfig {
+            padding: 1,
+            stride: 2,
+        };
+        let conv1 = Conv1D::load(cfg.n_mels, n_state, 3, cfg1, &format!("{p}.conv1"), vb)?;
+        let conv2 = Conv1D::load(n_state, n_state, 3, cfg2, &format!("{p}.conv2"), vb)?;
         let positional_embedding = sinusoids(cfg.n_audio_ctx, n_state)?.to_device(&vb.device)?;
         let blocks = (0..cfg.n_audio_layer)
             .map(|i| {

From 3aac1047fec43a4d756ae4e60a8ae82f7c3e636e Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Tue, 4 Jul 2023 10:52:34 +0100
Subject: [PATCH 09/36] Sketch the conv1d op.

---
 candle-core/src/backprop.rs           |  8 +++++++-
 candle-core/src/cpu_backend.rs        | 11 +++++++++++
 candle-core/src/cuda_backend.rs       | 11 +++++++++++
 candle-core/src/dummy_cuda_backend.rs | 11 +++++++++++
 candle-core/src/op.rs                 |  8 ++++++++
 candle-core/src/storage.rs            | 27 +++++++++++++++++++++++++++
 candle-core/src/tensor.rs             | 22 ++++++++++++++++++++++
 7 files changed, 97 insertions(+), 1 deletion(-)

diff --git a/candle-core/src/backprop.rs b/candle-core/src/backprop.rs
index 45448505..a44f732f 100644
--- a/candle-core/src/backprop.rs
+++ b/candle-core/src/backprop.rs
@@ -33,7 +33,12 @@ impl Tensor {
                         track_grad |= tg;
                         nodes
                     }
-                    Op::Add(lhs, rhs)
+                    Op::Conv1D {
+                        arg: lhs,
+                        kernel: rhs,
+                        ..
+                    }
+                    | Op::Add(lhs, rhs)
                     | Op::Mul(lhs, rhs)
                     | Op::Sub(lhs, rhs)
                     | Op::Div(lhs, rhs)
@@ -147,6 +152,7 @@ impl Tensor {
                         let f_grad = pred.where_cond(&zeros, &grad)?;
                         *f_sum_grad = f_sum_grad.add(&f_grad)?;
                     }
+                    Op::Conv1D { .. } => return Err(Error::BackwardNotSupported { op: "conv1d" }),
                     Op::Embedding(_lhs, _rhs) => {
                         return Err(Error::BackwardNotSupported { op: "embedding" })
                     }
diff --git a/candle-core/src/cpu_backend.rs b/candle-core/src/cpu_backend.rs
index 0871175f..ed3a5998 100644
--- a/candle-core/src/cpu_backend.rs
+++ b/candle-core/src/cpu_backend.rs
@@ -627,6 +627,17 @@ impl CpuStorage {
         WCond(pred, layout).map(t, t_l, f, f_l)
     }
 
+    pub(crate) fn conv1d(
+        &self,
+        _l: &Layout,
+        _kernel: &Self,
+        _kernel_l: &Layout,
+        _padding: usize,
+        _stride: usize,
+    ) -> Result<Self> {
+        todo!()
+    }
+
     pub(crate) fn embedding(&self, ids_l: &Layout, rhs: &Self, rhs_l: &Layout) -> Result<Self> {
         let ids = self.as_slice::<u32>()?;
         let (vocab_size, hidden_size) = rhs_l.shape().r2()?;
diff --git a/candle-core/src/cuda_backend.rs b/candle-core/src/cuda_backend.rs
index 0c87004b..ec69688c 100644
--- a/candle-core/src/cuda_backend.rs
+++ b/candle-core/src/cuda_backend.rs
@@ -801,6 +801,17 @@ impl CudaStorage {
         Ok(Self { slice, device })
     }
 
+    pub(crate) fn conv1d(
+        &self,
+        _l: &Layout,
+        _kernel: &Self,
+        _kernel_l: &Layout,
+        _padding: usize,
+        _stride: usize,
+    ) -> Result<Self> {
+        todo!()
+    }
+
     pub(crate) fn embedding(&self, layout: &Layout, rhs: &Self, rhs_l: &Layout) -> Result<Self> {
         let device = self.device().clone();
         let slice = Embedding(self, layout).map(&rhs.slice, &device, rhs_l)?;
diff --git a/candle-core/src/dummy_cuda_backend.rs b/candle-core/src/dummy_cuda_backend.rs
index b025eeab..eca5961b 100644
--- a/candle-core/src/dummy_cuda_backend.rs
+++ b/candle-core/src/dummy_cuda_backend.rs
@@ -100,6 +100,17 @@ impl CudaStorage {
         Err(Error::NotCompiledWithCudaSupport)
     }
 
+    pub(crate) fn conv1d(
+        &self,
+        _l: &Layout,
+        _kernel: &Self,
+        _kernel_l: &Layout,
+        _padding: usize,
+        _stride: usize,
+    ) -> Result<Self> {
+        Err(Error::NotCompiledWithCudaSupport)
+    }
+
     pub(crate) fn embedding(&self, _: &Layout, _: &Self, _: &Layout) -> Result<Self> {
         Err(Error::NotCompiledWithCudaSupport)
     }
diff --git a/candle-core/src/op.rs b/candle-core/src/op.rs
index 860be0b3..ee57b325 100644
--- a/candle-core/src/op.rs
+++ b/candle-core/src/op.rs
@@ -12,6 +12,14 @@ pub(crate) enum Op {
     Embedding(Tensor, Tensor),
     WhereCond(Tensor, Tensor, Tensor),
 
+    #[allow(dead_code)]
+    Conv1D {
+        arg: Tensor,
+        kernel: Tensor,
+        padding: usize,
+        stride: usize,
+    },
+
     Cat(Vec<Tensor>, usize),
 
     #[allow(dead_code)] // add is currently unused.
diff --git a/candle-core/src/storage.rs b/candle-core/src/storage.rs
index 4e630a58..235080c0 100644
--- a/candle-core/src/storage.rs
+++ b/candle-core/src/storage.rs
@@ -144,6 +144,33 @@ impl Storage {
         }
     }
 
+    pub(crate) fn conv1d(
+        &self,
+        l: &Layout,
+        kernel: &Self,
+        kernel_l: &Layout,
+        padding: usize,
+        stride: usize,
+    ) -> Result<Self> {
+        self.same_device(kernel, "conv1d")?;
+        self.same_dtype(kernel, "conv1d")?;
+        match (self, &kernel) {
+            (Storage::Cpu(inp), Storage::Cpu(kernel)) => {
+                let s = inp.conv1d(l, kernel, kernel_l, padding, stride)?;
+                Ok(Self::Cpu(s))
+            }
+            (Storage::Cuda(inp), Storage::Cuda(kernel)) => {
+                let s = inp.conv1d(l, kernel, kernel_l, padding, stride)?;
+                Ok(Self::Cuda(s))
+            }
+            (lhs, rhs) => Err(Error::DeviceMismatchBinaryOp {
+                lhs: lhs.device().location(),
+                rhs: rhs.device().location(),
+                op: "conv1d",
+            }),
+        }
+    }
+
     pub(crate) fn where_cond(
         &self,
         layout: &Layout,
diff --git a/candle-core/src/tensor.rs b/candle-core/src/tensor.rs
index a468d879..26d44718 100644
--- a/candle-core/src/tensor.rs
+++ b/candle-core/src/tensor.rs
@@ -432,6 +432,28 @@ impl Tensor {
         Ok(from_storage(storage, dims, op, false))
     }
 
+    pub fn conv1d(&self, kernel: &Self, padding: usize, stride: usize) -> Result<Self> {
+        let storage = self.storage.conv1d(
+            self.layout(),
+            &kernel.storage,
+            kernel.layout(),
+            padding,
+            stride,
+        )?;
+        let op = if self.track_op() || kernel.track_op() {
+            Some(Op::Conv1D {
+                arg: self.clone(),
+                kernel: kernel.clone(),
+                padding,
+                stride,
+            })
+        } else {
+            None
+        };
+        let dims = self.dims();
+        Ok(from_storage(storage, dims, op, false))
+    }
+
     pub fn matmul(&self, rhs: &Self) -> Result<Self> {
         let a_dims = self.shape().dims();
         let b_dims = rhs.shape().dims();

From a424d95473ea9268ffb1dde4d73ce0cff9904845 Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Tue, 4 Jul 2023 11:15:45 +0100
Subject: [PATCH 10/36] Add more of the conv1d op.

---
 candle-core/src/conv.rs                  | 24 +++++++++++++++++++++
 candle-core/src/cpu_backend.rs           |  3 +--
 candle-core/src/cuda_backend.rs          |  3 +--
 candle-core/src/dummy_cuda_backend.rs    |  3 +--
 candle-core/src/lib.rs                   |  1 +
 candle-core/src/storage.rs               |  7 +++---
 candle-core/src/tensor.rs                | 27 ++++++++++++++++++------
 candle-examples/examples/whisper/main.rs |  3 +--
 8 files changed, 52 insertions(+), 19 deletions(-)
 create mode 100644 candle-core/src/conv.rs

diff --git a/candle-core/src/conv.rs b/candle-core/src/conv.rs
new file mode 100644
index 00000000..90bb5229
--- /dev/null
+++ b/candle-core/src/conv.rs
@@ -0,0 +1,24 @@
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub(crate) struct ParamsConv1D {
+    pub(crate) b_size: Option<usize>,
+    pub(crate) c_out: usize,
+    pub(crate) c_in: usize,
+    pub(crate) k_size: usize,
+    pub(crate) padding: usize,
+    pub(crate) stride: usize,
+}
+
+impl ParamsConv1D {
+    pub(crate) fn l_out(&self, l_in: usize) -> usize {
+        let dilation = 1;
+        (l_in + 2 * self.padding - dilation * (self.k_size - 1) - 1) / self.stride + 1
+    }
+
+    pub(crate) fn out_dims(&self, l_in: usize) -> Vec<usize> {
+        let l_out = self.l_out(l_in);
+        match self.b_size {
+            None => vec![self.c_out, l_out],
+            Some(n) => vec![n, self.c_out, l_out],
+        }
+    }
+}
diff --git a/candle-core/src/cpu_backend.rs b/candle-core/src/cpu_backend.rs
index ed3a5998..54002184 100644
--- a/candle-core/src/cpu_backend.rs
+++ b/candle-core/src/cpu_backend.rs
@@ -632,8 +632,7 @@ impl CpuStorage {
         _l: &Layout,
         _kernel: &Self,
         _kernel_l: &Layout,
-        _padding: usize,
-        _stride: usize,
+        _params: &crate::conv::ParamsConv1D,
     ) -> Result<Self> {
         todo!()
     }
diff --git a/candle-core/src/cuda_backend.rs b/candle-core/src/cuda_backend.rs
index ec69688c..917655fc 100644
--- a/candle-core/src/cuda_backend.rs
+++ b/candle-core/src/cuda_backend.rs
@@ -806,8 +806,7 @@ impl CudaStorage {
         _l: &Layout,
         _kernel: &Self,
         _kernel_l: &Layout,
-        _padding: usize,
-        _stride: usize,
+        _params: &crate::conv::ParamsConv1D,
     ) -> Result<Self> {
         todo!()
     }
diff --git a/candle-core/src/dummy_cuda_backend.rs b/candle-core/src/dummy_cuda_backend.rs
index eca5961b..0dbd8d54 100644
--- a/candle-core/src/dummy_cuda_backend.rs
+++ b/candle-core/src/dummy_cuda_backend.rs
@@ -105,8 +105,7 @@ impl CudaStorage {
         _l: &Layout,
         _kernel: &Self,
         _kernel_l: &Layout,
-        _padding: usize,
-        _stride: usize,
+        _params: &crate::conv::ParamsConv1D,
     ) -> Result<Self> {
         Err(Error::NotCompiledWithCudaSupport)
     }
diff --git a/candle-core/src/lib.rs b/candle-core/src/lib.rs
index 0d4c2a8d..2365a34d 100644
--- a/candle-core/src/lib.rs
+++ b/candle-core/src/lib.rs
@@ -1,4 +1,5 @@
 mod backprop;
+mod conv;
 mod cpu_backend;
 #[cfg(feature = "cuda")]
 mod cuda_backend;
diff --git a/candle-core/src/storage.rs b/candle-core/src/storage.rs
index 235080c0..53ea1544 100644
--- a/candle-core/src/storage.rs
+++ b/candle-core/src/storage.rs
@@ -149,18 +149,17 @@ impl Storage {
         l: &Layout,
         kernel: &Self,
         kernel_l: &Layout,
-        padding: usize,
-        stride: usize,
+        params: &crate::conv::ParamsConv1D,
     ) -> Result<Self> {
         self.same_device(kernel, "conv1d")?;
         self.same_dtype(kernel, "conv1d")?;
         match (self, &kernel) {
             (Storage::Cpu(inp), Storage::Cpu(kernel)) => {
-                let s = inp.conv1d(l, kernel, kernel_l, padding, stride)?;
+                let s = inp.conv1d(l, kernel, kernel_l, params)?;
                 Ok(Self::Cpu(s))
             }
             (Storage::Cuda(inp), Storage::Cuda(kernel)) => {
-                let s = inp.conv1d(l, kernel, kernel_l, padding, stride)?;
+                let s = inp.conv1d(l, kernel, kernel_l, params)?;
                 Ok(Self::Cuda(s))
             }
             (lhs, rhs) => Err(Error::DeviceMismatchBinaryOp {
diff --git a/candle-core/src/tensor.rs b/candle-core/src/tensor.rs
index 26d44718..590b81c4 100644
--- a/candle-core/src/tensor.rs
+++ b/candle-core/src/tensor.rs
@@ -433,13 +433,26 @@ impl Tensor {
     }
 
     pub fn conv1d(&self, kernel: &Self, padding: usize, stride: usize) -> Result<Self> {
-        let storage = self.storage.conv1d(
-            self.layout(),
-            &kernel.storage,
-            kernel.layout(),
+        let (c_out, c_in_k, k_size) = kernel.shape().r3()?;
+        let (b_size, c_in, l_in) = match *self.dims() {
+            [b_size, c_in, l_in] => (Some(b_size), c_in, l_in),
+            [c_in, l_in] => (None, c_in, l_in),
+            _ => todo!("proper error message"),
+        };
+        if c_in != c_in_k {
+            todo!("proper error message")
+        }
+        let params = crate::conv::ParamsConv1D {
+            b_size,
+            c_out,
+            c_in,
+            k_size,
             padding,
             stride,
-        )?;
+        };
+        let storage =
+            self.storage
+                .conv1d(self.layout(), &kernel.storage, kernel.layout(), &params)?;
         let op = if self.track_op() || kernel.track_op() {
             Some(Op::Conv1D {
                 arg: self.clone(),
@@ -450,8 +463,8 @@ impl Tensor {
         } else {
             None
         };
-        let dims = self.dims();
-        Ok(from_storage(storage, dims, op, false))
+        let out_dims = params.out_dims(l_in);
+        Ok(from_storage(storage, out_dims, op, false))
     }
 
     pub fn matmul(&self, rhs: &Self) -> Result<Self> {
diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs
index 75ab2189..a380d30e 100644
--- a/candle-examples/examples/whisper/main.rs
+++ b/candle-examples/examples/whisper/main.rs
@@ -236,8 +236,7 @@ impl Conv1D {
     fn forward(&self, x: &Tensor) -> candle::Result<Tensor> {
         let (bsize, _, _) = x.shape().r3()?;
         let w = self.weight.broadcast_left(bsize)?.t()?;
-        // TODO: Add the conv1d operation
-        let x = x.matmul(&w)?;
+        let x = x.conv1d(&w, self.config.padding, self.config.stride)?;
         match &self.bias {
             None => Ok(x),
             Some(bias) => x.broadcast_add(bias),

From 950b4af49e56b640b87eb273e839b2fd466e1424 Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Tue, 4 Jul 2023 11:29:28 +0100
Subject: [PATCH 11/36] Proper conv1d dispatch.

---
 candle-core/src/conv.rs        | 11 +++++++----
 candle-core/src/cpu_backend.rs | 30 +++++++++++++++++++++++++-----
 candle-core/src/tensor.rs      |  3 ++-
 3 files changed, 34 insertions(+), 10 deletions(-)

diff --git a/candle-core/src/conv.rs b/candle-core/src/conv.rs
index 90bb5229..041bb6fb 100644
--- a/candle-core/src/conv.rs
+++ b/candle-core/src/conv.rs
@@ -1,6 +1,9 @@
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub(crate) struct ParamsConv1D {
     pub(crate) b_size: Option<usize>,
+    // Maybe we should have a version without l_in as this bit depends on the input and not only on
+    // the weights.
+    pub(crate) l_in: usize,
     pub(crate) c_out: usize,
     pub(crate) c_in: usize,
     pub(crate) k_size: usize,
@@ -9,13 +12,13 @@ pub(crate) struct ParamsConv1D {
 }
 
 impl ParamsConv1D {
-    pub(crate) fn l_out(&self, l_in: usize) -> usize {
+    pub(crate) fn l_out(&self) -> usize {
         let dilation = 1;
-        (l_in + 2 * self.padding - dilation * (self.k_size - 1) - 1) / self.stride + 1
+        (self.l_in + 2 * self.padding - dilation * (self.k_size - 1) - 1) / self.stride + 1
     }
 
-    pub(crate) fn out_dims(&self, l_in: usize) -> Vec<usize> {
-        let l_out = self.l_out(l_in);
+    pub(crate) fn out_dims(&self) -> Vec<usize> {
+        let l_out = self.l_out();
         match self.b_size {
             None => vec![self.c_out, l_out],
             Some(n) => vec![n, self.c_out, l_out],
diff --git a/candle-core/src/cpu_backend.rs b/candle-core/src/cpu_backend.rs
index 54002184..718b071c 100644
--- a/candle-core/src/cpu_backend.rs
+++ b/candle-core/src/cpu_backend.rs
@@ -202,6 +202,26 @@ fn copy_strided_src_<T: Copy + std::fmt::Display>(
     }
 }
 
+struct Conv1D<'a>(&'a crate::conv::ParamsConv1D);
+
+impl<'a> Map2 for Conv1D<'a> {
+    const OP: &'static str = "conv1d";
+    fn f<T: 'static + num_traits::Num + Copy>(
+        &self,
+        _inp: &[T],
+        _inp_l: &Layout,
+        _k: &[T],
+        _k_l: &Layout,
+    ) -> Result<Vec<T>> {
+        let p = self.0;
+        let l_out = p.l_out();
+        let out_elems = p.c_out * l_out * p.b_size.unwrap_or(1);
+        let dst = vec![T::zero(); out_elems];
+        // TODO: actually implement the ops.
+        Ok(dst)
+    }
+}
+
 struct MatMul((usize, usize, usize, usize));
 
 impl Map2 for MatMul {
@@ -629,12 +649,12 @@ impl CpuStorage {
 
     pub(crate) fn conv1d(
         &self,
-        _l: &Layout,
-        _kernel: &Self,
-        _kernel_l: &Layout,
-        _params: &crate::conv::ParamsConv1D,
+        l: &Layout,
+        kernel: &Self,
+        kernel_l: &Layout,
+        params: &crate::conv::ParamsConv1D,
     ) -> Result<Self> {
-        todo!()
+        Conv1D(params).map(self, l, kernel, kernel_l)
     }
 
     pub(crate) fn embedding(&self, ids_l: &Layout, rhs: &Self, rhs_l: &Layout) -> Result<Self> {
diff --git a/candle-core/src/tensor.rs b/candle-core/src/tensor.rs
index 590b81c4..25ab0a9b 100644
--- a/candle-core/src/tensor.rs
+++ b/candle-core/src/tensor.rs
@@ -444,6 +444,7 @@ impl Tensor {
         }
         let params = crate::conv::ParamsConv1D {
             b_size,
+            l_in,
             c_out,
             c_in,
             k_size,
@@ -463,7 +464,7 @@ impl Tensor {
         } else {
             None
         };
-        let out_dims = params.out_dims(l_in);
+        let out_dims = params.out_dims();
         Ok(from_storage(storage, out_dims, op, false))
     }
 

From aea090401ded5789e95f1f8efb7404a66b508356 Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Tue, 4 Jul 2023 12:03:28 +0100
Subject: [PATCH 12/36] Run the text decoding bit.

---
 .../examples/whisper/extract_weights.py       |  2 +-
 candle-examples/examples/whisper/main.rs      | 39 +++++++++++--------
 2 files changed, 24 insertions(+), 17 deletions(-)

diff --git a/candle-examples/examples/whisper/extract_weights.py b/candle-examples/examples/whisper/extract_weights.py
index d6ccffc6..65602703 100644
--- a/candle-examples/examples/whisper/extract_weights.py
+++ b/candle-examples/examples/whisper/extract_weights.py
@@ -8,6 +8,6 @@ data = torch.load("tiny.en.pt")
 weights = {}
 for k, v in data["model_state_dict"].items():
     weights[k] = v.contiguous()
-    print(k, v.shape)
+    print(k, v.shape, v.dtype)
 save_file(weights, "tiny.en.safetensors")
 print(data["dims"])
diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs
index a380d30e..c3a2769f 100644
--- a/candle-examples/examples/whisper/main.rs
+++ b/candle-examples/examples/whisper/main.rs
@@ -3,7 +3,7 @@
 // TODO:
 // - kv-cache support?
 
-use anyhow::{Error as E, Result};
+use anyhow::Result;
 use candle::{safetensors::SafeTensors, DType, Device, Shape, Tensor};
 use clap::Parser;
 use std::collections::HashMap;
@@ -97,7 +97,7 @@ struct Config {
 }
 
 impl Config {
-    fn tiny() -> Self {
+    fn tiny_en() -> Self {
         Self {
             n_mels: 80,
             n_vocab: 51864,
@@ -302,8 +302,8 @@ struct MultiHeadAttention {
 impl MultiHeadAttention {
     fn load(n_state: usize, n_head: usize, p: &str, vb: &VarBuilder) -> Result<Self> {
         let query = Linear::load(n_state, n_state, &format!("{p}.query"), vb)?;
-        let value = Linear::load_no_bias(n_state, n_state, &format!("{p}.value"), vb)?;
-        let key = Linear::load(n_state, n_state, &format!("{p}.key"), vb)?;
+        let value = Linear::load(n_state, n_state, &format!("{p}.value"), vb)?;
+        let key = Linear::load_no_bias(n_state, n_state, &format!("{p}.key"), vb)?;
         let out = Linear::load(n_state, n_state, &format!("{p}.out"), vb)?;
         Ok(Self {
             query,
@@ -500,12 +500,13 @@ impl TextDecoder {
         let last = x_dims[x_dims.len() - 1];
         let token_embedding = self.token_embedding.forward(x)?;
         let positional_embedding = self.positional_embedding.narrow(0, 0, last)?;
-        let mut x = (token_embedding + positional_embedding)?;
+        let mut x = token_embedding.broadcast_add(&positional_embedding)?;
         for block in self.blocks.iter() {
             x = block.forward(&x, Some(xa))?
         }
         let x = self.ln.forward(&x)?;
-        let logits = x.matmul(&self.token_embedding.embeddings.t()?)?;
+        let w = self.token_embedding.embeddings.broadcast_left(x_dims[0])?;
+        let logits = x.matmul(&w.t()?)?;
         Ok(logits)
     }
 }
@@ -517,9 +518,9 @@ struct Whisper {
 }
 
 impl Whisper {
-    fn load(p: &str, vb: &VarBuilder, cfg: &Config) -> Result<Self> {
-        let encoder = AudioEncoder::load(&format!("{p}.encoder"), vb, cfg)?;
-        let decoder = TextDecoder::load(&format!("{p}.decoder"), vb, cfg)?;
+    fn load(vb: &VarBuilder, cfg: &Config) -> Result<Self> {
+        let encoder = AudioEncoder::load("encoder", vb, cfg)?;
+        let decoder = TextDecoder::load("decoder", vb, cfg)?;
         Ok(Self { encoder, decoder })
     }
     fn forward(&self, mel: &Tensor, tokens: &Tensor) -> Result<Tensor> {
@@ -537,15 +538,13 @@ struct Args {
     cpu: bool,
 
     #[arg(long)]
-    tokenizer_config: String,
+    weights: String,
 
     #[arg(long)]
-    weights: String,
+    input: String,
 }
 
 fn main() -> Result<()> {
-    use tokenizers::Tokenizer;
-
     let args = Args::parse();
     let device = if args.cpu {
         Device::Cpu
@@ -553,11 +552,19 @@ fn main() -> Result<()> {
         Device::new_cuda(0)?
     };
 
-    let mut tokenizer = Tokenizer::from_file(args.tokenizer_config).map_err(E::msg)?;
-    let _tokenizer = tokenizer.with_padding(None).with_truncation(None);
+    let input = unsafe { candle::safetensors::MmapedFile::new(args.input)? };
+    let input = input.deserialize()?;
+    let x = input.tensor("x", &device)?.to_dtype(DType::U32)?;
+    let xa = input.tensor("xa", &device)?;
 
     let weights = unsafe { candle::safetensors::MmapedFile::new(args.weights)? };
     let weights = weights.deserialize()?;
-    let _vb = VarBuilder::from_safetensors(vec![weights], DTYPE, device);
+    let vb = VarBuilder::from_safetensors(vec![weights], DTYPE, device.clone());
+    let cfg = Config::tiny_en();
+
+    let model = Whisper::load(&vb, &cfg)?;
+    let logits = model.decoder.forward(&x, &xa)?;
+    println!("{logits}");
+    println!("python logits: {}", input.tensor("logits", &device)?);
     Ok(())
 }

From 0d99b4379224838265bf9bedd836dd655b94a001 Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Tue, 4 Jul 2023 12:27:05 +0100
Subject: [PATCH 13/36] Line up the textdecoder values with the python
 implementation.

---
 candle-examples/examples/whisper/main.rs | 36 ++++++++++++------------
 1 file changed, 18 insertions(+), 18 deletions(-)

diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs
index c3a2769f..6341c5ee 100644
--- a/candle-examples/examples/whisper/main.rs
+++ b/candle-examples/examples/whisper/main.rs
@@ -323,13 +323,18 @@ impl MultiHeadAttention {
         Ok(out)
     }
 
-    fn qkv_attention(&self, q: &Tensor, k: &Tensor, v: &Tensor) -> Result<Tensor> {
-        let (n_batch, n_ctx, n_state) = q.shape().r3()?;
+    fn reshape_head(&self, x: &Tensor) -> Result<Tensor> {
+        let (n_batch, n_ctx, n_state) = x.shape().r3()?;
         let target_dims = &[n_batch, n_ctx, self.n_head, n_state / self.n_head];
+        Ok(x.reshape(target_dims)?.transpose(1, 2)?)
+    }
+
+    fn qkv_attention(&self, q: &Tensor, k: &Tensor, v: &Tensor) -> Result<Tensor> {
+        let (_, _, n_state) = q.shape().r3()?;
         let scale = ((n_state / self.n_head) as f64).powf(-0.25);
-        let q = (q.reshape(target_dims)?.transpose(1, 2)? * scale)?;
-        let k = (k.reshape(target_dims)?.transpose(1, 2)?.transpose(2, 3)? * scale)?;
-        let v = v.reshape(target_dims)?.transpose(1, 2)?;
+        let q = (self.reshape_head(q)? * scale)?;
+        let k = (self.reshape_head(k)?.transpose(2, 3)? * scale)?;
+        let v = self.reshape_head(v)?.contiguous()?;
         let qk = q.matmul(&k)?;
         let w = qk.softmax(qk.rank() - 1)?;
         let wv = w.matmul(&v)?.transpose(1, 2)?.flatten(Some(2), None)?;
@@ -341,8 +346,7 @@ impl MultiHeadAttention {
 struct ResidualAttentionBlock {
     attn: MultiHeadAttention,
     attn_ln: LayerNorm,
-    cross_attn: Option<MultiHeadAttention>,
-    cross_attn_ln: Option<LayerNorm>,
+    cross_attn: Option<(MultiHeadAttention, LayerNorm)>,
     mlp_linear1: Linear,
     mlp_linear2: Linear,
     mlp_ln: LayerNorm,
@@ -352,13 +356,13 @@ impl ResidualAttentionBlock {
     fn load(n_state: usize, n_head: usize, ca: bool, p: &str, vb: &VarBuilder) -> Result<Self> {
         let attn = MultiHeadAttention::load(n_state, n_head, &format!("{p}.attn"), vb)?;
         let attn_ln = LayerNorm::load(n_state, &format!("{p}.attn_ln"), vb)?;
-        let (cross_attn, cross_attn_ln) = if ca {
+        let cross_attn = if ca {
             let cross_attn =
                 MultiHeadAttention::load(n_state, n_head, &format!("{p}.cross_attn"), vb)?;
             let cross_attn_ln = LayerNorm::load(n_state, &format!("{p}.cross_attn_ln"), vb)?;
-            (Some(cross_attn), Some(cross_attn_ln))
+            Some((cross_attn, cross_attn_ln))
         } else {
-            (None, None)
+            None
         };
         let n_mlp = n_state * 4;
         let mlp_linear1 = Linear::load(n_state, n_mlp, &format!("{p}.mlp.0"), vb)?;
@@ -368,7 +372,6 @@ impl ResidualAttentionBlock {
             attn,
             attn_ln,
             cross_attn,
-            cross_attn_ln,
             mlp_linear1,
             mlp_linear2,
             mlp_ln,
@@ -379,11 +382,8 @@ impl ResidualAttentionBlock {
         let attn = self.attn.forward(&self.attn_ln.forward(x)?, None)?;
         let mut x = (x + attn)?;
         // Cross-Attn
-        if let Some(cross_attn_ln) = &self.cross_attn_ln {
-            x = cross_attn_ln.forward(&x)?
-        }
-        if let Some(cross_attn) = &self.cross_attn {
-            x = cross_attn.forward(&x, xa)?
+        if let Some((attn, ln)) = &self.cross_attn {
+            x = (&x + attn.forward(&ln.forward(&x)?, xa)?)?;
         }
         // Mlp
         let mlp = self.mlp_linear2.forward(
@@ -482,7 +482,7 @@ impl TextDecoder {
             vb.get((n_ctx, n_state), &format!("{p}.positional_embedding"))?;
         let blocks = (0..cfg.n_text_layer)
             .map(|i| {
-                ResidualAttentionBlock::load(n_state, n_head, false, &format!("{p}.blocks.{i}"), vb)
+                ResidualAttentionBlock::load(n_state, n_head, true, &format!("{p}.blocks.{i}"), vb)
             })
             .collect::<Result<Vec<_>>>()?;
         let ln = LayerNorm::load(n_state, &format!("{p}.ln"), vb)?;
@@ -502,7 +502,7 @@ impl TextDecoder {
         let positional_embedding = self.positional_embedding.narrow(0, 0, last)?;
         let mut x = token_embedding.broadcast_add(&positional_embedding)?;
         for block in self.blocks.iter() {
-            x = block.forward(&x, Some(xa))?
+            x = block.forward(&x, Some(xa))?;
         }
         let x = self.ln.forward(&x)?;
         let w = self.token_embedding.embeddings.broadcast_left(x_dims[0])?;

From 599160605c0294c94c33f64aeca0ac9f388d03c7 Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Tue, 4 Jul 2023 13:13:28 +0100
Subject: [PATCH 14/36] Use the stored embeddings.

---
 candle-examples/examples/whisper/main.rs | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs
index 6341c5ee..1b6f4bfe 100644
--- a/candle-examples/examples/whisper/main.rs
+++ b/candle-examples/examples/whisper/main.rs
@@ -424,6 +424,7 @@ impl AudioEncoder {
     fn load(p: &str, vb: &VarBuilder, cfg: &Config) -> Result<Self> {
         let n_state = cfg.n_audio_state;
         let n_head = cfg.n_audio_head;
+        let n_ctx = cfg.n_audio_ctx;
         let cfg1 = ConvConfig {
             padding: 1,
             stride: 1,
@@ -434,7 +435,12 @@ impl AudioEncoder {
         };
         let conv1 = Conv1D::load(cfg.n_mels, n_state, 3, cfg1, &format!("{p}.conv1"), vb)?;
         let conv2 = Conv1D::load(n_state, n_state, 3, cfg2, &format!("{p}.conv2"), vb)?;
-        let positional_embedding = sinusoids(cfg.n_audio_ctx, n_state)?.to_device(&vb.device)?;
+        /* The positional embeddings could be regenerated via the following. */
+        let positional_embedding = if true {
+            vb.get((n_ctx, n_state), &format!("{p}.positional_embedding"))?
+        } else {
+            sinusoids(n_ctx, n_state)?.to_device(&vb.device)?
+        };
         let blocks = (0..cfg.n_audio_layer)
             .map(|i| {
                 ResidualAttentionBlock::load(n_state, n_head, false, &format!("{p}.blocks.{i}"), vb)

From 99b83773b51382915c94bd3b2298522036235f72 Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Tue, 4 Jul 2023 13:21:59 +0100
Subject: [PATCH 15/36] Small cleanup.

---
 candle-examples/examples/whisper/main.rs | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs
index 1b6f4bfe..71b03e72 100644
--- a/candle-examples/examples/whisper/main.rs
+++ b/candle-examples/examples/whisper/main.rs
@@ -435,10 +435,10 @@ impl AudioEncoder {
         };
         let conv1 = Conv1D::load(cfg.n_mels, n_state, 3, cfg1, &format!("{p}.conv1"), vb)?;
         let conv2 = Conv1D::load(n_state, n_state, 3, cfg2, &format!("{p}.conv2"), vb)?;
-        /* The positional embeddings could be regenerated via the following. */
         let positional_embedding = if true {
             vb.get((n_ctx, n_state), &format!("{p}.positional_embedding"))?
         } else {
+            /* The positional embeddings could be regenerated via the following. */
             sinusoids(n_ctx, n_state)?.to_device(&vb.device)?
         };
         let blocks = (0..cfg.n_audio_layer)
@@ -474,7 +474,6 @@ struct TextDecoder {
     positional_embedding: Tensor,
     blocks: Vec<ResidualAttentionBlock>,
     ln: LayerNorm,
-    mask: Tensor,
 }
 
 impl TextDecoder {
@@ -492,13 +491,11 @@ impl TextDecoder {
             })
             .collect::<Result<Vec<_>>>()?;
         let ln = LayerNorm::load(n_state, &format!("{p}.ln"), vb)?;
-        let mask = Tensor::new(&[0u32], &vb.device)?; // TODO
         Ok(Self {
             token_embedding,
             positional_embedding,
             blocks,
             ln,
-            mask,
         })
     }
     fn forward(&self, x: &Tensor, xa: &Tensor) -> Result<Tensor> {

From b3d4d0fd0f7c0115cf5c8cea60094abef5536e56 Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Tue, 4 Jul 2023 13:50:41 +0100
Subject: [PATCH 16/36] Very inefficient conv1d implementation.

---
 candle-core/src/cpu_backend.rs | 53 +++++++++++++++++++++++++++++-----
 1 file changed, 45 insertions(+), 8 deletions(-)

diff --git a/candle-core/src/cpu_backend.rs b/candle-core/src/cpu_backend.rs
index 718b071c..4eb57bc7 100644
--- a/candle-core/src/cpu_backend.rs
+++ b/candle-core/src/cpu_backend.rs
@@ -206,18 +206,55 @@ struct Conv1D<'a>(&'a crate::conv::ParamsConv1D);
 
 impl<'a> Map2 for Conv1D<'a> {
     const OP: &'static str = "conv1d";
-    fn f<T: 'static + num_traits::Num + Copy>(
+    fn f<T: 'static + num_traits::NumAssign + Copy>(
         &self,
-        _inp: &[T],
-        _inp_l: &Layout,
-        _k: &[T],
-        _k_l: &Layout,
+        inp: &[T],
+        inp_l: &Layout,
+        k: &[T],
+        k_l: &Layout,
     ) -> Result<Vec<T>> {
+        // TODO: Optimize this (proper algorithm, simd, multithread, remove bound checks, etc).
         let p = self.0;
+        let inp = &inp[inp_l.start_offset()..];
+        let k = &k[k_l.start_offset()..];
+        let inp_stride = inp_l.stride();
+        let (inp_stride0, inp_stride) = if inp_stride.len() == 3 {
+            (inp_stride[0], &inp_stride[1..])
+        } else {
+            (0, inp_stride) // This value never gets used anyway
+        };
+        let k_stride = k_l.stride();
+        let k_over_2 = p.k_size / 2;
         let l_out = p.l_out();
-        let out_elems = p.c_out * l_out * p.b_size.unwrap_or(1);
-        let dst = vec![T::zero(); out_elems];
-        // TODO: actually implement the ops.
+        let dst_elems = p.c_out * l_out * p.b_size.unwrap_or(1);
+        let mut dst = vec![T::zero(); dst_elems];
+        // The output shape is [b_size, c_out, l_out]
+        for b_idx in 0..p.b_size.unwrap_or(1) {
+            let inp_idx = b_idx * inp_stride0;
+            let dst_idx = b_idx * p.c_out * l_out;
+            for dst_c_idx in 0..p.c_out {
+                let dst_idx = dst_idx + dst_c_idx * l_out;
+                for dst_l in 0..l_out {
+                    let dst_idx = dst_idx + dst_l;
+                    let mut d = T::zero();
+                    for offset in 0..p.k_size {
+                        // inp[bidx, src_c_idx, dst_l + offset - k//2] * k[dst_c_idx, src_c_idx, offset]
+                        if k_over_2 <= dst_l + offset && dst_l + offset < k_over_2 + p.l_in {
+                            let src_l = dst_l + offset - k_over_2;
+                            for src_c_idx in 0..p.c_in {
+                                let inp_idx =
+                                    inp_idx + src_c_idx * inp_stride[0] + src_l * inp_stride[1];
+                                let k_idx = dst_c_idx * k_stride[0]
+                                    + src_c_idx * k_stride[1]
+                                    + offset * k_stride[2];
+                                d += inp[inp_idx] * k[k_idx]
+                            }
+                        }
+                    }
+                    dst[dst_idx] = d
+                }
+            }
+        }
         Ok(dst)
     }
 }

From c3739d001bfb1a1305fc1ff398761e380f87bfb1 Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Tue, 4 Jul 2023 14:06:09 +0100
Subject: [PATCH 17/36] Get the audio-encoder to return some values.

---
 candle-examples/examples/whisper/main.rs | 22 +++++++++++-----------
 1 file changed, 11 insertions(+), 11 deletions(-)

diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs
index 71b03e72..839dfc13 100644
--- a/candle-examples/examples/whisper/main.rs
+++ b/candle-examples/examples/whisper/main.rs
@@ -233,13 +233,15 @@ impl Conv1D {
         })
     }
 
-    fn forward(&self, x: &Tensor) -> candle::Result<Tensor> {
-        let (bsize, _, _) = x.shape().r3()?;
-        let w = self.weight.broadcast_left(bsize)?.t()?;
-        let x = x.conv1d(&w, self.config.padding, self.config.stride)?;
+    fn forward(&self, x: &Tensor) -> Result<Tensor> {
+        let x = x.conv1d(&self.weight, self.config.padding, self.config.stride)?;
         match &self.bias {
             None => Ok(x),
-            Some(bias) => x.broadcast_add(bias),
+            Some(bias) => {
+                let b = bias.shape().r1()?;
+                let bias = bias.reshape((1, b, 1))?;
+                Ok(x.broadcast_add(&bias)?)
+            }
         }
     }
 }
@@ -381,11 +383,9 @@ impl ResidualAttentionBlock {
     fn forward(&self, x: &Tensor, xa: Option<&Tensor>) -> Result<Tensor> {
         let attn = self.attn.forward(&self.attn_ln.forward(x)?, None)?;
         let mut x = (x + attn)?;
-        // Cross-Attn
         if let Some((attn, ln)) = &self.cross_attn {
             x = (&x + attn.forward(&ln.forward(&x)?, xa)?)?;
         }
-        // Mlp
         let mlp = self.mlp_linear2.forward(
             &self
                 .mlp_linear1
@@ -557,8 +557,8 @@ fn main() -> Result<()> {
 
     let input = unsafe { candle::safetensors::MmapedFile::new(args.input)? };
     let input = input.deserialize()?;
-    let x = input.tensor("x", &device)?.to_dtype(DType::U32)?;
-    let xa = input.tensor("xa", &device)?;
+    let tokens = input.tensor("tokens", &device)?.to_dtype(DType::U32)?;
+    let mel = input.tensor("mel", &device)?;
 
     let weights = unsafe { candle::safetensors::MmapedFile::new(args.weights)? };
     let weights = weights.deserialize()?;
@@ -566,8 +566,8 @@ fn main() -> Result<()> {
     let cfg = Config::tiny_en();
 
     let model = Whisper::load(&vb, &cfg)?;
-    let logits = model.decoder.forward(&x, &xa)?;
+    let logits = model.forward(&mel, &tokens)?;
     println!("{logits}");
-    println!("python logits: {}", input.tensor("logits", &device)?);
+    println!("python logits: {}", input.tensor("dec", &device)?);
     Ok(())
 }

From 459e2e1ae34c624ae9f54aac5430f9ed7e1666f3 Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Tue, 4 Jul 2023 15:05:04 +0100
Subject: [PATCH 18/36] Properly handle the stride in conv1d.

---
 candle-core/src/cpu_backend.rs           |  5 +++--
 candle-examples/examples/whisper/main.rs | 13 ++++++-------
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/candle-core/src/cpu_backend.rs b/candle-core/src/cpu_backend.rs
index 4eb57bc7..b2345756 100644
--- a/candle-core/src/cpu_backend.rs
+++ b/candle-core/src/cpu_backend.rs
@@ -238,9 +238,10 @@ impl<'a> Map2 for Conv1D<'a> {
                     let dst_idx = dst_idx + dst_l;
                     let mut d = T::zero();
                     for offset in 0..p.k_size {
+                        let src_l_plus = p.stride * dst_l + offset;
                         // inp[bidx, src_c_idx, dst_l + offset - k//2] * k[dst_c_idx, src_c_idx, offset]
-                        if k_over_2 <= dst_l + offset && dst_l + offset < k_over_2 + p.l_in {
-                            let src_l = dst_l + offset - k_over_2;
+                        if k_over_2 <= src_l_plus && src_l_plus < k_over_2 + p.l_in {
+                            let src_l = src_l_plus - k_over_2;
                             for src_c_idx in 0..p.c_in {
                                 let inp_idx =
                                     inp_idx + src_c_idx * inp_stride[0] + src_l * inp_stride[1];
diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs
index 839dfc13..d119b6a7 100644
--- a/candle-examples/examples/whisper/main.rs
+++ b/candle-examples/examples/whisper/main.rs
@@ -435,12 +435,7 @@ impl AudioEncoder {
         };
         let conv1 = Conv1D::load(cfg.n_mels, n_state, 3, cfg1, &format!("{p}.conv1"), vb)?;
         let conv2 = Conv1D::load(n_state, n_state, 3, cfg2, &format!("{p}.conv2"), vb)?;
-        let positional_embedding = if true {
-            vb.get((n_ctx, n_state), &format!("{p}.positional_embedding"))?
-        } else {
-            /* The positional embeddings could be regenerated via the following. */
-            sinusoids(n_ctx, n_state)?.to_device(&vb.device)?
-        };
+        let positional_embedding = sinusoids(n_ctx, n_state)?.to_device(&vb.device)?;
         let blocks = (0..cfg.n_audio_layer)
             .map(|i| {
                 ResidualAttentionBlock::load(n_state, n_head, false, &format!("{p}.blocks.{i}"), vb)
@@ -567,7 +562,11 @@ fn main() -> Result<()> {
 
     let model = Whisper::load(&vb, &cfg)?;
     let logits = model.forward(&mel, &tokens)?;
-    println!("{logits}");
+    println!("tokens\n{tokens}");
+    println!("logits:\n{logits}");
     println!("python logits: {}", input.tensor("dec", &device)?);
+    let enc = model.encoder.forward(&mel)?;
+    println!("encoder:\n{enc}");
+    println!("python enc: {}", input.tensor("enc", &device)?);
     Ok(())
 }

From 31663bc04f25a7432befb491893dbb691597c46f Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Tue, 4 Jul 2023 15:25:47 +0100
Subject: [PATCH 19/36] Add the causal mask in text decoding.

---
 candle-examples/examples/whisper/main.rs | 37 +++++++++++++++++-------
 1 file changed, 27 insertions(+), 10 deletions(-)

diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs
index d119b6a7..8f72fc19 100644
--- a/candle-examples/examples/whisper/main.rs
+++ b/candle-examples/examples/whisper/main.rs
@@ -316,11 +316,11 @@ impl MultiHeadAttention {
         })
     }
 
-    fn forward(&self, x: &Tensor, xa: Option<&Tensor>) -> Result<Tensor> {
+    fn forward(&self, x: &Tensor, xa: Option<&Tensor>, mask: Option<&Tensor>) -> Result<Tensor> {
         let q = self.query.forward(x)?;
         let k = self.key.forward(xa.unwrap_or(x))?;
         let v = self.value.forward(xa.unwrap_or(x))?;
-        let wv = self.qkv_attention(&q, &k, &v)?;
+        let wv = self.qkv_attention(&q, &k, &v, mask)?;
         let out = self.out.forward(&wv)?;
         Ok(out)
     }
@@ -331,13 +331,23 @@ impl MultiHeadAttention {
         Ok(x.reshape(target_dims)?.transpose(1, 2)?)
     }
 
-    fn qkv_attention(&self, q: &Tensor, k: &Tensor, v: &Tensor) -> Result<Tensor> {
-        let (_, _, n_state) = q.shape().r3()?;
+    fn qkv_attention(
+        &self,
+        q: &Tensor,
+        k: &Tensor,
+        v: &Tensor,
+        mask: Option<&Tensor>,
+    ) -> Result<Tensor> {
+        let (_, n_ctx, n_state) = q.shape().r3()?;
         let scale = ((n_state / self.n_head) as f64).powf(-0.25);
         let q = (self.reshape_head(q)? * scale)?;
         let k = (self.reshape_head(k)?.transpose(2, 3)? * scale)?;
         let v = self.reshape_head(v)?.contiguous()?;
-        let qk = q.matmul(&k)?;
+        let mut qk = q.matmul(&k)?;
+        if let Some(mask) = mask {
+            let mask = mask.narrow(0, 0, n_ctx)?.narrow(1, 0, n_ctx)?;
+            qk = qk.broadcast_add(&mask)?
+        }
         let w = qk.softmax(qk.rank() - 1)?;
         let wv = w.matmul(&v)?.transpose(1, 2)?.flatten(Some(2), None)?;
         Ok(wv)
@@ -380,11 +390,11 @@ impl ResidualAttentionBlock {
         })
     }
 
-    fn forward(&self, x: &Tensor, xa: Option<&Tensor>) -> Result<Tensor> {
-        let attn = self.attn.forward(&self.attn_ln.forward(x)?, None)?;
+    fn forward(&self, x: &Tensor, xa: Option<&Tensor>, mask: Option<&Tensor>) -> Result<Tensor> {
+        let attn = self.attn.forward(&self.attn_ln.forward(x)?, None, mask)?;
         let mut x = (x + attn)?;
         if let Some((attn, ln)) = &self.cross_attn {
-            x = (&x + attn.forward(&ln.forward(&x)?, xa)?)?;
+            x = (&x + attn.forward(&ln.forward(&x)?, xa, None)?)?;
         }
         let mlp = self.mlp_linear2.forward(
             &self
@@ -456,7 +466,7 @@ impl AudioEncoder {
         let x = x.transpose(1, 2)?;
         let mut x = x.broadcast_add(&self.positional_embedding)?;
         for block in self.blocks.iter() {
-            x = block.forward(&x, None)?
+            x = block.forward(&x, None, None)?
         }
         let x = self.ln_post.forward(&x)?;
         Ok(x)
@@ -469,6 +479,7 @@ struct TextDecoder {
     positional_embedding: Tensor,
     blocks: Vec<ResidualAttentionBlock>,
     ln: LayerNorm,
+    mask: Tensor,
 }
 
 impl TextDecoder {
@@ -486,11 +497,17 @@ impl TextDecoder {
             })
             .collect::<Result<Vec<_>>>()?;
         let ln = LayerNorm::load(n_state, &format!("{p}.ln"), vb)?;
+        let mask: Vec<_> = (0..n_ctx)
+            .flat_map(|i| (0..n_ctx).map(move |j| if j > i { f32::NEG_INFINITY } else { 0f32 }))
+            .collect();
+        let mask = Tensor::from_vec(mask, (n_ctx, n_ctx), &vb.device)?;
+
         Ok(Self {
             token_embedding,
             positional_embedding,
             blocks,
             ln,
+            mask,
         })
     }
     fn forward(&self, x: &Tensor, xa: &Tensor) -> Result<Tensor> {
@@ -500,7 +517,7 @@ impl TextDecoder {
         let positional_embedding = self.positional_embedding.narrow(0, 0, last)?;
         let mut x = token_embedding.broadcast_add(&positional_embedding)?;
         for block in self.blocks.iter() {
-            x = block.forward(&x, Some(xa))?;
+            x = block.forward(&x, Some(xa), Some(&self.mask))?;
         }
         let x = self.ln.forward(&x)?;
         let w = self.token_embedding.embeddings.broadcast_left(x_dims[0])?;

From babf3b40653e6ad2b0adbd80e169c8a5dc8651bf Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Tue, 4 Jul 2023 18:25:47 +0100
Subject: [PATCH 20/36] Sketch the decode step for whisper.

---
 candle-examples/examples/whisper/main.rs | 115 +++++++++++++++++++----
 1 file changed, 99 insertions(+), 16 deletions(-)

diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs
index 8f72fc19..5b4ab4ce 100644
--- a/candle-examples/examples/whisper/main.rs
+++ b/candle-examples/examples/whisper/main.rs
@@ -2,14 +2,32 @@
 // https://github.com/openai/whisper/blob/main/whisper/model.py
 // TODO:
 // - kv-cache support?
+// - language detection?
 
-use anyhow::Result;
+use anyhow::{Error as E, Result};
 use candle::{safetensors::SafeTensors, DType, Device, Shape, Tensor};
 use clap::Parser;
 use std::collections::HashMap;
 
 const DTYPE: DType = DType::F32;
 
+// Audio parameters.
+const SAMPLE_RATE: usize = 16000;
+const N_FFT: usize = 400;
+const N_MELS: usize = 80;
+const HOP_LENGTH: usize = 160;
+const CHUNK_LENGTH: usize = 30;
+const N_SAMPLES: usize = CHUNK_LENGTH * SAMPLE_RATE; // 480000 samples in a 30-second chunk
+const N_FRAMES: usize = N_SAMPLES / HOP_LENGTH; // 3000 frames in a mel spectrogram input
+const N_SAMPLES_PER_TOKEN: usize = HOP_LENGTH * 2; // the initial convolutions has stride 2
+const FRAMES_PER_SECOND: usize = SAMPLE_RATE / HOP_LENGTH; // 10ms per audio frame
+const TOKENS_PER_SECOND: usize = SAMPLE_RATE / N_SAMPLES_PER_TOKEN; // 20ms per audio token
+
+const NO_SPEECH_THRESHOLD: f64 = 0.6;
+const LOGPROB_THRESHOLD: f64 = -1.0;
+const TEMPERATURES: [f64; 6] = [0.0, 0.2, 0.4, 0.6, 0.8, 1.0];
+const COMPRESSION_RATIO_THRESHOLD: f64 = 2.4;
+
 struct VarBuilder<'a> {
     safetensors: Option<(HashMap<String, usize>, Vec<SafeTensors<'a>>)>,
     dtype: DType,
@@ -510,6 +528,7 @@ impl TextDecoder {
             mask,
         })
     }
+
     fn forward(&self, x: &Tensor, xa: &Tensor) -> Result<Tensor> {
         let x_dims = x.dims();
         let last = x_dims[x_dims.len() - 1];
@@ -530,13 +549,18 @@ impl TextDecoder {
 struct Whisper {
     encoder: AudioEncoder,
     decoder: TextDecoder,
+    config: Config,
 }
 
 impl Whisper {
-    fn load(vb: &VarBuilder, cfg: &Config) -> Result<Self> {
-        let encoder = AudioEncoder::load("encoder", vb, cfg)?;
-        let decoder = TextDecoder::load("decoder", vb, cfg)?;
-        Ok(Self { encoder, decoder })
+    fn load(vb: &VarBuilder, config: Config) -> Result<Self> {
+        let encoder = AudioEncoder::load("encoder", vb, &config)?;
+        let decoder = TextDecoder::load("decoder", vb, &config)?;
+        Ok(Self {
+            encoder,
+            decoder,
+            config,
+        })
     }
     fn forward(&self, mel: &Tensor, tokens: &Tensor) -> Result<Tensor> {
         let enc = self.encoder.forward(mel)?;
@@ -557,9 +581,60 @@ struct Args {
 
     #[arg(long)]
     input: String,
+
+    #[arg(long)]
+    tokenizer_config: String,
+}
+
+#[derive(Debug, Clone)]
+struct DecodingResult {
+    audio_features: Tensor,
+    tokens: Vec<usize>,
+    text: String,
+    avg_logprob: f64,
+    no_speech_prob: f64,
+    temperature: f64,
+    compression_ratio: f64,
+}
+
+#[derive(Debug, Clone)]
+struct Segment {
+    start: f64,
+    end: f64,
+    tokens: Vec<usize>,
+    result: DecodingResult,
+}
+
+fn decode(model: &Whisper, mel: &Tensor, t: f64) -> Result<DecodingResult> {
+    let audio_features = model.encoder.forward(&mel)?;
+    let sample_len = model.config.n_text_ctx / 2;
+    let mut tokens: Vec<u32> = vec![]; // TODO: get initial tokens
+    for i in 0..sample_len {
+        let tokens = Tensor::new(tokens.as_slice(), &mel.device())?;
+        let logits = model.decoder.forward(&tokens, mel)?;
+        // logits
+    }
+    todo!()
+}
+
+fn decode_with_fallback(model: &Whisper, segment: &Tensor) -> Result<DecodingResult> {
+    for (i, &t) in TEMPERATURES.iter().enumerate() {
+        let dr: DecodingResult = decode(model, segment, t)?;
+        if i == TEMPERATURES.len() - 1 {
+            return Ok(dr);
+        }
+        let needs_fallback = dr.compression_ratio > COMPRESSION_RATIO_THRESHOLD
+            || dr.avg_logprob < LOGPROB_THRESHOLD;
+        if !needs_fallback || dr.no_speech_prob > NO_SPEECH_THRESHOLD {
+            return Ok(dr);
+        }
+    }
+    unreachable!()
 }
 
 fn main() -> Result<()> {
+    use tokenizers::Tokenizer;
+
     let args = Args::parse();
     let device = if args.cpu {
         Device::Cpu
@@ -567,23 +642,31 @@ fn main() -> Result<()> {
         Device::new_cuda(0)?
     };
 
+    let tokenizer = Tokenizer::from_file(args.tokenizer_config).map_err(E::msg)?;
+
     let input = unsafe { candle::safetensors::MmapedFile::new(args.input)? };
     let input = input.deserialize()?;
-    let tokens = input.tensor("tokens", &device)?.to_dtype(DType::U32)?;
     let mel = input.tensor("mel", &device)?;
 
     let weights = unsafe { candle::safetensors::MmapedFile::new(args.weights)? };
     let weights = weights.deserialize()?;
-    let vb = VarBuilder::from_safetensors(vec![weights], DTYPE, device.clone());
-    let cfg = Config::tiny_en();
+    let vb = VarBuilder::from_safetensors(vec![weights], DTYPE, device);
+    let model = Whisper::load(&vb, Config::tiny_en())?;
 
-    let model = Whisper::load(&vb, &cfg)?;
-    let logits = model.forward(&mel, &tokens)?;
-    println!("tokens\n{tokens}");
-    println!("logits:\n{logits}");
-    println!("python logits: {}", input.tensor("dec", &device)?);
-    let enc = model.encoder.forward(&mel)?;
-    println!("encoder:\n{enc}");
-    println!("python enc: {}", input.tensor("enc", &device)?);
+    let (_, content_frames) = mel.shape().r2()?;
+    let content_frames = content_frames - N_SAMPLES;
+    let mut seek = 0;
+    while seek < content_frames {
+        let time_offset = (seek * HOP_LENGTH) as f64 / SAMPLE_RATE as f64;
+        let segment_size = usize::min(content_frames - seek, N_FRAMES);
+        let mel_segment = mel.narrow(1, seek, segment_size)?;
+        let segment_duration = (segment_size * HOP_LENGTH) as f64 / SAMPLE_RATE as f64;
+        let dr = decode_with_fallback(&model, &mel_segment)?;
+        if dr.no_speech_prob > NO_SPEECH_THRESHOLD && dr.avg_logprob < LOGPROB_THRESHOLD {
+            seek += segment_size;
+            continue;
+        }
+        //
+    }
     Ok(())
 }

From 80f25e6fbbef4b46876c254692ca08d581a61be1 Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Tue, 4 Jul 2023 21:43:20 +0100
Subject: [PATCH 21/36] More decoding.

---
 candle-examples/examples/whisper/main.rs | 126 +++++++++++++++++------
 1 file changed, 95 insertions(+), 31 deletions(-)

diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs
index 5b4ab4ce..45ea74a0 100644
--- a/candle-examples/examples/whisper/main.rs
+++ b/candle-examples/examples/whisper/main.rs
@@ -7,7 +7,9 @@
 use anyhow::{Error as E, Result};
 use candle::{safetensors::SafeTensors, DType, Device, Shape, Tensor};
 use clap::Parser;
+use rand::{distributions::Distribution, SeedableRng};
 use std::collections::HashMap;
+use tokenizers::Tokenizer;
 
 const DTYPE: DType = DType::F32;
 
@@ -584,12 +586,15 @@ struct Args {
 
     #[arg(long)]
     tokenizer_config: String,
+
+    /// The seed to use when generating random samples.
+    #[arg(long, default_value_t = 299792458)]
+    seed: u64,
 }
 
 #[derive(Debug, Clone)]
 struct DecodingResult {
-    audio_features: Tensor,
-    tokens: Vec<usize>,
+    tokens: Vec<u32>,
     text: String,
     avg_logprob: f64,
     no_speech_prob: f64,
@@ -600,73 +605,132 @@ struct DecodingResult {
 #[derive(Debug, Clone)]
 struct Segment {
     start: f64,
-    end: f64,
-    tokens: Vec<usize>,
-    result: DecodingResult,
+    duration: f64,
+    dr: DecodingResult,
 }
 
-fn decode(model: &Whisper, mel: &Tensor, t: f64) -> Result<DecodingResult> {
-    let audio_features = model.encoder.forward(&mel)?;
-    let sample_len = model.config.n_text_ctx / 2;
-    let mut tokens: Vec<u32> = vec![]; // TODO: get initial tokens
-    for i in 0..sample_len {
-        let tokens = Tensor::new(tokens.as_slice(), &mel.device())?;
-        let logits = model.decoder.forward(&tokens, mel)?;
-        // logits
-    }
-    todo!()
+struct Decode {
+    model: Whisper,
+    rng: rand::rngs::StdRng,
+    tokenizer: Tokenizer,
 }
 
-fn decode_with_fallback(model: &Whisper, segment: &Tensor) -> Result<DecodingResult> {
-    for (i, &t) in TEMPERATURES.iter().enumerate() {
-        let dr: DecodingResult = decode(model, segment, t)?;
-        if i == TEMPERATURES.len() - 1 {
-            return Ok(dr);
-        }
-        let needs_fallback = dr.compression_ratio > COMPRESSION_RATIO_THRESHOLD
-            || dr.avg_logprob < LOGPROB_THRESHOLD;
-        if !needs_fallback || dr.no_speech_prob > NO_SPEECH_THRESHOLD {
-            return Ok(dr);
+impl Decode {
+    fn decode(&mut self, mel: &Tensor, t: f64) -> Result<DecodingResult> {
+        let model = &self.model;
+        let audio_features = model.encoder.forward(mel)?;
+        let sample_len = model.config.n_text_ctx / 2;
+        let mut sum_logprob = 0f64;
+        let no_speech_prob = f64::NAN;
+        let mut tokens: Vec<u32> = vec![]; // TODO: get initial tokens
+        for _i in 0..sample_len {
+            let tokens_t = Tensor::new(tokens.as_slice(), &mel.device())?;
+            let logits = model.decoder.forward(&tokens_t, &audio_features)?;
+            let next_token = if t > 0f64 {
+                let prs = (&logits / t)?.softmax(logits.rank() - 1)?;
+                let logits_v: Vec<f32> = prs.to_vec1()?;
+                let distr = rand::distributions::WeightedIndex::new(&logits_v)?;
+                distr.sample(&mut self.rng) as u32
+            } else {
+                let logits_v: Vec<f32> = logits.to_vec1()?;
+                logits_v
+                    .iter()
+                    .enumerate()
+                    .max_by(|(_, u), (_, v)| u.total_cmp(v))
+                    .map(|(i, _)| i as u32)
+                    .unwrap()
+            };
+            tokens.push(next_token);
+            let prob = logits
+                .softmax(logits.rank() - 1)?
+                .get(next_token as usize)?
+                .to_scalar::<f32>()? as f64;
+            sum_logprob += prob.ln();
+            // 50256 is the eot token, TODO: parameterize this.
+            if next_token == 50256 || tokens.len() > model.config.n_text_ctx {
+                break;
+            }
         }
+        let text = self
+            .tokenizer
+            .decode(tokens.clone(), true)
+            .map_err(E::msg)?;
+        let avg_logprob = sum_logprob / tokens.len() as f64;
+
+        Ok(DecodingResult {
+            tokens,
+            text,
+            avg_logprob,
+            no_speech_prob,
+            temperature: t,
+            compression_ratio: f64::NAN,
+        })
+    }
+
+    fn decode_with_fallback(&mut self, segment: &Tensor) -> Result<DecodingResult> {
+        for (i, &t) in TEMPERATURES.iter().enumerate() {
+            let dr: DecodingResult = self.decode(segment, t)?;
+            if i == TEMPERATURES.len() - 1 {
+                return Ok(dr);
+            }
+            let needs_fallback = dr.compression_ratio > COMPRESSION_RATIO_THRESHOLD
+                || dr.avg_logprob < LOGPROB_THRESHOLD;
+            if !needs_fallback || dr.no_speech_prob > NO_SPEECH_THRESHOLD {
+                return Ok(dr);
+            }
+        }
+        unreachable!()
     }
-    unreachable!()
 }
 
 fn main() -> Result<()> {
-    use tokenizers::Tokenizer;
-
     let args = Args::parse();
     let device = if args.cpu {
         Device::Cpu
     } else {
         Device::new_cuda(0)?
     };
+    let rng = rand::rngs::StdRng::seed_from_u64(args.seed);
 
     let tokenizer = Tokenizer::from_file(args.tokenizer_config).map_err(E::msg)?;
 
     let input = unsafe { candle::safetensors::MmapedFile::new(args.input)? };
     let input = input.deserialize()?;
     let mel = input.tensor("mel", &device)?;
+    println!("loaded mel: {:?}", mel.dims());
+    let mel = if mel.rank() > 2 { mel.squeeze(0)? } else { mel };
 
     let weights = unsafe { candle::safetensors::MmapedFile::new(args.weights)? };
     let weights = weights.deserialize()?;
     let vb = VarBuilder::from_safetensors(vec![weights], DTYPE, device);
     let model = Whisper::load(&vb, Config::tiny_en())?;
+    let mut dc = Decode {
+        model,
+        rng,
+        tokenizer,
+    };
 
     let (_, content_frames) = mel.shape().r2()?;
     let content_frames = content_frames - N_SAMPLES;
     let mut seek = 0;
+    let mut segments = vec![];
     while seek < content_frames {
         let time_offset = (seek * HOP_LENGTH) as f64 / SAMPLE_RATE as f64;
         let segment_size = usize::min(content_frames - seek, N_FRAMES);
         let mel_segment = mel.narrow(1, seek, segment_size)?;
         let segment_duration = (segment_size * HOP_LENGTH) as f64 / SAMPLE_RATE as f64;
-        let dr = decode_with_fallback(&model, &mel_segment)?;
+        let dr = dc.decode_with_fallback(&mel_segment)?;
+        seek += segment_size;
         if dr.no_speech_prob > NO_SPEECH_THRESHOLD && dr.avg_logprob < LOGPROB_THRESHOLD {
-            seek += segment_size;
             continue;
         }
-        //
+        let segment = Segment {
+            start: time_offset,
+            duration: segment_duration,
+            dr,
+        };
+        println!("{seek} {segment:?}");
+        segments.push(segment)
     }
     Ok(())
 }

From 9fe7a428954e774f760830a063376e9fcf228e13 Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Tue, 4 Jul 2023 22:18:07 +0100
Subject: [PATCH 22/36] More whisper sampling.

---
 candle-examples/examples/whisper/main.rs | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs
index 45ea74a0..647ac6f9 100644
--- a/candle-examples/examples/whisper/main.rs
+++ b/candle-examples/examples/whisper/main.rs
@@ -619,13 +619,20 @@ impl Decode {
     fn decode(&mut self, mel: &Tensor, t: f64) -> Result<DecodingResult> {
         let model = &self.model;
         let audio_features = model.encoder.forward(mel)?;
+        println!("audio features: {:?}", audio_features.dims());
         let sample_len = model.config.n_text_ctx / 2;
         let mut sum_logprob = 0f64;
         let no_speech_prob = f64::NAN;
-        let mut tokens: Vec<u32> = vec![]; // TODO: get initial tokens
+        // TODO: 50257 is the start of transcipt token, be more principled about get initial tokens
+        let mut tokens: Vec<u32> = vec![50257];
         for _i in 0..sample_len {
             let tokens_t = Tensor::new(tokens.as_slice(), &mel.device())?;
+            // Insert a batch dim.
+            let tokens_t = tokens_t.unsqueeze(0)?;
             let logits = model.decoder.forward(&tokens_t, &audio_features)?;
+            let logits = logits.squeeze(0)?;
+            let (seq_len, _) = logits.shape().r2()?;
+            let logits = logits.get(seq_len - 1)?;
             let next_token = if t > 0f64 {
                 let prs = (&logits / t)?.softmax(logits.rank() - 1)?;
                 let logits_v: Vec<f32> = prs.to_vec1()?;
@@ -698,7 +705,6 @@ fn main() -> Result<()> {
     let input = input.deserialize()?;
     let mel = input.tensor("mel", &device)?;
     println!("loaded mel: {:?}", mel.dims());
-    let mel = if mel.rank() > 2 { mel.squeeze(0)? } else { mel };
 
     let weights = unsafe { candle::safetensors::MmapedFile::new(args.weights)? };
     let weights = weights.deserialize()?;
@@ -710,14 +716,13 @@ fn main() -> Result<()> {
         tokenizer,
     };
 
-    let (_, content_frames) = mel.shape().r2()?;
-    let content_frames = content_frames - N_SAMPLES;
+    let (_, _, content_frames) = mel.shape().r3()?;
     let mut seek = 0;
     let mut segments = vec![];
     while seek < content_frames {
         let time_offset = (seek * HOP_LENGTH) as f64 / SAMPLE_RATE as f64;
         let segment_size = usize::min(content_frames - seek, N_FRAMES);
-        let mel_segment = mel.narrow(1, seek, segment_size)?;
+        let mel_segment = mel.narrow(2, seek, segment_size)?;
         let segment_duration = (segment_size * HOP_LENGTH) as f64 / SAMPLE_RATE as f64;
         let dr = dc.decode_with_fallback(&mel_segment)?;
         seek += segment_size;

From 8cf803d1a3e32bc929754b3dc7e3a49906648fe8 Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Wed, 5 Jul 2023 05:46:53 +0100
Subject: [PATCH 23/36] Split the model in a separate file.

---
 candle-examples/examples/whisper/main.rs  | 546 +---------------------
 candle-examples/examples/whisper/model.rs | 545 +++++++++++++++++++++
 2 files changed, 548 insertions(+), 543 deletions(-)
 create mode 100644 candle-examples/examples/whisper/model.rs

diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs
index 647ac6f9..8bf31f34 100644
--- a/candle-examples/examples/whisper/main.rs
+++ b/candle-examples/examples/whisper/main.rs
@@ -5,11 +5,12 @@
 // - language detection?
 
 use anyhow::{Error as E, Result};
-use candle::{safetensors::SafeTensors, DType, Device, Shape, Tensor};
+use candle::{DType, Device, Tensor};
 use clap::Parser;
 use rand::{distributions::Distribution, SeedableRng};
-use std::collections::HashMap;
 use tokenizers::Tokenizer;
+mod model;
+use model::{Config, VarBuilder, Whisper};
 
 const DTYPE: DType = DType::F32;
 
@@ -30,547 +31,6 @@ const LOGPROB_THRESHOLD: f64 = -1.0;
 const TEMPERATURES: [f64; 6] = [0.0, 0.2, 0.4, 0.6, 0.8, 1.0];
 const COMPRESSION_RATIO_THRESHOLD: f64 = 2.4;
 
-struct VarBuilder<'a> {
-    safetensors: Option<(HashMap<String, usize>, Vec<SafeTensors<'a>>)>,
-    dtype: DType,
-    device: Device,
-}
-
-impl<'a> VarBuilder<'a> {
-    pub fn from_safetensors(
-        safetensors: Vec<SafeTensors<'a>>,
-        dtype: DType,
-        device: Device,
-    ) -> Self {
-        let mut routing = HashMap::new();
-        for (index, sf) in safetensors.iter().enumerate() {
-            for k in sf.names() {
-                routing.insert(k.to_string(), index);
-            }
-        }
-        Self {
-            safetensors: Some((routing, safetensors)),
-            device,
-            dtype,
-        }
-    }
-
-    pub fn zeros(dtype: DType, device: Device) -> Self {
-        Self {
-            safetensors: None,
-            device,
-            dtype,
-        }
-    }
-
-    pub fn get<S: Into<Shape>>(&self, s: S, tensor_name: &str) -> candle::Result<Tensor> {
-        let s: Shape = s.into();
-        match &self.safetensors {
-            None => Tensor::zeros(s, self.dtype, &self.device),
-            Some((routing, safetensors)) => {
-                // Unwrap or 0  just to let the proper error flow.
-                let index = routing.get(tensor_name).unwrap_or(&0);
-                let tensor = safetensors[*index]
-                    .tensor(tensor_name, &self.device)?
-                    .to_dtype(self.dtype)?;
-                if *tensor.shape() != s {
-                    let msg = format!("shape mismatch for {tensor_name}");
-                    Err(candle::Error::UnexpectedShape {
-                        msg,
-                        expected: s,
-                        got: tensor.shape().clone(),
-                    })?
-                }
-                Ok(tensor)
-            }
-        }
-    }
-}
-
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
-enum HiddenAct {
-    Gelu,
-    Relu,
-}
-
-impl HiddenAct {
-    fn forward(&self, xs: &Tensor) -> candle::Result<Tensor> {
-        match self {
-            Self::Gelu => xs.gelu(),
-            Self::Relu => xs.relu(),
-        }
-    }
-}
-
-#[derive(Debug, Clone, PartialEq)]
-struct Config {
-    n_mels: usize,
-    n_audio_ctx: usize,
-    n_audio_state: usize,
-    n_audio_head: usize,
-    n_audio_layer: usize,
-    n_vocab: usize,
-    n_text_ctx: usize,
-    n_text_state: usize,
-    n_text_head: usize,
-    n_text_layer: usize,
-}
-
-impl Config {
-    fn tiny_en() -> Self {
-        Self {
-            n_mels: 80,
-            n_vocab: 51864,
-            n_audio_ctx: 1500,
-            n_audio_state: 384,
-            n_audio_head: 6,
-            n_audio_layer: 4,
-            n_text_ctx: 448,
-            n_text_state: 384,
-            n_text_head: 6,
-            n_text_layer: 4,
-        }
-    }
-}
-
-struct Embedding {
-    embeddings: Tensor,
-    hidden_size: usize,
-}
-
-impl Embedding {
-    fn new(embeddings: Tensor, hidden_size: usize) -> Self {
-        Self {
-            embeddings,
-            hidden_size,
-        }
-    }
-
-    fn load(vocab_size: usize, hidden_size: usize, p: &str, vb: &VarBuilder) -> Result<Self> {
-        let embeddings = vb.get((vocab_size, hidden_size), &format!("{p}.weight"))?;
-        Ok(Self::new(embeddings, hidden_size))
-    }
-
-    fn forward(&self, indexes: &Tensor) -> Result<Tensor> {
-        let mut final_dims = indexes.dims().to_vec();
-        final_dims.push(self.hidden_size);
-        let indexes = indexes.flatten_all()?;
-        let values = Tensor::embedding(&indexes, &self.embeddings)?;
-        let values = values.reshape(final_dims)?;
-        Ok(values)
-    }
-}
-
-struct Linear {
-    weight: Tensor,
-    bias: Option<Tensor>,
-}
-
-impl Linear {
-    fn load(size1: usize, size2: usize, p: &str, vb: &VarBuilder) -> Result<Self> {
-        let weight = vb.get((size2, size1), &format!("{p}.weight"))?;
-        let bias = vb.get(size2, &format!("{p}.bias"))?;
-        Ok(Self {
-            weight,
-            bias: Some(bias),
-        })
-    }
-
-    fn load_no_bias(size1: usize, size2: usize, p: &str, vb: &VarBuilder) -> Result<Self> {
-        let weight = vb.get((size2, size1), &format!("{p}.weight"))?;
-        Ok(Self { weight, bias: None })
-    }
-
-    fn forward(&self, x: &Tensor) -> candle::Result<Tensor> {
-        let (bsize, _, _) = x.shape().r3()?;
-        let w = self.weight.broadcast_left(bsize)?.t()?;
-        let x = x.matmul(&w)?;
-        match &self.bias {
-            None => Ok(x),
-            Some(bias) => x.broadcast_add(bias),
-        }
-    }
-}
-
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
-struct ConvConfig {
-    padding: usize,
-    stride: usize,
-}
-
-impl Default for ConvConfig {
-    fn default() -> Self {
-        Self {
-            padding: 0,
-            stride: 1,
-        }
-    }
-}
-
-struct Conv1D {
-    weight: Tensor,
-    bias: Option<Tensor>,
-    config: ConvConfig,
-}
-
-impl Conv1D {
-    fn load(
-        in_channels: usize,
-        out_channels: usize,
-        kernel_size: usize,
-        config: ConvConfig,
-        p: &str,
-        vb: &VarBuilder,
-    ) -> Result<Self> {
-        let weight = vb.get(
-            (out_channels, in_channels, kernel_size),
-            &format!("{p}.weight"),
-        )?;
-        let bias = vb.get(out_channels, &format!("{p}.bias"))?;
-        Ok(Self {
-            weight,
-            bias: Some(bias),
-            config,
-        })
-    }
-
-    fn load_no_bias(
-        in_channels: usize,
-        out_channels: usize,
-        kernel_size: usize,
-        config: ConvConfig,
-        p: &str,
-        vb: &VarBuilder,
-    ) -> Result<Self> {
-        let weight = vb.get(
-            (out_channels, in_channels, kernel_size),
-            &format!("{p}.weight"),
-        )?;
-        Ok(Self {
-            weight,
-            bias: None,
-            config,
-        })
-    }
-
-    fn forward(&self, x: &Tensor) -> Result<Tensor> {
-        let x = x.conv1d(&self.weight, self.config.padding, self.config.stride)?;
-        match &self.bias {
-            None => Ok(x),
-            Some(bias) => {
-                let b = bias.shape().r1()?;
-                let bias = bias.reshape((1, b, 1))?;
-                Ok(x.broadcast_add(&bias)?)
-            }
-        }
-    }
-}
-
-struct Dropout {
-    pr: f64,
-}
-
-impl Dropout {
-    fn new(pr: f64) -> Self {
-        Self { pr }
-    }
-
-    fn forward(&self, x: &Tensor) -> Result<Tensor> {
-        // TODO
-        Ok(x.clone())
-    }
-}
-
-// This layer norm version handles both weight and bias so removes the mean.
-struct LayerNorm {
-    weight: Tensor,
-    bias: Tensor,
-    eps: f64,
-}
-
-impl LayerNorm {
-    fn load(size: usize, p: &str, vb: &VarBuilder) -> Result<Self> {
-        let weight = vb.get(size, &format!("{p}.weight"))?;
-        let bias = vb.get(size, &format!("{p}.bias"))?;
-        Ok(Self {
-            weight,
-            bias,
-            eps: 1e-5,
-        })
-    }
-
-    fn forward(&self, x: &Tensor) -> Result<Tensor> {
-        let (_bsize, _seq_len, hidden_size) = x.shape().r3()?;
-        let mean_x = (x.sum(&[2])? / hidden_size as f64)?;
-        let x = x.broadcast_sub(&mean_x)?;
-        let norm_x = ((&x * &x)?.sum(&[2])? / hidden_size as f64)?;
-        let x_normed = x.broadcast_div(&(norm_x + self.eps)?.sqrt()?)?;
-        let x = x_normed
-            .broadcast_mul(&self.weight)?
-            .broadcast_add(&self.bias)?;
-        Ok(x)
-    }
-}
-
-// https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L62
-struct MultiHeadAttention {
-    query: Linear,
-    key: Linear,
-    value: Linear,
-    out: Linear,
-    n_head: usize,
-}
-
-impl MultiHeadAttention {
-    fn load(n_state: usize, n_head: usize, p: &str, vb: &VarBuilder) -> Result<Self> {
-        let query = Linear::load(n_state, n_state, &format!("{p}.query"), vb)?;
-        let value = Linear::load(n_state, n_state, &format!("{p}.value"), vb)?;
-        let key = Linear::load_no_bias(n_state, n_state, &format!("{p}.key"), vb)?;
-        let out = Linear::load(n_state, n_state, &format!("{p}.out"), vb)?;
-        Ok(Self {
-            query,
-            key,
-            value,
-            out,
-            n_head,
-        })
-    }
-
-    fn forward(&self, x: &Tensor, xa: Option<&Tensor>, mask: Option<&Tensor>) -> Result<Tensor> {
-        let q = self.query.forward(x)?;
-        let k = self.key.forward(xa.unwrap_or(x))?;
-        let v = self.value.forward(xa.unwrap_or(x))?;
-        let wv = self.qkv_attention(&q, &k, &v, mask)?;
-        let out = self.out.forward(&wv)?;
-        Ok(out)
-    }
-
-    fn reshape_head(&self, x: &Tensor) -> Result<Tensor> {
-        let (n_batch, n_ctx, n_state) = x.shape().r3()?;
-        let target_dims = &[n_batch, n_ctx, self.n_head, n_state / self.n_head];
-        Ok(x.reshape(target_dims)?.transpose(1, 2)?)
-    }
-
-    fn qkv_attention(
-        &self,
-        q: &Tensor,
-        k: &Tensor,
-        v: &Tensor,
-        mask: Option<&Tensor>,
-    ) -> Result<Tensor> {
-        let (_, n_ctx, n_state) = q.shape().r3()?;
-        let scale = ((n_state / self.n_head) as f64).powf(-0.25);
-        let q = (self.reshape_head(q)? * scale)?;
-        let k = (self.reshape_head(k)?.transpose(2, 3)? * scale)?;
-        let v = self.reshape_head(v)?.contiguous()?;
-        let mut qk = q.matmul(&k)?;
-        if let Some(mask) = mask {
-            let mask = mask.narrow(0, 0, n_ctx)?.narrow(1, 0, n_ctx)?;
-            qk = qk.broadcast_add(&mask)?
-        }
-        let w = qk.softmax(qk.rank() - 1)?;
-        let wv = w.matmul(&v)?.transpose(1, 2)?.flatten(Some(2), None)?;
-        Ok(wv)
-    }
-}
-
-// https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L111
-struct ResidualAttentionBlock {
-    attn: MultiHeadAttention,
-    attn_ln: LayerNorm,
-    cross_attn: Option<(MultiHeadAttention, LayerNorm)>,
-    mlp_linear1: Linear,
-    mlp_linear2: Linear,
-    mlp_ln: LayerNorm,
-}
-
-impl ResidualAttentionBlock {
-    fn load(n_state: usize, n_head: usize, ca: bool, p: &str, vb: &VarBuilder) -> Result<Self> {
-        let attn = MultiHeadAttention::load(n_state, n_head, &format!("{p}.attn"), vb)?;
-        let attn_ln = LayerNorm::load(n_state, &format!("{p}.attn_ln"), vb)?;
-        let cross_attn = if ca {
-            let cross_attn =
-                MultiHeadAttention::load(n_state, n_head, &format!("{p}.cross_attn"), vb)?;
-            let cross_attn_ln = LayerNorm::load(n_state, &format!("{p}.cross_attn_ln"), vb)?;
-            Some((cross_attn, cross_attn_ln))
-        } else {
-            None
-        };
-        let n_mlp = n_state * 4;
-        let mlp_linear1 = Linear::load(n_state, n_mlp, &format!("{p}.mlp.0"), vb)?;
-        let mlp_linear2 = Linear::load(n_mlp, n_state, &format!("{p}.mlp.2"), vb)?;
-        let mlp_ln = LayerNorm::load(n_state, &format!("{p}.mlp_ln"), vb)?;
-        Ok(Self {
-            attn,
-            attn_ln,
-            cross_attn,
-            mlp_linear1,
-            mlp_linear2,
-            mlp_ln,
-        })
-    }
-
-    fn forward(&self, x: &Tensor, xa: Option<&Tensor>, mask: Option<&Tensor>) -> Result<Tensor> {
-        let attn = self.attn.forward(&self.attn_ln.forward(x)?, None, mask)?;
-        let mut x = (x + attn)?;
-        if let Some((attn, ln)) = &self.cross_attn {
-            x = (&x + attn.forward(&ln.forward(&x)?, xa, None)?)?;
-        }
-        let mlp = self.mlp_linear2.forward(
-            &self
-                .mlp_linear1
-                .forward(&self.mlp_ln.forward(&x)?)?
-                .gelu()?,
-        )?;
-        Ok((x + mlp)?)
-    }
-}
-
-fn sinusoids(length: usize, channels: usize) -> Result<Tensor> {
-    let max_timescale = 10000f32;
-    let log_timescale_increment = max_timescale.ln() / (channels / 2 - 1) as f32;
-    let inv_timescales: Vec<_> = (0..channels / 2)
-        .map(|i| (i as f32 * (-log_timescale_increment)).exp())
-        .collect();
-    let arange: Vec<_> = (0..length).map(|c| c as f32).collect();
-    let inv_timescales = Tensor::new(inv_timescales.as_slice(), &Device::Cpu)?.unsqueeze(0)?;
-    let arange = Tensor::new(arange.as_slice(), &Device::Cpu)?.unsqueeze(1)?;
-    let sh = (length, channels / 2);
-    let scaled_time = (arange.broadcast_as(sh)? * inv_timescales.broadcast_as(sh)?)?;
-    let sincos = Tensor::cat(&[scaled_time.sin()?, scaled_time.cos()?], 1)?;
-    Ok(sincos)
-}
-
-// https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L143
-struct AudioEncoder {
-    conv1: Conv1D,
-    conv2: Conv1D,
-    positional_embedding: Tensor,
-    blocks: Vec<ResidualAttentionBlock>,
-    ln_post: LayerNorm,
-}
-
-impl AudioEncoder {
-    fn load(p: &str, vb: &VarBuilder, cfg: &Config) -> Result<Self> {
-        let n_state = cfg.n_audio_state;
-        let n_head = cfg.n_audio_head;
-        let n_ctx = cfg.n_audio_ctx;
-        let cfg1 = ConvConfig {
-            padding: 1,
-            stride: 1,
-        };
-        let cfg2 = ConvConfig {
-            padding: 1,
-            stride: 2,
-        };
-        let conv1 = Conv1D::load(cfg.n_mels, n_state, 3, cfg1, &format!("{p}.conv1"), vb)?;
-        let conv2 = Conv1D::load(n_state, n_state, 3, cfg2, &format!("{p}.conv2"), vb)?;
-        let positional_embedding = sinusoids(n_ctx, n_state)?.to_device(&vb.device)?;
-        let blocks = (0..cfg.n_audio_layer)
-            .map(|i| {
-                ResidualAttentionBlock::load(n_state, n_head, false, &format!("{p}.blocks.{i}"), vb)
-            })
-            .collect::<Result<Vec<_>>>()?;
-        let ln_post = LayerNorm::load(n_state, &format!("{p}.ln_post"), vb)?;
-        Ok(Self {
-            conv1,
-            conv2,
-            positional_embedding,
-            blocks,
-            ln_post,
-        })
-    }
-    fn forward(&self, x: &Tensor) -> Result<Tensor> {
-        let x = self.conv1.forward(x)?.gelu()?;
-        let x = self.conv2.forward(&x)?.gelu()?;
-        let x = x.transpose(1, 2)?;
-        let mut x = x.broadcast_add(&self.positional_embedding)?;
-        for block in self.blocks.iter() {
-            x = block.forward(&x, None, None)?
-        }
-        let x = self.ln_post.forward(&x)?;
-        Ok(x)
-    }
-}
-
-// https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L176
-struct TextDecoder {
-    token_embedding: Embedding,
-    positional_embedding: Tensor,
-    blocks: Vec<ResidualAttentionBlock>,
-    ln: LayerNorm,
-    mask: Tensor,
-}
-
-impl TextDecoder {
-    fn load(p: &str, vb: &VarBuilder, cfg: &Config) -> Result<Self> {
-        let n_state = cfg.n_text_state;
-        let n_head = cfg.n_text_head;
-        let n_ctx = cfg.n_text_ctx;
-        let token_embedding =
-            Embedding::load(cfg.n_vocab, n_state, &format!("{p}.token_embedding"), vb)?;
-        let positional_embedding =
-            vb.get((n_ctx, n_state), &format!("{p}.positional_embedding"))?;
-        let blocks = (0..cfg.n_text_layer)
-            .map(|i| {
-                ResidualAttentionBlock::load(n_state, n_head, true, &format!("{p}.blocks.{i}"), vb)
-            })
-            .collect::<Result<Vec<_>>>()?;
-        let ln = LayerNorm::load(n_state, &format!("{p}.ln"), vb)?;
-        let mask: Vec<_> = (0..n_ctx)
-            .flat_map(|i| (0..n_ctx).map(move |j| if j > i { f32::NEG_INFINITY } else { 0f32 }))
-            .collect();
-        let mask = Tensor::from_vec(mask, (n_ctx, n_ctx), &vb.device)?;
-
-        Ok(Self {
-            token_embedding,
-            positional_embedding,
-            blocks,
-            ln,
-            mask,
-        })
-    }
-
-    fn forward(&self, x: &Tensor, xa: &Tensor) -> Result<Tensor> {
-        let x_dims = x.dims();
-        let last = x_dims[x_dims.len() - 1];
-        let token_embedding = self.token_embedding.forward(x)?;
-        let positional_embedding = self.positional_embedding.narrow(0, 0, last)?;
-        let mut x = token_embedding.broadcast_add(&positional_embedding)?;
-        for block in self.blocks.iter() {
-            x = block.forward(&x, Some(xa), Some(&self.mask))?;
-        }
-        let x = self.ln.forward(&x)?;
-        let w = self.token_embedding.embeddings.broadcast_left(x_dims[0])?;
-        let logits = x.matmul(&w.t()?)?;
-        Ok(logits)
-    }
-}
-
-// https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L221
-struct Whisper {
-    encoder: AudioEncoder,
-    decoder: TextDecoder,
-    config: Config,
-}
-
-impl Whisper {
-    fn load(vb: &VarBuilder, config: Config) -> Result<Self> {
-        let encoder = AudioEncoder::load("encoder", vb, &config)?;
-        let decoder = TextDecoder::load("decoder", vb, &config)?;
-        Ok(Self {
-            encoder,
-            decoder,
-            config,
-        })
-    }
-    fn forward(&self, mel: &Tensor, tokens: &Tensor) -> Result<Tensor> {
-        let enc = self.encoder.forward(mel)?;
-        let dec = self.decoder.forward(tokens, &enc)?;
-        Ok(dec)
-    }
-}
-
 #[derive(Parser, Debug)]
 #[command(author, version, about, long_about = None)]
 struct Args {
diff --git a/candle-examples/examples/whisper/model.rs b/candle-examples/examples/whisper/model.rs
new file mode 100644
index 00000000..94607934
--- /dev/null
+++ b/candle-examples/examples/whisper/model.rs
@@ -0,0 +1,545 @@
+use anyhow::Result;
+use candle::{safetensors::SafeTensors, DType, Device, Shape, Tensor};
+use std::collections::HashMap;
+
+pub struct VarBuilder<'a> {
+    safetensors: Option<(HashMap<String, usize>, Vec<SafeTensors<'a>>)>,
+    dtype: DType,
+    device: Device,
+}
+
+impl<'a> VarBuilder<'a> {
+    pub fn from_safetensors(
+        safetensors: Vec<SafeTensors<'a>>,
+        dtype: DType,
+        device: Device,
+    ) -> Self {
+        let mut routing = HashMap::new();
+        for (index, sf) in safetensors.iter().enumerate() {
+            for k in sf.names() {
+                routing.insert(k.to_string(), index);
+            }
+        }
+        Self {
+            safetensors: Some((routing, safetensors)),
+            device,
+            dtype,
+        }
+    }
+
+    pub fn zeros(dtype: DType, device: Device) -> Self {
+        Self {
+            safetensors: None,
+            device,
+            dtype,
+        }
+    }
+
+    pub fn get<S: Into<Shape>>(&self, s: S, tensor_name: &str) -> candle::Result<Tensor> {
+        let s: Shape = s.into();
+        match &self.safetensors {
+            None => Tensor::zeros(s, self.dtype, &self.device),
+            Some((routing, safetensors)) => {
+                // Unwrap or 0  just to let the proper error flow.
+                let index = routing.get(tensor_name).unwrap_or(&0);
+                let tensor = safetensors[*index]
+                    .tensor(tensor_name, &self.device)?
+                    .to_dtype(self.dtype)?;
+                if *tensor.shape() != s {
+                    let msg = format!("shape mismatch for {tensor_name}");
+                    Err(candle::Error::UnexpectedShape {
+                        msg,
+                        expected: s,
+                        got: tensor.shape().clone(),
+                    })?
+                }
+                Ok(tensor)
+            }
+        }
+    }
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum HiddenAct {
+    Gelu,
+    Relu,
+}
+
+impl HiddenAct {
+    fn forward(&self, xs: &Tensor) -> candle::Result<Tensor> {
+        match self {
+            Self::Gelu => xs.gelu(),
+            Self::Relu => xs.relu(),
+        }
+    }
+}
+
+#[derive(Debug, Clone, PartialEq)]
+pub struct Config {
+    pub n_mels: usize,
+    pub n_audio_ctx: usize,
+    pub n_audio_state: usize,
+    pub n_audio_head: usize,
+    pub n_audio_layer: usize,
+    pub n_vocab: usize,
+    pub n_text_ctx: usize,
+    pub n_text_state: usize,
+    pub n_text_head: usize,
+    pub n_text_layer: usize,
+}
+
+impl Config {
+    pub fn tiny_en() -> Self {
+        Self {
+            n_mels: 80,
+            n_vocab: 51864,
+            n_audio_ctx: 1500,
+            n_audio_state: 384,
+            n_audio_head: 6,
+            n_audio_layer: 4,
+            n_text_ctx: 448,
+            n_text_state: 384,
+            n_text_head: 6,
+            n_text_layer: 4,
+        }
+    }
+}
+
+struct Embedding {
+    embeddings: Tensor,
+    hidden_size: usize,
+}
+
+impl Embedding {
+    fn new(embeddings: Tensor, hidden_size: usize) -> Self {
+        Self {
+            embeddings,
+            hidden_size,
+        }
+    }
+
+    fn load(vocab_size: usize, hidden_size: usize, p: &str, vb: &VarBuilder) -> Result<Self> {
+        let embeddings = vb.get((vocab_size, hidden_size), &format!("{p}.weight"))?;
+        Ok(Self::new(embeddings, hidden_size))
+    }
+
+    fn forward(&self, indexes: &Tensor) -> Result<Tensor> {
+        let mut final_dims = indexes.dims().to_vec();
+        final_dims.push(self.hidden_size);
+        let indexes = indexes.flatten_all()?;
+        let values = Tensor::embedding(&indexes, &self.embeddings)?;
+        let values = values.reshape(final_dims)?;
+        Ok(values)
+    }
+}
+
+struct Linear {
+    weight: Tensor,
+    bias: Option<Tensor>,
+}
+
+impl Linear {
+    fn load(size1: usize, size2: usize, p: &str, vb: &VarBuilder) -> Result<Self> {
+        let weight = vb.get((size2, size1), &format!("{p}.weight"))?;
+        let bias = vb.get(size2, &format!("{p}.bias"))?;
+        Ok(Self {
+            weight,
+            bias: Some(bias),
+        })
+    }
+
+    fn load_no_bias(size1: usize, size2: usize, p: &str, vb: &VarBuilder) -> Result<Self> {
+        let weight = vb.get((size2, size1), &format!("{p}.weight"))?;
+        Ok(Self { weight, bias: None })
+    }
+
+    fn forward(&self, x: &Tensor) -> candle::Result<Tensor> {
+        let (bsize, _, _) = x.shape().r3()?;
+        let w = self.weight.broadcast_left(bsize)?.t()?;
+        let x = x.matmul(&w)?;
+        match &self.bias {
+            None => Ok(x),
+            Some(bias) => x.broadcast_add(bias),
+        }
+    }
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+struct ConvConfig {
+    padding: usize,
+    stride: usize,
+}
+
+impl Default for ConvConfig {
+    fn default() -> Self {
+        Self {
+            padding: 0,
+            stride: 1,
+        }
+    }
+}
+
+struct Conv1D {
+    weight: Tensor,
+    bias: Option<Tensor>,
+    config: ConvConfig,
+}
+
+impl Conv1D {
+    fn load(
+        in_channels: usize,
+        out_channels: usize,
+        kernel_size: usize,
+        config: ConvConfig,
+        p: &str,
+        vb: &VarBuilder,
+    ) -> Result<Self> {
+        let weight = vb.get(
+            (out_channels, in_channels, kernel_size),
+            &format!("{p}.weight"),
+        )?;
+        let bias = vb.get(out_channels, &format!("{p}.bias"))?;
+        Ok(Self {
+            weight,
+            bias: Some(bias),
+            config,
+        })
+    }
+
+    fn load_no_bias(
+        in_channels: usize,
+        out_channels: usize,
+        kernel_size: usize,
+        config: ConvConfig,
+        p: &str,
+        vb: &VarBuilder,
+    ) -> Result<Self> {
+        let weight = vb.get(
+            (out_channels, in_channels, kernel_size),
+            &format!("{p}.weight"),
+        )?;
+        Ok(Self {
+            weight,
+            bias: None,
+            config,
+        })
+    }
+
+    fn forward(&self, x: &Tensor) -> Result<Tensor> {
+        let x = x.conv1d(&self.weight, self.config.padding, self.config.stride)?;
+        match &self.bias {
+            None => Ok(x),
+            Some(bias) => {
+                let b = bias.shape().r1()?;
+                let bias = bias.reshape((1, b, 1))?;
+                Ok(x.broadcast_add(&bias)?)
+            }
+        }
+    }
+}
+
+struct Dropout {
+    pr: f64,
+}
+
+impl Dropout {
+    fn new(pr: f64) -> Self {
+        Self { pr }
+    }
+
+    fn forward(&self, x: &Tensor) -> Result<Tensor> {
+        // TODO
+        Ok(x.clone())
+    }
+}
+
+// This layer norm version handles both weight and bias so removes the mean.
+struct LayerNorm {
+    weight: Tensor,
+    bias: Tensor,
+    eps: f64,
+}
+
+impl LayerNorm {
+    fn load(size: usize, p: &str, vb: &VarBuilder) -> Result<Self> {
+        let weight = vb.get(size, &format!("{p}.weight"))?;
+        let bias = vb.get(size, &format!("{p}.bias"))?;
+        Ok(Self {
+            weight,
+            bias,
+            eps: 1e-5,
+        })
+    }
+
+    fn forward(&self, x: &Tensor) -> Result<Tensor> {
+        let (_bsize, _seq_len, hidden_size) = x.shape().r3()?;
+        let mean_x = (x.sum(&[2])? / hidden_size as f64)?;
+        let x = x.broadcast_sub(&mean_x)?;
+        let norm_x = ((&x * &x)?.sum(&[2])? / hidden_size as f64)?;
+        let x_normed = x.broadcast_div(&(norm_x + self.eps)?.sqrt()?)?;
+        let x = x_normed
+            .broadcast_mul(&self.weight)?
+            .broadcast_add(&self.bias)?;
+        Ok(x)
+    }
+}
+
+// https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L62
+struct MultiHeadAttention {
+    query: Linear,
+    key: Linear,
+    value: Linear,
+    out: Linear,
+    n_head: usize,
+}
+
+impl MultiHeadAttention {
+    fn load(n_state: usize, n_head: usize, p: &str, vb: &VarBuilder) -> Result<Self> {
+        let query = Linear::load(n_state, n_state, &format!("{p}.query"), vb)?;
+        let value = Linear::load(n_state, n_state, &format!("{p}.value"), vb)?;
+        let key = Linear::load_no_bias(n_state, n_state, &format!("{p}.key"), vb)?;
+        let out = Linear::load(n_state, n_state, &format!("{p}.out"), vb)?;
+        Ok(Self {
+            query,
+            key,
+            value,
+            out,
+            n_head,
+        })
+    }
+
+    fn forward(&self, x: &Tensor, xa: Option<&Tensor>, mask: Option<&Tensor>) -> Result<Tensor> {
+        let q = self.query.forward(x)?;
+        let k = self.key.forward(xa.unwrap_or(x))?;
+        let v = self.value.forward(xa.unwrap_or(x))?;
+        let wv = self.qkv_attention(&q, &k, &v, mask)?;
+        let out = self.out.forward(&wv)?;
+        Ok(out)
+    }
+
+    fn reshape_head(&self, x: &Tensor) -> Result<Tensor> {
+        let (n_batch, n_ctx, n_state) = x.shape().r3()?;
+        let target_dims = &[n_batch, n_ctx, self.n_head, n_state / self.n_head];
+        Ok(x.reshape(target_dims)?.transpose(1, 2)?)
+    }
+
+    fn qkv_attention(
+        &self,
+        q: &Tensor,
+        k: &Tensor,
+        v: &Tensor,
+        mask: Option<&Tensor>,
+    ) -> Result<Tensor> {
+        let (_, n_ctx, n_state) = q.shape().r3()?;
+        let scale = ((n_state / self.n_head) as f64).powf(-0.25);
+        let q = (self.reshape_head(q)? * scale)?;
+        let k = (self.reshape_head(k)?.transpose(2, 3)? * scale)?;
+        let v = self.reshape_head(v)?.contiguous()?;
+        let mut qk = q.matmul(&k)?;
+        if let Some(mask) = mask {
+            let mask = mask.narrow(0, 0, n_ctx)?.narrow(1, 0, n_ctx)?;
+            qk = qk.broadcast_add(&mask)?
+        }
+        let w = qk.softmax(qk.rank() - 1)?;
+        let wv = w.matmul(&v)?.transpose(1, 2)?.flatten(Some(2), None)?;
+        Ok(wv)
+    }
+}
+
+// https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L111
+struct ResidualAttentionBlock {
+    attn: MultiHeadAttention,
+    attn_ln: LayerNorm,
+    cross_attn: Option<(MultiHeadAttention, LayerNorm)>,
+    mlp_linear1: Linear,
+    mlp_linear2: Linear,
+    mlp_ln: LayerNorm,
+}
+
+impl ResidualAttentionBlock {
+    fn load(n_state: usize, n_head: usize, ca: bool, p: &str, vb: &VarBuilder) -> Result<Self> {
+        let attn = MultiHeadAttention::load(n_state, n_head, &format!("{p}.attn"), vb)?;
+        let attn_ln = LayerNorm::load(n_state, &format!("{p}.attn_ln"), vb)?;
+        let cross_attn = if ca {
+            let cross_attn =
+                MultiHeadAttention::load(n_state, n_head, &format!("{p}.cross_attn"), vb)?;
+            let cross_attn_ln = LayerNorm::load(n_state, &format!("{p}.cross_attn_ln"), vb)?;
+            Some((cross_attn, cross_attn_ln))
+        } else {
+            None
+        };
+        let n_mlp = n_state * 4;
+        let mlp_linear1 = Linear::load(n_state, n_mlp, &format!("{p}.mlp.0"), vb)?;
+        let mlp_linear2 = Linear::load(n_mlp, n_state, &format!("{p}.mlp.2"), vb)?;
+        let mlp_ln = LayerNorm::load(n_state, &format!("{p}.mlp_ln"), vb)?;
+        Ok(Self {
+            attn,
+            attn_ln,
+            cross_attn,
+            mlp_linear1,
+            mlp_linear2,
+            mlp_ln,
+        })
+    }
+
+    fn forward(&self, x: &Tensor, xa: Option<&Tensor>, mask: Option<&Tensor>) -> Result<Tensor> {
+        let attn = self.attn.forward(&self.attn_ln.forward(x)?, None, mask)?;
+        let mut x = (x + attn)?;
+        if let Some((attn, ln)) = &self.cross_attn {
+            x = (&x + attn.forward(&ln.forward(&x)?, xa, None)?)?;
+        }
+        let mlp = self.mlp_linear2.forward(
+            &self
+                .mlp_linear1
+                .forward(&self.mlp_ln.forward(&x)?)?
+                .gelu()?,
+        )?;
+        Ok((x + mlp)?)
+    }
+}
+
+fn sinusoids(length: usize, channels: usize) -> Result<Tensor> {
+    let max_timescale = 10000f32;
+    let log_timescale_increment = max_timescale.ln() / (channels / 2 - 1) as f32;
+    let inv_timescales: Vec<_> = (0..channels / 2)
+        .map(|i| (i as f32 * (-log_timescale_increment)).exp())
+        .collect();
+    let arange: Vec<_> = (0..length).map(|c| c as f32).collect();
+    let inv_timescales = Tensor::new(inv_timescales.as_slice(), &Device::Cpu)?.unsqueeze(0)?;
+    let arange = Tensor::new(arange.as_slice(), &Device::Cpu)?.unsqueeze(1)?;
+    let sh = (length, channels / 2);
+    let scaled_time = (arange.broadcast_as(sh)? * inv_timescales.broadcast_as(sh)?)?;
+    let sincos = Tensor::cat(&[scaled_time.sin()?, scaled_time.cos()?], 1)?;
+    Ok(sincos)
+}
+
+// https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L143
+pub struct AudioEncoder {
+    conv1: Conv1D,
+    conv2: Conv1D,
+    positional_embedding: Tensor,
+    blocks: Vec<ResidualAttentionBlock>,
+    ln_post: LayerNorm,
+}
+
+impl AudioEncoder {
+    fn load(p: &str, vb: &VarBuilder, cfg: &Config) -> Result<Self> {
+        let n_state = cfg.n_audio_state;
+        let n_head = cfg.n_audio_head;
+        let n_ctx = cfg.n_audio_ctx;
+        let cfg1 = ConvConfig {
+            padding: 1,
+            stride: 1,
+        };
+        let cfg2 = ConvConfig {
+            padding: 1,
+            stride: 2,
+        };
+        let conv1 = Conv1D::load(cfg.n_mels, n_state, 3, cfg1, &format!("{p}.conv1"), vb)?;
+        let conv2 = Conv1D::load(n_state, n_state, 3, cfg2, &format!("{p}.conv2"), vb)?;
+        let positional_embedding = sinusoids(n_ctx, n_state)?.to_device(&vb.device)?;
+        let blocks = (0..cfg.n_audio_layer)
+            .map(|i| {
+                ResidualAttentionBlock::load(n_state, n_head, false, &format!("{p}.blocks.{i}"), vb)
+            })
+            .collect::<Result<Vec<_>>>()?;
+        let ln_post = LayerNorm::load(n_state, &format!("{p}.ln_post"), vb)?;
+        Ok(Self {
+            conv1,
+            conv2,
+            positional_embedding,
+            blocks,
+            ln_post,
+        })
+    }
+    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
+        let x = self.conv1.forward(x)?.gelu()?;
+        let x = self.conv2.forward(&x)?.gelu()?;
+        let x = x.transpose(1, 2)?;
+        let mut x = x.broadcast_add(&self.positional_embedding)?;
+        for block in self.blocks.iter() {
+            x = block.forward(&x, None, None)?
+        }
+        let x = self.ln_post.forward(&x)?;
+        Ok(x)
+    }
+}
+
+// https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L176
+pub struct TextDecoder {
+    token_embedding: Embedding,
+    positional_embedding: Tensor,
+    blocks: Vec<ResidualAttentionBlock>,
+    ln: LayerNorm,
+    mask: Tensor,
+}
+
+impl TextDecoder {
+    fn load(p: &str, vb: &VarBuilder, cfg: &Config) -> Result<Self> {
+        let n_state = cfg.n_text_state;
+        let n_head = cfg.n_text_head;
+        let n_ctx = cfg.n_text_ctx;
+        let token_embedding =
+            Embedding::load(cfg.n_vocab, n_state, &format!("{p}.token_embedding"), vb)?;
+        let positional_embedding =
+            vb.get((n_ctx, n_state), &format!("{p}.positional_embedding"))?;
+        let blocks = (0..cfg.n_text_layer)
+            .map(|i| {
+                ResidualAttentionBlock::load(n_state, n_head, true, &format!("{p}.blocks.{i}"), vb)
+            })
+            .collect::<Result<Vec<_>>>()?;
+        let ln = LayerNorm::load(n_state, &format!("{p}.ln"), vb)?;
+        let mask: Vec<_> = (0..n_ctx)
+            .flat_map(|i| (0..n_ctx).map(move |j| if j > i { f32::NEG_INFINITY } else { 0f32 }))
+            .collect();
+        let mask = Tensor::from_vec(mask, (n_ctx, n_ctx), &vb.device)?;
+
+        Ok(Self {
+            token_embedding,
+            positional_embedding,
+            blocks,
+            ln,
+            mask,
+        })
+    }
+
+    pub fn forward(&self, x: &Tensor, xa: &Tensor) -> Result<Tensor> {
+        let x_dims = x.dims();
+        let last = x_dims[x_dims.len() - 1];
+        let token_embedding = self.token_embedding.forward(x)?;
+        let positional_embedding = self.positional_embedding.narrow(0, 0, last)?;
+        let mut x = token_embedding.broadcast_add(&positional_embedding)?;
+        for block in self.blocks.iter() {
+            x = block.forward(&x, Some(xa), Some(&self.mask))?;
+        }
+        let x = self.ln.forward(&x)?;
+        let w = self.token_embedding.embeddings.broadcast_left(x_dims[0])?;
+        let logits = x.matmul(&w.t()?)?;
+        Ok(logits)
+    }
+}
+
+// https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L221
+pub struct Whisper {
+    pub encoder: AudioEncoder,
+    pub decoder: TextDecoder,
+    pub config: Config,
+}
+
+impl Whisper {
+    pub fn load(vb: &VarBuilder, config: Config) -> Result<Self> {
+        let encoder = AudioEncoder::load("encoder", vb, &config)?;
+        let decoder = TextDecoder::load("decoder", vb, &config)?;
+        Ok(Self {
+            encoder,
+            decoder,
+            config,
+        })
+    }
+
+    pub fn forward(&self, mel: &Tensor, tokens: &Tensor) -> Result<Tensor> {
+        let enc = self.encoder.forward(mel)?;
+        let dec = self.decoder.forward(tokens, &enc)?;
+        Ok(dec)
+    }
+}

From 3ba4bfc501acafbdff526118b7f6be1ac99f3db5 Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Wed, 5 Jul 2023 05:50:33 +0100
Subject: [PATCH 24/36] More pretty printing.

---
 candle-examples/examples/whisper/main.rs | 1 +
 1 file changed, 1 insertion(+)

diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs
index 8bf31f34..36caa183 100644
--- a/candle-examples/examples/whisper/main.rs
+++ b/candle-examples/examples/whisper/main.rs
@@ -187,6 +187,7 @@ fn main() -> Result<()> {
         let dr = dc.decode_with_fallback(&mel_segment)?;
         seek += segment_size;
         if dr.no_speech_prob > NO_SPEECH_THRESHOLD && dr.avg_logprob < LOGPROB_THRESHOLD {
+            println!("no speech detected, skipping {seek} {dr:?}");
             continue;
         }
         let segment = Segment {

From 6d1e79d3782c93c7f6a097ce71a91a9a277e52ed Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Wed, 5 Jul 2023 06:42:29 +0100
Subject: [PATCH 25/36] Bugfix for to_scalar (use the proper start offset).

---
 candle-core/src/tensor.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/candle-core/src/tensor.rs b/candle-core/src/tensor.rs
index 25ab0a9b..95f663f0 100644
--- a/candle-core/src/tensor.rs
+++ b/candle-core/src/tensor.rs
@@ -326,7 +326,7 @@ impl Tensor {
         }
         let from_cpu_storage = |cpu_storage: &crate::CpuStorage| {
             let data = S::cpu_storage_as_slice(cpu_storage)?;
-            Ok::<_, Error>(data[0])
+            Ok::<_, Error>(data[self.layout().start_offset()])
         };
         match self.storage.as_ref() {
             Storage::Cpu(cpu_storage) => from_cpu_storage(cpu_storage),

From 9694e35db08b70ca3a9f60b54cf9f0701425edbd Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Wed, 5 Jul 2023 08:37:26 +0100
Subject: [PATCH 26/36] Clean the decode loop of the whisper example.

---
 candle-examples/examples/whisper/main.rs  | 36 ++++++++++++++---------
 candle-examples/examples/whisper/model.rs |  2 ++
 2 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs
index 36caa183..1446e067 100644
--- a/candle-examples/examples/whisper/main.rs
+++ b/candle-examples/examples/whisper/main.rs
@@ -2,7 +2,8 @@
 // https://github.com/openai/whisper/blob/main/whisper/model.py
 // TODO:
 // - kv-cache support?
-// - language detection?
+// - Language detection?
+// - Batch size greater than 1.
 
 use anyhow::{Error as E, Result};
 use candle::{DType, Device, Tensor};
@@ -31,6 +32,10 @@ const LOGPROB_THRESHOLD: f64 = -1.0;
 const TEMPERATURES: [f64; 6] = [0.0, 0.2, 0.4, 0.6, 0.8, 1.0];
 const COMPRESSION_RATIO_THRESHOLD: f64 = 2.4;
 
+// Tokenizer dependent bits.
+const SOT_TOKEN: u32 = 50257;
+const EOT_TOKEN: u32 = 50256;
+
 #[derive(Parser, Debug)]
 #[command(author, version, about, long_about = None)]
 struct Args {
@@ -83,11 +88,12 @@ impl Decode {
         let sample_len = model.config.n_text_ctx / 2;
         let mut sum_logprob = 0f64;
         let no_speech_prob = f64::NAN;
-        // TODO: 50257 is the start of transcipt token, be more principled about get initial tokens
-        let mut tokens: Vec<u32> = vec![50257];
+        let mut tokens = vec![SOT_TOKEN];
         for _i in 0..sample_len {
             let tokens_t = Tensor::new(tokens.as_slice(), &mel.device())?;
-            // Insert a batch dim.
+
+            // The model expects a batch dim but this inference loop does not handle
+            // it so we add it at this point.
             let tokens_t = tokens_t.unsqueeze(0)?;
             let logits = model.decoder.forward(&tokens_t, &audio_features)?;
             let logits = logits.squeeze(0)?;
@@ -112,11 +118,10 @@ impl Decode {
                 .softmax(logits.rank() - 1)?
                 .get(next_token as usize)?
                 .to_scalar::<f32>()? as f64;
-            sum_logprob += prob.ln();
-            // 50256 is the eot token, TODO: parameterize this.
-            if next_token == 50256 || tokens.len() > model.config.n_text_ctx {
+            if next_token == EOT_TOKEN || tokens.len() > model.config.n_text_ctx {
                 break;
             }
+            sum_logprob += prob.ln();
         }
         let text = self
             .tokenizer
@@ -136,14 +141,17 @@ impl Decode {
 
     fn decode_with_fallback(&mut self, segment: &Tensor) -> Result<DecodingResult> {
         for (i, &t) in TEMPERATURES.iter().enumerate() {
-            let dr: DecodingResult = self.decode(segment, t)?;
+            let dr: Result<DecodingResult> = self.decode(segment, t);
             if i == TEMPERATURES.len() - 1 {
-                return Ok(dr);
+                return dr;
             }
-            let needs_fallback = dr.compression_ratio > COMPRESSION_RATIO_THRESHOLD
-                || dr.avg_logprob < LOGPROB_THRESHOLD;
-            if !needs_fallback || dr.no_speech_prob > NO_SPEECH_THRESHOLD {
-                return Ok(dr);
+            // On errors, we try again with a different temperature.
+            if let Ok(dr) = dr {
+                let needs_fallback = dr.compression_ratio > COMPRESSION_RATIO_THRESHOLD
+                    || dr.avg_logprob < LOGPROB_THRESHOLD;
+                if !needs_fallback || dr.no_speech_prob > NO_SPEECH_THRESHOLD {
+                    return Ok(dr);
+                }
             }
         }
         unreachable!()
@@ -195,7 +203,7 @@ fn main() -> Result<()> {
             duration: segment_duration,
             dr,
         };
-        println!("{seek} {segment:?}");
+        println!("{seek}: {segment:?}");
         segments.push(segment)
     }
     Ok(())
diff --git a/candle-examples/examples/whisper/model.rs b/candle-examples/examples/whisper/model.rs
index 94607934..53ee6a90 100644
--- a/candle-examples/examples/whisper/model.rs
+++ b/candle-examples/examples/whisper/model.rs
@@ -1,3 +1,5 @@
+// We use anyhow rather than candle errors as it provides better support for getting the backtrace
+// back when using RUST_LIB_BACKTRACE=1.
 use anyhow::Result;
 use candle::{safetensors::SafeTensors, DType, Device, Shape, Tensor};
 use std::collections::HashMap;

From a824c5c3e3f758874817514e87105c22f6352ded Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Wed, 5 Jul 2023 08:54:04 +0100
Subject: [PATCH 27/36] Populate the no-speech probability.

---
 candle-examples/examples/whisper/main.rs | 34 ++++++++++++++++++------
 1 file changed, 26 insertions(+), 8 deletions(-)

diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs
index 1446e067..c77118f6 100644
--- a/candle-examples/examples/whisper/main.rs
+++ b/candle-examples/examples/whisper/main.rs
@@ -35,6 +35,8 @@ const COMPRESSION_RATIO_THRESHOLD: f64 = 2.4;
 // Tokenizer dependent bits.
 const SOT_TOKEN: u32 = 50257;
 const EOT_TOKEN: u32 = 50256;
+const NO_SPEECH_TOKEN: u32 = 50361;
+const NO_TIMESTAMP_TOKEN: u32 = 50362;
 
 #[derive(Parser, Debug)]
 #[command(author, version, about, long_about = None)]
@@ -87,9 +89,9 @@ impl Decode {
         println!("audio features: {:?}", audio_features.dims());
         let sample_len = model.config.n_text_ctx / 2;
         let mut sum_logprob = 0f64;
-        let no_speech_prob = f64::NAN;
+        let mut no_speech_prob = f64::NAN;
         let mut tokens = vec![SOT_TOKEN];
-        for _i in 0..sample_len {
+        for i in 0..sample_len {
             let tokens_t = Tensor::new(tokens.as_slice(), &mel.device())?;
 
             // The model expects a batch dim but this inference loop does not handle
@@ -97,10 +99,21 @@ impl Decode {
             let tokens_t = tokens_t.unsqueeze(0)?;
             let logits = model.decoder.forward(&tokens_t, &audio_features)?;
             let logits = logits.squeeze(0)?;
+
+            // Extract the no speech probability on the first iteration by looking at the first
+            // token logits and the probability for the according token.
+            if i == 0 {
+                no_speech_prob = logits
+                    .get(0)?
+                    .softmax(0)?
+                    .get(NO_SPEECH_TOKEN as usize)?
+                    .to_scalar::<f32>()? as f64;
+            }
+
             let (seq_len, _) = logits.shape().r2()?;
             let logits = logits.get(seq_len - 1)?;
             let next_token = if t > 0f64 {
-                let prs = (&logits / t)?.softmax(logits.rank() - 1)?;
+                let prs = (&logits / t)?.softmax(0)?;
                 let logits_v: Vec<f32> = prs.to_vec1()?;
                 let distr = rand::distributions::WeightedIndex::new(&logits_v)?;
                 distr.sample(&mut self.rng) as u32
@@ -146,11 +159,16 @@ impl Decode {
                 return dr;
             }
             // On errors, we try again with a different temperature.
-            if let Ok(dr) = dr {
-                let needs_fallback = dr.compression_ratio > COMPRESSION_RATIO_THRESHOLD
-                    || dr.avg_logprob < LOGPROB_THRESHOLD;
-                if !needs_fallback || dr.no_speech_prob > NO_SPEECH_THRESHOLD {
-                    return Ok(dr);
+            match dr {
+                Ok(dr) => {
+                    let needs_fallback = dr.compression_ratio > COMPRESSION_RATIO_THRESHOLD
+                        || dr.avg_logprob < LOGPROB_THRESHOLD;
+                    if !needs_fallback || dr.no_speech_prob > NO_SPEECH_THRESHOLD {
+                        return Ok(dr);
+                    }
+                }
+                Err(err) => {
+                    println!("Error running at {t}: {err}")
                 }
             }
         }

From 7a6bc6d2dcdc9b1fb485987444998dde5ecc4a5b Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Wed, 5 Jul 2023 09:54:12 +0100
Subject: [PATCH 28/36] Mel spectogram computation (fft bits).

---
 candle-examples/Cargo.toml                |  1 +
 candle-examples/examples/whisper/audio.rs | 79 +++++++++++++++++++++++
 candle-examples/examples/whisper/main.rs  |  2 +
 3 files changed, 82 insertions(+)
 create mode 100644 candle-examples/examples/whisper/audio.rs

diff --git a/candle-examples/Cargo.toml b/candle-examples/Cargo.toml
index a71ca17b..ca910441 100644
--- a/candle-examples/Cargo.toml
+++ b/candle-examples/Cargo.toml
@@ -12,6 +12,7 @@ readme = "README.md"
 
 [dependencies]
 candle = { path = "../candle-core", default-features=false }
+num-traits = "0.2.15"
 
 [dev-dependencies]
 anyhow = { version = "1", features = ["backtrace"] }
diff --git a/candle-examples/examples/whisper/audio.rs b/candle-examples/examples/whisper/audio.rs
new file mode 100644
index 00000000..7ae5dd39
--- /dev/null
+++ b/candle-examples/examples/whisper/audio.rs
@@ -0,0 +1,79 @@
+// Audio processing code, adapted from whisper.cpp
+// https://github.com/ggerganov/whisper.cpp
+
+trait Float: num_traits::Float + num_traits::FloatConst + num_traits::NumAssign {}
+
+// https://github.com/ggerganov/whisper.cpp/blob/4774d2feb01a772a15de81ffc34b34a1f294f020/whisper.cpp#L2357
+fn fft<T: Float>(inp: &[T]) -> Vec<T> {
+    let n = inp.len();
+    let zero = T::zero();
+    if n == 1 {
+        return vec![inp[0], zero];
+    }
+    if n % 2 == 1 {
+        return dft(inp);
+    }
+    let mut out = vec![zero; n * 2];
+
+    let mut even = vec![];
+    even.reserve(n / 2);
+    let mut odd = vec![];
+    odd.reserve(n / 2);
+
+    for (i, &inp) in inp.iter().enumerate() {
+        if i % 2 == 0 {
+            even.push(inp)
+        } else {
+            odd.push(inp);
+        }
+    }
+
+    let even_fft = fft(&even);
+    let odd_fft = fft(&odd);
+
+    let two_pi = T::PI() + T::PI();
+    let n_t = T::from(n).unwrap();
+    for k in 0..n / 2 {
+        let k_t = T::from(k).unwrap();
+        let theta = two_pi * k_t / n_t;
+        let re = theta.cos();
+        let im = -theta.sin();
+
+        let re_odd = odd_fft[2 * k];
+        let im_odd = odd_fft[2 * k + 1];
+
+        out[2 * k] = even_fft[2 * k] + re * re_odd - im * im_odd;
+        out[2 * k + 1] = even_fft[2 * k + 1] + re * im_odd + im * re_odd;
+
+        out[2 * (k + n / 2)] = even_fft[2 * k] - re * re_odd + im * im_odd;
+        out[2 * (k + n / 2) + 1] = even_fft[2 * k + 1] - re * im_odd - im * re_odd;
+    }
+    out
+}
+
+// https://github.com/ggerganov/whisper.cpp/blob/4774d2feb01a772a15de81ffc34b34a1f294f020/whisper.cpp#L2337
+fn dft<T: Float>(inp: &[T]) -> Vec<T> {
+    let zero = T::zero();
+    let n = inp.len();
+    let two_pi = T::PI() + T::PI();
+
+    let mut out = Vec::new();
+    out.reserve(2 * n);
+    let n_t = T::from(n).unwrap();
+    for k in 0..n {
+        let k_t = T::from(k).unwrap();
+        let mut re = zero;
+        let mut im = zero;
+
+        for (j, &inp) in inp.iter().enumerate() {
+            let j_t = T::from(j).unwrap();
+            let angle = two_pi * k_t * j_t / n_t;
+            re += inp * angle.cos();
+            im -= inp * angle.sin();
+        }
+
+        out.push(re);
+        out.push(im);
+    }
+    out
+}
diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs
index c77118f6..f37f7ac7 100644
--- a/candle-examples/examples/whisper/main.rs
+++ b/candle-examples/examples/whisper/main.rs
@@ -10,6 +10,8 @@ use candle::{DType, Device, Tensor};
 use clap::Parser;
 use rand::{distributions::Distribution, SeedableRng};
 use tokenizers::Tokenizer;
+
+mod audio;
 mod model;
 use model::{Config, VarBuilder, Whisper};
 

From f4c8a196a8bc3a43fd9c51dca4182994d6af5912 Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Wed, 5 Jul 2023 10:14:20 +0100
Subject: [PATCH 29/36] Mel spectogram.

---
 candle-examples/examples/whisper/audio.rs | 66 +++++++++++++++++++++++
 1 file changed, 66 insertions(+)

diff --git a/candle-examples/examples/whisper/audio.rs b/candle-examples/examples/whisper/audio.rs
index 7ae5dd39..aca4607f 100644
--- a/candle-examples/examples/whisper/audio.rs
+++ b/candle-examples/examples/whisper/audio.rs
@@ -77,3 +77,69 @@ fn dft<T: Float>(inp: &[T]) -> Vec<T> {
     }
     out
 }
+
+#[allow(clippy::too_many_arguments)]
+// https://github.com/ggerganov/whisper.cpp/blob/4774d2feb01a772a15de81ffc34b34a1f294f020/whisper.cpp#L2414
+fn log_mel_spectrogram_w<T: Float>(
+    ith: usize,
+    hann: &[T],
+    samples: &[T],
+    filters: &[T],
+    fft_size: usize,
+    fft_step: usize,
+    speed_up: bool,
+    n_len: usize,
+    n_mel: usize,
+    n_threads: usize,
+) {
+    let n_fft = if speed_up {
+        1 + fft_size / 4
+    } else {
+        1 + fft_size / 2
+    };
+
+    let zero = T::zero();
+    let half = T::from(0.5).unwrap();
+    let mut fft_in = vec![zero; fft_size];
+    let mut mel = vec![zero; n_len * n_mel];
+
+    for i in (ith..n_len).step_by(n_threads) {
+        let offset = i * fft_step;
+
+        // apply Hanning window
+        for j in 0..fft_size {
+            fft_in[j] = if offset + j < samples.len() {
+                hann[j] * samples[offset + j]
+            } else {
+                zero
+            }
+        }
+
+        // FFT -> mag^2
+        let mut fft_out: Vec<T> = fft(&fft_in);
+
+        for j in 0..fft_size {
+            fft_out[j] = fft_out[2 * j] * fft_out[2 * j] + fft_out[2 * j + 1] * fft_out[2 * j + 1];
+        }
+        for j in 1..fft_size / 2 {
+            let v = fft_out[fft_size - j];
+            fft_out[j] += v;
+        }
+
+        if speed_up {
+            // scale down in the frequency domain results in a speed up in the time domain
+            for j in 0..n_fft {
+                fft_out[j] = half * (fft_out[2 * j] + fft_out[2 * j + 1]);
+            }
+        }
+
+        // mel spectrogram
+        for j in 0..n_mel {
+            let mut sum = zero;
+            for k in 0..n_fft {
+                sum += fft_out[k] * filters[j * n_fft + k];
+            }
+            mel[j * n_len + i] = T::max(sum, T::from(1e-10).unwrap()).log10();
+        }
+    }
+}

From dd1d55f5c7e95256fa784c4c4c75607f387c3c02 Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Wed, 5 Jul 2023 10:49:37 +0100
Subject: [PATCH 30/36] Mel spectogram computation.

---
 candle-examples/examples/whisper/audio.rs | 57 ++++++++++++++++++++++-
 1 file changed, 56 insertions(+), 1 deletion(-)

diff --git a/candle-examples/examples/whisper/audio.rs b/candle-examples/examples/whisper/audio.rs
index aca4607f..4b5ed181 100644
--- a/candle-examples/examples/whisper/audio.rs
+++ b/candle-examples/examples/whisper/audio.rs
@@ -1,6 +1,8 @@
 // Audio processing code, adapted from whisper.cpp
 // https://github.com/ggerganov/whisper.cpp
 
+const WHISPER_CHUNK_SIZE: usize = 30;
+
 trait Float: num_traits::Float + num_traits::FloatConst + num_traits::NumAssign {}
 
 // https://github.com/ggerganov/whisper.cpp/blob/4774d2feb01a772a15de81ffc34b34a1f294f020/whisper.cpp#L2357
@@ -91,7 +93,7 @@ fn log_mel_spectrogram_w<T: Float>(
     n_len: usize,
     n_mel: usize,
     n_threads: usize,
-) {
+) -> Vec<T> {
     let n_fft = if speed_up {
         1 + fft_size / 4
     } else {
@@ -142,4 +144,57 @@ fn log_mel_spectrogram_w<T: Float>(
             mel[j * n_len + i] = T::max(sum, T::from(1e-10).unwrap()).log10();
         }
     }
+    mel
+}
+
+fn log_mel_spectrogram<T: Float>(
+    samples: &[T],
+    filters: &[T],
+    fft_size: usize,
+    fft_step: usize,
+    n_mel: usize,
+    speed_up: bool,
+) -> Vec<T> {
+    let zero = T::zero();
+    let two_pi = T::PI() + T::PI();
+    let half = T::from(0.5).unwrap();
+    let one = T::from(1.0).unwrap();
+    let four = T::from(4.0).unwrap();
+    let fft_size_t = T::from(fft_size).unwrap();
+
+    let hann: Vec<T> = (0..fft_size)
+        .map(|i| half * (one - ((two_pi * T::from(i).unwrap()) / fft_size_t).cos()))
+        .collect();
+    let n_len = samples.len() / fft_step;
+
+    // pad audio with at least one extra chunk of zeros
+    let pad = 100 * WHISPER_CHUNK_SIZE / 2;
+    let n_len = if n_len % pad != 0 {
+        (n_len / pad + 1) * pad
+    } else {
+        n_len
+    };
+    let n_len = n_len + pad;
+    let samples = {
+        let mut samples_padded = samples.to_vec();
+        let to_add = n_len * fft_step - samples.len();
+        samples_padded.extend(std::iter::repeat(zero).take(to_add));
+        samples_padded
+    };
+
+    // Use a single thread for now.
+    let mut mel = log_mel_spectrogram_w(
+        0, &hann, &samples, filters, fft_size, fft_step, speed_up, n_len, n_mel, 1,
+    );
+    let mmax = mel
+        .iter()
+        .max_by(|&u, &v| u.partial_cmp(v).unwrap_or(std::cmp::Ordering::Greater))
+        .copied()
+        .unwrap_or(zero)
+        - T::from(8).unwrap();
+    for m in mel.iter_mut() {
+        let v = T::max(*m, mmax);
+        *m = v / four + one
+    }
+    mel
 }

From 648d1511d5b077a503caa9d8c6a947f19ce1ea20 Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Wed, 5 Jul 2023 11:02:49 +0100
Subject: [PATCH 31/36] PCM conversion.

---
 candle-examples/examples/whisper/audio.rs | 26 ++++++++++++++++++++++-
 1 file changed, 25 insertions(+), 1 deletion(-)

diff --git a/candle-examples/examples/whisper/audio.rs b/candle-examples/examples/whisper/audio.rs
index 4b5ed181..64858141 100644
--- a/candle-examples/examples/whisper/audio.rs
+++ b/candle-examples/examples/whisper/audio.rs
@@ -1,6 +1,10 @@
 // Audio processing code, adapted from whisper.cpp
 // https://github.com/ggerganov/whisper.cpp
 
+const WHISPER_SAMPLE_RATE: usize = 16000;
+const WHISPER_N_FFT: usize = 400;
+const WHISPER_N_MEL: usize = 80;
+const WHISPER_HOP_LENGTH: usize = 160;
 const WHISPER_CHUNK_SIZE: usize = 30;
 
 trait Float: num_traits::Float + num_traits::FloatConst + num_traits::NumAssign {}
@@ -147,7 +151,7 @@ fn log_mel_spectrogram_w<T: Float>(
     mel
 }
 
-fn log_mel_spectrogram<T: Float>(
+fn log_mel_spectrogram_<T: Float>(
     samples: &[T],
     filters: &[T],
     fft_size: usize,
@@ -198,3 +202,23 @@ fn log_mel_spectrogram<T: Float>(
     }
     mel
 }
+
+fn pcm_to_mel<T: Float>(samples: &[T], filters: &[T]) -> anyhow::Result<Vec<T>> {
+    if filters.len() != WHISPER_N_MEL * WHISPER_N_FFT {
+        anyhow::bail!(
+            "unexpected filter length {} (n_mel: {}, n_fft: {})",
+            filters.len(),
+            WHISPER_N_MEL,
+            WHISPER_N_FFT
+        )
+    }
+    let mel = log_mel_spectrogram_(
+        samples,
+        filters,
+        WHISPER_N_FFT,
+        WHISPER_HOP_LENGTH,
+        WHISPER_N_MEL,
+        false,
+    );
+    Ok(mel)
+}

From c701ee33a7a457cfa46f27a0cbf618186993e281 Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Wed, 5 Jul 2023 11:05:08 +0100
Subject: [PATCH 32/36] Add the mel filters.

---
 .../examples/whisper/mel_filters.safetensors    | Bin 0 -> 64400 bytes
 1 file changed, 0 insertions(+), 0 deletions(-)
 create mode 100644 candle-examples/examples/whisper/mel_filters.safetensors

diff --git a/candle-examples/examples/whisper/mel_filters.safetensors b/candle-examples/examples/whisper/mel_filters.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..98f3af4498d37fc529b1dbcf17044ae4ae187f2c
GIT binary patch
literal 64400
zcmeI*dsNT&9|!Qnax0g}?I%$#Z9;t`^?tsM6|<RZD)-BmN<VVRH<ukMQq!`PTa4OV
zDpK|fg`eC*$z{{56LZ@LP1vl>HT<&u{?xnve?L2ib9`Q((>brlc|Onk^ZoWc`N-^_
zuSNDV1A-^HJKKvz_Wq%t&Izy=_8+<$?49gQLB==#8t3lpWN>zw;AHP_3^h)g9T;c|
z2sKq+cXo2?>1uGE5Ej<H^6SlaOWAo={_|b=Hw7s0HUZN@b0kbS$ez3nhu0|Ze*x3N
zNGxxZ!JaS?6i^84%szonk8Wg76h<6Tpz^s(PM;IYo-hg&Pzc1gaFF{wL)jCB5l0l*
zne8Pvri^4y7zGL_1WtSV%Ww-<_C#UC5d|vW?~=wKYxaaupnyVPY5r>Y`Py4h+`<tB
zv=@kPkt~Oe-pZb6p9BJ^1ilH#l1}cP?1>7AGwlT`f6gdxzX})HLlA%h3W1#3Kk_xa
z!iOUYyb;)#a*;h@6eyq&2n~rAU-+J2PZUNRQNYP%0KzS@*b_#90t$f(KPF@L{;lkZ
z!iXaZ)N6Vh#-P>g38O#(g}|21&E&||k?e`Wh$9LdIq6#YnU$&R38O#(g}_f!d}Q!r
zPxeG%#1RFan9ZyFY?nQI!YEKcArSk0tTg1*VowxC98q9HyTr=RK<4r@kc<Ka6avp)
zr^_z_qSzCK5l0mG*>SIYx_AtG!YEKcA<(<&9$AqZ%$_KWIHJJ0d)s9<%Ov)MQJ{cA
zz%OLIEU~!6o+yksqCj-XQn~IDf7c3=KmnCN|9!*d_LzSB`H2dNGwlVg#kG`)V@>Rd
z_DLXsN?_UjEIcS(E9*wP^EHYJiZiVROqplI;#1MGQKxeDMC&AQ`RW96oO{bt_xkhM
zt2!vowHCO1G!Ic3%VqZMIG!n5CxOeW1cpBUOf16)dAETb&lD9DXW9!G9+to~$t2xN
z;(4ZMp9TU{C2;!M0g>t|<-B2?WUZn?kE-tDbqeS%kUVTR=FSX}vFnO>rsy69?ynT+
zv9+$e*yUq6vrZ(tqLkuLX8}KBV|;hDvwUZTgS>LdhF#G)4BTBQFnV1W{Bqo7>D~s?
zB{7r#--A+%L!AZU%fE($Utj6z{S3c)MX@V7hk?5*1rEMQL(&Q<JM{kzKCctm6{Qx3
zIt$$ImyE&+7uojRKT&Y*9=oD*7`VGqAb)>6=C|q~T~p0vKO0xRwpVI#sI$QChemv2
z*HCuv*hC5}Wmj|#19w*nltx>lp7%|p#dMHid71d++EhM=DaAO{Rbc1_A2D`B8cx;y
zmwXX&3*Xzg@odpG4BWd~fyCUA$m{q3e@(KIi)STaR_Pp`E!C>=1|0?J?AapLzu1CN
zr>x|WnK$5a$A@Q&jzQqgH4<35uq%>>=D`@zRJKpwhu7)N*_Rs0@h05_mN=~Qc<nV6
z`M+Goh$m)pRMOYz;dq<RV!8!^`_@RnbMjna`Zf&HjW=Ms;~`cbUxhY77ulB@>G3Ar
z1P%mj5QR%4FxBxgdOf&<dL=6{dqx@iqFWHSZ;b>B^Wucv^I**Qr~pTgmg1511`IrB
z&DSqA(&J6K3FMysKwLZXF;1=8jh6m@U_fFTN>ciuVz4967~R6aeXA9i@Z8|>(xo-t
z+qedPi9f-q))8cmF~Mp|5&KfD9&gZ5pu*yiX!e;AzEOM7>UlA0XMKe$#ReSfJc6Ck
zF%aBYDKLCnk%v!sb3Aldj_i6T@#K6CqIxcYy<0=vYQ2lU`$eh8p^gF-6QhLrQv*DI
zUWKUsXYk;38tm%_pvdfza9q@vozXEA+*v8`!}cBG!hk-wdSx{(G|$I}?Ki{X$T+;1
z|4@9JpTN#2^*Gc~V2kZe@gUy~QO8!Gne8FW@3#qigN;}`x(-@2wdLy>9YevLl>*)K
zR*Cp|_NZHLG5VfQ$DsTaG#GC}#Jje*``}A<Mybc4jskry_7%_Xnxo?tUpP)%i`?)`
z^h{WeL8JO&T)i72GUg?p&vXn1cdk-kyXzQ_>!-8D)?QsuY!M2J;LR{h*?<<8f>2P_
z4o_e05ak9J_NGcf4s;ZVS=&{VG$|3eT|E%fWeJ=QBx6<i3Y3S8#EFVVaBQ3)5(ekG
z>uAHBDNvO_-TIS-Z(g}*b`p4z5{5r=H=+1!G`blEBjcPUyuaHh0^3aHc~cc2uWKnF
zr%n|6di^R&lU(rUH?y$hLIV0`M&o+gaP%p+fmpOp{AFj(_XV^}2bZo&U{_;jkzi3E
z(#>sAKhhW3X3KFfCKl#d6A|y@h~D>aiA@t{iRjMT`K(qIBCl&Fkhe9*qj2zg@#y#0
zB4$uee3cXgp9OIUm>Y^GtvvDh(Z=X?E=Ob;oOtGFmkuth5(szhA+AUKR~*i=L|m6X
zn2<OVwhl2knrejCJFdv{uY)x)8Nw~TExV(F<V-7pSf>l_DNBZkD`CgP9$PCcsPKTi
z5{Q9kBM}!g7Ar*$thTI!&WCpk?_4(#ncAMuY+5CR%T^^YYFUwc^qqcUcwCnFk56s%
zy6J>Nv;E-tG6GKyk4Nw0J+bhlCFaH#h}&IF;(W+9zQ<M-Ca-HJu<cnB@px{6X!Gxz
zVpC>I?0)ZkoJjRY@!de&yYT^D6}LsF&A*AkD^o?OS#R-6R73VhyOeO@w+T2MkMY=4
z<}G{<?h?gjPeu1+JN!Pt3wuffVQb@yCT9#Vk7$CjX;;M$h3mzf1)ar>G57fF_BKpj
z(@tPWae>Fh`$6LM{C#5Wf~Uep+G4VN9~p8I8V(x*t1r8u(OE0Bv%V_q;*x|{b0ONf
z1oQmSE-762Z31%OF^?%z#)>t=)5YW2cSXiPYm6^;Mx$xN5Pp3KR?c%qvSl+&A6O<<
R&f6uX51di?r|j?k{12c*uoM6Q

literal 0
HcmV?d00001


From 26d1a7803fd143bf92b56c5470208ecc029be979 Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Wed, 5 Jul 2023 11:20:33 +0100
Subject: [PATCH 33/36] Load the mel filters.

---
 candle-examples/examples/whisper/main.rs | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs
index f37f7ac7..66bde0e8 100644
--- a/candle-examples/examples/whisper/main.rs
+++ b/candle-examples/examples/whisper/main.rs
@@ -59,6 +59,12 @@ struct Args {
     /// The seed to use when generating random samples.
     #[arg(long, default_value_t = 299792458)]
     seed: u64,
+
+    #[arg(
+        long,
+        default_value = "candle-examples/examples/whisper/mel_filters.safetensors"
+    )]
+    filters: String,
 }
 
 #[derive(Debug, Clone)]
@@ -189,6 +195,11 @@ fn main() -> Result<()> {
 
     let tokenizer = Tokenizer::from_file(args.tokenizer_config).map_err(E::msg)?;
 
+    let filters = unsafe { candle::safetensors::MmapedFile::new(args.filters)? };
+    let filters = filters.deserialize()?;
+    let filters = filters.tensor("mel_80", &device)?;
+    println!("loaded mel filters {:?}", filters.shape());
+
     let input = unsafe { candle::safetensors::MmapedFile::new(args.input)? };
     let input = input.deserialize()?;
     let mel = input.tensor("mel", &device)?;

From 95f378ebb41bdf5abbdb71167e3df9eab14789df Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Wed, 5 Jul 2023 11:53:58 +0100
Subject: [PATCH 34/36] Read wav files.

---
 candle-examples/Cargo.toml                |  1 +
 candle-examples/examples/whisper/audio.rs | 17 +++--
 candle-examples/examples/whisper/main.rs  | 84 ++++++++++++++---------
 3 files changed, 61 insertions(+), 41 deletions(-)

diff --git a/candle-examples/Cargo.toml b/candle-examples/Cargo.toml
index ca910441..a3e64a17 100644
--- a/candle-examples/Cargo.toml
+++ b/candle-examples/Cargo.toml
@@ -21,6 +21,7 @@ clap = { version = "4.2.4", features = ["derive"] }
 rand = "0.8.5"
 tokenizers = { version = "0.13.3", default-features=false, features=["onig"] }
 tokio = { version = "1.28.2", features = ["macros", "rt-multi-thread"] }
+wav = "1.0.0"
 
 [features]
 default = ["cuda"]
diff --git a/candle-examples/examples/whisper/audio.rs b/candle-examples/examples/whisper/audio.rs
index 64858141..dbf5d2c0 100644
--- a/candle-examples/examples/whisper/audio.rs
+++ b/candle-examples/examples/whisper/audio.rs
@@ -1,13 +1,16 @@
 // Audio processing code, adapted from whisper.cpp
 // https://github.com/ggerganov/whisper.cpp
 
-const WHISPER_SAMPLE_RATE: usize = 16000;
-const WHISPER_N_FFT: usize = 400;
-const WHISPER_N_MEL: usize = 80;
-const WHISPER_HOP_LENGTH: usize = 160;
-const WHISPER_CHUNK_SIZE: usize = 30;
+pub const WHISPER_SAMPLE_RATE: usize = 16000;
+pub const WHISPER_N_FFT: usize = 400;
+pub const WHISPER_N_MEL: usize = 80;
+pub const WHISPER_HOP_LENGTH: usize = 160;
+pub const WHISPER_CHUNK_SIZE: usize = 30;
 
-trait Float: num_traits::Float + num_traits::FloatConst + num_traits::NumAssign {}
+pub trait Float: num_traits::Float + num_traits::FloatConst + num_traits::NumAssign {}
+
+impl Float for f32 {}
+impl Float for f64 {}
 
 // https://github.com/ggerganov/whisper.cpp/blob/4774d2feb01a772a15de81ffc34b34a1f294f020/whisper.cpp#L2357
 fn fft<T: Float>(inp: &[T]) -> Vec<T> {
@@ -203,7 +206,7 @@ fn log_mel_spectrogram_<T: Float>(
     mel
 }
 
-fn pcm_to_mel<T: Float>(samples: &[T], filters: &[T]) -> anyhow::Result<Vec<T>> {
+pub fn pcm_to_mel<T: Float>(samples: &[T], filters: &[T]) -> anyhow::Result<Vec<T>> {
     if filters.len() != WHISPER_N_MEL * WHISPER_N_FFT {
         anyhow::bail!(
             "unexpected filter length {} (n_mel: {}, n_fft: {})",
diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs
index 66bde0e8..231feac8 100644
--- a/candle-examples/examples/whisper/main.rs
+++ b/candle-examples/examples/whisper/main.rs
@@ -40,33 +40,6 @@ const EOT_TOKEN: u32 = 50256;
 const NO_SPEECH_TOKEN: u32 = 50361;
 const NO_TIMESTAMP_TOKEN: u32 = 50362;
 
-#[derive(Parser, Debug)]
-#[command(author, version, about, long_about = None)]
-struct Args {
-    /// Run on CPU rather than on GPU.
-    #[arg(long)]
-    cpu: bool,
-
-    #[arg(long)]
-    weights: String,
-
-    #[arg(long)]
-    input: String,
-
-    #[arg(long)]
-    tokenizer_config: String,
-
-    /// The seed to use when generating random samples.
-    #[arg(long, default_value_t = 299792458)]
-    seed: u64,
-
-    #[arg(
-        long,
-        default_value = "candle-examples/examples/whisper/mel_filters.safetensors"
-    )]
-    filters: String,
-}
-
 #[derive(Debug, Clone)]
 struct DecodingResult {
     tokens: Vec<u32>,
@@ -184,6 +157,35 @@ impl Decode {
     }
 }
 
+#[derive(Parser, Debug)]
+#[command(author, version, about, long_about = None)]
+struct Args {
+    /// Run on CPU rather than on GPU.
+    #[arg(long)]
+    cpu: bool,
+
+    #[arg(long)]
+    weights: String,
+
+    /// The input to be processed, in wav formats.
+    #[arg(long)]
+    input: String,
+
+    #[arg(long)]
+    tokenizer_config: String,
+
+    /// The seed to use when generating random samples.
+    #[arg(long, default_value_t = 299792458)]
+    seed: u64,
+
+    /// The mel filters in safetensors format.
+    #[arg(
+        long,
+        default_value = "candle-examples/examples/whisper/mel_filters.safetensors"
+    )]
+    filters: String,
+}
+
 fn main() -> Result<()> {
     let args = Args::parse();
     let device = if args.cpu {
@@ -195,14 +197,28 @@ fn main() -> Result<()> {
 
     let tokenizer = Tokenizer::from_file(args.tokenizer_config).map_err(E::msg)?;
 
-    let filters = unsafe { candle::safetensors::MmapedFile::new(args.filters)? };
-    let filters = filters.deserialize()?;
-    let filters = filters.tensor("mel_80", &device)?;
-    println!("loaded mel filters {:?}", filters.shape());
+    let mel_filters = unsafe { candle::safetensors::MmapedFile::new(args.filters)? };
+    let mel_filters = mel_filters.deserialize()?;
+    let mel_filters = mel_filters.tensor("mel_80", &device)?;
+    println!("loaded mel filters {:?}", mel_filters.shape());
+    let mel_filters = mel_filters.flatten_all()?.to_vec1::<f32>()?;
 
-    let input = unsafe { candle::safetensors::MmapedFile::new(args.input)? };
-    let input = input.deserialize()?;
-    let mel = input.tensor("mel", &device)?;
+    let mut input = std::fs::File::open(args.input)?;
+    let (header, data) = wav::read(&mut input)?;
+    println!("loaded wav data: {header:?}");
+    if header.sampling_rate != audio::WHISPER_SAMPLE_RATE as u32 {
+        anyhow::bail!(
+            "wav file must have a {} sampling rate",
+            audio::WHISPER_SAMPLE_RATE
+        )
+    }
+    let data = data.as_sixteen().expect("expected 16 bit wav file");
+    let pcm_data: Vec<_> = data[..data.len() / header.channel_count as usize]
+        .iter()
+        .map(|v| *v as f32 / 32768.)
+        .collect();
+    let mel = audio::pcm_to_mel(&pcm_data, &mel_filters)?;
+    let mel = Tensor::new(&mel[..], &device)?;
     println!("loaded mel: {:?}", mel.dims());
 
     let weights = unsafe { candle::safetensors::MmapedFile::new(args.weights)? };

From 63e5a266bf749933fbc0f7f907d2d48f8fe1e1ba Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Wed, 5 Jul 2023 12:19:21 +0100
Subject: [PATCH 35/36] Put everything together.

---
 candle-examples/examples/whisper/audio.rs | 30 +++++++++--------------
 candle-examples/examples/whisper/main.rs  | 13 +++++-----
 2 files changed, 17 insertions(+), 26 deletions(-)

diff --git a/candle-examples/examples/whisper/audio.rs b/candle-examples/examples/whisper/audio.rs
index dbf5d2c0..d50b7923 100644
--- a/candle-examples/examples/whisper/audio.rs
+++ b/candle-examples/examples/whisper/audio.rs
@@ -1,12 +1,6 @@
 // Audio processing code, adapted from whisper.cpp
 // https://github.com/ggerganov/whisper.cpp
 
-pub const WHISPER_SAMPLE_RATE: usize = 16000;
-pub const WHISPER_N_FFT: usize = 400;
-pub const WHISPER_N_MEL: usize = 80;
-pub const WHISPER_HOP_LENGTH: usize = 160;
-pub const WHISPER_CHUNK_SIZE: usize = 30;
-
 pub trait Float: num_traits::Float + num_traits::FloatConst + num_traits::NumAssign {}
 
 impl Float for f32 {}
@@ -175,7 +169,7 @@ fn log_mel_spectrogram_<T: Float>(
     let n_len = samples.len() / fft_step;
 
     // pad audio with at least one extra chunk of zeros
-    let pad = 100 * WHISPER_CHUNK_SIZE / 2;
+    let pad = 100 * super::CHUNK_LENGTH / 2;
     let n_len = if n_len % pad != 0 {
         (n_len / pad + 1) * pad
     } else {
@@ -206,22 +200,20 @@ fn log_mel_spectrogram_<T: Float>(
     mel
 }
 
-pub fn pcm_to_mel<T: Float>(samples: &[T], filters: &[T]) -> anyhow::Result<Vec<T>> {
-    if filters.len() != WHISPER_N_MEL * WHISPER_N_FFT {
+pub fn pcm_to_mel<T: Float>(
+    samples: &[T],
+    filters: &[T],
+    n_mel: usize,
+    n_fft: usize,
+) -> anyhow::Result<Vec<T>> {
+    if filters.len() != n_mel * n_fft {
         anyhow::bail!(
             "unexpected filter length {} (n_mel: {}, n_fft: {})",
             filters.len(),
-            WHISPER_N_MEL,
-            WHISPER_N_FFT
+            n_mel,
+            n_fft
         )
     }
-    let mel = log_mel_spectrogram_(
-        samples,
-        filters,
-        WHISPER_N_FFT,
-        WHISPER_HOP_LENGTH,
-        WHISPER_N_MEL,
-        false,
-    );
+    let mel = log_mel_spectrogram_(samples, filters, n_fft, super::HOP_LENGTH, n_mel, false);
     Ok(mel)
 }
diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs
index 231feac8..6e15fa8a 100644
--- a/candle-examples/examples/whisper/main.rs
+++ b/candle-examples/examples/whisper/main.rs
@@ -201,24 +201,23 @@ fn main() -> Result<()> {
     let mel_filters = mel_filters.deserialize()?;
     let mel_filters = mel_filters.tensor("mel_80", &device)?;
     println!("loaded mel filters {:?}", mel_filters.shape());
+    let (n_mel, n_fft) = mel_filters.shape().r2()?;
     let mel_filters = mel_filters.flatten_all()?.to_vec1::<f32>()?;
 
     let mut input = std::fs::File::open(args.input)?;
     let (header, data) = wav::read(&mut input)?;
     println!("loaded wav data: {header:?}");
-    if header.sampling_rate != audio::WHISPER_SAMPLE_RATE as u32 {
-        anyhow::bail!(
-            "wav file must have a {} sampling rate",
-            audio::WHISPER_SAMPLE_RATE
-        )
+    if header.sampling_rate != SAMPLE_RATE as u32 {
+        anyhow::bail!("wav file must have a {} sampling rate", SAMPLE_RATE)
     }
     let data = data.as_sixteen().expect("expected 16 bit wav file");
     let pcm_data: Vec<_> = data[..data.len() / header.channel_count as usize]
         .iter()
         .map(|v| *v as f32 / 32768.)
         .collect();
-    let mel = audio::pcm_to_mel(&pcm_data, &mel_filters)?;
-    let mel = Tensor::new(&mel[..], &device)?;
+    let mel = audio::pcm_to_mel(&pcm_data, &mel_filters, n_mel, n_fft)?;
+    let mel_len = mel.len();
+    let mel = Tensor::from_vec(mel, (1, n_mel, mel_len / n_mel), &device)?;
     println!("loaded mel: {:?}", mel.dims());
 
     let weights = unsafe { candle::safetensors::MmapedFile::new(args.weights)? };

From d560855c2a9abb0556077fcd62b6bd773e22da5c Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Wed, 5 Jul 2023 12:56:04 +0100
Subject: [PATCH 36/36] Bugfix for the mel filters.

---
 candle-examples/examples/whisper/audio.rs | 23 ++++++++++-------------
 candle-examples/examples/whisper/main.rs  |  6 +++---
 2 files changed, 13 insertions(+), 16 deletions(-)

diff --git a/candle-examples/examples/whisper/audio.rs b/candle-examples/examples/whisper/audio.rs
index d50b7923..d095e239 100644
--- a/candle-examples/examples/whisper/audio.rs
+++ b/candle-examples/examples/whisper/audio.rs
@@ -148,7 +148,7 @@ fn log_mel_spectrogram_w<T: Float>(
     mel
 }
 
-fn log_mel_spectrogram_<T: Float>(
+fn log_mel_spectrogram_<T: Float + std::fmt::Display>(
     samples: &[T],
     filters: &[T],
     fft_size: usize,
@@ -200,20 +200,17 @@ fn log_mel_spectrogram_<T: Float>(
     mel
 }
 
-pub fn pcm_to_mel<T: Float>(
+pub fn pcm_to_mel<T: Float + std::fmt::Display>(
     samples: &[T],
     filters: &[T],
-    n_mel: usize,
-    n_fft: usize,
 ) -> anyhow::Result<Vec<T>> {
-    if filters.len() != n_mel * n_fft {
-        anyhow::bail!(
-            "unexpected filter length {} (n_mel: {}, n_fft: {})",
-            filters.len(),
-            n_mel,
-            n_fft
-        )
-    }
-    let mel = log_mel_spectrogram_(samples, filters, n_fft, super::HOP_LENGTH, n_mel, false);
+    let mel = log_mel_spectrogram_(
+        samples,
+        filters,
+        super::N_FFT,
+        super::HOP_LENGTH,
+        super::N_MELS,
+        false,
+    );
     Ok(mel)
 }
diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs
index 6e15fa8a..6ea3e536 100644
--- a/candle-examples/examples/whisper/main.rs
+++ b/candle-examples/examples/whisper/main.rs
@@ -201,7 +201,6 @@ fn main() -> Result<()> {
     let mel_filters = mel_filters.deserialize()?;
     let mel_filters = mel_filters.tensor("mel_80", &device)?;
     println!("loaded mel filters {:?}", mel_filters.shape());
-    let (n_mel, n_fft) = mel_filters.shape().r2()?;
     let mel_filters = mel_filters.flatten_all()?.to_vec1::<f32>()?;
 
     let mut input = std::fs::File::open(args.input)?;
@@ -215,9 +214,10 @@ fn main() -> Result<()> {
         .iter()
         .map(|v| *v as f32 / 32768.)
         .collect();
-    let mel = audio::pcm_to_mel(&pcm_data, &mel_filters, n_mel, n_fft)?;
+    println!("pcm data loaded {}", pcm_data.len());
+    let mel = audio::pcm_to_mel(&pcm_data, &mel_filters)?;
     let mel_len = mel.len();
-    let mel = Tensor::from_vec(mel, (1, n_mel, mel_len / n_mel), &device)?;
+    let mel = Tensor::from_vec(mel, (1, N_MELS, mel_len / N_MELS), &device)?;
     println!("loaded mel: {:?}", mel.dims());
 
     let weights = unsafe { candle::safetensors::MmapedFile::new(args.weights)? };