src/make_report.py

"""
Final cross-model report: results table, quality-vs-diversity frontier chart,
and qualitative examples. Reads outputs/eval/results.csv.
"""
from __future__ import annotations

import argparse
import csv
import json
import sys
from pathlib import Path

import numpy as np

import logbook

ROOT = Path(__file__).resolve().parent.parent

ORDER = ["base", "E0-baseline", "E1-div-individual", "E2-div-group",
         "E3-multipos", "E4-divpo-emb", "E4-divpo-prob"]

PRETTY = {
    "base": "Base (Qwen3-4B-Instruct)",
    "E0-baseline": "E0 · GRPO, quality only",
    "E1-div-individual": "E1 · div-grpo-individual (d_i)",
    "E2-div-group": "E2 · div-grpo-group (d_i + m_i)",
    "E3-multipos": "E3 · multi-positive weighted DPO",
    "E4-divpo-emb": "E4a · DivPO (embedding)",
    "E4-divpo-prob": "E4b · DivPO (probability)",
}


def load(path: Path) -> list[dict]:
    rows = list(csv.DictReader(open(path)))
    for r in rows:
        for k, v in r.items():
            if k != "model":
                try:
                    r[k] = float(v)
                except (TypeError, ValueError):
                    pass
    key = {m: i for i, m in enumerate(ORDER)}
    rows.sort(key=lambda r: key.get(r["model"], 99))
    return rows


def frontier_chart(rows, out_path):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    fig, ax = plt.subplots(1, 3, figsize=(17, 5.2))
    cols = plt.cm.tab10(np.linspace(0, 1, 10))

    panels = [
        ("pairwise", "Mean pairwise embedding distance", 0),
        ("logdet", "Log-det volume (semantic coverage)", 1),
        ("eff_rank", "Effective rank (mode count)", 2),
    ]
    base = next((r for r in rows if r["model"] == "base"), None)

    for (xk, xlabel, i) in panels:
        a = ax[i]
        for j, r in enumerate(rows):
            a.scatter(r[xk], r["quality"], s=170, color=cols[j % 10],
                      edgecolor="black", zorder=3,
                      label=PRETTY.get(r["model"], r["model"]))
            a.annotate(r["model"].replace("-div", "").replace("divpo-", ""),
                       (r[xk], r["quality"]), fontsize=7,
                       xytext=(5, 5), textcoords="offset points")
        if base:
            a.axhline(base["quality"], ls=":", c="gray", lw=1)
            a.axvline(base[xk], ls=":", c="gray", lw=1)
        a.set_xlabel(xlabel); a.set_ylabel("Judge quality (0-10)")
        a.grid(alpha=.3)
        a.set_title(f"Quality vs {xlabel.split('(')[0].strip()}", fontsize=10)

    handles, labels = ax[0].get_legend_handles_labels()
    fig.legend(handles, labels, loc="lower center", ncol=4, fontsize=8,
               frameon=False, bbox_to_anchor=(0.5, -0.06))
    fig.suptitle("Quality–diversity frontier (dotted lines = base model)", fontsize=13)
    plt.tight_layout()
    plt.savefig(out_path, dpi=140, bbox_inches="tight")
    plt.close()


def delta_table(rows):
    base = next((r for r in rows if r["model"] == "base"), None)
    if not base:
        return []
    keys = ["quality", "pairwise", "logdet", "distinct4", "self_bleu",
            "eff_rank", "n_clusters", "ends_cleanly"]
    out = []
    for r in rows:
        d = {"model": PRETTY.get(r["model"], r["model"])}
        for k in keys:
            base_v = base.get(k, 0.0)
            d[f"Δ{k}"] = round(r.get(k, 0.0) - base_v, 4)
        out.append(d)
    return out


QUALITY_TOLERANCE = 0.3   # judge points we are willing to lose for diversity


def recommend(rows) -> str:
    """Pick the arm to scale, from the data rather than from vibes.

    Rule, stated before looking at results: among arms whose quality is within
    QUALITY_TOLERANCE of the BASE model (not of the best arm -- we care about not
    degrading the model, not about winning on quality), take the largest gain in
    effective rank. Effective rank is the primary diversity measure because it is
    the one metric shown to separate collapse from spread (see 02_setup). Ties
    and near-ties are reported rather than hidden.
    """
    base = next((r for r in rows if r["model"] == "base"), None)
    if not base:
        return "_No base row; cannot compute deltas._"

    cand = []
    for r in rows:
        if r["model"] == "base":
            continue
        dq = r["quality"] - base["quality"]
        der = r.get("eff_rank", 0.0) - base.get("eff_rank", 0.0)
        dld = r.get("logdet", 0.0) - base.get("logdet", 0.0)
        cand.append({"model": PRETTY.get(r["model"], r["model"]),
                     "raw": r["model"], "Δquality": round(dq, 3),
                     "Δeff_rank": round(der, 3), "Δlogdet": round(dld, 3),
                     "ends_cleanly": round(r.get("ends_cleanly", 0.0), 3),
                     "eligible": dq >= -QUALITY_TOLERANCE and r.get("ends_cleanly", 0) > 0.9})

    ok = [c for c in cand if c["eligible"]]
    ranked = sorted(ok, key=lambda c: -c["Δeff_rank"])

    if not ranked:
        return (f"**No arm qualifies.** Every arm either lost more than "
                f"{QUALITY_TOLERANCE} judge points against base or fell below a "
                f"90% clean-completion rate. The honest recommendation is to fix "
                f"the objective before scaling anything to 8B.\n\n"
                + logbook.table(cand, ["model", "Δquality", "Δeff_rank",
                                       "Δlogdet", "ends_cleanly", "eligible"]))

    win = ranked[0]
    runner = ranked[1] if len(ranked) > 1 else None
    margin = (win["Δeff_rank"] - runner["Δeff_rank"]) if runner else None
    close = margin is not None and margin < 0.25

    txt = (f"**Scale `{win['model']}` to Qwen3-8B.** It gains "
           f"{win['Δeff_rank']:+.2f} effective rank and {win['Δlogdet']:+.2f} "
           f"log-det against base while holding quality at "
           f"{win['Δquality']:+.2f} and completion at "
           f"{100*win['ends_cleanly']:.0f}%.\n\n")
    if runner:
        txt += (f"Runner-up is `{runner['model']}` "
                f"({runner['Δeff_rank']:+.2f} effective rank). ")
        txt += ("The margin is **{:.2f}, which is small enough that this "
                "ordering should not be treated as settled** on 50 eval prompts "
                "— both are worth carrying forward.\n\n".format(margin)
                if close else
                "The margin ({:.2f}) is clear.\n\n".format(margin))
    excluded = [c for c in cand if not c["eligible"]]
    if excluded:
        txt += ("Excluded for quality or completeness regression: "
                + ", ".join(f"`{c['model']}` (Δq {c['Δquality']:+.2f}, "
                            f"complete {100*c['ends_cleanly']:.0f}%)"
                            for c in excluded) + ".\n\n")
    txt += logbook.table(cand, ["model", "Δquality", "Δeff_rank", "Δlogdet",
                                "ends_cleanly", "eligible"])
    return txt


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--results", default="outputs/eval/results.csv")
    ap.add_argument("--samples", default="outputs/eval/samples")
    args = ap.parse_args()

    rows = load(ROOT / args.results)
    fig = logbook.FIGS / "frontier.png"
    fig.parent.mkdir(parents=True, exist_ok=True)
    frontier_chart(rows, fig)

    main_cols = ["model", "quality", "ends_cleanly", "pairwise", "logdet",
                 "distinct4", "self_bleu", "eff_rank", "n_clusters", "tok_entropy", "words"]
    tbl = [{k: (PRETTY.get(r["model"], r["model"]) if k == "model"
                else round(r.get(k, 0.0), 4)) for k in main_cols} for r in rows]

    # qualitative examples
    qual = ""
    sdir = ROOT / args.samples
    for r in rows:
        p = sdir / f"{r['model']}_examples.json"
        if not p.exists():
            continue
        ex = json.load(open(p))
        if not ex:
            continue
        e = ex[0]
        qual += f"\n### {PRETTY.get(r['model'], r['model'])}\n\n"
        qual += f"*Prompt:* {e['prompt'][:200]}\n\n"
        qual += (f"*Set stats:* pairwise {e['pairwise']:.3f} · logdet {e['logdet']:.2f} "
                 f"· eff_rank {e['eff_rank']:.2f} · clusters {e['n_clusters']} "
                 f"· self-BLEU {e['self_bleu']:.3f}\n\n")
        for i, t in enumerate(e["texts"][:3]):
            first = t.strip().split("\n")[0][:180]
            qual += f"{i+1}. {first}…\n"
    body = f"""# Results — diversity-aware post-training for creative story generation

## Main table

{logbook.table(tbl, main_cols)}

`self_bleu` is inverted in meaning: **lower is more diverse**. `tok_entropy` is
monitor-only and was never optimized.

## Change vs. base model

{logbook.table(delta_table(rows))}

## Quality–diversity frontier

![frontier](figures/frontier.png)

## Recommendation — which method to scale to Qwen3-8B

Selection rule, fixed before results were seen: among arms whose judge quality
is within {QUALITY_TOLERANCE} points of the **base** model and whose clean-completion
rate stays above 90%, take the largest gain in **effective rank**. Quality is
measured against base rather than against the best arm because the goal is to
add diversity without degrading the model, not to win on quality. Effective rank
is the primary diversity axis because it is the one mode metric demonstrated to
separate collapse from spread (`02_setup_and_deviations.md`, §11).

{recommend(rows)}

## Qualitative examples (first 3 of 16 samples, same prompt)
{qual}
"""
    (logbook.LOGS.parent / "report.md").write_text(body)
    print(body[:2500])
    print("\nwrote report.md and", fig)
    return 0


if __name__ == "__main__":
    sys.exit(main())