diff options
Diffstat (limited to 'src/eval.py')
| -rw-r--r-- | src/eval.py | 7 |
1 files changed, 5 insertions, 2 deletions
diff --git a/src/eval.py b/src/eval.py index 0e769ac..7dcd289 100644 --- a/src/eval.py +++ b/src/eval.py @@ -132,7 +132,8 @@ def halting_report(model, cfg: Config) -> dict: else: code = "H4" # intermediate but no positive gap correlation — noisy/unused return {"code": code, "mean_steps": mean, "corr_gap": rho, - "min_steps": cfg.min_steps, "max_steps": cfg.max_steps} + "min_steps": cfg.min_steps, "max_steps": cfg.max_steps, + "gaps": gaps, "steps_by_n": steps} def _load(out_dir: str, ckpt: str, cfg: Config): @@ -146,8 +147,10 @@ def main() -> None: ap = argparse.ArgumentParser(description="prime-grokking eval") ap.add_argument("model") ap.add_argument("seed") + ap.add_argument("--ckpt", default="last.pt") + ap.add_argument("--runs-dir", default="runs") a = ap.parse_args() - out_dir = os.path.join("runs", a.model, f"seed{a.seed}") + out_dir = os.path.join(a.runs_dir, a.model, f"seed{a.seed}") cfg = Config.load(os.path.join(out_dir, "config.json")) _, val_in = get_splits(cfg) |