memra-engine 0.82.2

//! Teacher-forced replay acceptance driver (hqmtp MTP-heal protocol; see
//! HybridModel::replay_acceptance for the metric definition). Walks a FIXED corpus text and
//! scores the MTP head's draft chain against the trunk's own teacher-forced greedy picks at
//! sampled positions — no generation, so degenerate self-generated loops cannot inflate
//! acceptance, and two arms (bf16 full-prec ceiling vs NVFP4) score on IDENTICAL contexts.
//!
//! Run: replay-acceptance <model.gguf|hf_dir>   (corpus text via MEMRA_PROMPT_FILE)
//! Env:
//!   MEMRA_REPLAY_K=4        draft chain length per eval position
//!   MEMRA_REPLAY_STRIDE=16  eval every Nth corpus position
//!   MEMRA_REPLAY_CHUNK=512  forced-pass chunk (logits buffer = chunk x n_vocab f32)
//!   MEMRA_REPLAY_T=0        cap corpus tokens (0 = all)
//!   MEMRA_REPLAY_DUMP=f.jsonl  per-position rows: {"pos","drafts","targets","hits"}
//!   MEMRA_REPLAY_GATE=1     re-run a 2048-token prefix at chunk=64 and require identical
//!                          greedy track + drafts (chunk-boundary correctness gate)
//!   MEMRA_FULL_PREC=1       arm A (bf16 ST ceiling) — same knob as run-spec

use memra_engine::Engine;
use memra_engine::hybrid::HybridModel;
use memra_gguf::GgufFile;
use std::io::Write;

const BUCKETS: &[(usize, usize)] = &[
    (0, 512), (512, 2048), (2048, 8192), (8192, 16384),
    (16384, 32768), (32768, 65536), (65536, usize::MAX),
];

fn env_usize(name: &str, default: usize) -> usize {
    std::env::var(name).ok().and_then(|v| v.parse().ok()).unwrap_or(default)
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let path = std::env::args().nth(1)
        .expect("usage: replay-acceptance <model.gguf|hf_dir>  (corpus via MEMRA_PROMPT_FILE)");
    let e = Engine::new(0)?;
    let is_dir = std::path::Path::new(&path).is_dir();
    let g: Option<GgufFile> = if is_dir { None } else { Some(GgufFile::open(&path)?) };
    let model = match &g {
        Some(g) => HybridModel::load(&e, g)?,
        None => {
            let st = memra_gguf::source::SafetensorsSource::open(std::path::Path::new(&path))?;
            HybridModel::load_from_source(&e, &st)?
        }
    };
    if model.mtp.is_none() {
        eprintln!("ERROR: model has no MTP/NextN head — replay acceptance is undefined.");
        std::process::exit(2);
    }

    let prompt_path = std::env::var("MEMRA_PROMPT_FILE")
        .expect("replay-acceptance needs MEMRA_PROMPT_FILE (corpus text file OR directory)");
    let tok = match &g {
        Some(g) => memra_tokenizer::Tokenizer::from_gguf(g)?,
        None => memra_tokenizer::Tokenizer::from_hf_dir(std::path::Path::new(&path))
            .map_err(|err| format!("HF tokenizer init failed: {err}"))?,
    };

    // DIRECTORY mode (bulk distillation extraction): iterate every file in the dir with ONE
    // resident model — the per-invocation model load dominated bulk extraction otherwise.
    // Per file <stem>: hiddens -> $MEMRA_REPLAY_HDUMP/<stem>.hiddens (+ .meta.json). Draft
    // chains are usually skipped here via a huge MEMRA_REPLAY_STRIDE (prefill-speed pass).
    if std::path::Path::new(&prompt_path).is_dir() {
        let hdir = std::env::var("MEMRA_REPLAY_HDUMP")
            .expect("directory mode needs MEMRA_REPLAY_HDUMP=<output dir>");
        std::fs::create_dir_all(&hdir)?;
        let k = env_usize("MEMRA_REPLAY_K", 4);
        let stride = env_usize("MEMRA_REPLAY_STRIDE", 16);
        let chunk = env_usize("MEMRA_REPLAY_CHUNK", 512);
        let cap = env_usize("MEMRA_REPLAY_T", 65536);
        let mut files: Vec<_> = std::fs::read_dir(&prompt_path)?
            .filter_map(|d| d.ok().map(|d| d.path()))
            .filter(|p| p.is_file()).collect();
        files.sort();
        for fp in files {
            let stem = fp.file_stem().unwrap().to_string_lossy().to_string();
            let out = format!("{hdir}/{stem}.hiddens");
            if std::path::Path::new(&format!("{out}.meta.json")).exists() {
                println!("skip {stem} (done)");
                continue;
            }
            let text = std::fs::read_to_string(&fp)?;
            let mut ids = tok.encode(&text, true);
            if cap > 0 && ids.len() > cap { ids.truncate(cap); }
            if ids.len() < 64 { println!("skip {stem} (too short)"); continue; }
            let mut hf = std::fs::File::create(&out)?;
            let t0 = std::time::Instant::now();
            let (_rows, bg) = model.replay_acceptance(&e, &ids, k, stride, chunk, Some(&mut hf))?;
            use std::io::Write;
            let mut mf = std::fs::File::create(format!("{out}.meta.json"))?;
            write!(mf, "{{\"n_tokens\":{},\"n_embd\":{},\"dtype\":\"bf16\",\"bg\":{:?}}}",
                   ids.len(), model.cfg.n_embd, bg)?;
            println!("extracted {stem}: {} tokens in {:.1}s", ids.len(),
                     t0.elapsed().as_secs_f64());
        }
        println!("DIR-EXTRACT-DONE");
        return Ok(());
    }
    let text = std::fs::read_to_string(&prompt_path).expect("MEMRA_PROMPT_FILE unreadable");
    let mut ids = tok.encode(&text, true);
    let cap = env_usize("MEMRA_REPLAY_T", 0);
    if cap > 0 && ids.len() > cap { ids.truncate(cap); }
    let k = env_usize("MEMRA_REPLAY_K", 4);
    let stride = env_usize("MEMRA_REPLAY_STRIDE", 16);
    let chunk = env_usize("MEMRA_REPLAY_CHUNK", 512);
    println!("corpus: {} chars -> {} tokens | K={k} stride={stride} chunk={chunk} full_prec={}",
             text.len(), ids.len(), std::env::var("MEMRA_FULL_PREC").is_ok());

    // MEMRA_REPLAY_HDUMP=<path>: stream every position's pre-norm trunk hidden (f32 LE,
    // [n_tokens, n_embd]) + write <path>.meta.json {n_tokens, n_embd, bg} — the distillation
    // training extraction (engine = source of truth for hiddens).
    let hdump_path = std::env::var("MEMRA_REPLAY_HDUMP").ok();
    let mut hdump_file = match &hdump_path {
        Some(p) => Some(std::fs::File::create(p)?),
        None => None,
    };
    let t0 = std::time::Instant::now();
    let (rows, bg) = model.replay_acceptance(&e, &ids, k, stride, chunk, hdump_file.as_mut())?;
    let dt = t0.elapsed().as_secs_f64();
    if let Some(p) = &hdump_path {
        use std::io::Write;
        let n_embd = model.cfg.n_embd;
        let mut mf = std::fs::File::create(format!("{p}.meta.json"))?;
        write!(mf, "{{\"n_tokens\":{},\"n_embd\":{},\"dtype\":\"bf16\",\"bg\":{:?}}}",
               ids.len(), n_embd, bg)?;
        println!("hdump: {} tokens x {} f32 -> {p} (+.meta.json)", ids.len(), n_embd);
    }
    println!("replay: {} eval positions in {dt:.1}s ({:.1} corpus tok/s incl. chains)",
             rows.len(), ids.len() as f64 / dt);

    // Chunk-boundary gate: an independent pass over a prefix with a DIFFERENT chunk size.
    // The greedy track must be EXACT (a position/rope/fill offset bug breaks it hard). Drafts
    // are chunk-FP-order sensitive (batched fill reduce order -> ULP -> deep-chain argmax flips
    // on close calls; measured 2/127 positions, slot 3 only, same-chunk reruns bit-identical) —
    // allow <=1% slot disagreement and report it. Cross-arm runs must use ONE chunk size.
    if std::env::var("MEMRA_REPLAY_GATE").is_ok() {
        let n_gate = 2048.min(ids.len());
        let (rows_g, bg_g) = model.replay_acceptance(&e, &ids[..n_gate], k, stride, 64, None)?;
        let mut bg_bad = 0usize;
        for i in 1..n_gate.saturating_sub(1) {
            if bg[i] != bg_g[i] { bg_bad += 1; if bg_bad <= 3 {
                eprintln!("[gate] bg mismatch at pos {i}: {} vs {}", bg[i], bg_g[i]); } }
        }
        let main_prefix: Vec<_> = rows.iter().filter(|r| r.0 + k <= n_gate).collect();
        let mut slots_diff = 0usize;
        let mut slots_all = 0usize;
        let mut pos_bad = 0usize;
        for (a, b) in main_prefix.iter().zip(rows_g.iter()) {
            if a.0 != b.0 { pos_bad += 1; continue; }
            slots_all += k;
            slots_diff += a.1.iter().zip(b.1.iter()).filter(|(x, y)| x != y).count();
        }
        let frac = if slots_all > 0 { slots_diff as f64 / slots_all as f64 } else { 0.0 };
        if bg_bad > 0 || pos_bad > 0 {
            eprintln!("[gate] FAIL: bg_bad={bg_bad} pos_bad={pos_bad} (chunk={chunk} vs 64) — offset bug");
            std::process::exit(3);
        }
        // slot_diff is NOT a failure: it is the chunk-FP-order noise floor of THIS arm's draft
        // chain (quant arms measure higher — tighter head margins under NVFP4 hiddens; 9B-st-ct
        // measured 2.0-2.8% vs 0.2% GGUF/bf16). Deltas below this floor are noise.
        println!("[gate] PASS: bg exact; draft slot noise floor {slots_diff}/{slots_all}={frac:.4} (chunk={chunk} vs 64)");
    }

    // aggregate per context-depth bucket
    let slot_hdr: String = (0..k).map(|j| format!("  slot{j}")).collect();
    println!("{:>14} {:>7}{slot_hdr}   chain", "ctx-bucket", "n");
    for &(lo, hi) in BUCKETS {
        let sel: Vec<_> = rows.iter().filter(|r| r.0 >= lo && r.0 < hi).collect();
        if sel.is_empty() { continue; }
        let n = sel.len();
        let mut slot_hit = vec![0usize; k];
        let mut chain_len = 0usize;
        for (_, d, t) in &sel {
            for j in 0..k { if d[j] == t[j] { slot_hit[j] += 1; } }
            let mut c = 0usize;
            while c < k && d[c] == t[c] { c += 1; }
            chain_len += c;
        }
        let slots: String = (0..k)
            .map(|j| format!(" {:6.3}", slot_hit[j] as f64 / n as f64)).collect();
        let hi_s = if hi == usize::MAX { "inf".into() } else { hi.to_string() };
        println!("{:>14} {:>7}{slots}   {:5.3}", format!("{lo}-{hi_s}"), n,
                 chain_len as f64 / (n * k) as f64);
    }

    if let Ok(dump) = std::env::var("MEMRA_REPLAY_DUMP") {
        let mut f = std::fs::File::create(&dump)?;
        for (p, d, t) in &rows {
            let hits: Vec<bool> = d.iter().zip(t.iter()).map(|(a, b)| a == b).collect();
            writeln!(f, "{{\"pos\":{p},\"drafts\":{d:?},\"targets\":{t:?},\"hits\":{hits:?}}}")?;
        }
        println!("wrote {} rows -> {dump}", rows.len());
    }
    Ok(())
}