1 files changed, 10 insertions, 4 deletions
diff --git a/scripts/plot.py b/scripts/plot.py
index 14536e5..6a71c3b 100644
--- a/scripts/plot.py
+++ b/scripts/plot.py
@@ -1,4 +1,5 @@
-"""Plot train/val curves from a metrics.csv. Usage: python -m scripts.plot <model> <seed>"""
+"""Plot train/val curves from a metrics.csv. Usage: python -m scripts.plot <model> <seed> [--runs-dir DIR]"""
+import argparse
 import csv
 import sys
 
@@ -8,8 +9,14 @@ import matplotlib.pyplot as plt
 
 
 def main() -> None:
-    model, seed = sys.argv[1], sys.argv[2]
-    path = f"runs/{model}/seed{seed}/metrics.csv"
+    ap = argparse.ArgumentParser()
+    ap.add_argument("model")
+    ap.add_argument("seed")
+    ap.add_argument("--runs-dir", default="runs")
+    a = ap.parse_args()
+    model, seed = a.model, a.seed
+    path = f"{a.runs_dir}/{model}/seed{seed}/metrics.csv"
+    out = f"{a.runs_dir}/{model}/seed{seed}/curves.png"
     with open(path) as fh:
         rows = list(csv.DictReader(fh))
     steps = [int(r["step"]) for r in rows]
@@ -40,7 +47,6 @@ def main() -> None:
     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}")