cera 0.5.5

//! Cross-implementation correctness gate for the dense text models cera serves
//! through `LlamaModel` — NEOX-rope (Qwen2, Qwen3) and NORM-rope (LLaMA, Mistral,
//! Granite 3.x) — against golden fixtures generated from upstream llama.cpp on the
//! *same* quantized GGUF (see `scripts/oracle/`). Granite additionally exercises
//! the embedding/residual/attention/logit scalar multipliers (folded into the
//! gated `l_out-{i}`, `embd`, and `result_output` sums).
//!
//! Iterates every fixture set under `tests/fixtures/oracle/<model>/`; each
//! `index.json` names its `model_file`, looked up under `target/oracle/models/`
//! (override the dir with `CERA_ORACLE_MODELS_DIR`). The whole suite is inert
//! unless `CERA_ORACLE=1` is set, which no CI workflow does, so CI never reaches
//! it regardless of which fixtures happen to be on the runner. Once that opt-in
//! IS set, an absent GGUF is a hard failure rather than a skip: see the
//! assertion in `text_models_match_llama_cpp_oracle`.
//!
//! Two gates per prompt, plus one informational signal:
//!   1. tokenizer parity   — cera's `encode` matches llama.cpp's input tokens
//!   2. per-layer sums      — cera's activation `sum` checksums match llama.cpp's,
//!      keyed by (node name, op) so repeated node names don't collide. Covers the
//!      embedding and every layer's residual-stream output (`l_out-{i}` — captures
//!      that layer's attention/rope/bias/FFN). Deterministic and localizes any
//!      math bug to the first diverging layer. The final-logit `result_output`
//!      sum is NOT gated: it sums ~10^5 partially-cancelling logits, so its
//!      relative diff is both noisy (Q8_0 accumulation doesn't average out) and
//!      insensitive (a wrong rope convention barely moves it) — reported as info.
//!      • greedy continuation (informational) — cera's `--temp 0` argmax decode vs
//!      llama.cpp's greedy text, reported MATCH / DIVERGES but never gated:
//!      greedy decode flips at near-tied logits, and Q8_0 noise tips those ties
//!      into a different-but-coherent continuation that is not a bug.
//!
//! Gated behind `CERA_ORACLE=1` and `#[ignore]`. Run:
//!   CERA_ORACLE=1 cargo test -p cera --release --test oracle_text -- --ignored --nocapture

#![cfg(feature = "mmap")]

use std::collections::{BTreeSet, HashMap};
use std::path::PathBuf;

use cera::gguf::GgufFile;
use cera::kv_cache::{InferenceState, KvCompression};
use cera::model::Model;
use cera::model::llama::LlamaModel;
use cera::model::transformer::oracle_dump;
use cera::sampler::argmax;
use cera::tokenizer::BpeTokenizer;

/// Relative difference, robust near zero. Catches gross math bugs (sign flip,
/// wrong layout → diffs ≫ 1) while tolerating Q8_0 accumulation-order noise
/// between cera's and llama.cpp's CPU kernels.
fn rel_diff(a: f64, b: f64) -> f64 {
    (a - b).abs() / (a.abs() + b.abs() + 1e-9)
}

/// Per-node relative tolerance for the sum gate. Sized to clear the observed
/// Q8_0 cross-implementation accumulation noise (≤ ~3.5% on early small-sum
/// residuals across Qwen2 + Qwen3, all prompts) with margin, while still
/// catching real math bugs — a sign flip / wrong layout / misapplied bias
/// shifts a node's sum by tens of percent to >100%, far above this.
const SUM_REL_TOL: f64 = 0.05;

/// Below this absolute sum, cancellation makes the *relative* diff meaningless,
/// so the node is reported but not gated. A real bug propagates to the many
/// large-magnitude nodes downstream, so this loses no coverage.
const SUM_MAG_FLOOR: f64 = 10.0;

fn fixtures_root() -> PathBuf {
    PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/oracle")
}

fn models_dir() -> PathBuf {
    if let Ok(d) = std::env::var("CERA_ORACLE_MODELS_DIR") {
        return PathBuf::from(d);
    }
    PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../target/oracle/models")
}

/// Encode a prompt the same way the oracle did: byte-level BPE, with a leading
/// BOS to match llama.cpp's special-token prefixing.
///
/// cera's `encode()` does not prepend BOS — that's the Session / chat-template
/// layer's job — so the harness normalizes that prefix here to isolate BPE-merge
/// parity. We can't read it off `tokenizer.ggml.add_bos_token`: that key is
/// *absent* on some BPE GGUFs (e.g. `llama-bpe`), yet llama.cpp still prepends
/// BOS by vocab-type default.
///
/// Decide from the *whole* golden sequence rather than `want_tokens.first()`
/// alone: only prepend when the golden is exactly `[bos] ++ encode(prompt)`.
/// Keying on the full body is unambiguous even when the prompt's first *content*
/// token legitimately equals `bos_id` and the model does NOT add BOS — e.g. Qwen
/// (`bos_id == eos_id == <|endoftext|>`, `add_bos_token = false`) on a prompt
/// that literally starts with "<|endoftext|>": a first-token heuristic would
/// double-prepend and false-fail the tokenizer gate. A genuine BPE divergence in
/// the body still surfaces — the `[1..] == base` check fails, no BOS is added,
/// and the mismatch is reported.
fn encode_with_bos(tok: &BpeTokenizer, want_tokens: &[u32], prompt: &str) -> Vec<u32> {
    let base = tok.encode(prompt);
    match tok.bos_token() {
        Some(bos)
            if want_tokens.first() == Some(&bos)
                && want_tokens.len() == base.len() + 1
                && want_tokens[1..] == base[..] =>
        {
            let mut tokens = Vec::with_capacity(base.len() + 1);
            tokens.push(bos);
            tokens.extend_from_slice(&base);
            tokens
        }
        _ => base,
    }
}

/// Outcome of validating one model's fixture set.
enum ModelOutcome {
    /// The GGUF was present and every prompt ran. Carries the gate failures,
    /// empty meaning pass.
    Checked(Vec<String>),
    /// The GGUF was absent, so nothing in this fixture set was verified. Fatal
    /// at the call site rather than here, so one message can name all of them.
    /// The path already ends in the file name, so it carries both.
    Absent { path: PathBuf },
}

/// Validate one model's fixture set.
fn check_model(fixture_dir: &std::path::Path) -> ModelOutcome {
    // Committed fixtures must parse; corruption is a hard failure.
    let index_path = fixture_dir.join("index.json");
    let index: serde_json::Value = serde_json::from_str(
        &std::fs::read_to_string(&index_path)
            .unwrap_or_else(|e| panic!("read {}: {e}", index_path.display())),
    )
    .unwrap_or_else(|e| panic!("parse {}: {e}", index_path.display()));
    let model_file = index["model_file"].as_str().unwrap();
    let mp = models_dir().join(model_file);
    if !mp.exists() {
        return ModelOutcome::Absent { path: mp };
    }
    eprintln!("=== oracle model: {} ===", mp.display());

    let gguf = GgufFile::open(&mp).expect("open gguf");
    let tokenizer = BpeTokenizer::from_gguf(&gguf).expect("tokenizer");
    let model = LlamaModel::from_gguf(GgufFile::open(&mp).expect("open gguf"), 8192)
        .expect("load LlamaModel");
    let n_layers = model.config().n_layers;
    let last = n_layers - 1;

    // Each cera-recorded node maps to exactly one llama.cpp graph op, so the
    // oracle lookup is keyed by (name, op): names like `Qcur-{i}` repeat across
    // ops (MUL_MAT/ADD/RESHAPE/ROPE) in the fixture, and a name-only key would
    // silently collide. cera records the embedding, each layer's residual-stream
    // output (`l_out-{i}`), the final norm, and the logits.
    //
    // Why the residual stream and not finer sub-steps: `l_out-{i}` is the layer's
    // output, so any bug in that layer's attention/rope/bias/FFN surfaces here
    // and localizes to the layer. The finer post-rope Q/K sums were evaluated as
    // gate nodes and rejected — rope rotation makes them cancel toward zero, so
    // their relative sum diff is a noisy checksum (>10% under pure Q8_0 noise).
    //
    // Which op a node maps to is per-arch DETERMINISTIC, so derive the exact
    // expected op from the model's own scalar metadata rather than accepting an
    // OR-list of "acceptable" ops. Granite scales the embeddings and the logits,
    // so its "embd" callback fires on a SCALE node (not GET_ROWS) and
    // "result_output" on SCALE (not MUL_MAT); plain archs never scale, so they
    // stay GET_ROWS / MUL_MAT. An exact per-arch op keeps the gate a precise
    // contract — it catches a genuinely wrong mapping (e.g. an arch that should
    // scale `embd` but emits GET_ROWS), which a first-present-wins list could not.
    let scalars = model.config().scalars;
    let expected_op = |name: &str| -> &'static str {
        if name == "embd" {
            if scalars.embedding != 1.0 {
                "SCALE"
            } else {
                "GET_ROWS"
            }
        } else if name.starts_with("l_out-") {
            "ADD"
        } else if name == "result_norm" {
            "MUL"
        } else if name == "result_output" {
            if scalars.logit != 1.0 {
                "SCALE"
            } else {
                "MUL_MAT"
            }
        } else {
            panic!("unmapped oracle node name {name:?}")
        }
    };

    // llama.cpp prunes its graph so only the LAST token flows past the final
    // layer: every node of layer `last` (and the result_* nodes) covers one
    // position; earlier layers cover all positions. Fold cera occurrences to
    // match: sum all-position nodes over tokens, take the last occurrence for
    // last-position nodes.
    let last_pos_only =
        |name: &str| name.starts_with("result_") || name.ends_with(&format!("-{last}"));
    // `l_out-{last}`, `result_norm`, and `result_output` are single-token sums
    // whose *relative* diff is a weak cross-impl checksum:
    //   - `l_out-{last}` / `result_norm` are post-residual sums that nearly
    //     cancel (≈ 0), so tiny Q8_0 accumulation noise blows up their rel diff.
    //   - `result_output` sums ~10^5 partially-cancelling logits (128k vocab on
    //     Llama-3), so Q8_0 accumulation-order differences don't average out in
    //     the sum either — its rel diff drifts past tol on some short prompts
    //     even when the argmax (the value users consume) is identical (verified:
    //     cera's greedy decode matched llama.cpp byte-for-byte on the prompt that
    //     tripped the old gate). It's also insensitive (a wrong rope convention
    //     barely moves it), so it's a poor gate at any tolerance.
    // The sensitive, localizing gate is the per-layer residual sums (all gated);
    // `embd` is exact (0.0000) and covers the tied output-projection weights. The
    // greedy continuation below is an additional informational cross-check.
    // Report these three sums informationally rather than gating them.
    let informational = |name: &str| {
        name == "result_norm" || name == "result_output" || name == format!("l_out-{last}")
    };

    let mut failures = Vec::new();
    for entry in index["prompts"].as_array().unwrap() {
        let fname = entry["fixture"].as_str().unwrap();
        let fx: serde_json::Value =
            serde_json::from_str(&std::fs::read_to_string(fixture_dir.join(fname)).unwrap())
                .unwrap();
        let prompt = fx["prompt"].as_str().unwrap();
        let want_tokens: Vec<u32> = fx["input_tokens"]
            .as_array()
            .unwrap()
            .iter()
            .map(|v| v.as_u64().unwrap() as u32)
            .collect();

        // Gate 1 — tokenizer parity.
        let got_tokens = encode_with_bos(&tokenizer, &want_tokens, prompt);
        if got_tokens != want_tokens {
            failures.push(format!(
                "[{fname}] tokenizer mismatch:\n    cera: {got_tokens:?}\n    llama:{want_tokens:?}"
            ));
            continue; // tokenization drives the forward pass; sum gate moot
        }

        // Gate 2 — per-substep sum checksums. Prefill with the dump active.
        let mut state =
            InferenceState::from_config_with_compression(model.config(), &KvCompression::None)
                .unwrap();
        oracle_dump::begin();
        let _ = model.forward_prefill(&got_tokens, 0, &mut state);
        let occ = oracle_dump::take();

        let mut cera: HashMap<String, f64> = HashMap::new();
        for (name, sum) in occ {
            if last_pos_only(&name) {
                cera.insert(name, sum); // last wins
            } else {
                *cera.entry(name).or_insert(0.0) += sum;
            }
        }
        // Oracle sums keyed by (name, op) so per-substep nodes don't collide.
        let mut want: HashMap<(&str, &str), f64> = HashMap::new();
        for node in fx["nodes"].as_array().unwrap() {
            want.insert(
                (node["name"].as_str().unwrap(), node["op"].as_str().unwrap()),
                node["sum"].as_f64().unwrap(),
            );
        }

        let mut worst = 0.0f64;
        let mut checked = 0usize;
        for (name, &got) in &cera {
            let op = expected_op(name);
            let Some(&exp) = want.get(&(name.as_str(), op)) else {
                failures.push(format!(
                    "[{fname}] oracle has no node {name:?} with op {op:?}"
                ));
                continue;
            };
            let d = rel_diff(got, exp);
            // Skip nodes that can't be a reliable checksum: explicitly noisy
            // ones, and any whose oracle sum is near zero (cancellation makes the
            // relative diff meaningless). Real bugs still surface on the many
            // large-magnitude nodes downstream.
            if informational(name) || exp.abs() < SUM_MAG_FLOOR {
                eprintln!("[{fname}] (info) {name}/{op}: cera={got:.4} llama={exp:.4} rel={d:.4}");
                continue;
            }
            checked += 1;
            worst = worst.max(d);
            if d > SUM_REL_TOL {
                failures.push(format!(
                    "[{fname}] sum mismatch at {name}/{op}: cera={got:.4} llama={exp:.4} rel={d:.4}"
                ));
            }
        }
        // Most per-layer residuals should be gated (a handful of near-zero-sum
        // ones are legitimately skipped by the magnitude floor). A much smaller
        // count means the instrumentation or fixtures drifted.
        assert!(
            checked >= n_layers / 2,
            "[{fname}] only {checked} nodes checked — instrumentation/fixture drift"
        );
        eprintln!("[{fname}] sums OK — {checked} gated nodes, worst rel diff {worst:.5}");

        // Greedy continuation (end-to-end argmax) — INFORMATIONAL, not gated.
        // This is the human-meaningful final-output signal, but it cannot be a
        // hard gate: greedy decode flips whenever two tokens are near-tied, and
        // tiny Q8_0 cross-impl noise tips those ties either way. Empirically a
        // few prompts diverge into a *different but equally coherent*
        // continuation (e.g. two valid Spanish replies that split at token 0),
        // which is not a bug. So decode cera's continuation, compare to
        // llama.cpp's greedy text, and report MATCH / DIVERGES without failing —
        // a MATCH confirms exact argmax-path agreement, a DIVERGES is a prompt
        // to eyeball (a real projection bug yields garbage, not a plausible
        // alternative). The gated per-layer sums + `embd` (exact, and the tied
        // output-projection weights) remain the actual correctness gate.
        let n_predict = index["n_predict"].as_u64().unwrap_or(16) as usize;
        let want_text = fx["greedy_text"].as_str().unwrap().trim_end();
        let mut gstate =
            InferenceState::from_config_with_compression(model.config(), &KvCompression::None)
                .unwrap();
        let mut logits = model.forward_prefill(&got_tokens, 0, &mut gstate);
        let mut out_tokens: Vec<u32> = Vec::new();
        for _ in 0..n_predict {
            let next = argmax(&logits);
            // llama.cpp's greedy decode stops at EOS and does not render it, so
            // match that: break before appending the EOS token.
            if tokenizer.eos_token() == Some(next) {
                break;
            }
            out_tokens.push(next);
            logits = model.forward(&[next], gstate.seq_len, &mut gstate);
        }
        let got_text = tokenizer.decode(&out_tokens);
        let got_text = got_text.trim_end();
        if got_text == want_text {
            eprintln!("[{fname}] greedy MATCH — {got_text:?}");
        } else {
            eprintln!(
                "[{fname}] greedy DIVERGES (tie-flip; not gated):\n    cera: {got_text:?}\n    llama:{want_text:?}"
            );
        }
    }
    // Prefix each failure with the fixture set it came from. The sets share most
    // of their prompt file names (`1-2-3.json`, `the-capital-of-france-is.json`,
    // …), so an unprefixed line cannot be attributed to a model once the caller
    // is looking at failures from more than one.
    let set_name = fixture_dir
        .file_name()
        .map(|n| n.to_string_lossy().into_owned())
        .unwrap_or_else(|| fixture_dir.display().to_string());
    ModelOutcome::Checked(
        failures
            .into_iter()
            .map(|f| format!("{set_name} {f}"))
            .collect(),
    )
}

#[test]
#[ignore] // run with --ignored + CERA_ORACLE=1
fn text_models_match_llama_cpp_oracle() {
    if std::env::var("CERA_ORACLE").as_deref() != Ok("1") {
        eprintln!("skipping: CERA_ORACLE=1 not set");
        return;
    }

    let mut dirs: Vec<PathBuf> = std::fs::read_dir(fixtures_root())
        .expect("read fixtures dir")
        .filter_map(|e| e.ok().map(|e| e.path()))
        .filter(|p| p.is_dir())
        .collect();
    dirs.sort();
    assert!(!dirs.is_empty(), "no oracle fixture sets found");

    let mut all_failures = Vec::new();
    // A BTreeSet, not a Vec: `llama-3_2-1b` and `llama-3_2-1b-long` are two
    // fixture sets over one GGUF, so a Vec would name that file twice and
    // inflate the count. Ordering by path keeps the report deterministic.
    let mut missing: BTreeSet<PathBuf> = BTreeSet::new();
    for dir in &dirs {
        match check_model(dir) {
            ModelOutcome::Checked(failures) => all_failures.extend(failures),
            ModelOutcome::Absent { path } => {
                missing.insert(path);
            }
        }
    }

    // An absent GGUF is fatal here, unlike the `CERA_REQUIRE_MODEL` opt-in the
    // sibling suites use. CI sets that variable in the `gpu-tests` job, whose
    // wgpu LFM2 step needs only `core`-tier fixtures, and pointedly leaves it
    // unset in `batched-prefill-parity`, which fetches only `core` on a pull
    // request while granite/llama3/qwen2/qwen3 are `arch` tier, so absence is a
    // legitimate skip there. (Cited by job name, not line: ci.yml's line numbers
    // shift often enough that a line citation goes stale between branches.)
    // This suite is different: no workflow sets `CERA_ORACLE`, so it runs only
    // when a human asks, and absence then means the caller wanted coverage and
    // silently did not get it. Skipping reports the identical green as a real
    // run, which is how the whole fixture set went unnoticed after macOS purged
    // `target/` (cargo marks it `CACHEDIR.TAG`, and the fetch script's default
    // dest lives inside it).
    //
    // Both conditions are reported together: gate failures are only collected,
    // never printed as they happen, so failing on `missing` alone would discard
    // a real regression found on the models that WERE present and send the
    // caller off to download several GB before they ever saw it.
    let mut report = Vec::new();
    if !missing.is_empty() {
        report.push(format!(
            "{} required oracle model(s) absent:\n  {}\n\
             Fetch them with `scripts/fetch_test_models.sh --set all --dest {}`, \
             or point CERA_ORACLE_MODELS_DIR at a directory that already has them.",
            missing.len(),
            missing
                .iter()
                .map(|p| p.display().to_string())
                .collect::<Vec<_>>()
                .join("\n  "),
            models_dir().display(),
        ));
    }
    if !all_failures.is_empty() {
        report.push(format!(
            "gate failures on the model(s) that were present:\n{}",
            all_failures.join("\n"),
        ));
    }
    assert!(
        report.is_empty(),
        "CERA_ORACLE=1 oracle run did not pass:\n\n{}",
        report.join("\n\n"),
    );
}