Use the binary decoder for llama2.c. (#230)

* Use the binary decoder for llama2.c.

* Add the temperature.

* Formatting tweak.

* Fix the rotary embeddings.

candle-examples/examples/llama2-c/main.rs

@@ -13,6 +13,7 @@ use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 use candle::{DType, Device, Error, IndexOp, Layout, Shape, Tensor};
 use candle_nn::{Embedding, Linear, VarBuilder};
 use candle_transformers::generation::LogitsProcessor;
+use std::io::Write;
 
 use model::{Config, Llama};
 
@@ -38,21 +39,33 @@ struct TransformerWeights {
     freq_cis_imag: Tensor, // (seq_len, head_size/2)
 }
 
-impl Config {
+fn read_i32<R: std::io::Read>(r: &mut R) -> Result<i32> {
-    fn read_i32<R: std::io::Read>(r: &mut R) -> Result<i32> {
     let mut buf = [0u8; 4];
     r.read_exact(&mut buf)?;
     Ok(i32::from_le_bytes(buf))
-    }
+}
 
+fn read_tensor<R: std::io::Read, S: Into<Shape>>(
+    r: &mut R,
+    shape: S,
+    dev: &Device,
+) -> Result<Tensor> {
+    let shape = shape.into();
+    let mut data_t = vec![0f32; shape.elem_count()];
+    r.read_f32_into::<LittleEndian>(&mut data_t)?;
+    let tensor = Tensor::from_vec(data_t, shape, dev)?;
+    Ok(tensor)
+}
+
+impl Config {
     fn from_reader<R: std::io::Read>(r: &mut R) -> Result<Self> {
-        let dim = Self::read_i32(r)? as usize;
+        let dim = read_i32(r)? as usize;
-        let hidden_dim = Self::read_i32(r)? as usize;
+        let hidden_dim = read_i32(r)? as usize;
-        let n_layers = Self::read_i32(r)? as usize;
+        let n_layers = read_i32(r)? as usize;
-        let n_heads = Self::read_i32(r)? as usize;
+        let n_heads = read_i32(r)? as usize;
-        let n_kv_heads = Self::read_i32(r)? as usize;
+        let n_kv_heads = read_i32(r)? as usize;
-        let vocab_size = Self::read_i32(r)? as usize;
+        let vocab_size = read_i32(r)? as usize;
-        let seq_len = Self::read_i32(r)? as usize;
+        let seq_len = read_i32(r)? as usize;
         Ok(Self {
             dim,
             hidden_dim,
@@ -71,33 +84,21 @@ impl Config {
 }
 
 impl TransformerWeights {
-    fn read_tensor<R: std::io::Read, S: Into<Shape>>(
-        r: &mut R,
-        shape: S,
-        dev: &Device,
-    ) -> Result<Tensor> {
-        let shape = shape.into();
-        let mut data_t = vec![0f32; shape.elem_count()];
-        r.read_f32_into::<LittleEndian>(&mut data_t)?;
-        let tensor = Tensor::from_vec(data_t, shape, dev)?;
-        Ok(tensor)
-    }
-
     fn from_reader<R: std::io::Read>(r: &mut R, c: &Config, dev: &Device) -> Result<Self> {
-        let token_embedding_table = Self::read_tensor(r, (c.vocab_size, c.dim), dev)?;
+        let token_embedding_table = read_tensor(r, (c.vocab_size, c.dim), dev)?;
-        let rms_att_weight = Self::read_tensor(r, (c.n_layers, c.dim), dev)?;
+        let rms_att_weight = read_tensor(r, (c.n_layers, c.dim), dev)?;
-        let wq = Self::read_tensor(r, (c.n_layers, c.dim, c.dim), dev)?;
+        let wq = read_tensor(r, (c.n_layers, c.dim, c.dim), dev)?;
-        let wk = Self::read_tensor(r, (c.n_layers, c.dim, c.dim), dev)?;
+        let wk = read_tensor(r, (c.n_layers, c.dim, c.dim), dev)?;
-        let wv = Self::read_tensor(r, (c.n_layers, c.dim, c.dim), dev)?;
+        let wv = read_tensor(r, (c.n_layers, c.dim, c.dim), dev)?;
-        let wo = Self::read_tensor(r, (c.n_layers, c.dim, c.dim), dev)?;
+        let wo = read_tensor(r, (c.n_layers, c.dim, c.dim), dev)?;
-        let rms_ffn_weight = Self::read_tensor(r, (c.n_layers, c.dim), dev)?;
+        let rms_ffn_weight = read_tensor(r, (c.n_layers, c.dim), dev)?;
-        let w1 = Self::read_tensor(r, (c.n_layers, c.hidden_dim, c.dim), dev)?;
+        let w1 = read_tensor(r, (c.n_layers, c.hidden_dim, c.dim), dev)?;
-        let w2 = Self::read_tensor(r, (c.n_layers, c.dim, c.hidden_dim), dev)?;
+        let w2 = read_tensor(r, (c.n_layers, c.dim, c.hidden_dim), dev)?;
-        let w3 = Self::read_tensor(r, (c.n_layers, c.hidden_dim, c.dim), dev)?;
+        let w3 = read_tensor(r, (c.n_layers, c.hidden_dim, c.dim), dev)?;
-        let rms_final_weight = Self::read_tensor(r, c.dim, dev)?;
+        let rms_final_weight = read_tensor(r, c.dim, dev)?;
         let head_size = c.head_size();
-        let freq_cis_real = Self::read_tensor(r, (c.seq_len, head_size / 2), dev)?;
+        let freq_cis_real = read_tensor(r, (c.seq_len, head_size / 2), dev)?;
-        let freq_cis_imag = Self::read_tensor(r, (c.seq_len, head_size / 2), dev)?;
+        let freq_cis_imag = read_tensor(r, (c.seq_len, head_size / 2), dev)?;
         Ok(Self {
             token_embedding_table,
             rms_att_weight,
@@ -181,13 +182,36 @@ struct Args {
     /// Config file in binary format.
     #[arg(long)]
     config: String,
+
+    /// Tokenizer config file in binary format.
+    #[arg(long)]
+    tokenizer: String,
+
+    /// The temperature used to generate samples.
+    #[arg(long)]
+    temperature: Option<f64>,
+}
+
+struct Tokenizer {
+    tokens: Vec<String>,
+}
+
+impl Tokenizer {
+    fn from_reader<R: std::io::Read>(r: &mut R, c: &Config) -> Result<Self> {
+        let mut tokens = Vec::with_capacity(c.vocab_size);
+        for _token_index in 0..c.vocab_size {
+            let token_len = read_i32(r)?;
+            let mut token = vec![0u8; token_len as usize];
+            r.read_exact(&mut token);
+            tokens.push(String::from_utf8_lossy(&token).into_owned())
+        }
+        Ok(Self { tokens })
+    }
 }
 
 fn main() -> anyhow::Result<()> {
     let args = Args::parse();
     let device = candle_examples::device(args.cpu)?;
-    let t = Tensor::arange(0f32, 14f32, &device)?.reshape((2, 7))?;
-    println!("{t}");
     let mut file = std::fs::File::open(&args.config)?;
     let config = Config::from_reader(&mut file)?;
     println!("config: {config:?}");
@@ -196,13 +220,15 @@ fn main() -> anyhow::Result<()> {
     let cache = model::Cache::new(true, &config, vb.pp("rot"))?;
     let model = Llama::load(vb, &cache, &config)?;
 
+    let mut file = std::fs::File::open(&args.tokenizer)?;
+    let tokenizer = Tokenizer::from_reader(&mut file, &config)?;
+
     println!("starting the inference loop");
-    let mut logits_processor = LogitsProcessor::new(299792458, None);
+    let mut logits_processor = LogitsProcessor::new(299792458, args.temperature);
-    let mut new_tokens: Vec<u32> = vec![];
-    let start_gen = std::time::Instant::now();
     let mut index_pos = 0;
     let mut tokens = vec![1u32];
 
+    let start_gen = std::time::Instant::now();
     for index in 0..config.seq_len - 10 {
         let start_gen = std::time::Instant::now();
         let context_size = if cache.use_kv_cache && index > 0 {
@@ -218,23 +244,14 @@ fn main() -> anyhow::Result<()> {
 
         let next_token = logits_processor.sample(&logits)?;
         tokens.push(next_token);
-        new_tokens.push(next_token);
+        print!("{}", tokenizer.tokens[next_token as usize]);
-        println!("> {:?}", start_gen.elapsed());
+        std::io::stdout().flush()?;
-        println!(
-            "{} token: {} '{}'",
-            index + 1,
-            next_token,
-            0,
-            // tokenizer.decode(vec![next_token], true).map_err(E::msg)?
-        );
     }
     let dt = start_gen.elapsed();
     println!(
-        "{} tokens generated ({} token/s)\n----\n{}\n----",
+        "\n{} tokens generated ({:.2} token/s)\n",
-        config.seq_len,
+        tokens.len(),
-        config.seq_len as f64 / dt.as_secs_f64(),
+        tokens.len() as f64 / dt.as_secs_f64(),
-        0,
-        // tokenizer.decode(new_tokens, true).map_err(E::msg)?
     );
     Ok(())
 }

candle-examples/examples/llama2-c/model.rs

@@ -30,12 +30,14 @@ impl Cache {
     pub fn new(use_kv_cache: bool, cfg: &Config, vb: VarBuilder) -> Result<Self> {
         let freq_cis_real = vb.get((cfg.seq_len, cfg.head_size() / 2), "freq_cis_real")?;
         let freq_cis_imag = vb.get((cfg.seq_len, cfg.head_size() / 2), "freq_cis_imag")?;
+        let cos = freq_cis_real.reshape((cfg.seq_len, cfg.head_size() / 2, 1))?;
+        let sin = freq_cis_imag.reshape((cfg.seq_len, cfg.head_size() / 2, 1))?;
         Ok(Self {
             masks: Arc::new(Mutex::new(HashMap::new())),
             use_kv_cache,
             kvs: Arc::new(Mutex::new(vec![None; cfg.n_layers])),
-            cos: freq_cis_real,
+            cos,
-            sin: freq_cis_imag,
+            sin,
             device: vb.device().clone(),
         })
     }
@@ -110,16 +112,17 @@ struct CausalSelfAttention {
 
 impl CausalSelfAttention {
     fn apply_rotary_emb(&self, x: &Tensor, index_pos: usize) -> Result<Tensor> {
-        let (b_sz, _, seq_len, n_embd) = x.dims4()?;
+        let (b_sz, seq_len, h, n_embd) = x.dims4()?;
         let cos = self.cache.cos.narrow(0, index_pos, seq_len)?;
         let sin = self.cache.sin.narrow(0, index_pos, seq_len)?;
-        let cos = cos.broadcast_as((b_sz, 1, seq_len, n_embd / 2))?;
+        let cos = cos.broadcast_as((b_sz, seq_len, 1, n_embd / 2, 1))?;
-        let sin = sin.broadcast_as((b_sz, 1, seq_len, n_embd / 2))?;
+        let sin = sin.broadcast_as((b_sz, seq_len, 1, n_embd / 2, 1))?;
-        let x0 = x.narrow(D::Minus1, 0, n_embd / 2)?;
+        let x = x.reshape((b_sz, seq_len, h, n_embd / 2, 2))?;
-        let x1 = x.narrow(D::Minus1, n_embd / 2, n_embd / 2)?;
+        let x0 = x.narrow(D::Minus1, 0, 1)?;
+        let x1 = x.narrow(D::Minus1, 1, 1)?;
         let dst0 = (x0.broadcast_mul(&cos)? - x1.broadcast_mul(&sin)?)?;
         let dst1 = (x0.broadcast_mul(&sin)? + x1.broadcast_mul(&cos)?)?;
-        let rope = Tensor::cat(&[&dst0, &dst1], D::Minus1)?;
+        let rope = Tensor::cat(&[&dst0, &dst1], D::Minus1)?.reshape((b_sz, seq_len, h, n_embd))?;
         Ok(rope)
     }