```
use candle_core::{DType, Device, Result, Tensor, D};
use candle_nn::{loss, ops, Activation, Embedding, Linear, Module, Optimizer, VarBuilder, VarMap};

const VOCAB_SIZE: u8 = 12;
const PAD: u32 = 0;
const UNK: u32 = 1;
const VOCAB_START: u32 = 2;
const ASCII_MIN: u8 = 32;
const ASCII_MAX: u8 = 126;
const EMB_DIM: usize = 1;
const HIDDEN_DIM: usize = 1;
const BATCH_SIZE: usize = 64;
const MAX_SEQ_LEN: usize = 2048;
const EPOCHS: usize = 10;
const LR: f64 = 5e-5; // Lower LR for stability with 7 layers

fn tokenize(input: char) -> u32 {
    assert!(input.is_ascii());

    match input {
        'a' => 1,
        'b' => 2,
        'c' => 3,
        'd' => 4,
        'e' => 5,
        'f' => 6,
        'g' => 7,
        'h' => 8,
        'j' => 9,
        'k' => 10,

        _ => 11,
    }
}

fn split<T: Clone>(dataset: Vec<T>) -> (Vec<T>, Vec<T>) {
    let x = (dataset.len() as f32 * 0.7).round() as usize;
    let train = dataset[..x].into();
    let test = dataset[x..].into();

    (train, test)
}

fn dataset() -> Vec<(char, char)> {
    vec![
        ('a', 'b'), // 1
        ('b', 'c'), // 2
        ('c', 'd'), // 3
        ('d', 'e'), // 4
        ('e', 'f'), // 5
        ('f', 'g'), // 6
        ('g', 'h'), // 7
        ('h', 'i'), // 8
        ('i', 'j'), // 9
        ('j', 'k'), // 10
    ]
}

struct TextSet {
    train_x: Tensor,
    train_y: Tensor,
    test_x: Tensor,
    test_y: Tensor,
}

// 7-LAYER STRONK MLP
struct TheNn {
    embed: Embedding,
    l1: Activation,
    l2: Activation,
    l3: Activation,
    l4: Activation,
    l5: Activation,
    l6: Activation,
    l7: Activation,
    output: Linear,
}

impl TheNn {
    fn new(vs: VarBuilder) -> Result<Self> {
        let embed = Embedding::new(
            vs.get((VOCAB_SIZE as usize, EMB_DIM), "embed.weight")?,
            VOCAB_SIZE as _,
        );

        let [l1, l2, l3, l4, l5, l6, l7] = [Activation::Swiglu; 7];

        let output = Linear::new(
            vs.get((HIDDEN_DIM, VOCAB_SIZE as usize), "output.weight")?,
            Some(vs.get(VOCAB_SIZE as usize, "output.bias")?),
        );

        Ok(Self {
            embed,
            l1,
            l2,
            l3,
            l4,
            l5,
            l6,
            l7,
            output,
        })
    }

    fn forward(&self, ids: &Tensor) -> Result<Tensor> {
        let x = self.embed.forward(ids)?; // (b, s, emb)
        let x = x.mean(D::Minus2)?; // global average: (b, emb)

        let x = self.l1.forward(&x)?;
        let x = self.l2.forward(&x)?;
        let x = self.l3.forward(&x)?;
        let x = self.l4.forward(&x)?;
        let x = self.l5.forward(&x)?;
        let x = self.l6.forward(&x)?;
        let x = self.l7.forward(&x)?;

        self.output.forward(&x)
    }
}

fn train_epoch(model: &TheNn, opt: &mut candle_nn::AdamW, x: &Tensor, y: &Tensor) -> Result<f32> {
    let mut sum_loss = 0f32;
    let logits = model.forward(x)?;
    let probs = ops::log_softmax(&logits, D::Minus1)?;
    let loss = loss::nll(&probs, y)?;
    opt.backward_step(&loss)?;

    sum_loss += loss.to_vec0::<f32>()?;

    Ok(sum_loss)
}

fn accuracy(model: &TheNn, x: &Tensor, y: &Tensor) -> Result<f32> {
    let logits = model.forward(x)?;
    let preds = logits.argmax(D::Minus1)?;
    preds
        .eq(y)?
        .to_dtype(DType::F32)?
        .mean_all()?
        .to_scalar::<f32>()
}

fn ready_dataset(dev: &Device) -> TextSet {
    let (train, test) = split(dataset());

    TextSet {
        train_x: Tensor::new(
            &*train.iter().map(|t| tokenize(t.0)).collect::<Vec<_>>(),
            dev,
        )
        .unwrap(),
        train_y: Tensor::new(
            &*train.iter().map(|t| tokenize(t.1)).collect::<Vec<_>>(),
            dev,
        )
        .unwrap(),

        test_x: Tensor::new(
            &*test.iter().map(|t| tokenize(t.0)).collect::<Vec<_>>(),
            dev,
        )
        .unwrap(),
        test_y: Tensor::new(
            &*test.iter().map(|t| tokenize(t.1)).collect::<Vec<_>>(),
            dev,
        )
        .unwrap(),
    }
}

fn main() -> Result<()> {
    let data_dir = "./data";

    println!("🐒 STRONK 7-LAYER MLP INITIALIZING 🐒");
    println!(
        "EMB_DIM: {}, HIDDEN_DIM: {}, LAYERS: 7",
        EMB_DIM, HIDDEN_DIM
    );

    let dev = Device::cuda_if_available(0)?;
    println!(
        "Using device: {}",
        if dev.is_cuda() { "CUDA" } else { "CPU" }
    );

    let ds = ready_dataset(&dev);

    let vars = VarMap::new();
    let vs = VarBuilder::from_varmap(&vars, DType::F32, &dev);
    let model = TheNn::new(vs.clone())?;

    let mut opt = candle_nn::AdamW::new(
        vars.all_vars(),
        candle_nn::ParamsAdamW {
            lr: LR,
            weight_decay: 0.0,
            ..Default::default()
        },
    )?;

    println!("🐒 TRAINING STARTING - {} EPOCHS 🐒", EPOCHS);
    println!("Epoch    Loss      Accuracy");
    println!("-------------------------");

    for epoch in 1..=EPOCHS {
        let train_loss = train_epoch(&model, &mut opt, &ds.train_x, &ds.train_y)?;
        let acc = accuracy(&model, &ds.test_x, &ds.test_y)?;
        println!("{:3}    {:8.5}   {:5.2}%", epoch, train_loss, 100.0 * acc);
    }

    println!("🐒 TRAINING COMPLETE 🐒");

    // Optional: Save the model
    println!("Saving model to stronk_7layer.safetensors");
    vars.save("stronk_7layer.safetensors")?;

    Ok(())
}
```

🐒 STRONK 7-LAYER MLP INITIALIZING 🐒
EMB_DIM: 1, HIDDEN_DIM: 1, LAYERS: 7
Using device: CPU
🐒 TRAINING STARTING - 10 EPOCHS 🐒
Epoch    Loss      Accuracy
-------------------------
Error: shape mismatch in reshape, lhs: [7, 1], rhs: [7, 12]
   0: candle_core::error::Error::bt
             at /home/u/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/candle-core-0.8.4/src/error.rs:255:25
   1: candle_core::tensor::Tensor::reshape
             at /home/u/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/candle-core-0.8.4/src/tensor.rs:2203:24
   2: <candle_nn::embedding::Embedding as candle_core::Module>::forward
             at /home/u/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/candle-nn-0.8.4/src/embedding.rs:34:22
   3: ml_exps::TheNn::forward
             at ./src/main.rs:107:17
   4: ml_exps::train_epoch
             at ./src/main.rs:124:18
   5: ml_exps::main
             at ./src/main.rs:207:26
   6: core::ops::function::FnOnce::call_once
             at /build/rustc-1.85.0-src/library/core/src/ops/function.rs:250:5
   7: std::sys::backtrace::__rust_begin_short_backtrace
             at /build/rustc-1.85.0-src/library/std/src/sys/backtrace.rs:152:18
   8: std::rt::lang_start::{{closure}}
             at /build/rustc-1.85.0-src/library/std/src/rt.rs:195:18
   9: std::rt::lang_start_internal
  10: std::rt::lang_start
             at /build/rustc-1.85.0-src/library/std/src/rt.rs:194:17
  11: main
  12: __libc_start_call_main
  13: __libc_start_main_alias_2
  14: _start