mod adapter;
mod config;
mod data;
mod loader;
mod lora;
mod model;
mod train;

use anyhow::Result;
use burn::backend::ndarray::NdArrayDevice;
use burn::backend::{Autodiff, NdArray};
use burn::module::AutodiffModule;
use burn::prelude::*;
use burn::tensor::TensorData;
use clap::{Parser, Subcommand};
use std::path::PathBuf;
use tokenizers::Tokenizer;

use config::{HfConfig, LoraConfig};
use model::Llama;

const MODEL_DIR: &str = "models/smollm2-135m-instruct";
const DATA: &str = "data/dolly-15k/databricks-dolly-15k.jsonl";

type Back = Autodiff<NdArray<f32>>;
type Inner = NdArray<f32>;

#[derive(Parser)]
#[command(about = "LoRA fine-tuning for Llama models, written in Rust with burn")]
struct Cli {
    /// Hugging Face model directory (config.json, tokenizer.json, model.safetensors).
    #[arg(long, global = true, default_value = MODEL_DIR)]
    model: PathBuf,
    #[command(subcommand)]
    cmd: Cmd,
}

#[derive(Subcommand)]
enum Cmd {
    /// Load the weights and generate text, to check the model is wired correctly.
    Check {
        #[arg(default_value = "The capital of France is")]
        prompt: String,
        #[arg(long, default_value_t = 12)]
        tokens: usize,
        /// Directory of a LoRA adapter to apply on top of the base model.
        #[arg(long)]
        adapter: Option<PathBuf>,
    },
    /// Train a LoRA adapter on dolly-15k.
    Train {
        /// dolly-15k JSONL file.
        #[arg(long, default_value = DATA)]
        data: PathBuf,
        #[arg(long, default_value_t = 20)]
        steps: usize,
        #[arg(long, default_value_t = 1)]
        batch: usize,
        #[arg(long, default_value_t = 1e-4)]
        lr: f64,
        #[arg(long, default_value_t = 256)]
        max_len: usize,
        #[arg(long, default_value_t = 256)]
        examples: usize,
        /// Where to save the trained adapter.
        #[arg(long, default_value = "adapters/dolly")]
        out: PathBuf,
    },
}

fn generate(
    model: &Llama<Inner>,
    tok: &Tokenizer,
    cfg: &HfConfig,
    prompt: &str,
    n: usize,
    device: &NdArrayDevice,
) -> Result<()> {
    let mut ids: Vec<i64> = tok
        .encode(prompt, false)
        .map_err(|e| anyhow::anyhow!("{e}"))?
        .get_ids()
        .iter()
        .map(|&i| i as i64)
        .collect();

    print!("{prompt}");
    use std::io::Write;
    for _ in 0..n {
        let len = ids.len();
        let input =
            Tensor::<Inner, 2, Int>::from_data(TensorData::new(ids.clone(), [1, len]), device);
        let next = model
            .forward(input)
            .slice([0..1, len - 1..len, 0..cfg.vocab_size])
            .reshape([cfg.vocab_size])
            .argmax(0)
            .into_scalar() as i64;
        // SmolLM2: 0 = <|endoftext|>, 2 = <|im_end|>
        if next == 2 || next == 0 {
            break;
        }
        ids.push(next);
        print!(
            "{}",
            tok.decode(&[next as u32], false)
                .map_err(|e| anyhow::anyhow!("{e}"))?
        );
        std::io::stdout().flush().ok();
    }
    println!("\n");
    Ok(())
}

fn main() -> Result<()> {
    let cli = Cli::parse();
    let device = Default::default();

    let model_dir = cli.model.as_path();
    let cfg = HfConfig::load(model_dir)?;
    let tok = Tokenizer::from_file(model_dir.join("tokenizer.json"))
        .map_err(|e| anyhow::anyhow!("{e}"))?;

    println!(
        "model : {} layers, hidden {}, {} heads / {} kv, head_dim {}, vocab {}",
        cfg.num_hidden_layers,
        cfg.hidden_size,
        cfg.num_attention_heads,
        cfg.num_key_value_heads,
        cfg.head_dim(),
        cfg.vocab_size
    );
    let describe = |l: &LoraConfig| {
        println!(
            "lora  : rank {}, alpha {}, scaling {:.2}, targets {:?}",
            l.rank,
            l.alpha,
            l.scaling(),
            l.targets
        );
    };

    match cli.cmd {
        Cmd::Check {
            prompt,
            tokens,
            adapter: adapter_dir,
        } => {
            // The adapter config wins: it sets the rank and targets,
            // and therefore the shape of the tensors to create.
            let lora = match &adapter_dir {
                Some(dir) => adapter::load_config(dir)?,
                None => LoraConfig::default(),
            };
            describe(&lora);

            let mut model = Llama::<Inner>::new(&cfg, &lora, cfg.max_seq(), &device);
            loader::load_hf_weights(&mut model, model_dir)?;

            if let Some(dir) = &adapter_dir {
                let n = adapter::load(&mut model, dir)?;
                println!(
                    "adapter        : {n} tensors from {}  |B| = {:.6}",
                    dir.display(),
                    model.lora_norm()
                );
            }
            println!();
            generate(&model, &tok, &cfg, &prompt, tokens, &device)?;
        }
        Cmd::Train {
            data,
            steps,
            batch,
            lr,
            max_len,
            examples,
            out,
        } => {
            let lora = LoraConfig::default();
            describe(&lora);
            let mut model = Llama::<Back>::new(&cfg, &lora, cfg.max_seq(), &device);
            loader::load_hf_weights(&mut model, model_dir)?;
            // freeze AFTER loading: the base weights stop requiring gradients
            let model = model.freeze_base();

            let data = data::load_dolly(&data, &tok, max_len, examples)?;
            println!("examples loaded: {}\n", data.len());

            let cfg_t = train::TrainCfg {
                steps,
                batch_size: batch,
                lr,
            };
            let model = train::train(model, &data, cfg_t, &device)?;

            adapter::save(&model, &lora, &out)?;
            println!(
                "\nadapter saved to {}/  (|B| = {:.6})",
                out.display(),
                model.lora_norm()
            );

            println!("\n--- sample after training ---");
            let infer = model.valid();
            generate(
                &infer,
                &tok,
                &cfg,
                "<|im_start|>user\nWhat is the capital of France?<|im_end|>\n<|im_start|>assistant\n",
                40,
                &device,
            )?;
        }
    }
    Ok(())
}