inferencelayer 0.2.4

//! `gguf-quantize` — safetensors checkpoint → quantized GGUF, with ZERO external tools: the
//! in-repo quantizer family + GGUF writer replace `convert_hf_to_gguf.py` + `llama-quantize`
//! in the artifact pipeline (they remain only as dev-time oracles in gated tests).
//!
//! ```text
//! gguf-quantize --model <safetensors dir> --out <file.gguf> [--preset q4_k_m]
//!               [--tensor-type <substr>=<type> ...] [--tokenizer-pre <name>]
//! ```
//!
//! Presets: f16, q8_0, q4_0, q5_0, q4_k_m (Q6_K on attn_v/ffn_down, Q4_K elsewhere),
//! q4_k_s, q5_k_m, q6_k. `--tensor-type` overrides match by substring (`ffn_up=q4_0` hits every
//! layer's up projection) — the hook the calibrated dynamic search drives. K-quants fall back
//! per llama.cpp convention when a row length is not a multiple of 256 (Q4_K/Q5_K→Q5_0,
//! Q6_K→Q8_0).
//!
//! Arch support mirrors `gguf::load_gguf`: llama (q/k rows PERMUTED to the ecosystem's
//! interleaved-rotary convention — the exact inverse of the loader's un-permute) and qwen2
//! (HF order + fused-bias tensors). The tokenizer (BPE vocab + merges + special tokens) embeds
//! from `tokenizer.json`, so llama.cpp/ollama can run the artifact stand-alone.

use anyhow::{Context, Result, bail, ensure};
use inferencelayer::gguf_write::{GgufWriter, WVal, encode_ggml_im};

fn type_id(name: &str) -> Result<u32> {
    Ok(match name.to_ascii_lowercase().as_str() {
        "f32" => 0,
        "f16" => 1,
        "q4_0" => 2,
        "q5_0" => 6,
        "q8_0" => 8,
        "q4_k" => 12,
        "q5_k" => 13,
        "q6_k" => 14,
        "iq4_nl" => 20,
        "iq4_xs" => 23,
        "iq2_xs" => 17,
        "iq3_xxs" => 18,
        "iq3_s" => 21,
        other => bail!("unknown type {other}"),
    })
}

fn type_name(id: u32) -> &'static str {
    match id {
        0 => "F32",
        1 => "F16",
        2 => "Q4_0",
        6 => "Q5_0",
        8 => "Q8_0",
        12 => "Q4_K",
        13 => "Q5_K",
        14 => "Q6_K",
        17 => "IQ2_XS",
        18 => "IQ3_XXS",
        21 => "IQ3_S",
        20 => "IQ4_NL",
        23 => "IQ4_XS",
        _ => "?",
    }
}

/// The preset's base type for a 2-D weight tensor, before overrides and fallbacks.
fn preset_type(preset: &str, gguf_name: &str) -> Result<u32> {
    Ok(match preset {
        "f16" => 1,
        "q8_0" => 8,
        "q4_0" => 2,
        "q5_0" => 6,
        "q4_k_s" => 12,
        "q6_k" => 14,
        "iq4_xs" => 23,
        "iq3_xxs" => 18,
        "iq2_xs" => 17,
        "q5_k_m" => {
            if gguf_name.contains("attn_v") || gguf_name.contains("ffn_down") {
                14
            } else {
                13
            }
        }
        // the stock M recipe shape: bit-hungry tensors up a rung
        "q4_k_m" => {
            if gguf_name.contains("attn_v") || gguf_name.contains("ffn_down") {
                14
            } else {
                12
            }
        }
        other => bail!("unknown preset {other}"),
    })
}

/// llama.cpp's ne0%256 fallbacks for K-quants; IQ4_XS falls back to IQ4_NL (same LUT, 32-blocks).
fn fallback(ty: u32, cols: usize) -> u32 {
    if matches!(ty, 12 | 13 | 14 | 17 | 18 | 21 | 23) && cols % 256 != 0 {
        match ty {
            14 => 8,
            17 | 18 | 21 | 23 => 20,
            _ => 6,
        }
    } else {
        ty
    }
}

fn main() -> Result<()> {
    let mut model = String::new();
    let mut out = String::new();
    let mut preset = "q4_k_m".to_string();
    let mut overrides: Vec<(String, u32)> = Vec::new();
    let mut imatrix_path = String::new();
    let mut tok_pre = "gpt-2".to_string();
    let mut it = std::env::args().skip(1);
    while let Some(a) = it.next() {
        match a.as_str() {
            "--model" => model = it.next().context("--model value")?,
            "--out" => out = it.next().context("--out value")?,
            "--preset" => preset = it.next().context("--preset value")?,
            "--tokenizer-pre" => tok_pre = it.next().context("--tokenizer-pre value")?,
            "--imatrix" => imatrix_path = it.next().context("--imatrix value")?,
            "--tensor-type" => {
                let kv = it.next().context("--tensor-type value")?;
                let (k, v) = kv.split_once('=').context("--tensor-type substr=type")?;
                overrides.push((k.to_string(), type_id(v)?));
            }
            other => bail!("unknown argument {other}"),
        }
    }
    ensure!(!model.is_empty() && !out.is_empty(), "--model and --out required");
    // Multi-quant batch (Unsloth's multi-export): comma-separated presets share one
    // safetensors read path; each emits <out stem>-<preset>.gguf.
    if preset.contains(',') {
        let presets: Vec<String> = preset.split(',').map(|x| x.trim().to_string()).collect();
        let base = std::path::Path::new(&out);
        let stem = base.file_stem().and_then(|x| x.to_str()).unwrap_or("out");
        let dir = base.parent().unwrap_or(std::path::Path::new("."));
        for p in &presets {
            let po = dir.join(format!("{stem}-{p}.gguf"));
            let mut args = vec![
                std::env::args().next().unwrap_or_default(),
                "--model".into(), model.clone(),
                "--out".into(), po.to_string_lossy().into_owned(),
                "--preset".into(), p.clone(),
                "--tokenizer-pre".into(), tok_pre.clone(),
            ];
            if !imatrix_path.is_empty() {
                args.push("--imatrix".into());
                args.push(imatrix_path.clone());
            }
            for (k, v) in &overrides {
                args.push("--tensor-type".into());
                args.push(format!("{k}={}", type_name(*v).to_ascii_lowercase()));
            }
            let st = std::process::Command::new(std::env::current_exe()?)
                .args(&args[1..])
                .status()?;
            ensure!(st.success(), "preset {p} failed");
        }
        return Ok(());
    }
    // Per-COLUMN importance (imatrix diagonal) by GGUF tensor name; JSON from
    // scripts/make_imatrix.py. Column space = the tensor's INPUT dim, so the llama q/k row
    // permute below does not touch it.
    let imatrix: std::collections::HashMap<String, Vec<f32>> = if imatrix_path.is_empty() {
        Default::default()
    } else {
        serde_json::from_slice(&std::fs::read(&imatrix_path).context("read --imatrix")?)
            .context("parse --imatrix json")?
    };
    let dir = std::path::Path::new(&model);

    // ── config + arch ───────────────────────────────────────────────────────────────────────
    let cfg_raw: serde_json::Value =
        serde_json::from_slice(&std::fs::read(dir.join("config.json"))?)?;
    let cfg = inferencelayer::weights::Lfm2Config::from_json(&std::fs::read(
        dir.join("config.json"),
    )?)?;
    use inferencelayer::weights::Arch;
    let (arch_name, is_llama) = match cfg.arch {
        Arch::Llama => ("llama", true),
        Arch::Qwen2 => ("qwen2", false),
        other => bail!("gguf-quantize handles llama/qwen2 checkpoints (got {other:?})"),
    };
    let st = inferencelayer::weights::lazyst_open_for_tools(dir)?;
    let (h, hd, nh, nkv) = (cfg.hidden, cfg.head_dim, cfg.n_heads, cfg.n_kv_heads);

    let mut g = GgufWriter::new();
    g.kv("general.architecture", WVal::Str(arch_name.into()));
    g.kv("general.name", WVal::Str("inferencelayer".into()));
    let k = |s: &str| format!("{arch_name}.{s}");
    g.kv(&k("embedding_length"), WVal::U32(h as u32));
    g.kv(&k("block_count"), WVal::U32(cfg.n_layers as u32));
    g.kv(&k("attention.head_count"), WVal::U32(nh as u32));
    g.kv(&k("attention.head_count_kv"), WVal::U32(nkv as u32));
    g.kv(&k("feed_forward_length"), WVal::U32(cfg.intermediate as u32));
    g.kv(&k("attention.layer_norm_rms_epsilon"), WVal::F32(cfg.eps));
    g.kv(&k("rope.freq_base"), WVal::F32(cfg.rope_theta));
    g.kv(&k("rope.dimension_count"), WVal::U32(if cfg.rotary_dim > 0 { cfg.rotary_dim } else { hd } as u32));
    g.kv(&k("attention.key_length"), WVal::U32(hd as u32));
    g.kv(&k("attention.value_length"), WVal::U32(hd as u32));
    g.kv(
        &k("context_length"),
        WVal::U32(cfg_raw["max_position_embeddings"].as_u64().unwrap_or(4096) as u32),
    );

    // ── tokenizer (BPE, gpt2-style) from tokenizer.json ─────────────────────────────────────
    let tj: serde_json::Value =
        serde_json::from_slice(&std::fs::read(dir.join("tokenizer.json"))?)
            .context("tokenizer.json")?;
    let vocab = tj["model"]["vocab"]
        .as_object()
        .context("tokenizer.json model.vocab")?;
    let mut tokens: Vec<(u32, String)> = vocab
        .iter()
        .map(|(t, id)| (id.as_u64().unwrap() as u32, t.clone()))
        .collect();
    let mut special_ids = std::collections::HashSet::new();
    if let Some(added) = tj["added_tokens"].as_array() {
        for a in added {
            let id = a["id"].as_u64().unwrap() as u32;
            if a["special"].as_bool().unwrap_or(false) {
                special_ids.insert(id);
            }
            if !vocab.contains_key(a["content"].as_str().unwrap_or("")) {
                tokens.push((id, a["content"].as_str().unwrap_or("").to_string()));
            }
        }
    }
    tokens.sort_by_key(|(id, _)| *id);
    // dense id space check — GGUF token arrays are positional
    for (i, (id, _)) in tokens.iter().enumerate() {
        ensure!(*id as usize == i, "non-dense vocab at {i} (id {id})");
    }
    let mut token_type: Vec<i32> = tokens
        .iter()
        .map(|(id, _)| if special_ids.contains(id) { 3 } else { 1 })
        .collect();
    let mut token_strs: Vec<String> = tokens.iter().map(|(_, t)| t.clone()).collect();
    // Checkpoints often PAD the embedding matrix past the real vocab (qwen2: 151936 rows vs
    // 151665 tokens); llama.cpp sizes token_embd from the TOKEN LIST, so pad it with UNUSED
    // entries exactly like convert_hf_to_gguf — without this our qwen2 artifacts load in the
    // engine but are REFUSED by llama.cpp (found by the chat-template interop check).
    let n_emb_rows = st.shape("model.embed_tokens.weight")?[0];
    for i in token_strs.len()..n_emb_rows {
        token_strs.push(format!("[PAD{i}]"));
        token_type.push(5); // UNUSED
    }
    let merges: Vec<String> = tj["model"]["merges"]
        .as_array()
        .map(|a| {
            a.iter()
                .map(|m| {
                    if let Some(s) = m.as_str() {
                        s.to_string() // "a b" form
                    } else {
                        let p = m.as_array().unwrap();
                        format!("{} {}", p[0].as_str().unwrap(), p[1].as_str().unwrap())
                    }
                })
                .collect()
        })
        .unwrap_or_default();
    g.kv("tokenizer.ggml.model", WVal::Str("gpt2".into()));
    g.kv("tokenizer.ggml.pre", WVal::Str(tok_pre));
    g.kv("tokenizer.ggml.tokens", WVal::StrArray(token_strs));
    g.kv("tokenizer.ggml.token_type", WVal::I32Array(token_type));
    g.kv("tokenizer.ggml.merges", WVal::StrArray(merges));
    if let Some(b) = cfg_raw["bos_token_id"].as_u64() {
        g.kv("tokenizer.ggml.bos_token_id", WVal::U32(b as u32));
    }
    if let Some(e) = cfg_raw["eos_token_id"].as_u64() {
        g.kv("tokenizer.ggml.eos_token_id", WVal::U32(e as u32));
    }
    // Instruct checkpoints: embed the chat template so llama.cpp/ollama render turns without
    // a hand-written TEMPLATE stanza (Unsloth's GGUF export does the same).
    if let Ok(tc) = std::fs::read(std::path::Path::new(&model).join("tokenizer_config.json")) {
        if let Ok(v) = serde_json::from_slice::<serde_json::Value>(&tc) {
            if let Some(t) = v["chat_template"].as_str() {
                g.kv("tokenizer.chat_template", WVal::Str(t.to_string()));
                eprintln!("  embedded chat_template ({} chars)", t.len());
            }
        }
    }

    // ── tensors ─────────────────────────────────────────────────────────────────────────────
    // llama arch: PERMUTE q/k rows to the ecosystem's interleaved-rotary convention (inverse of
    // the loader's un-permute: HF row f = a·(hd/2)+b → gguf row g = 2b+a, per head).
    let permute = |w: &[f32], n_head: usize| -> Vec<f32> {
        let mut outv = vec![0f32; w.len()];
        for hh in 0..n_head {
            for f in 0..hd {
                let (a, b) = (f / (hd / 2), f % (hd / 2));
                let gr = 2 * b + a;
                outv[(hh * hd + gr) * h..(hh * hd + gr) * h + h]
                    .copy_from_slice(&w[(hh * hd + f) * h..(hh * hd + f) * h + h]);
            }
        }
        outv
    };
    let pick = |gguf_name: &str, cols: usize| -> Result<u32> {
        let mut ty = preset_type(&preset, gguf_name)?;
        for (pat, t) in &overrides {
            if gguf_name.contains(pat.as_str()) {
                ty = *t;
            }
        }
        Ok(fallback(ty, cols))
    };
    let mut n_im_used = 0usize;
    let mut add_w = |g: &mut GgufWriter, gguf_name: &str, w: &[f32], rows: usize, cols: usize| -> Result<()> {
        let ty = pick(gguf_name, cols)?;
        let im = imatrix.get(gguf_name).map(|v| v.as_slice());
        if im.is_some() {
            n_im_used += 1;
        }
        let bytes = encode_ggml_im(w, cols, ty, im)?;
        eprintln!(
            "  {gguf_name:34} [{rows}x{cols}] {}{}",
            type_name(ty),
            if im.is_some() { " +imatrix" } else { "" }
        );
        g.tensor(gguf_name, &[rows, cols], ty, bytes);
        Ok(())
    };
    let f32_tensor = |g: &mut GgufWriter, gguf_name: &str, v: &[f32]| {
        let mut b = Vec::with_capacity(v.len() * 4);
        for x in v {
            b.extend_from_slice(&x.to_le_bytes());
        }
        g.tensor(gguf_name, &[v.len()], 0, b);
    };

    let tens = |name: &str| -> Result<Vec<f32>> { st.tensor_f32(name) };
    let shape = |name: &str| -> Result<Vec<usize>> { Ok(st.shape(name)?.to_vec()) };

    let emb = tens("model.embed_tokens.weight")?;
    let esh = shape("model.embed_tokens.weight")?;
    add_w(&mut g, "token_embd.weight", &emb, esh[0], esh[1])?;
    if st.has("lm_head.weight") {
        let hw = tens("lm_head.weight")?;
        let hs = shape("lm_head.weight")?;
        add_w(&mut g, "output.weight", &hw, hs[0], hs[1])?;
    }
    f32_tensor(&mut g, "output_norm.weight", &tens("model.norm.weight")?);
    for i in 0..cfg.n_layers {
        let p = format!("model.layers.{i}");
        let b = format!("blk.{i}");
        f32_tensor(&mut g, &format!("{b}.attn_norm.weight"), &tens(&format!("{p}.input_layernorm.weight"))?);
        f32_tensor(&mut g, &format!("{b}.ffn_norm.weight"), &tens(&format!("{p}.post_attention_layernorm.weight"))?);
        for (hf, gg, nheads, perm) in [
            ("self_attn.q_proj", "attn_q", nh, is_llama),
            ("self_attn.k_proj", "attn_k", nkv, is_llama),
            ("self_attn.v_proj", "attn_v", nkv, false),
            ("self_attn.o_proj", "attn_output", 0, false),
            ("mlp.gate_proj", "ffn_gate", 0, false),
            ("mlp.up_proj", "ffn_up", 0, false),
            ("mlp.down_proj", "ffn_down", 0, false),
        ] {
            let w = tens(&format!("{p}.{hf}.weight"))?;
            let sh = shape(&format!("{p}.{hf}.weight"))?;
            let w = if perm { permute(&w, nheads) } else { w };
            add_w(&mut g, &format!("{b}.{gg}.weight"), &w, sh[0], sh[1])?;
        }
        if !is_llama {
            for (hf, gg) in [
                ("self_attn.q_proj", "attn_q"),
                ("self_attn.k_proj", "attn_k"),
                ("self_attn.v_proj", "attn_v"),
            ] {
                let bias = tens(&format!("{p}.{hf}.bias"))?;
                f32_tensor(&mut g, &format!("{b}.{gg}.bias"), &bias);
            }
        }
    }
    g.write_to(std::path::Path::new(&out))?;
    let sz = std::fs::metadata(&out)?.len();
    println!(
        "wrote {out} ({:.1} MB, preset {preset}, {} overrides, {} imatrix-weighted tensors)",
        sz as f64 / 1e6,
        overrides.len(),
        n_im_used
    );
    Ok(())
}