social-data-attribution-w2/scripts/analysis/unlearn_report.py

"""
Generate a self-contained interactive HTML report for single-topic unlearning results.

Reads per-topic eval_fast_results.json, forget_ppl_log.json, ppl_stopping_threshold.json,
plus the baseline_eval_fast_results.json. Produces one HTML file with Plotly charts.

Usage:
    uv run python scripts/analysis/unlearn_report.py
    uv run python scripts/analysis/unlearn_report.py --data-dir artifacts/unlearn_report/data --output report.html
"""
import argparse
import json
import math
import sys
from pathlib import Path

TOPICS = [
    "adult_content", "art_and_design", "crime_and_law", "education_and_jobs",
    "electronics_and_hardware", "entertainment", "fashion_and_beauty",
    "finance_and_business", "food_and_dining", "games", "health",
    "history_and_geography", "home_and_hobbies", "industrial", "literature",
    "politics", "religion", "science_math_and_technology", "social_life",
    "software", "software_development", "sports_and_fitness", "transportation",
    "travel_and_tourism",
]
ACCURACY_BENCHMARKS = ["gsm8k", "mmlu_stem", "mmlu_social_science", "socialiqa"]
ALL_BENCHMARKS = ACCURACY_BENCHMARKS + ["wikitext"]
BENCH_LABELS = {
    "gsm8k": "GSM8K", "mmlu_stem": "MMLU-STEM",
    "mmlu_social_science": "MMLU-SocSci", "socialiqa": "SocialIQA",
    "wikitext": "Wikitext PPL",
}
PLOTLY_CDN = "https://cdn.plot.ly/plotly-2.35.0.min.js"


def fmt(t: str) -> str:
    return t.replace("_", " ").title()


def safe(v) -> bool:
    return v is not None and not math.isnan(v)


def nanmean(vals: list[float]) -> float:
    clean = [v for v in vals if safe(v)]
    return sum(clean) / len(clean) if clean else float("nan")


def score_for(metrics: dict, bench: str) -> float:
    m = metrics.get(bench)
    if m is None:
        return float("nan")
    key = "word_perplexity" if bench == "wikitext" else "accuracy"
    return m.get(key, float("nan"))


def find_stop_step(ppl_log: dict, threshold: float) -> int:
    for step_str in sorted(ppl_log, key=lambda s: int(s)):
        if ppl_log[step_str] >= threshold:
            return int(step_str)
    return max((int(s) for s in ppl_log), default=0)


def load_data(data_dir: Path):
    evals, ppl_logs, thresholds = {}, {}, {}
    for topic in TOPICS:
        td = data_dir / topic / "f1000_r9000"
        ep = td / "eval_fast_results.json"
        pp = td / "forget_ppl_log.json"
        tp = td / "ppl_stopping_threshold.json"
        if ep.exists():
            evals[topic] = json.loads(ep.read_text())
        if pp.exists():
            ppl_logs[topic] = json.loads(pp.read_text())
        if tp.exists():
            thresholds[topic] = json.loads(tp.read_text())["threshold"]
    bp = data_dir / "baseline_eval_fast_results.json"
    if bp.exists():
        evals["_baseline"] = json.loads(bp.read_text())
    return evals, ppl_logs, thresholds


def build_rows(evals, ppl_logs, thresholds):
    bl_metrics = evals.get("_baseline", {}).get("metrics", {})
    baseline = {b: score_for(bl_metrics, b) for b in ALL_BENCHMARKS}
    completed = [t for t in TOPICS if t in evals]
    rows = []
    for topic in completed:
        metrics = evals[topic]["metrics"]
        r = {"topic": topic, "label": fmt(topic)}
        for b in ALL_BENCHMARKS:
            s = score_for(metrics, b)
            bl = baseline.get(b, float("nan"))
            r[f"{b}_score"] = s
            r[f"{b}_bl"] = bl
            if safe(s) and safe(bl) and bl != 0:
                r[f"{b}_gamma"] = (s - bl) / abs(bl)
            else:
                r[f"{b}_gamma"] = float("nan")
        if topic in thresholds:
            r["threshold"] = thresholds[topic]
        if topic in ppl_logs:
            log = ppl_logs[topic]
            r["stop_step"] = find_stop_step(log, thresholds.get(topic, 1e30))
            steps = sorted(log, key=lambda s: int(s))
            r["final_ppl"] = log[steps[-1]] if steps else float("nan")
        rows.append(r)
    return rows, baseline, completed


def generate_report(evals, ppl_logs, thresholds, output: Path):
    rows, baseline, completed = build_rows(evals, ppl_logs, thresholds)
    if not rows:
        print("No completed topics found.")
        sys.exit(1)

    stored_bl_siq = evals[completed[0]]["metrics"].get("socialiqa", {}).get("baseline", float("nan"))
    fast_bl_siq = baseline.get("socialiqa", float("nan"))

    parts = [_head(), "<body>"]
    parts.append("<h1>Single-Topic Unlearning Report</h1>")
    parts.append(f"<p class='sub'>OLMo-3-1025-7B &middot; NGDiff &middot; "
                 f"1,000 forget / 9,000 retain docs &middot; "
                 f"{len(completed)} of {len(TOPICS)} topics</p>")

    parts.append(_section1(rows, baseline, stored_bl_siq, fast_bl_siq))
    parts.append(_section2(rows, baseline))
    parts.append(_section3(rows, baseline))
    parts.append(_section4(rows, ppl_logs, thresholds))
    parts.append(_section5(rows))
    parts.append(_section6(rows, completed, stored_bl_siq, fast_bl_siq))

    parts.append("</body></html>")
    output.parent.mkdir(parents=True, exist_ok=True)
    output.write_text("\n".join(parts))
    print(f"Report: {output} ({len(completed)} topics)")


def _head():
    return f"""<!DOCTYPE html>
<html lang="en"><head><meta charset="UTF-8">
<title>Single-Topic Unlearning Report</title>
<script src="{PLOTLY_CDN}"></script>
<style>
body {{ font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', system-ui, sans-serif;
  max-width: 1200px; margin: 0 auto; padding: 20px; color: #1a1a1a; line-height: 1.6; }}
h1 {{ border-bottom: 2px solid #333; padding-bottom: 8px; }}
h2 {{ margin-top: 48px; border-bottom: 1px solid #ccc; padding-bottom: 4px; }}
h3 {{ margin-top: 28px; color: #444; }}
.sub {{ color: #666; font-size: 0.92em; }}
.chart {{ margin: 20px 0; }}
table {{ border-collapse: collapse; width: 100%; margin: 16px 0; font-size: 0.88em; }}
th, td {{ border: 1px solid #ddd; padding: 5px 10px; text-align: right; }}
th {{ background: #f5f5f5; text-align: center; font-weight: 600; }}
td:first-child, th:first-child {{ text-align: left; }}
tr:nth-child(even) {{ background: #fafafa; }}
.note {{ background: #f8f8f8; border-left: 3px solid #999; padding: 12px 16px;
  margin: 16px 0; font-size: 0.93em; }}
.warn {{ background: #fff8e1; border-left-color: #f9a825; }}
.neg {{ color: #c62828; }} .pos {{ color: #2e7d32; }} .dim {{ color: #888; }}
.legend-box {{ display: inline-block; width: 14px; height: 14px;
  vertical-align: middle; margin-right: 4px; border: 1px solid #ccc; }}
</style></head>"""


def _section1(rows, baseline, stored_bl_siq, fast_bl_siq):
    parts = ["<h2>1. Experiment Overview</h2>"]
    parts.append("""<table>
<tr><th>Parameter</th><th>Value</th></tr>
<tr><td>Base model</td><td>allenai/OLMo-3-1025-7B</td></tr>
<tr><td>Unlearning method</td><td>NGDiff (gradient difference with LoRA adapter)</td></tr>
<tr><td>Forget docs per topic</td><td>1,000</td></tr>
<tr><td>Retain docs per topic</td><td>9,000</td></tr>
<tr><td>Topics</td><td>24 WebOrganizer topic bins</td></tr>
<tr><td>Evaluation</td><td>200-sample fast eval (GSM8K, MMLU-STEM, MMLU-SocSci, SocialIQA, Wikitext-2)</td></tr>
<tr><td>Baseline comparison</td><td>Fast-eval run on base model (no adapter)</td></tr>
</table>""")
    parts.append(f"""<div class="note">
<p>Gamma values throughout this report are computed against the <strong>fast-eval baseline</strong>:
the same 200-sample evaluation run on the unmodified base model. This is the apples-to-apples
comparison. The stored baselines in each topic JSON come from a separate full-benchmark evaluation
and differ on some benchmarks.</p></div>""")

    siq_scores = [r["socialiqa_score"] for r in rows if safe(r.get("socialiqa_score"))]
    if safe(stored_bl_siq) and safe(fast_bl_siq) and abs(stored_bl_siq - fast_bl_siq) > 0.05:
        parts.append(f"""<div class="note warn">
<p><strong>SocialIQA baseline discrepancy.</strong> The stored full-benchmark baseline is
{stored_bl_siq*100:.1f}%, but the fast-eval baseline (200 samples) is {fast_bl_siq*100:.1f}%.
Unlearned models score {min(siq_scores)*100:.0f}%-{max(siq_scores)*100:.0f}% on the same subset.
The large gamma values in the stored JSON (e.g. -40%) are artifacts of comparing against 80.3%
instead of {fast_bl_siq*100:.0f}%. SocialIQA results from this 200-sample eval should be
interpreted with caution.</p></div>""")
    return "\n".join(parts)


def _section2(rows, baseline):
    parts = ["<h2>2. Benchmark Impact Heatmap</h2>"]

    sorted_rows = sorted(rows, key=lambda r: nanmean(
        [r.get(f"{b}_gamma", float("nan")) for b in ACCURACY_BENCHMARKS]))
    labels = [r["label"] for r in sorted_rows]
    bench_labels = [BENCH_LABELS[b] for b in ACCURACY_BENCHMARKS]
    z, text = [], []
    for r in sorted_rows:
        zrow, trow = [], []
        for b in ACCURACY_BENCHMARKS:
            g = r.get(f"{b}_gamma", float("nan"))
            zrow.append(g * 100 if safe(g) else None)
            trow.append(f"{g*100:.1f}%" if safe(g) else "n/a")
        z.append(zrow)
        text.append(trow)

    parts.append(f"""<div id="heatmap" class="chart"></div>
<script>
Plotly.newPlot('heatmap', [{{
  z: {json.dumps(z)}, x: {json.dumps(bench_labels)}, y: {json.dumps(labels)},
  text: {json.dumps(text)}, texttemplate: '%{{text}}',
  type: 'heatmap', name: 'Gamma',
  colorscale: [[0,'#c62828'],[0.5,'#ffffff'],[1,'#2e7d32']],
  zmid: 0, zmin: -50, zmax: 10,
  colorbar: {{title: 'Change (%)', titleside: 'right'}}
}}], {{
  title: 'Accuracy Change vs Fast-Eval Baseline (%)',
  height: {max(520, len(labels) * 26)},
  margin: {{l: 220, r: 80, t: 50, b: 60}},
  yaxis: {{autorange: 'reversed'}}, xaxis: {{side: 'top'}}
}});</script>""")

    parts.append("""<div class="note">
<p><strong>Color legend:</strong>
<span class="legend-box" style="background:#c62828"></span> Red = accuracy dropped vs baseline.
<span class="legend-box" style="background:#ffffff"></span> White = no change.
<span class="legend-box" style="background:#2e7d32"></span> Green = accuracy improved.
Cell values are percentage-point change relative to the fast-eval baseline score.</p></div>""")

    worst_pairs = []
    for r in sorted_rows:
        for b in ACCURACY_BENCHMARKS:
            g = r.get(f"{b}_gamma", float("nan"))
            if safe(g):
                worst_pairs.append((r["label"], BENCH_LABELS[b], g))
    worst_pairs.sort(key=lambda x: x[2])
    top5 = worst_pairs[:5]
    wp_str = "; ".join(f"{t}/{b} ({g*100:.1f}%)" for t, b, g in top5)
    parts.append(f"""<div class="note">
<p>Largest accuracy drops: {wp_str}.</p></div>""")
    return "\n".join(parts)


def _grouped_bar(chart_id, rows, bench, baseline_val):
    labels = [r["label"] for r in rows]
    bl_pct = baseline_val * 100 if safe(baseline_val) else 0
    unlearned = []
    colors = []
    for r in rows:
        s = r.get(f"{bench}_score", float("nan"))
        pct = s * 100 if safe(s) else 0
        unlearned.append(pct)
        if not safe(s):
            colors.append("#999")
        elif pct < bl_pct - 5:
            colors.append("#c62828")
        elif pct < bl_pct - 1:
            colors.append("#f9a825")
        else:
            colors.append("#2e7d32")

    return f"""<div id="{chart_id}" class="chart"></div>
<script>
Plotly.newPlot('{chart_id}', [
  {{x: {json.dumps(labels)}, y: {json.dumps([bl_pct]*len(labels))},
    type: 'bar', name: 'Baseline ({bl_pct:.1f}%)',
    marker: {{color: '#1565c0', opacity: 0.35}}}},
  {{x: {json.dumps(labels)}, y: {json.dumps(unlearned)},
    type: 'bar', name: 'Unlearned',
    marker: {{color: {json.dumps(colors)}}},
    text: {json.dumps([f"{v:.1f}" for v in unlearned])},
    textposition: 'outside', textfont: {{size: 9}}}}
], {{
  title: '{BENCH_LABELS[bench]} Accuracy: Baseline vs Unlearned',
  barmode: 'group', height: 420,
  margin: {{l: 60, r: 20, t: 50, b: 130}},
  xaxis: {{tickangle: -45}},
  yaxis: {{title: 'Accuracy (%)', range: [0, 100]}},
  legend: {{x: 0.01, y: 0.99}}
}});</script>"""


def _bench_narrative(rows, bench, baseline_val):
    pairs = [(r["label"], r.get(f"{bench}_gamma", float("nan"))) for r in rows]
    valid = [(t, g) for t, g in pairs if safe(g)]
    valid.sort(key=lambda x: x[1])
    if not valid:
        return '<div class="note"><p>No valid data.</p></div>'
    worst3 = ", ".join(f"{t} ({g*100:.1f}%)" for t, g in valid[:3])
    best3 = ", ".join(f"{t} ({g*100:+.1f}%)" for t, g in valid[-3:])
    avg = nanmean([g for _, g in valid])
    bl_str = f"{baseline_val*100:.1f}%" if safe(baseline_val) else "n/a"

    special = ""
    if bench == "gsm8k":
        zero_topics = [t for t, g in valid if g <= -0.99]
        if zero_topics:
            names = ", ".join(zero_topics)
            special = (f" {names} scored 0.0% (complete failure). "
                       "This is due to PPL overshoot: the saved adapter was at a step "
                       "where PPL had already diverged by orders of magnitude past the "
                       "threshold. A checkpoint-1500 re-evaluation is pending.")

    return f"""<div class="note">
<p>Fast-eval baseline: {bl_str}. Average change across topics: {avg*100:.1f}%.{special}</p>
<p>3 worst drops: {worst3}.</p>
<p>3 least affected: {best3}.</p></div>"""


def _section3(rows, baseline):
    parts = ["<h2>3. Per-Benchmark Deep Dive</h2>"]
    parts.append("""<div class="note">
<p><strong>Bar color key:</strong>
<span class="legend-box" style="background:#c62828"></span> Dropped &gt;5% from baseline.
<span class="legend-box" style="background:#f9a825"></span> Dropped 1-5%.
<span class="legend-box" style="background:#2e7d32"></span> Within 1% or improved.
<span class="legend-box" style="background:#1565c0;opacity:0.35"></span> Baseline (translucent blue).</p>
</div>""")

    for bench in ACCURACY_BENCHMARKS:
        bl = baseline.get(bench, float("nan"))
        parts.append(f"<h3>{BENCH_LABELS[bench]}</h3>")
        parts.append(_grouped_bar(f"bar_{bench}", rows, bench, bl))
        parts.append(_bench_narrative(rows, bench, bl))

    parts.append(f"<h3>{BENCH_LABELS['wikitext']}</h3>")
    parts.append(_wikitext_chart(rows, baseline))
    parts.append(_wikitext_narrative(rows, baseline))
    return "\n".join(parts)


def _wikitext_chart(rows, baseline):
    bl = baseline.get("wikitext", float("nan"))
    labels = [r["label"] for r in rows]
    ppls = [r.get("wikitext_score", float("nan")) for r in rows]
    safe_ppls = [p if safe(p) else 1 for p in ppls]
    bl_val = bl if safe(bl) else 1
    colors = []
    for p in safe_ppls:
        if p > bl_val * 10:
            colors.append("#c62828")
        elif p > bl_val * 2:
            colors.append("#f9a825")
        else:
            colors.append("#2e7d32")
    hover = [f"PPL: {p:.2f}" if p < 1e4 else f"PPL: {p:.2e}" for p in safe_ppls]
    bar_text = [f"{p:.1f}" if p < 1000 else f"{p:.1e}" for p in safe_ppls]

    return f"""<div id="bar_wikitext" class="chart"></div>
<script>
Plotly.newPlot('bar_wikitext', [
  {{x: {json.dumps(labels)}, y: {json.dumps(safe_ppls)},
    type: 'bar', name: 'Unlearned PPL',
    marker: {{color: {json.dumps(colors)}}},
    text: {json.dumps(bar_text)}, textposition: 'outside', textfont: {{size: 8}},
    hovertext: {json.dumps(hover)}}},
  {{x: ['{labels[0]}','{labels[-1]}'], y: [{bl_val},{bl_val}],
    type: 'scatter', mode: 'lines', name: 'Baseline ({bl_val:.1f})',
    line: {{color: '#1565c0', dash: 'dash', width: 2}}}}
], {{
  title: 'Wikitext-2 Perplexity (log scale)',
  height: 420, margin: {{l: 70, r: 20, t: 50, b: 130}},
  xaxis: {{tickangle: -45}},
  yaxis: {{title: 'Word Perplexity', type: 'log'}},
  legend: {{x: 0.01, y: 0.99}}
}});</script>"""


def _wikitext_narrative(rows, baseline):
    bl = baseline.get("wikitext", float("nan"))
    pairs = [(r["label"], r.get("wikitext_score", float("nan"))) for r in rows]
    valid = [(t, p) for t, p in pairs if safe(p)]
    valid.sort(key=lambda x: x[1], reverse=True)
    worst3 = ", ".join(
        f"{t} ({p:.0f})" if p < 1e4 else f"{t} ({p:.1e})" for t, p in valid[:3])
    best3 = ", ".join(f"{t} ({p:.1f})" for t, p in valid[-3:])
    bl_str = f"{bl:.1f}" if safe(bl) else "n/a"
    return f"""<div class="note">
<p>Baseline PPL: {bl_str}. Wikitext perplexity measures general language modeling quality;
lower is better.</p>
<p>Most degraded: {worst3}.</p>
<p>Least degraded: {best3}.</p>
<p>The variation spans several orders of magnitude (from near-baseline to >10<sup>6</sup>),
reflecting the PPL overshoot problem. The adapter is saved at the step where forget-set PPL
crossed the threshold, but the 250-step check interval means general LM quality may already be
severely damaged by that point.</p></div>"""


def _section4(rows, ppl_logs, thresholds):
    parts = ["<h2>4. PPL Trajectory Analysis</h2>"]

    traces_json = []
    annotations_json = []
    for topic in sorted(ppl_logs):
        log = ppl_logs[topic]
        steps = sorted(log, key=lambda s: int(s))
        xs = [int(s) for s in steps]
        ys = [min(log[s], 1e6) for s in steps]
        traces_json.append({
            "x": xs, "y": ys, "mode": "lines+markers",
            "name": fmt(topic), "line": {"width": 1.5}, "marker": {"size": 3},
        })
    for topic, thresh in thresholds.items():
        if thresh < 1e6:
            annotations_json.append({
                "y": math.log10(thresh), "yref": "y",
                "x": 0.98, "xref": "paper",
                "text": f"{fmt(topic)[:12]} thr={thresh:.0f}",
                "showarrow": False, "font": {"size": 7, "color": "#888"},
            })

    parts.append(f"""<div id="ppl_traj" class="chart"></div>
<script>
Plotly.newPlot('ppl_traj', {json.dumps(traces_json)}, {{
  title: 'Forget-Set PPL Over Training (capped at 1e6)',
  height: 520, margin: {{l: 80, r: 40, t: 50, b: 60}},
  xaxis: {{title: 'Optimizer Step'}},
  yaxis: {{title: 'Forget PPL', type: 'log'}},
  showlegend: true, legend: {{font: {{size: 8}}, x: 1.02, y: 1}}
}});</script>""")

    tbl_rows = []
    for r in sorted(rows, key=lambda x: x.get("stop_step", 9999)):
        stop = r.get("stop_step", "?")
        thresh = r.get("threshold", float("nan"))
        final = r.get("final_ppl", float("nan"))
        if safe(thresh) and safe(final) and thresh > 0:
            overshoot = final / thresh
        else:
            overshoot = float("nan")
        t_s = f"{thresh:.1f}" if safe(thresh) else "n/a"
        f_s = f"{final:.1f}" if (safe(final) and final < 1e5) else (
            f"{final:.2e}" if safe(final) else "n/a")
        o_s = f"{overshoot:.1f}x" if (safe(overshoot) and overshoot < 1e5) else (
            f"{overshoot:.2e}x" if safe(overshoot) else "n/a")
        cls = ' class="neg"' if (safe(overshoot) and overshoot > 100) else ""
        tbl_rows.append(
            f"<tr><td>{r['label']}</td><td>{stop}</td><td>{t_s}</td>"
            f"<td{cls}>{f_s}</td><td{cls}>{o_s}</td></tr>")

    parts.append(f"""<table>
<tr><th>Topic</th><th>Stop Step</th><th>Threshold</th><th>Final PPL</th><th>Overshoot</th></tr>
{"".join(tbl_rows)}</table>""")

    fast = [r["label"] for r in rows if r.get("stop_step", 9999) <= 1100]
    med = [r["label"] for r in rows if 1100 < r.get("stop_step", 9999) <= 1300]
    std = [r["label"] for r in rows if 1300 < r.get("stop_step", 9999) <= 1600]
    slow = [r["label"] for r in rows if r.get("stop_step", 9999) > 1600]
    parts.append(f"""<div class="note">
<p>All topics follow a shared pattern: forget-set PPL stays nearly flat for 750-1000 optimizer
steps, then diverges rapidly. This phase-transition behavior means the 250-step check interval
often catches the adapter well past the intended threshold.</p>
<ul>
<li><strong>Fast stop (~1000 steps):</strong> {', '.join(fast) or 'none'}</li>
<li><strong>Medium (~1250 steps):</strong> {', '.join(med) or 'none'}</li>
<li><strong>Standard (~1500 steps):</strong> {', '.join(std) or 'none'}</li>
<li><strong>Slow (&gt;1600 steps):</strong> {', '.join(slow) or 'none'}</li>
</ul>
<p>Topics with lower thresholds (where the base model had stronger knowledge of the forget set)
generally required more steps to unlearn. The coarse check interval is the primary driver of the
overshoot problem, not the learning rate or adapter capacity.</p></div>""")
    return "\n".join(parts)


def _section5(rows):
    parts = ["<h2>5. Cross-Benchmark Correlations</h2>"]

    for bench, bid in [("gsm8k", "scatter_gsm_wiki"), ("mmlu_stem", "scatter_mmlu_wiki")]:
        xs, ys, labels = [], [], []
        for r in rows:
            wg = r.get("wikitext_gamma", float("nan"))
            bg = r.get(f"{bench}_gamma", float("nan"))
            if not safe(wg) or not safe(bg):
                continue
            xs.append(wg * 100)
            ys.append(bg * 100)
            labels.append(r["label"])
        trace = {
            "x": xs, "y": ys, "mode": "markers+text",
            "name": BENCH_LABELS[bench] + " vs Wikitext",
            "text": labels, "textposition": "top center",
            "textfont": {"size": 7}, "marker": {"size": 9},
        }
        parts.append(f"""<div id="{bid}" class="chart"></div>
<script>
Plotly.newPlot('{bid}', [{json.dumps(trace)}], {{
  title: '{BENCH_LABELS[bench]} Change vs Wikitext PPL Change (%)',
  height: 480, margin: {{l: 60, r: 20, t: 50, b: 60}},
  xaxis: {{title: 'Wikitext PPL gamma (%)'}},
  yaxis: {{title: '{BENCH_LABELS[bench]} gamma (%)'}},
  showlegend: false
}});</script>""")

    parts.append("""<div class="note">
<p>These scatter plots test whether topics with worse Wikitext PPL degradation also show worse
benchmark drops. If general language-model damage were the sole driver of accuracy loss,
we would expect a clear negative correlation (more PPL damage = lower accuracy).</p>
<p>GSM8K shows some correlation: topics with higher wikitext degradation (transportation, industrial,
finance) also tend to have larger GSM8K drops. This could reflect genuine content overlap between
those topics and numerical reasoning tasks, or it could simply mean more aggressive unlearning
damages both metrics.</p>
<p>MMLU-STEM is more uniform: most topics cluster in a narrow band regardless of wikitext impact,
suggesting MMLU knowledge is distributed broadly across the training data and not easily disrupted
by removing 1,000 documents from any single topic.</p></div>""")
    return "\n".join(parts)


def _section6(rows, completed, stored_bl_siq, fast_bl_siq):
    parts = ["<h2>6. Discussion</h2>"]
    pending = [fmt(t) for t in TOPICS if t not in completed]
    pending_str = ", ".join(pending) if pending else "none"

    gsm_zeros = [r["label"] for r in rows
                 if safe(r.get("gsm8k_score")) and r["gsm8k_score"] == 0.0]
    extreme_overshoot = [r for r in rows
                         if safe(r.get("final_ppl")) and safe(r.get("threshold"))
                         and r["threshold"] > 0 and r["final_ppl"] / r["threshold"] > 1e4]
    gsm_zero_note = ""
    if gsm_zeros:
        names = ", ".join(gsm_zeros)
        gsm_zero_note = f"""
<p><strong>GSM8K 0.0% scores ({names}).</strong> This is the extreme case of PPL overshoot.
The adapter was saved at a step where forget-set PPL had diverged by orders of magnitude past the
threshold. A checkpoint saved before the divergence is expected to produce normal GSM8K scores.
Re-evaluation with checkpoint-1500 is pending.</p>"""
    if extreme_overshoot:
        eo_details = []
        for r in extreme_overshoot:
            ratio = r["final_ppl"] / r["threshold"]
            eo_details.append(f"{r['label']} (final PPL {r['final_ppl']:.2e} vs "
                              f"threshold {r['threshold']:.0f}, {ratio:.0e}x overshoot)")
        gsm_zero_note += f"""
<p><strong>Extreme PPL overshoot cases:</strong> {'; '.join(eo_details)}.
For science_math_and_technology specifically, the adapter saved at the stop step (PPL 18.3M) has
degenerated well past the forgetting target (threshold 76). The checkpoint at step 1500
(PPL 9.8, still below threshold) is expected to retain normal benchmark scores. Re-evaluation
with checkpoint-1500 is pending.</p>"""

    siq_note = ""
    if safe(stored_bl_siq) and safe(fast_bl_siq) and abs(stored_bl_siq - fast_bl_siq) > 0.05:
        siq_scores = [r["socialiqa_score"] for r in rows if safe(r.get("socialiqa_score"))]
        siq_note = f"""
<p><strong>SocialIQA baseline problem.</strong> The 200-sample subset used in fast eval produces a
baseline of {fast_bl_siq*100:.0f}%, far below the stored full-eval baseline of {stored_bl_siq*100:.1f}%.
Unlearned models score {min(siq_scores)*100:.0f}%-{max(siq_scores)*100:.0f}% on the same subset. The
apparent uniformity of SocialIQA drops may simply reflect that this particular 200-sample slice is
not representative of the full benchmark. Until a full 1,954-sample SocialIQA eval is run, these
results are unreliable for drawing conclusions about social-reasoning impact.</p>"""

    parts.append(f"""<div class="note">
<p><strong>Observations.</strong></p>
<ul>
<li>GSM8K shows the most topic-dependent variation among accuracy benchmarks. Topics plausibly
related to numerical content (finance, industrial, transportation) show larger drops, while
topics like entertainment, fashion, and social life show minimal change.</li>
<li>MMLU-STEM and MMLU-SocSci are relatively stable across all topics, with most drops under 5%.
This is consistent with MMLU knowledge being distributed across many training topics.</li>
<li>Wikitext PPL varies by orders of magnitude across topics, driven primarily by the PPL overshoot
problem rather than by the topic content itself.</li>
</ul>
{siq_note}
{gsm_zero_note}
<p><strong>PPL overshoot problem.</strong> The 250-step check interval is too coarse for the
phase-transition dynamics observed in forget-set PPL. PPL stays flat for ~1000 steps then diverges
exponentially. By the time the check catches the crossing, the adapter may be 10<sup>5</sup>x past
the threshold. Halving the check interval (to 125 steps) or implementing a binary-search rollback
would reduce overshoot without adding significant compute cost.</p>

<p><strong>Pending topics:</strong> {pending_str}.</p>

<p><strong>Known limitations:</strong></p>
<ul>
<li>200-sample fast eval introduces sampling variance; a few percentage points of noise is expected</li>
<li>Single random seed (42) for document selection</li>
<li>1,000 forget docs out of ~5.5M total is a small fraction of training data</li>
<li>WebOrganizer topic labels are argmax from a fastText classifier; label noise is not quantified</li>
<li>An or-chain bug in eval_harness.py may affect multi-choice scoring on some benchmarks</li>
<li>PPL overshoot means the saved adapter captures a state past the intended forgetting point</li>
</ul></div>""")
    return "\n".join(parts)


def main():
    parser = argparse.ArgumentParser(description="Generate unlearning report")
    parser.add_argument("--data-dir", type=Path,
                        default=Path("artifacts/unlearn_report/data"))
    parser.add_argument("--output", type=Path,
                        default=Path("artifacts/unlearn_report/unlearn_single_topic_report.html"))
    args = parser.parse_args()

    evals, ppl_logs, thresholds = load_data(args.data_dir)
    topic_count = len([t for t in evals if t != "_baseline"])
    print(f"Loaded {topic_count} topic evals, {len(ppl_logs)} PPL logs, "
          f"{len(thresholds)} thresholds")
    if "_baseline" in evals:
        print("Baseline eval loaded")
    generate_report(evals, ppl_logs, thresholds, args.output)


if __name__ == "__main__":
    main()