diff options
Diffstat (limited to 'scripts/build_report.py')
| -rw-r--r-- | scripts/build_report.py | 21 |
1 files changed, 15 insertions, 6 deletions
diff --git a/scripts/build_report.py b/scripts/build_report.py index 4b74686..90849a2 100644 --- a/scripts/build_report.py +++ b/scripts/build_report.py @@ -17,7 +17,7 @@ import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt -from src.eval import FLAGGED_SIEVE_PREDS +from src.eval import _sieve_rank_signature CODE_LEGEND = { "O1": "Sharp transition — grokking-like", @@ -63,20 +63,29 @@ def img_b64(path: str) -> str: def probe_fig(result: dict, model: str) -> str: errors = result.get("errors", []) - flagged = [e for e in errors if e["pred"] in FLAGGED_SIEVE_PREDS] + if not errors: + return "" + probe = result.get("probe", {}) + rank = probe.get("sieve_rank") + lo = min(e["n"] for e in errors) - 5 if errors else 100 + hi = max(e["n"] for e in errors) + 5 if errors else 200 + flagged = set() + if rank is not None: + flagged = _sieve_rank_signature(rank, lo, hi) fig, ax = plt.subplots(figsize=(10, 3.2)) xs, ys, cs, ls = [], [], [], [] for e in errors: xs.append(e["n"]); ys.append(1.0) - cs.append("tab:red" if e["pred"] in FLAGGED_SIEVE_PREDS else "tab:blue") + cs.append("tab:red" if e["pred"] in flagged else "tab:blue") ls.append(f"n={e['n']}\npred {e['pred']}\ntarget {e['target']}") ax.scatter(xs, ys, c=cs, s=28, zorder=3) - ax.set_xlim(100, 201) + ax.set_xlim(lo - 1, hi + 1) ax.set_ylim(0.6, 1.4) ax.set_yticks([1.0]); ax.set_yticklabels(["wrong"]) ax.set_xlabel("input n") - ax.set_title(f"{model}: probe errors (red = predicted a no-small-divisor composite " - f"{sorted(FLAGGED_SIEVE_PREDS)})") + rank_str = f" rank-{rank}" if rank is not None else "" + ax.set_title(f"{model}: probe errors (red = sieve{rank_str} signature " + f"{sorted(flagged) if flagged else 'none'})") for x, y, l in zip(xs, ys, ls): ax.annotate(l, (x, y), textcoords="offset points", xytext=(0, 10), fontsize=6, ha="center", rotation=90, va="bottom") |