From 898a0570dfe619ec2bcf330b07dce9b23cb63d54 Mon Sep 17 00:00:00 2001 From: Void Agent Date: Fri, 14 Aug 2026 13:11:02 +0100 Subject: implement Experiment 1: data pipeline, tied-RNN w/ ACT halting, transformer baseline, train/eval/plot, 16 tests --- scripts/plot.py | 49 +++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 49 insertions(+) create mode 100644 scripts/plot.py (limited to 'scripts/plot.py') diff --git a/scripts/plot.py b/scripts/plot.py new file mode 100644 index 0000000..14536e5 --- /dev/null +++ b/scripts/plot.py @@ -0,0 +1,49 @@ +"""Plot train/val curves from a metrics.csv. Usage: python -m scripts.plot """ +import csv +import sys + +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt + + +def main() -> None: + model, seed = sys.argv[1], sys.argv[2] + path = f"runs/{model}/seed{seed}/metrics.csv" + with open(path) as fh: + rows = list(csv.DictReader(fh)) + steps = [int(r["step"]) for r in rows] + train_loss = [float(r["train_loss"]) for r in rows] + train_em = [float(r["train_em"]) for r in rows] + val_em = [float(r["val_em"]) for r in rows] + train_tok = [float(r["train_token_acc"]) for r in rows] + val_tok = [float(r["val_token_acc"]) for r in rows] + halt = [float(r["mean_halt_steps"]) for r in rows] + + fig, axes = plt.subplots(2, 2, figsize=(12, 8)) + axes[0, 0].plot(steps, train_loss, label="train loss", color="tab:blue") + axes[0, 0].set_title("Train loss (token CE + halt penalty)") + axes[0, 0].set_xlabel("step") + axes[0, 1].plot(steps, train_em, label="train EM", color="tab:orange") + axes[0, 1].plot(steps, val_em, label="val EM", color="tab:green") + axes[0, 1].axhline(0.9, ls="--", c="gray", lw=0.7) + axes[0, 1].set_title(f"Exact-match (model={model}, seed={seed})") + axes[0, 1].set_ylim(-0.05, 1.05) + axes[0, 1].legend() + axes[1, 0].plot(steps, train_tok, label="train tok acc", color="tab:red") + axes[1, 0].plot(steps, val_tok, label="val tok acc", color="tab:purple") + axes[1, 0].set_title("Token accuracy") + axes[1, 0].set_ylim(-0.05, 1.05) + axes[1, 0].legend() + axes[1, 1].plot(steps, halt, label="mean halt steps", color="tab:brown") + axes[1, 1].set_title("RNN halting (mean steps used)") + axes[1, 1].set_xlabel("step") + fig.suptitle(f"prime-grokking — {model} seed {seed}") + plt.tight_layout() + out = f"runs/{model}/seed{seed}/curves.png" + plt.savefig(out, dpi=110) + print(f"saved {out}") + + +if __name__ == "__main__": + main() -- cgit v1.2.3