code/analyze_results.py

#!/usr/bin/env python3
"""
Turn the five eval JSONs into a publication-ready verdict with statistics.

Adds the rigor a 124-sample benchmark needs:
  * 95% bootstrap CI on each model's Jaccard (10k resamples)
  * PAIRED bootstrap on (LoRA - zero-shot) per-sample delta — the correct test
    for "is the adapter's gain real?", since both run on the same samples
  * stratification by prompt length (short < 2000 chars vs long >= 2000) and by
    timeout, to quantify DiffusionGemma's long-context limitation honestly
  * a markdown comparison table + a plain verdict against the targets

Usage (run anywhere the JSONs + test.jsonl are readable):
  python3 analyze_results.py --dir ./data --out report.md
"""
import argparse, json, statistics
from pathlib import Path

# deterministic bootstrap without numpy (LCG) so results are reproducible everywhere
class _Rng:
    def __init__(self, seed=12345): self.s = seed & 0xFFFFFFFF
    def randint(self, n):
        self.s = (1103515245 * self.s + 12345) & 0x7FFFFFFF
        return self.s % n

FREQ_FLOOR = {"jaccard": 0.474, "exact": 0.113, "precision": 0.600, "recall": 0.631, "top1": 0.750}
MODELS = [
    ("DiffusionGemma zero-shot", "eval_zeroshot_clean"),
    ("DiffusionGemma + our LoRA", "eval_lora"),
    ("Gemma-4-26B AR sibling", "eval_gemma4_ar"),
    ("Qwen3.6-35B-A3B", "eval_qwen36_35b"),
    ("Qwen3.6-27B", "eval_qwen36_27b"),
]


def bootstrap_ci(vals, reps=10000, seed=12345):
    if not vals:
        return (0.0, 0.0, 0.0)
    rng = _Rng(seed)
    n = len(vals)
    means = []
    for _ in range(reps):
        means.append(sum(vals[rng.randint(n)] for _ in range(n)) / n)
    means.sort()
    return (sum(vals) / n, means[int(0.025 * reps)], means[int(0.975 * reps)])


def paired_bootstrap(deltas, reps=10000, seed=999):
    """P(mean delta > 0) and 95% CI of the mean per-sample delta."""
    if not deltas:
        return (0.0, 0.0, 0.0, 0.0)
    rng = _Rng(seed)
    n = len(deltas)
    means = []
    for _ in range(reps):
        means.append(sum(deltas[rng.randint(n)] for _ in range(n)) / n)
    means.sort()
    p_win = sum(1 for m in means if m > 0) / reps
    return (sum(deltas) / n, means[int(0.025 * reps)], means[int(0.975 * reps)], p_win)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--dir", required=True)
    ap.add_argument("--test", default=None, help="test.jsonl for prompt-length stratification")
    ap.add_argument("--out", default=None)
    args = ap.parse_args()
    d = Path(args.dir)

    # prompt lengths for stratification
    plen = {}
    test_path = Path(args.test) if args.test else d / "test.jsonl"
    if test_path.exists():
        for i, line in enumerate(open(test_path)):
            plen[i] = len(json.loads(line)["prompt"])

    loaded = {}
    for name, fname in MODELS:
        p = d / f"{fname}.json"
        if p.exists():
            try:
                loaded[name] = json.load(open(p))
            except Exception:
                pass

    out = []
    out.append("# DiffusionGemma tool-selection — benchmark report\n")
    out.append("Clean leak-free test, 124 held-out tasks, identical harness/parser/metrics. "
               "Greedy decoding (temp 0). CIs are 10k-resample bootstrap.\n")

    # main table with CI on Jaccard. Two views: "practical" (timeouts = 0-score
    # failures, the real-world number) and "completed-only" (accuracy when the model
    # actually answers within the cap) — fair separation of accuracy vs convergence.
    out.append("| Model | Jaccard [95% CI] (all) | Completed | Jaccard (completed-only) | Top-1 | Timeouts |")
    out.append("|---|---|---|---|---|---|")
    out.append(f"| _freq top-3 floor_ | _{FREQ_FLOOR['jaccard']:.3f}_ | _124/124_ | _{FREQ_FLOOR['jaccard']:.3f}_ | _{FREQ_FLOOR['top1']:.3f}_ | — |")
    jac_vals = {}
    comp_jac = {}
    for name, _ in MODELS:
        if name not in loaded:
            out.append(f"| {name} | _pending_ | | | | |")
            continue
        a = loaded[name]["aggregate"]
        samples = loaded[name]["samples"]
        jv = [s["jaccard"] for s in samples]
        jac_vals[name] = jv
        comp = [s["jaccard"] for s in samples if not s.get("timed_out")]
        comp_jac[name] = comp
        mean, lo, hi = bootstrap_ci(jv)
        to = a.get("timeouts", 0)
        cmean = sum(comp) / len(comp) if comp else 0.0
        out.append(f"| **{name}** | {mean:.3f} [{lo:.3f}, {hi:.3f}] | {len(comp)}/{len(samples)} | "
                   f"{cmean:.3f} | {a['top1']:.3f} | {to} |")

    # convergence finding: fine-tuning sharpens the output distribution, so the
    # entropy-bounded diffusion sampler early-stops instead of running to the cap.
    zs, lora = "DiffusionGemma zero-shot", "DiffusionGemma + our LoRA"
    if zs in loaded and lora in loaded:
        zs_to = loaded[zs]["aggregate"].get("timeouts", 0)
        lo_to = loaded[lora]["aggregate"].get("timeouts", 0)
        out.append("\n## Fine-tuning makes generation CONVERGE (not just more accurate)\n")
        out.append(f"- Zero-shot base times out on **{zs_to}/124** prompts; the LoRA on **{lo_to}/124**.")
        out.append("- DiffusionGemma's sampler uses entropy-bounded early stopping. The uncertain "
                   "base model never converges and runs the full denoising budget (→ timeout); the "
                   "fine-tuned model is confident, early-stops, and answers in a few seconds. "
                   "Fine-tuning buys both accuracy AND a large generation speedup.")

    # paired significance: LoRA vs zero-shot, on samples BOTH actually completed (fair)
    if zs in loaded and lora in loaded:
        zs_done = {s["i"]: s["jaccard"] for s in loaded[zs]["samples"] if not s.get("timed_out")}
        lo_s = {s["i"]: s["jaccard"] for s in loaded[lora]["samples"]}
        common = sorted(set(zs_done) & set(lo_s))
        deltas = [lo_s[i] - zs_done[i] for i in common]
        md, dlo, dhi, pwin = paired_bootstrap(deltas)
        sig = "**significant**" if (dlo > 0 or dhi < 0) else "NOT significant"
        out.append(f"\n## Fine-tuning effect on accuracy (paired, on the {len(common)} samples zero-shot completed)\n")
        out.append(f"- Mean per-sample Jaccard delta (LoRA − zero-shot): **{md:+.3f}** "
                   f"[95% CI {dlo:+.3f}, {dhi:+.3f}] — {sig}")
        out.append(f"- P(LoRA > zero-shot) = **{pwin:.3f}**")
        improved = sum(1 for x in deltas if x > 1e-9)
        worsened = sum(1 for x in deltas if x < -1e-9)
        out.append(f"- Per-sample: improved {improved}, worsened {worsened}, unchanged {len(deltas)-improved-worsened}")

    # stratification by prompt length (DiffusionGemma rows only — where it matters)
    if plen:
        out.append("\n## By prompt length (the long-context limitation)\n")
        out.append("| Model | short (<2000c) Jaccard | long (≥2000c) Jaccard | long timeouts |")
        out.append("|---|---|---|---|")
        for name in (zs, lora):
            if name not in loaded:
                continue
            sshort, slong, to_long = [], [], 0
            for s in loaded[name]["samples"]:
                L = plen.get(s["i"], 0)
                (slong if L >= 2000 else sshort).append(s["jaccard"])
                if L >= 2000 and s.get("timed_out"):
                    to_long += 1
            ms = sum(sshort)/len(sshort) if sshort else 0
            ml = sum(slong)/len(slong) if slong else 0
            out.append(f"| {name} | {ms:.3f} (n={len(sshort)}) | {ml:.3f} (n={len(slong)}) | {to_long} |")

    # verdict
    out.append("\n## Verdict\n")
    if lora in jac_vals:
        lj = sum(jac_vals[lora]) / len(jac_vals[lora])
        zj = sum(jac_vals[zs]) / len(jac_vals[zs]) if zs in jac_vals else None
        def line(target, label):
            if target is None: return f"- vs {label}: _baseline pending_"
            verdict = "✅ BEATS" if lj > target else "❌ below"
            return f"- vs {label} ({target:.3f}): {verdict} (LoRA {lj:.3f})"
        out.append(line(zj, "zero-shot self"))
        out.append(line(FREQ_FLOOR["jaccard"], "frequency floor — minimum to be interesting"))
        ar = "Gemma-4-26B AR sibling"
        if ar in jac_vals:
            out.append(line(sum(jac_vals[ar])/len(jac_vals[ar]), "AR sibling — the diffusion-vs-AR question"))
        q = "Qwen3.6-35B-A3B"
        if q in jac_vals:
            out.append(line(sum(jac_vals[q])/len(jac_vals[q]), "Qwen3.6-35B zero-shot — specialist-beats-bigger-generalist"))
    else:
        out.append("- _LoRA eval pending_")

    report = "\n".join(out)
    print(report)
    if args.out:
        Path(args.out).write_text(report)
        print(f"\n[wrote {args.out}]")


if __name__ == "__main__":
    main()