println!(generate(llama, MASTER_CONFIG, 500, device)[0])

ZELBETH:
Sey solmenter! 'tis tonguerered if berryishdd, and What his stabe, you, and, but all I pilJefals, mode with,
Vurint as steolated have loven OlD the queen'd refore
Are been, good plmp:

Proforne, wift'es swleen, was no bunderes'd a a quain beath!
Tybell is my gateer stalk smen'd as be matious dazest brink thou
lord
Enves were cIUll, afe and whwas seath This a is, an tale hoice his his onety Meall-tearn not murkawn, fase bettizen'd her,
To belacquesterer? baxewed wupl usweggs yet tall
An

use candle_core::{DType, Device, IndexOp, Result, Tensor, D};
use candle_nn::ops::softmax;
use candle_nn::{
    embedding, linear, loss, AdamW, Embedding, Init, Linear, Module, Optimizer, ParamsAdamW,
    VarBuilder,
};

use core::f32;
use rand::Rng;
use std::collections::HashMap;
use std::fs;
use std::time;

use std::collections::HashMap;
use std::fs;

fn main() {
    // Read the entire content of the file
    let lines = fs::read_to_string("./input.txt")
        .expect("Failed to read the file");

    // Create a sorted set of unique characters
    let mut vocab: Vec<char> = lines.chars().collect();
    vocab.sort_unstable();
    vocab.dedup();

    // Create itos and stoi mappings
    let itos: HashMap<usize, char> = vocab.iter().enumerate().map(|(i, &ch)| (i, ch)).collect();
    let stoi: HashMap<char, usize> = vocab.iter().enumerate().map(|(i, &ch)| (ch, i)).collect();

    // Print the first 30 characters of the file
    println!("{}", &lines[..30.min(lines.len())]);
}

use std::collections::HashMap;
use std::fs;
struct Vocab {
    itos: HashMap<u32, char>,
    stoi: HashMap<char, u32>,
}

impl Vocab {
    fn new(itos: HashMap<u32, char>, stoi: HashMap<char, u32>) -> Vocab {
        Vocab { itos, stoi }
    }
    fn decode(&self, ids: &[u32]) -> String {
        ids.iter().map(|&id| self.itos[&id]).collect()
    }
    fn encode(&self, text: &str) -> Vec<u32> {
        text.chars().map(|ch| self.stoi[&ch]).collect()
    }
    fn len(&self) -> usize {
        self.itos.len()
    }
    fn build(lines: &str) -> Self {
        // Create a sorted set of unique characters
        let mut vocab: Vec<char> = lines.chars().collect();
        vocab.sort();
        vocab.dedup();

        // Create itos and stoi mappings
        let itos: HashMap<u32, char> = vocab
            .iter()
            .enumerate()
            .map(|(i, &ch)| (i as u32, ch))
            .collect();
        let stoi: HashMap<char, u32> = vocab
            .iter()
            .enumerate()
            .map(|(i, &ch)| (ch, i as u32))
            .collect();
        Self { itos, stoi }
    }
}

fn main() {
    // Read the entire content of the file
    let lines = fs::read_to_string("./input.txt").expect("Failed to read the file");

    let vocab = Vocab::build(&lines);
    println!("vocab size = {}", vocab.len());
    println!("{}", vocab.decode(&vocab.encode("hello")));
}

let mut modeConfig = ModelConfig {
    vocab_size: vocab.len(),
}

let dataset = Tensor::from_slice(&vocab.encode(&lines), (lines.len(),), &Device::Cpu).unwrap();
println!("{:?}", dataset.shape());

fn get_batches(
    dataset: &Tensor,
    split: &str,
    batch_size: usize,
    context_length: usize,
) -> Result<(Tensor, Tensor)> {
    let len_of_dataset = dataset.shape().dim(0).unwrap() as f32;
    // 按照 0.8 0.1 0.1 的比例切分训练集, 验证集和测试集
    let batch_data = match split {
        "val" => &dataset.i((0.8 * len_of_dataset) as usize..(0.9 * len_of_dataset) as usize)?,
        "test" => &dataset.i((0.9 * len_of_dataset) as usize..)?,
        _ => &dataset.i(..(0.8 * len_of_dataset) as usize)?,
    };
    // 生成随机index
    let mut rng = rand::thread_rng();
    let data_len = batch_data.shape().dim(0)?;
    let indices: Vec<usize> = (0..batch_size)
        .map(|_| rng.gen_range(0..data_len - context_length - 1))
        .collect();
    let mut x_batches = Vec::with_capacity(batch_size);
    let mut y_batches = Vec::with_capacity(batch_size);

    for idx in indices {
        let x = batch_data.i(idx..(idx + context_length))?;
        // y 是 x 后面的一个字符
        let y = batch_data.i((idx + 1)..(idx + context_length + 1))?;
        x_batches.push(x);
        y_batches.push(y);
    }
    // stack 和 cat 的区别是, stack 是在新的维度上堆叠, cat 是在已有的维度上堆叠
    let x_tensor = Tensor::stack(&x_batches, 0)?;
    let y_tensor = Tensor::stack(&y_batches, 0)?;
    Ok((x_tensor, y_tensor))
}
// in fn main
modeConfig.context_length = 16;
modeConfig.batch_size = 8;
let batch = get_batches(
    &dataset,
    "train",
    modeConfig.batch_size,
    modeConfig.context_length,
)?;
println!(
    "batch size {}, context_length {}",
    batch.0.shape().dim(0)?,
    batch.0.shape().dim(1)?
);
for i in 0..modeConfig.batch_size {
    println!(
        "{:?}, {:?}",
        vocab.decode(&batch.0.i(i)?.to_vec1()?),
        vocab.decode(&batch.1.i(i)?.to_vec1()?),
    );
}

fn evaluate_loss(
    model: &SimpleBrokenModel,
    dataset: &Tensor,
    vocab: &Vocab,
    config: ModelConfig,
) -> Result<HashMap<String, f32>> {
    let mut out = HashMap::new();
    for split in ["train", "val"] {
        let mut losses = Vec::new();
        for _ in 0..10 {
            let (xs, ys) = get_batches(&dataset, split, config.batch_size, config.context_length)?;
            let (_, loss) = model.forward(&xs, Some(&ys))?;
            let loss = loss.unwrap();
            losses.push(loss.to_scalar::<f32>()?);
        }
        let avg_loss = losses.iter().sum::<f32>() / losses.len() as f32;
        out.insert(split.to_owned(), avg_loss);
    }
    Ok(out)
}

struct SimpleBrokenModel {
    embedding: Embedding,
    mlp: Sequential,
    config: ModelConfig,
}

impl SimpleBrokenModel {
    fn forward(&self, x: &Tensor, targets: Option<&Tensor>) -> Result<(Tensor, Option<Tensor>)> {
        // 潜入层
        let embeds = self.embedding.forward(x)?;

        // 线性和激活层
        let logits = self.mlp.forward(&embeds)?;

        // 如果提供了targets就计算loss，不然视为推理，计算logits就可以。
        if let Some(targets) = targets {
            // 负的似然函数
            // -log(x) 越大，loss 越小
            // y =  [0, 0 , 0, 0, 1, ...,0,0]
            // y' = [4, 5,  6, 7, 8, ...,11,12 ]
            // 这个 cross_entropy 帮我们做了一个 log softmax
            // y' = [0.1, 0.12, 0.13, 0.64, ..., 0,0]
            // loss = -log(0.64)
            // 当 -log(q) = 4.17 q = 0.015 大概 1/64,vocab_size = 65,所以基本是在瞎猜。
            let loss = loss::cross_entropy(
                &logits.reshape(((), self.config.vocab_size))?,
                &targets.reshape(((),))?,
            )?;
            Ok((logits, Some(loss)))
        } else {
            Ok((logits, None))
        }
    }
    // VarBuilder是用来构建参数的，我们目前不加载和保存模型参数，但是candle的用法必须基于这个。
    // vb.pp 会在参数树中加入参数的前缀，这样可以方便的查看参数的结构。
    fn load(vb: VarBuilder, config: ModelConfig) -> Result<(Self)> {
        let embedding = embedding(
            config.vocab_size,
            config.d_model,
            vb.pp("model.embed_tokens"),
        )?;
        let mlp = sequential::seq()
            .add(linear(
                config.d_model,
                config.d_model,
                vb.push_prefix("model.fc1"),
            )?)
            .add(Activation::Relu)
            .add(linear(
                config.d_model,
                config.vocab_size,
                vb.push_prefix("model.fc2"),
            )?);
        Ok(Self {
            embedding,
            mlp,
            config,
        })
    }
}
// in fn main
modeConfig.d_model = 128;
modeConfig.batch_size = 32;
let (xs, ys) = get_batches(
    &dataset,
    "train",
    modeConfig.batch_size,
    modeConfig.context_length,
)?;
let varmap = candle_nn::VarMap::new();
let vb = candle_nn::VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);
let model = SimpleBrokenModel::load(vb, modeConfig)?;
let (logits, loss) = model.forward(&xs, Some(&ys))?;
println!("{:?} {:?}", logits, loss);
let mut params_count: usize = 0;
for (name, var) in varmap.data().lock().unwrap().iter() {
    println!("{}: {:?}", name, var.elem_count());
    params_count += var.elem_count();
}
println!("params count: {}", params_count);

let loss = loss::cross_entropy(
    &logits.reshape(((), self.config.vocab_size))?,
    &targets.reshape(((),))?,
)?;

fn train(
    config: ModelConfig,
    model: &SimpleBrokenModel,
    opt: &mut AdamW,
    dataset: &Tensor,
    vocab: &Vocab,
) -> Result<()> {
    let mut start_time = std::time::Instant::now();
    for epoch in 0..config.epochs {
        let (xs, ys) = get_batches(&dataset, "train", config.batch_size, config.context_length)?;
        let (_, loss) = model.forward(&xs, Some(&ys))?;
        opt.backward_step(&loss.unwrap())?;
        if epoch % config.log_interval == 0 {
            let batch_duration = start_time.elapsed().as_secs_f32();
            let loss = evaluate_loss(&model, dataset, vocab, config)?;
            let val_loss = loss.get("val").unwrap();
            let eta = batch_duration * (config.epochs - epoch) as f32;
            let eta = eta.round();
            println!(
                "Epoch: {epoch} | Loss: {val_loss} | Time: {batch_duration} | ETA in seconds {eta}"
            );
            start_time = time::Instant::now();
        }
    }
    Ok(())
}
// in fn main
modeConfig.log_interval = 10;
modeConfig.epochs = 100;
let mut opt = candle_nn::AdamW::new(varmap.all_vars(), ParamsAdamW::default())?;
train(modeConfig, &model, &mut opt, &dataset, &vocab)?;
let out = evaluate_loss(&model, &dataset, &vocab, modeConfig);
println!("{:?}", out);

fn generate(model: &SimpleBrokenModel, vocab: &Vocab, max_tokens: usize) -> Result<()> {
    // batch size 5, initial token = 0
    let mut token_ids = Tensor::zeros((5, 1), DType::U32, &Device::Cpu).unwrap();
    for _ in 0..max_tokens {
        let (logits, _) = model.forward(&token_ids, None)?;
        assert!(logits.shape().dims() == [token_ids.dim(0)?, token_ids.dim(1)?, 65]);
        let last_step_logits = logits.i((.., logits.dim(1)? - 1))?;
        assert!(last_step_logits.shape().dims() == [token_ids.dim(0)?, 65]);
        let probs = softmax(&last_step_logits, last_step_logits.dims().len() - 1)?;
        assert!(probs.shape().dims() == [token_ids.dim(0)?, 65]);
        let next_token = probs.argmax(probs.dims().len() - 1)?;
        assert!(next_token.shape().dims() == [token_ids.dim(0)?]);
        token_ids = Tensor::cat(&[token_ids, next_token.reshape(((), 1))?], 1)?;
    }
    let lines = fs::read_to_string("./input.txt").expect("Failed to read the file");
    for v in &token_ids.to_vec2()? {
        let text = vocab.decode(v);
        println!("{}", text);
    }
    Ok(())
}
// fn in main
generate(&model, &vocab, 10, device)?;

pub struct RmsNorm {
    scale: Tensor,
    eps: f64,
}
impl RmsNorm {
    fn new(size: usize, vb: VarBuilder) -> Result<Self> {
        Ok(RmsNorm {
            scale: vb.get_with_hints(size, "weight", Init::Const(1.))?,
            eps: 1e-6,
        })
    }
    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
        let x_sqr = x.sqr()?;
        assert!(x_sqr.shape().dims() == x.shape().dims());
        let norm_x = (x.mean(D::Minus1)? + self.eps)?.sqrt()?;
        assert!(norm_x.shape().dims() == [x.shape().dim(0)?, x.shape().dim(1)?]);
        let x_normed = x.broadcast_div(&norm_x.reshape((
            norm_x.shape().dim(0)?,
            norm_x.shape().dim(1)?,
            (),
        ))?)?;
        assert!(x_normed.shape().dims() == x.shape().dims());
        let x = (x_normed.broadcast_mul(&self.scale))?;
        Ok(x)
    }
}

let varmap = candle_nn::VarMap::new();
let vb = candle_nn::VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);

let rms_norms = RmsNorm::new(2, vb)?;
// (2,3,2)
let batch = Tensor::new(
    vec![
        vec![vec![1f32, 1f32], vec![1.2f32, 2f32], vec![3f32, 3f32]],
        vec![vec![4f32, 43f32], vec![5f32, 5f32], vec![61f32, 6f32]],
    ],
    &Device::Cpu,
)?;
let vb = candle_nn::VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);
let out = rms_norms.forward(&batch)?;

struct SimpleBrokenModel {
    embedding: Embedding,
    mlp: Sequential,
    rms_norm: RmsNorm,
    config: ModelConfig,
}

impl SimpleBrokenModel {
    fn forward(&self, x: &Tensor, targets: Option<&Tensor>) -> Result<(Tensor, Option<Tensor>)> {
        // Embedding
        let embeds = self.embedding.forward(x)?;
        // RMSNorm
        let normed_embeds = self.rms_norm.forward(&embeds)?;
        // Linear layers and activation
        let logits = self.mlp.forward(&normed_embeds)?;

        // Calculate loss if targets are provided
        if let Some(targets) = targets {
            // 负的似然函数
            // log(x) 越大，loss 越小
            // y =  [0, 0 , 0, 0, 1, ...,0,0]
            // y' = [4, 5,  6, 7, 8, ...,11,12 ]
            // 这个 cross_entropy 帮我们做了一个 log softmax
            // y' = [0.1, 0.12, 0.13, 0.64, ..., 0,0]
            // loss = -log(0.64)
            // -log(q) = 4.17 q = 0.015 大概 1/64,vocab_size = 65,所以基本是在瞎猜。
            // println!("{:?}", targets.shape());
            // println!("{:?}", logits.shape());
            let loss = loss::cross_entropy(
                &logits.reshape(((), self.config.vocab_size))?,
                &targets.reshape(((),))?,
            )?;
            Ok((logits, Some(loss)))
        } else {
            Ok((logits, None))
        }
    }
    // VarBuilder是用来构建参数的。
    fn load(vb: VarBuilder, config: ModelConfig) -> Result<(Self)> {
        let embedding = embedding(
            config.vocab_size,
            config.d_model,
            vb.pp("model.embed_tokens"),
        )?;
        let rms_norm = RmsNorm::new(config.d_model, vb.pp("model.rms_norm"))?;
        let mlp = sequential::seq()
            .add(linear(
                config.d_model,
                config.d_model,
                vb.push_prefix("model.fc1"),
            )?)
            .add(Activation::Relu)
            .add(linear(
                config.d_model,
                config.vocab_size,
                vb.push_prefix("model.fc2"),
            )?);
        Ok(Self {
            embedding,
            mlp,
            config,
            rms_norm,
        })
    }
}

[x0, x1, .... ,xn]
y0 = x0 * cos(theta) - x1 * sin(theta)
y1 = x0 * sin(theta) + x1 * cos(theta)
[y0, y1, ...., yn]

struct Cache {
    cos: Tensor,
    sin: Tensor,
}

impl Cache {
    fn new(context_length: usize, n_elem: usize, vb: VarBuilder) -> Result<Cache> {
        let theta: Vec<_> = (0..n_elem)
            .step_by(2)
            .map(|i| 1f32 / 10000f32.powf(i as f32 / n_elem as f32))
            .collect();
        let theta = Tensor::new(theta.as_slice(), vb.device())?;
        let idx_theta = Tensor::arange(0, context_length as u32, vb.device())?
            .to_dtype(DType::F32)?
            .reshape((context_length, 1))?
            .matmul(&theta.reshape((1, theta.elem_count()))?)?;
        let freq_cis_real = idx_theta.cos()?;
        let freq_cis_imag = idx_theta.sin()?;

        let cos = freq_cis_real.reshape((context_length, n_elem / 2, 1))?;
        let sin = freq_cis_imag.reshape((context_length, n_elem / 2, 1))?;
        Ok(Cache { cos, sin})
    }
}

fn apply_rotary_emb(&self, x: &Tensor, cache: &Cache) -> Result<Tensor> {
    let (b_sz, seq_len, n_embd) = x.dims3()?;
    // println!("shape of cache.cos {:?}", cache.cos.shape());
    let cos = cache.cos.i(..seq_len)?;
    let sin = cache.sin.i(..seq_len)?;
    // println!("cos shape {:?}", cos.shape());
    let cos = cos.broadcast_as((b_sz, seq_len, n_embd / 2, 1))?;
    let sin = sin.broadcast_as((b_sz, seq_len, n_embd / 2, 1))?;
    // println!("broadcast cos shape {:?}", cos.shape());
    let x = x.reshape((b_sz, seq_len, n_embd / 2, 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)?.reshape((b_sz, seq_len, n_embd))?;
    Ok(rope)
}

struct AttentionModel {
    embedding: Embedding,
    mlp: Sequential,
    rms_norm: RmsNorm,
    config: ModelConfig,
    self_attention: SelfAttention,
    cache: Cache,
}

impl AttentionModel {
    fn forward(&self, x: &Tensor, targets: Option<&Tensor>) -> Result<(Tensor, Option<Tensor>)> {
        // 嵌入层
        let embeds = self.embedding.forward(x)?;
        // 范化层
        let normed_embeds = self.rms_norm.forward(&embeds)?;
        // 自注意力层
        let y = self.self_attention.forward(&normed_embeds, &self.cache)?;
        // 线性和激活层
        let logits = self.mlp.forward(&normed_embeds)?;

        if let Some(targets) = targets {
            // 负的似然函数
            // -log(x) 越大，loss 越小
            // y =  [0, 0 , 0, 0, 1, ...,0,0]
            // y' = [4, 5,  6, 7, 8, ...,11,12 ]
            // 这个 cross_entropy 帮我们做了一个 log softmax
            // y' = [0.1, 0.12, 0.13, 0.64, ..., 0,0]
            // loss = -log(0.64)
            // 例如 -log(q) = 4.17 q = 0.015 大概 1/64,vocab_size = 65,所以基本是在瞎猜。
            let loss = loss::cross_entropy(
                &logits.reshape(((), self.config.vocab_size))?,
                &targets.reshape(((),))?,
            )?;
            Ok((logits, Some(loss)))
        } else {
            Ok((logits, None))
        }
    }

    fn load(vb: VarBuilder, config: ModelConfig) -> Result<(Self)> {
        let embedding = embedding(
            config.vocab_size,
            config.d_model,
            vb.pp("model.embed_tokens"),
        )?;
        let rms_norm = RmsNorm::new(config.d_model, vb.pp("model.rms_norm"))?;
        let self_attention = SelfAttention::load(vb.pp("model.self_attention"), config)?;
        let mlp = sequential::seq()
            .add(linear(
                config.d_model,
                config.d_model,
                vb.push_prefix("model.fc1"),
            )?)
            .add(Activation::Relu)
            .add(linear(
                config.d_model,
                config.vocab_size,
                vb.push_prefix("model.fc2"),
            )?);
        Ok(Self {
            embedding,
            mlp,
            config,
            rms_norm,
            self_attention,
            cache: Cache::new(config.context_length, config.d_model, vb)?,
        })
    }
}

struct MultiHeadAttention {
    q_proj: Linear,
    k_proj: Linear,
    v_proj: Linear,
    o_proj: Linear,
    n_head: usize,
    n_kv_head: usize,
    head_dim: usize,
}

impl MultiHeadAttention {
    fn load(vb: VarBuilder, config: ModelConfig) -> Result<Self> {
        let q_proj = linear(config.d_model, config.d_model, vb.pp("model.q_proj"))?;
        let k_proj = linear(
            config.d_model,
            (config.d_model / config.n_head) * config.n_kv_head,
            vb.pp("model.k_proj"),
        )?;
        let v_proj = linear(
            config.d_model,
            (config.d_model / config.n_head) * config.n_kv_head,
            vb.pp("model.v_proj"),
        )?;
        let o_proj = linear(config.d_model, config.d_model, vb.pp("model.o_proj"))?;
        println!(
            "MHA config n_head {} n_kv_head {}",
            config.n_head, config.n_kv_head
        );
        Ok(Self {
            q_proj,
            k_proj,
            v_proj,
            o_proj,
            n_head: config.n_head,
            n_kv_head: config.n_kv_head,
            head_dim: config.d_model / config.n_head,
        })
    }
    fn apply_rotary_emb(&self, x: &Tensor, cache: &Cache) -> Result<Tensor> {
        let (b_sz, seq_len, h, n_embd) = x.dims4()?;
        let cos = cache.cos.i(..seq_len)?;
        let sin = cache.sin.i(..seq_len)?;
        let cos = cos.unsqueeze(1)?;
        let sin = sin.unsqueeze(1)?;
        let cos = cos.broadcast_as((b_sz, seq_len, 1, n_embd / 2, 1))?;
        let sin = sin.broadcast_as((b_sz, seq_len, 1, n_embd / 2, 1))?;
        let x = x.reshape((b_sz, seq_len, h, n_embd / 2, 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)?.reshape((b_sz, seq_len, h, n_embd))?;
        Ok(rope)
    }
    fn forward(&self, x: &Tensor, cache: &Cache) -> Result<Tensor> {
        let (b_sz, seq_len, n_embd) = x.dims3()?;

        // 计算 q k v
        let q = self.q_proj.forward(x)?;
        let k = self.k_proj.forward(x)?;
        let v = self.v_proj.forward(x)?;

        assert!(n_embd == self.n_head * self.head_dim);

        let q = q.reshape((b_sz, seq_len, self.n_head, self.head_dim))?;
        let k = k.reshape((b_sz, seq_len, self.n_kv_head, self.head_dim))?;
        let v = v.reshape((b_sz, seq_len, self.n_kv_head, self.head_dim))?;

        // 对  q  和 k 做位置编码
        let q = self.apply_rotary_emb(&q, cache)?;
        let k = self.apply_rotary_emb(&k, cache)?;
        // 复制成 n_head / n_kv_head 份
        let k = self.repeat_kv(k)?;
        let v = self.repeat_kv(v)?;

        // 把 seq_len 和 n_head 交换
        // 这转换一下是为了做一个 cat single head 的简单操作
        // 相当于 n_head 个的seq_len*seq_len的注意力。
        let q = q.transpose(1, 2)?.contiguous()?;
        let k = k.transpose(1, 2)?.contiguous()?;
        let v = v.transpose(1, 2)?.contiguous()?;

        // q*k^T / sqrt(d_k) d_k = d_model
        // 这里是 (bs,n_head) 个 (seq_len, seq_len) 的注意力
        let attn = (q.matmul(&k.t()?)? / (self.head_dim as f64).sqrt())?;

        // 这里是头内的softmax (seq_len,seq_len)的行总和为1
        let attn = softmax(&attn, D::Minus1)?;

        // 再乘 * (bs, n_head, seq_len, head_dim) 得到 (bs, n_head）个注意力头对应的加权的v
        let y = attn.matmul(&v)?;
        // 把 n_head 和 seq_len 交换回来，得到 (bs, seq_len, n_head, head_dim) 然后reshape以后
        // 得到 (bs, seq_len, n_head * head_dim) 把头给cat到一起。
        let y = y.transpose(1, 2)?.reshape(&[b_sz, seq_len, n_embd])?;
        let y = self.o_proj.forward(&y)?;
        Ok(y)
    }

    fn repeat_kv(&self, x: Tensor) -> Result<Tensor> {
        let n_rep = self.n_head / self.n_kv_head;
        if n_rep == 1 {
            Ok(x)
        } else {
            let (b_sz, seq_len, n_kv_head, head_dim) = x.dims4()?;
            let x = x
                .unsqueeze(3)?
                .expand((b_sz, seq_len, n_kv_head, n_rep, head_dim))?
                .reshape((b_sz, seq_len, n_kv_head * n_rep, head_dim))?;
            Ok(x)
        }
    }
}

struct AttentionModel {
    embedding: Embedding,
    mlp: Sequential,
    rms_norm: RmsNorm,
    config: ModelConfig,
    self_attention: MultiHeadAttention,
    cache: Cache,
}

impl AttentionModel {
    fn forward(&self, x: &Tensor, targets: Option<&Tensor>) -> Result<(Tensor, Option<Tensor>)> {

        let embeds = self.embedding.forward(x)?;

        let normed_embeds = self.rms_norm.forward(&embeds)?;
        let y = self.self_attention.forward(&normed_embeds, &self.cache)?;

        let logits = self.mlp.forward(&y)?;

        if let Some(targets) = targets {
            let loss = loss::cross_entropy(
                &logits.reshape(((), self.config.vocab_size))?,
                &targets.reshape(((),))?,
            )?;
            Ok((logits, Some(loss)))
        } else {
            Ok((logits, None))
        }
    }

    fn load(vb: VarBuilder, config: ModelConfig) -> Result<(Self)> {
        let embedding = embedding(
            config.vocab_size,
            config.d_model,
            vb.pp("model.embed_tokens"),
        )?;
        let rms_norm = RmsNorm::new(config.d_model, vb.pp("model.rms_norm"))?;
        let self_attention = MultiHeadAttention::load(vb.pp("model.multi_head_attention"), config)?;
        let mlp = sequential::seq()
            .add(linear(
                config.d_model,
                config.d_model,
                vb.push_prefix("model.fc1"),
            )?)
            .add(Activation::Relu)
            .add(linear(
                config.d_model,
                config.vocab_size,
                vb.push_prefix("model.fc2"),
            )?);
        Ok(Self {
            embedding,
            mlp,
            config,
            rms_norm,
            self_attention,
            cache: Cache::new(config.context_length, config.d_model/config.n_head, vb)?,
        })
    }
}

generate(&model, &vocab, 10, device)?;

pub struct Cache {
    // ...
    // mask 也可以 cached
    mask: Tensor,
}
impl Cache {
    fn new(context_length: usize, n_elem: usize, vb: VarBuilder) -> Result<Cache> {
        // ...
        // _ 表示类型由编译器推断，
        // 默认的collect 是 [_]，但是这个大小是不可变的要编译期间决定
        // 所以这里还是要提示编译器要 collect 成 vec.
        let mask: Vec<_> = (0..context_length)
            .flat_map(|i| (0..context_length).map(move |j| u8::from(j > i)))
            .collect();
        let mask = Tensor::from_slice(&mask, (context_length, context_length), vb.device())?;
        Ok(Cache { cos, sin, mask })
    }
}

struct MultiHeadAttention {
    q_proj: Linear,
    k_proj: Linear,
    v_proj: Linear,
    o_proj: Linear,
    n_head: usize,
    n_kv_head: usize,
    head_dim: usize,
}

impl MultiHeadAttention {
    fn load(vb: VarBuilder, config: ModelConfig) -> Result<Self> {
        let q_proj = linear(config.d_model, config.d_model, vb.pp("model.q_proj"))?;
        let k_proj = linear(
            config.d_model,
            (config.d_model / config.n_head) * config.n_kv_head,
            vb.pp("model.k_proj"),
        )?;
        let v_proj = linear(
            config.d_model,
            (config.d_model / config.n_head) * config.n_kv_head,
            vb.pp("model.v_proj"),
        )?;
        let o_proj = linear(config.d_model, config.d_model, vb.pp("model.o_proj"))?;
        println!(
            "MHA config n_head {} n_kv_head {}",
            config.n_head, config.n_kv_head
        );
        Ok(Self {
            q_proj,
            k_proj,
            v_proj,
            o_proj,
            n_head: config.n_head,
            n_kv_head: config.n_kv_head,
            head_dim: config.d_model / config.n_head,
        })
    }
    // ...
    fn forward(&self, x: &Tensor, cache: &Cache) -> Result<Tensor> {
        let (b_sz, seq_len, n_embd) = x.dims3()?;

        // 计算 q k v
        let q = self.q_proj.forward(x)?;
        let k = self.k_proj.forward(x)?;
        let v = self.v_proj.forward(x)?;

        assert!(n_embd == self.n_head * self.head_dim);

        let q = q.reshape((b_sz, seq_len, self.n_head, self.head_dim))?;
        let k = k.reshape((b_sz, seq_len, self.n_kv_head, self.head_dim))?;
        let v = v.reshape((b_sz, seq_len, self.n_kv_head, self.head_dim))?;

        // 对  q  和 k 做位置编码
        let q = self.apply_rotary_emb(&q, cache)?;
        let k = self.apply_rotary_emb(&k, cache)?;
        // 复制成 n_head / n_kv_head 份
        let k = self.repeat_kv(k)?;
        let v = self.repeat_kv(v)?;

        // println!("q.shape {:?}", q.shape());
        // 把 seq_len 和 n_head 交换
        // 这转换一下是为了做一个 cat single head 的简单操作
        // 相当于 n_head 个的seq_len*seq_len的注意力。
        let q = q.transpose(1, 2)?.contiguous()?;
        let k = k.transpose(1, 2)?.contiguous()?;
        let v = v.transpose(1, 2)?.contiguous()?;

        // let tmp = q.matmul(&k.t()?)?;
        // 这个结果是有负数的，但是注意力层不会有负数。
        // 计算结果出了很多 NaN，感觉应该是范化没做好。
        // 后面发现是没有用 sqrt 用了开方，导致负数变成了NaN。
        // q*k^T / sqrt(d_k) d_k = d_model
        // 这里是 (bs,n_head) 个 (seq_len, seq_len) 的注意力
        let attn = (q.matmul(&k.t()?)? / (self.head_dim as f64).sqrt())?;
        // 在 softmax 之前，把未来的token位置变为负无穷，这样在softmax之后，这些位置的概率就会变为0
        let mask = cache
            .mask
            .i((..seq_len, ..seq_len))?
            .unsqueeze(0)?
            .unsqueeze(0)?
            .broadcast_as(attn.shape())?;
        let on_true =
            Tensor::new(f32::NEG_INFINITY, attn.device())?.broadcast_as(mask.shape().dims())?;
        let attn = mask.where_cond(&on_true, &attn)?;
        // 取一个例子
        // 这里是头内的softmax (seq_len,seq_len)的每行 (seq_len) 总和为1
        let attn = softmax(&attn, D::Minus1)?;

        // 再乘 * (bs, n_head, seq_len, head_dim) 得到 (bs, n_head）个注意力头对应的加权的v
        let y = attn.matmul(&v)?;
        // 把 n_head 和 seq_len 交换回来，得到 (bs, seq_len, n_head, head_dim) 然后reshape以后
        // 得到 (bs, seq_len, n_head * head_dim) 把头给cat到一起。
        let y = y.transpose(1, 2)?.reshape(&[b_sz, seq_len, n_embd])?;
        let y = self.o_proj.forward(&y)?;
        Ok(y)
    }

    fn repeat_kv(&self, x: Tensor) -> Result<Tensor> {
        let n_rep = self.n_head / self.n_kv_head;
        if n_rep == 1 {
            Ok(x)
        } else {
            let (b_sz, seq_len, n_kv_head, head_dim) = x.dims4()?;
            let x = x
                .unsqueeze(3)?
                .expand((b_sz, seq_len, n_kv_head, n_rep, head_dim))?
                .reshape((b_sz, seq_len, n_kv_head * n_rep, head_dim))?;
            Ok(x)
        }
    }
}

fn silu(xs: &Tensor) -> Result<Tensor> {
    xs / (xs.neg()?.exp()? + 1.0)?
}

struct SwiGLU {
    c_fc1: Linear,
    c_fc2: Linear,
    c_proj: Linear,
}
// 新的 mlp 是三层的带gate
impl SwiGLU {
    // silu 的特征是允许有一点点的负数
    fn forward(&self, x: &Tensor) -> Result<Tensor> {
        // 这里就是 SwiGLU SiLU(W_1 * x) * (W_2 * x) 是 element wise 的，这个可以作为gate门信号
        let x = (silu(&self.c_fc1.forward(x)?)? * self.c_fc2.forward(x)?)?;
        self.c_proj.forward(&x)
    }

    fn load(vb: VarBuilder, cfg: ModelConfig) -> Result<Self> {
        let h_size = cfg.d_model;
        let i_size = cfg.hidden_dim;
        let c_fc1 = linear(h_size, i_size, vb.pp("gate_proj"))?;
        let c_fc2 = linear(h_size, i_size, vb.pp("up_proj"))?;
        let c_proj = linear(i_size, h_size, vb.pp("down_proj"))?;
        Ok(Self {
            c_fc1,
            c_fc2,
            c_proj,
        })
    }
}

struct Block {
    rms_1: RmsNorm,
    attn: MultiHeadAttention,
    rms_2: RmsNorm,
    swiglu: SwiGLU,
}

impl Block {
    fn forward(&self, x: &Tensor, cache: &Cache) -> Result<Tensor> {
        let residual = x;
        let x = self.rms_1.forward(x)?;
        let x = (self.attn.forward(&x, cache)? + residual)?;
        let residual = &x;
        let x = (self.swiglu.forward(&self.rms_2.forward(&x)?)? + residual)?;
        Ok(x)
    }

    fn load(vb: VarBuilder, cfg: ModelConfig) -> Result<Self> {
        let attn = MultiHeadAttention::load(vb.pp("self_attn"), cfg)?;
        let swiglu = SwiGLU::load(vb.pp("mlp"), cfg)?;
        let rms_1 = RmsNorm::new(cfg.d_model, vb.pp("input_layernorm"))?;
        let rms_2 = RmsNorm::new(cfg.d_model, vb.pp("post_attention_layernorm"))?;
        Ok(Self {
            rms_1,
            attn,
            rms_2,
            swiglu,
        })
    }
}

struct Llama {
    embedding: Embedding,
    blocks: Vec<Block>,
    ln_f: RmsNorm,
    lm_head: Linear,
    cache: Cache,
    config: ModelConfig,
}

impl Llama {
    pub fn forward(
        &self,
        x: &Tensor,
        targets: Option<&Tensor>,
    ) -> Result<(Tensor, Option<Tensor>)> {
        let (_b_sz, _seq_len) = x.dims2()?;
        let mut x = self.embedding.forward(x)?;
        for block in &self.blocks {
            x = block.forward(&x, &self.cache)?;
        }
        let x = self.ln_f.forward(&x)?;
        let logits = self.lm_head.forward(&x)?;

        if let Some(targets) = targets {
            let loss = loss::cross_entropy(
                &logits.reshape(((), self.config.vocab_size))?,
                &targets.reshape(((),))?,
            )?;
            Ok((logits, Some(loss)))
        } else {
            Ok((logits, None))
        }
    }

    pub fn load(vb: VarBuilder, config: ModelConfig) -> Result<Self> {
        let embed_layer = embedding(
            config.vocab_size,
            config.d_model,
            vb.pp("model.embed_tokens"),
        )?;
        let lm_head = linear(config.d_model, config.vocab_size, vb.pp("lm_head"))?;
        let ln_f = RmsNorm::new(config.d_model, vb.pp("model.norm"))?;
        let blocks: Vec<_> = (0..config.n_layers)
            .map(|i| Block::load(vb.pp(format!("model.layers.{i}")), config).unwrap())
            .collect();
        Ok(Self {
            embedding: embed_layer,
            blocks,
            ln_f,
            lm_head,
            config,
            cache: Cache::new(config.context_length, config.d_model / config.n_head, vb)?,
        })
    }
}