hf2q 0.1.1

//! ADR-033 P3 — convert orchestration scaffolding.
//!
//! Wires `StandardPolicy::target_for` + `GgmlQuantizer::quantize` +
//! `GgufWriter` into a single self-contained pipeline driver, so the
//! integration shape can be exercised end-to-end on synthetic tensors
//! before the per-arch safetensors mappers (P4+) come online.
//!
//! This module is intentionally **not** wired into `cmd_convert` (the
//! legacy two-pass pipeline at `src/main.rs::cmd_convert`); the legacy
//! pipeline stays load-bearing until P6 deletes it. Per
//! [[feedback-no-backwards-compat-2026-05-18]]: this is new code — no
//! migration shims.
//!
//! Per [[feedback-no-loop-suppression-2026-05-17]]: the orchestrator
//! returns typed errors; the only place where a tensor escapes the
//! `StandardPolicy` → `GgmlQuantizer` pipeline is the
//! vision/audio-pattern gate in [`crate::quantize::ggml_quants::vision`].
//! Matched tensors emit F16 directly; unmatched tensors that fail the
//! policy / quantizer surface as `OrchestratorError::*` — never silent
//! F16 demotion.

pub mod arch;
pub mod cli_driver;
pub mod model_card;
pub mod orchestrator;
pub mod quant_selector;
pub mod receipt;
pub mod sleef_expf;
pub mod source_dtype;
pub mod source_reader;
pub mod tokenizer;

pub use cli_driver::{run_convert, ConvertArgs, ConvertError};
pub use orchestrator::{
    ConvertOrchestrator, OrchestratorError, PlanEntry, PlannedSizeSummary, PlannedTypeSize,
    StreamingWriter,
};
pub use quant_selector::{approximate_for_apex, QuantSelector, QuantSelectorError};
pub use receipt::{ReceiptError, RemoteConversionSource};
pub use source_reader::{HfModelSource, HfTensor, SourceError, TensorMeta};
pub use tokenizer::{build_tokenizer_metadata, TokenizerError};

#[cfg(test)]
mod tests {
    //! End-to-end acceptance test for the new Llama3 convert path:
    //! synthesize a tiny Llama-3-shaped safetensors directory + config,
    //! run it through `HfModelSource::open` + `map_tensor_name` +
    //! `ConvertOrchestrator`, then re-parse via `mlx_native::gguf::GgufFile`
    //! and verify every tensor name + shape round-trips.

    use super::*;
    use crate::convert::arch::llama3;
    use crate::quantize::ggml_quants::standard_policy::HParams;
    use crate::quantize::ggml_quants::{ArchName, LlamaFtype};
    use safetensors::tensor::{Dtype, TensorView};
    use std::collections::HashMap;
    use std::fs;
    use std::path::Path;

    /// Tiny synthetic Llama-3-shaped model.
    ///
    /// Hyperparameters chosen so EVERY weight's innermost dim (GGUF
    /// `n_per_row` = PyTorch `in_dim`) is a multiple of 32, which lets us
    /// quantize to Q8_0 with no shape fallback. Q8_0 keeps the
    /// quantization step cheap (no K-quant superblock work) so the test
    /// is fast while still exercising the real `Quantizer` dispatch
    /// path.
    ///
    /// Rationale for Q8_0 instead of AllF32 (the mission spec's
    /// suggestion): `quantizer_for(GgmlType::F32)` returns
    /// `Err(NoQuantizerForType)` today — F32 is a pass-through type with
    /// no `Quantizer` impl by design (see ADR-033 §"Quantizer trait"
    /// hard-error contract + `quantizer.rs:103-117`). Modifying the
    /// orchestrator to pass F32 through is outside the mission's "DO
    /// NOT modify any file outside src/convert/" guardrail. Q8_0 keeps
    /// the test fast and the path-mapping logic isolated — every
    /// shape only multiplies through the trivial scale-max kernel.
    ///
    /// Synthetic shape table (iter-23+ workers extending to other arches
    /// can copy this pattern; just substitute the per-arch tensor
    /// inventory):
    ///
    /// - n_layers          = 2
    /// - hidden_size       = 32
    /// - ffn_size          = 64
    /// - n_heads           = 2
    /// - n_kv_heads        = 2  (MHA — keeps k/v shapes square at hidden)
    /// - head_dim          = hidden / n_heads = 16
    /// - vocab_size        = 64
    /// - max_pos           = 8192
    /// - rms_eps           = 1e-5
    /// - rope_theta        = 10000.0
    ///
    /// Per-tensor PyTorch shapes (HF convention: `[out, in]` for
    /// Linear; `[vocab, hidden]` for embeddings; `[hidden]` for
    /// RMSNorm; weights in F32 source-dtype for byte-equality):
    ///
    /// | HF name                                              | shape          |
    /// |------------------------------------------------------|----------------|
    /// | `model.embed_tokens.weight`                          | `[64, 32]`     |
    /// | `model.norm.weight`                                  | `[32]`         |
    /// | `lm_head.weight`                                     | `[64, 32]`     |
    /// | `model.layers.<i>.input_layernorm.weight`            | `[32]`         |
    /// | `model.layers.<i>.post_attention_layernorm.weight`   | `[32]`         |
    /// | `model.layers.<i>.self_attn.q_proj.weight`           | `[32, 32]`     |
    /// | `model.layers.<i>.self_attn.k_proj.weight`           | `[32, 32]`     |
    /// | `model.layers.<i>.self_attn.v_proj.weight`           | `[32, 32]`     |
    /// | `model.layers.<i>.self_attn.o_proj.weight`           | `[32, 32]`     |
    /// | `model.layers.<i>.mlp.gate_proj.weight`              | `[64, 32]`     |
    /// | `model.layers.<i>.mlp.up_proj.weight`                | `[64, 32]`     |
    /// | `model.layers.<i>.mlp.down_proj.weight`              | `[32, 64]`     |
    ///
    /// Total: 3 globals + 9 per-block × 2 layers = 21 tensors. Of
    /// those, 16 are 2-D (token_embd + lm_head + 7 per-block × 2
    /// layers); the other 5 are 1-D norms that we skip from the
    /// orchestrator feed (see the policy/norm caveat above).
    ///
    /// **Caveat for the 1-D RMSNorm weights.** GGUF requires
    /// `n_per_row` (innermost dim) to be a multiple of the block size
    /// for the picked `GgmlType`. A 1-D `[32]` norm tensor has
    /// `n_per_row=32` which IS a multiple of 32 → Q8_0 OK. But the
    /// `StandardPolicy` keeps norm tensors as F32 anyway (per
    /// `target_for`'s norm-tensor branch); since `quantizer_for(F32)`
    /// rejects, the test would fail if we tried to run them through.
    /// Three options: (a) skip norms in the harness (they don't
    /// exercise tensor-name mapping in any interesting way), (b) move
    /// the orchestrator's F32 pass-through to the dispatcher (out of
    /// mission scope), (c) push norms through the policy via a
    /// non-norm name. We pick (a) — the test's job is the convert
    /// path-mapping, not the policy's norm-tensor handling.
    fn synthesize_tiny_llama3(dir: &Path) {
        const HIDDEN: usize = 32;
        const FFN: usize = 64;
        const VOCAB: usize = 64;
        const LAYERS: usize = 2;

        let mut tensors: Vec<(String, Vec<usize>, Vec<u8>)> = Vec::new();

        let mk_f32_bytes = |numel: usize, seed: u32| -> Vec<u8> {
            (0..numel)
                .flat_map(|i| {
                    let x = ((i as u32).wrapping_mul(2654435761).wrapping_add(seed)) as i32;
                    let f = (x as f32) / (i32::MAX as f32);
                    f.to_le_bytes()
                })
                .collect()
        };

        // Globals (3). We deliberately INCLUDE `model.norm.weight` even
        // though we'll exclude it from the orchestrator feed (norm
        // tensors hit the policy's F32 passthrough — see the
        // module-level caveat above). Including it in the safetensors
        // file exercises `HfModelSource::open` for 1-D shapes too.
        tensors.push((
            "model.embed_tokens.weight".into(),
            vec![VOCAB, HIDDEN],
            mk_f32_bytes(VOCAB * HIDDEN, 1),
        ));
        tensors.push((
            "model.norm.weight".into(),
            vec![HIDDEN],
            mk_f32_bytes(HIDDEN, 2),
        ));
        tensors.push((
            "lm_head.weight".into(),
            vec![VOCAB, HIDDEN],
            mk_f32_bytes(VOCAB * HIDDEN, 3),
        ));

        // Per-block (11 × 2 = 22).
        for li in 0..LAYERS {
            let s = (li as u32) * 100;
            tensors.push((
                format!("model.layers.{li}.input_layernorm.weight"),
                vec![HIDDEN],
                mk_f32_bytes(HIDDEN, s + 10),
            ));
            tensors.push((
                format!("model.layers.{li}.post_attention_layernorm.weight"),
                vec![HIDDEN],
                mk_f32_bytes(HIDDEN, s + 11),
            ));
            tensors.push((
                format!("model.layers.{li}.self_attn.q_proj.weight"),
                vec![HIDDEN, HIDDEN],
                mk_f32_bytes(HIDDEN * HIDDEN, s + 12),
            ));
            tensors.push((
                format!("model.layers.{li}.self_attn.k_proj.weight"),
                vec![HIDDEN, HIDDEN],
                mk_f32_bytes(HIDDEN * HIDDEN, s + 13),
            ));
            tensors.push((
                format!("model.layers.{li}.self_attn.v_proj.weight"),
                vec![HIDDEN, HIDDEN],
                mk_f32_bytes(HIDDEN * HIDDEN, s + 14),
            ));
            tensors.push((
                format!("model.layers.{li}.self_attn.o_proj.weight"),
                vec![HIDDEN, HIDDEN],
                mk_f32_bytes(HIDDEN * HIDDEN, s + 15),
            ));
            tensors.push((
                format!("model.layers.{li}.mlp.gate_proj.weight"),
                vec![FFN, HIDDEN],
                mk_f32_bytes(FFN * HIDDEN, s + 16),
            ));
            tensors.push((
                format!("model.layers.{li}.mlp.up_proj.weight"),
                vec![FFN, HIDDEN],
                mk_f32_bytes(FFN * HIDDEN, s + 17),
            ));
            tensors.push((
                format!("model.layers.{li}.mlp.down_proj.weight"),
                vec![HIDDEN, FFN],
                mk_f32_bytes(HIDDEN * FFN, s + 18),
            ));
        }

        // Wrap in TensorViews + serialize. Views must borrow from the
        // owned byte vectors; keep them alive across the call.
        let views: Vec<(String, TensorView<'_>)> = tensors
            .iter()
            .map(|(n, sh, b)| {
                let v = TensorView::new(Dtype::F32, sh.clone(), b).expect("TensorView");
                (n.clone(), v)
            })
            .collect();
        let view_refs: Vec<(String, &TensorView<'_>)> =
            views.iter().map(|(n, v)| (n.clone(), v)).collect();
        let st_bytes =
            safetensors::tensor::serialize(view_refs, None).expect("serialize safetensors");
        fs::write(dir.join("model.safetensors"), st_bytes).expect("write safetensors");

        // Minimal Llama-3 config.json. Every field referenced by
        // `llama3::build_metadata` is present; nothing else (the mapper
        // tolerates missing optional keys but not missing required
        // ones).
        let cfg = serde_json::json!({
            "_name_or_path": "synthetic/Llama-3-Tiny-Test",
            "model_type": "llama",
            "hidden_size": HIDDEN,
            "num_hidden_layers": LAYERS,
            "intermediate_size": FFN,
            "num_attention_heads": 2,
            "num_key_value_heads": 2,
            "max_position_embeddings": 8192,
            "rms_norm_eps": 1.0e-5,
            "rope_theta": 10000.0,
            "vocab_size": VOCAB,
        });
        fs::write(
            dir.join("config.json"),
            serde_json::to_string_pretty(&cfg).unwrap(),
        )
        .expect("write config.json");
    }

    /// Acceptance test 3 — end-to-end safetensors → GGUF round-trip.
    /// See module doc for the synthetic-model shape table.
    #[test]
    fn llama3_end_to_end_tiny_safetensors_round_trip() {
        // ----- 1. Synthesize the tiny safetensors directory ------------
        let dir = tempfile::tempdir().unwrap();
        synthesize_tiny_llama3(dir.path());

        // ----- 2. Open via HfModelSource (streaming-mmap, no buffered Vec) ----
        let src = HfModelSource::open(dir.path()).expect("HfModelSource::open");
        assert_eq!(
            src.tensor_count(),
            3 + 9 * 2,
            "expected 21 HF tensors in the synthetic model (3 globals + 9 per-block × 2 layers)"
        );
        // Materialize them into a vec for the test's downstream
        // assertions. (Real driver code streams one at a time; this test
        // checks the path-mapping, so it doesn't matter that the test
        // itself collects.)
        let src_tensors: Vec<HfTensor> = src
            .iter_tensors()
            .collect::<Result<Vec<_>, _>>()
            .expect("stream");

        // ----- 3. Build the orchestrator at MostlyQ8_0 + Llama3 ------
        // Q8_0 has block_size=32; every weight's innermost dim is 32 or
        // 64, both multiples — no shape fallback hit. HParams are read
        // from config.json (mirrors the eventual cmd_convert wiring).
        let hidden = src.config["hidden_size"].as_u64().unwrap() as u32;
        let n_head = src.config["num_attention_heads"].as_u64().unwrap() as u32;
        let n_head_kv = src.config["num_key_value_heads"].as_u64().unwrap() as u32;
        let _ = hidden; // silence warning if HParams ever drops it

        let mut orch = ConvertOrchestrator::new(
            LlamaFtype::MostlyQ8_0,
            ArchName::Llama3,
            HParams {
                n_expert: 0,
                n_head,
                n_head_kv,
                n_layer: 2,
                n_mtp_layers: 0,
            },
        );

        // ----- 4. Emit metadata -------------------------------------
        for (k, v) in llama3::build_metadata(&src.config, 7 /* Q8_0 */, None, None, None) {
            orch.add_metadata(k, v);
        }

        // ----- 5. Map + plan every tensor ---------------------------
        // Strategy:
        //   - Map HF name → GGUF name via `llama3::map_tensor_name`.
        //   - Skip 1-D RMSNorm weights (policy → F32 → no Quantizer;
        //     see module-level caveat). Both `attn_norm` (post-map of
        //     `input_layernorm`) and `ffn_norm` are 1-D; the global
        //     `output_norm.weight` is also 1-D. We still assert these
        //     three name kinds map correctly via map_tensor_name
        //     above-band (test 1).
        //   - For 2-D tensors: reverse PyTorch shape `[out, in]` to
        //     GGUF shape `[in, out]` and feed through the orchestrator.
        //   - `layer_index` derived from the GGUF name (`blk.<N>.*`)
        //     since the orchestrator uses it for QsState counters.
        let mut expected_gguf_names: Vec<String> = Vec::new();
        let mut expected_shapes: HashMap<String, Vec<usize>> = HashMap::new();
        let mut plan: Vec<PlanEntry> = Vec::new();
        let mut datas: Vec<Vec<f32>> = Vec::new();
        for ht in &src_tensors {
            let gguf_name = llama3::map_tensor_name(&ht.name)
                .unwrap_or_else(|| panic!("unmapped HF tensor `{}`", ht.name));

            // Skip 1-D norm tensors — they'd hit the policy's F32
            // passthrough which has no Quantizer impl.
            if ht.shape.len() == 1 {
                continue;
            }

            // PyTorch [out, in] → GGUF [in, out]. The orchestrator
            // requires GGUF order (innermost-first).
            let gguf_shape: Vec<usize> = ht.shape.iter().rev().copied().collect();

            // `mlx_native::gguf::GgufFile` REVERSES on parse — so the
            // shape we'll see back on disk is the original PyTorch
            // order. Capture both for the round-trip assert.
            expected_shapes.insert(gguf_name.clone(), ht.shape.clone());
            expected_gguf_names.push(gguf_name.clone());

            // Layer index parsed from the canonical GGUF name. Globals
            // (`token_embd.weight`, `output.weight`) carry `None`.
            let layer_index = gguf_name
                .strip_prefix("blk.")
                .and_then(|s| s.split('.').next())
                .and_then(|s| s.parse::<usize>().ok());

            plan.push(PlanEntry {
                name: gguf_name,
                shape: gguf_shape,
                source_dtype: ht.source_dtype,
                layer_index,
            });
            datas.push(ht.data.clone());
        }

        // Sanity: 2-D tensor count.
        // Globals (2-D): embed, lm_head = 2. `output_norm` is 1-D → skipped.
        // Per-block (2-D): q, k, v, o, gate, up, down = 7 weights.
        // Total: 2 + 7×2 = 16.
        assert_eq!(
            expected_gguf_names.len(),
            16,
            "expected 16 quantizable 2-D tensors (2 globals + 7 per-block × 2 layers)"
        );

        orch.plan_tensors(plan).expect("plan");

        // ----- 6. Write to disk + reparse ---------------------------
        let out = tempfile::NamedTempFile::new().unwrap();
        {
            let f = fs::File::create(out.path()).unwrap();
            let mut sw = orch.begin_write(f).expect("begin_write");
            for (idx, d) in datas.iter().enumerate() {
                sw.stream_tensor(idx, d).expect("stream");
            }
            sw.finalize().expect("finalize");
        }

        let gguf = mlx_native::gguf::GgufFile::open(out.path()).expect("parse output GGUF");

        // ----- 7. Assertions ----------------------------------------
        // Tensor count: 16 (2 globals + 7×2 per-block).
        assert_eq!(gguf.tensor_count(), 16);

        // Metadata round-trip: 18 Llama3 KV pairs present (architecture,
        // type, name, basename, size_label, 12 llama.*,
        // quantization_version, file_type).
        assert_eq!(gguf.metadata_count(), 18);
        assert_eq!(gguf.metadata_string("general.architecture"), Some("llama"));
        assert_eq!(gguf.metadata_u32("llama.embedding_length"), Some(32));
        assert_eq!(gguf.metadata_u32("llama.block_count"), Some(2));
        assert_eq!(gguf.metadata_u32("llama.attention.head_count"), Some(2));
        assert_eq!(gguf.metadata_u32("general.file_type"), Some(7));

        // Every staged GGUF name appears with the right shape. The
        // reader reverses on parse — we wrote GGUF order `[in, out]`,
        // it returns PyTorch order `[out, in]`. We compare against the
        // ORIGINAL PyTorch shape (captured in `expected_shapes`).
        for name in &expected_gguf_names {
            let info = gguf
                .tensor_info(name)
                .unwrap_or_else(|| panic!("missing GGUF tensor `{name}`"));
            let want = &expected_shapes[name];
            assert_eq!(&info.shape, want, "shape round-trip failed for `{name}`");
            // Q8_0 has wire-position 3 in mlx_native's positional enum
            // (F32=0, F16=1, Q4_0=2, Q8_0=3, ...). Token_embd + output
            // route through the OUTPUT-branch which bumps to Q6_K
            // (position 6) at Q8_0 ftype only if `new_type != Q8_0`;
            // since new_type IS Q8_0 for `MostlyQ8_0` ftype, the bump
            // at C:432-435 does NOT fire — token_embd + output stay
            // Q8_0. Sanity that the policy did its job.
            assert_eq!(info.ggml_type as u32, 3, "{name} → Q8_0 (positional 3)");
        }
    }
}