summaryrefslogtreecommitdiff
path: root/src
diff options
context:
space:
mode:
Diffstat (limited to 'src')
-rw-r--r--src/__init__.py0
-rw-r--r--src/config.py101
-rw-r--r--src/data.py79
-rw-r--r--src/eval.py151
-rw-r--r--src/model_api.py38
-rw-r--r--src/models/__init__.py0
-rw-r--r--src/models/rnn.py98
-rw-r--r--src/models/transformer.py61
-rw-r--r--src/train.py154
9 files changed, 682 insertions, 0 deletions
diff --git a/src/__init__.py b/src/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/src/__init__.py
diff --git a/src/config.py b/src/config.py
new file mode 100644
index 0000000..74cb4d3
--- /dev/null
+++ b/src/config.py
@@ -0,0 +1,101 @@
+"""Experiment configuration. Defaults = Experiment 1 (design/preregistration.md)."""
+import argparse
+import json
+from dataclasses import dataclass, fields
+
+
+@dataclass
+class Config:
+ # data
+ vocab_mode: str = "digits" # "digits" | "integers"
+ range_start: int = 2
+ range_end: int = 100 # inclusive
+ holdout_frac: float = 0.30
+ seed: int = 0
+ # model
+ model: str = "rnn" # "rnn" | "transformer"
+ d_model: int = 128
+ max_steps: int = 20 # K tied iterations (RNN)
+ halting: bool = True # False -> fixed-K ablation
+ halt_eps: float = 0.05 # ACT cumulative threshold (documented; K small so no early break)
+ halt_penalty: float = 0.01 # lambda on mean steps
+ halt_warmup_steps: int = 1000 # penalty off before this (anti-collapse)
+ min_steps: int = 2 # ACT: halt prob forced to 0 before this step
+ n_layers: int = 2
+ n_heads: int = 4
+ # training
+ lr: float = 1e-3
+ weight_decay: float = 1.0
+ max_train_steps: int = 200_000
+ eval_every: int = 200
+ early_stop_em: float = 1.0
+ early_stop_patience: int = 5
+ batch_size: int = 32
+ max_out_len: int = 6
+ log_n_examples: int = 10
+ out_dir: str = "runs"
+
+ @property
+ def vocab(self) -> int:
+ """Token count. Integers mode: 0..range_end+1 (next prime can exceed range_end)."""
+ if self.vocab_mode == "integers":
+ return self.range_end + 2
+ return 11 # digits 0-9 + EOS
+
+ @property
+ def eos_id(self) -> int:
+ return self.vocab - 1
+
+ @property
+ def pad_id(self) -> int:
+ return self.vocab # one extra embedding row reserved for pad
+
+ def to_json(self) -> dict:
+ d = {f.name: getattr(self, f.name) for f in fields(self)}
+ d["vocab"] = self.vocab
+ d["eos_id"] = self.eos_id
+ d["pad_id"] = self.pad_id
+ return d
+
+ @classmethod
+ def from_json(cls, d: dict) -> "Config":
+ cfg = cls()
+ for f in fields(cls):
+ if f.name in d:
+ setattr(cfg, f.name, d[f.name])
+ return cfg
+
+ def save(self, path: str) -> None:
+ with open(path, "w") as fh:
+ json.dump(self.to_json(), fh, indent=2)
+
+ @classmethod
+ def load(cls, path: str) -> "Config":
+ with open(path) as fh:
+ return cls.from_json(json.load(fh))
+
+
+def _bool_arg(s: str) -> bool:
+ return s.lower() in ("1", "true", "yes", "on")
+
+
+def parse_args(argv=None) -> Config:
+ cfg = Config()
+ p = argparse.ArgumentParser(description="prime-grokking train")
+ p.add_argument("model_pos", nargs="?", default=None, help="model: rnn | transformer")
+ p.add_argument("seed_pos", nargs="?", default=None, help="seed (int)")
+ for f in fields(Config):
+ if f.type is bool:
+ p.add_argument(f"--{f.name}", default=None, type=_bool_arg)
+ elif f.type in (int, float, str):
+ p.add_argument(f"--{f.name}", default=None, type=f.type)
+ a = p.parse_args(argv)
+ if a.model_pos is not None:
+ cfg.model = a.model_pos
+ if a.seed_pos is not None:
+ cfg.seed = int(a.seed_pos)
+ for f in fields(Config):
+ v = getattr(a, f.name, None)
+ if v is not None:
+ setattr(cfg, f.name, v)
+ return cfg
diff --git a/src/data.py b/src/data.py
new file mode 100644
index 0000000..35682d3
--- /dev/null
+++ b/src/data.py
@@ -0,0 +1,79 @@
+"""Prime dataset: n -> next prime, digit-tokenized; splits and batching."""
+import random
+
+import torch
+
+from src.config import Config
+
+
+def sieve_primes(limit: int) -> list[int]:
+ """All primes <= limit (inclusive)."""
+ if limit < 2:
+ return []
+ is_prime = [True] * (limit + 1)
+ is_prime[0] = is_prime[1] = False
+ for p in range(2, int(limit ** 0.5) + 1):
+ if is_prime[p]:
+ for m in range(p * p, limit + 1, p):
+ is_prime[m] = False
+ return [i for i in range(2, limit + 1) if is_prime[i]]
+
+
+def next_prime(n: int, primes: list[int]) -> int:
+ for p in primes:
+ if p > n:
+ return p
+ raise ValueError(f"no prime > {n} in supplied list")
+
+
+def encode_int(n: int, cfg: Config) -> list[int]:
+ if cfg.vocab_mode == "integers":
+ return [n]
+ return [int(d) for d in str(n)]
+
+
+def decode_tokens(ts, cfg: Config) -> int:
+ """Decode a token sequence, stopping at EOS. -1 if nothing decodable."""
+ if cfg.vocab_mode == "integers":
+ return int(ts[0]) if len(ts) else -1
+ digits = []
+ for t in ts:
+ if t == cfg.eos_id:
+ break
+ digits.append(int(t))
+ return int("".join(map(str, digits))) if digits else -1
+
+
+def get_splits(cfg: Config) -> tuple[list[int], list[int]]:
+ """(train, val) input lists, seeded shuffle, no overlap."""
+ rng = random.Random(cfg.seed)
+ inputs = list(range(cfg.range_start, cfg.range_end + 1))
+ rng.shuffle(inputs)
+ n_val = max(1, round(len(inputs) * cfg.holdout_frac))
+ return sorted(inputs[n_val:]), sorted(inputs[:n_val])
+
+
+def build_examples(inputs: list[int], cfg: Config) -> list[tuple[list[int], list[int]]]:
+ """[(input_tokens, target_tokens+EOS), ...]"""
+ primes = sieve_primes(cfg.range_end + 100)
+ out = []
+ for n in inputs:
+ p = next_prime(n, primes)
+ out.append((encode_int(n, cfg), encode_int(p, cfg) + [cfg.eos_id]))
+ return out
+
+
+def make_batch(examples, cfg: Config) -> dict[str, torch.Tensor]:
+ xs, ys = zip(*examples)
+ T_in = max(len(x) for x in xs)
+ T_out = max(len(y) for y in ys)
+ B = len(examples)
+ x = torch.full((B, T_in), cfg.pad_id, dtype=torch.long)
+ y = torch.full((B, T_out), cfg.pad_id, dtype=torch.long)
+ for i, (xi, yi) in enumerate(examples):
+ x[i, : len(xi)] = torch.tensor(xi, dtype=torch.long)
+ y[i, : len(yi)] = torch.tensor(yi, dtype=torch.long)
+ # teacher-forced decoder input: BOS(=EOS reuse) then shifted y
+ y_in = torch.cat([torch.full((B, 1), cfg.eos_id, dtype=torch.long), y[:, :-1]], dim=1)
+ y_mask = y != cfg.pad_id
+ return {"x": x, "y": y, "y_in": y_in, "y_mask": y_mask}
diff --git a/src/eval.py b/src/eval.py
new file mode 100644
index 0000000..9d55437
--- /dev/null
+++ b/src/eval.py
@@ -0,0 +1,151 @@
+"""Post-run analysis: final metrics, [101,200] probe with sieve diagnostic, grokking signature.
+
+Interpretation codes are locked in design/preregistration.md — this module only MEASURES
+and classifies against those definitions (O1-O4, H1-H4, P1-P4).
+"""
+import argparse
+import csv
+import json
+import os
+
+import numpy as np
+import torch
+
+from src.config import Config
+from src.data import build_examples, decode_tokens, encode_int, get_splits, next_prime, sieve_primes
+from src.model_api import build_model, greedy_decode
+from src.train import evaluate
+
+FLAGGED_COMPOSITES = {121, 143, 169, 187} # need divisors 11, 13 — beyond the {2,3,5,7} sieve
+
+
+def probe_report(model, cfg: Config, lo: int = 101, hi: int = 200) -> dict:
+ primes = sieve_primes(hi + 200)
+ correct = 0
+ errors = []
+ easy_misses = 0
+ for n in range(lo, hi + 1):
+ x = torch.tensor(encode_int(n, cfg), dtype=torch.long).unsqueeze(0)
+ gen = greedy_decode(model, x, cfg)[0].tolist()
+ pred = decode_tokens(gen, cfg)
+ target = next_prime(n, primes)
+ if pred == target:
+ correct += 1
+ else:
+ errors.append({"n": n, "target": target, "pred": pred})
+ if n % 2 == 0 or n % 5 == 0:
+ easy_misses += 1
+ total = hi - lo + 1
+ acc = correct / total
+ flagged = [e for e in errors if e["n"] in FLAGGED_COMPOSITES]
+ # P-code classification (see preregistration.md)
+ if acc >= 0.85:
+ code = "P3"
+ elif easy_misses >= 5:
+ code = "P4"
+ elif errors and all(e["n"] in FLAGGED_COMPOSITES for e in errors) and len(errors) <= 6:
+ code = "P1"
+ else:
+ code = "P2"
+ return {
+ "code": code, "acc": acc, "correct": correct, "total": total,
+ "errors": errors, "flagged_errors": flagged, "easy_misses": easy_misses,
+ }
+
+
+def grokking_signature(metrics_path: str) -> dict:
+ """Classify the training curve against preregistered codes O1-O4."""
+ with open(metrics_path) as fh:
+ rows = list(csv.DictReader(fh))
+ if not rows:
+ return {"code": "O4", "note": "no eval rows"}
+ train = [float(r["train_em"]) for r in rows]
+ val = [float(r["val_em"]) for r in rows]
+ n = len(rows)
+ saturated = any(all(t >= 0.95 for t in train[i:i + 10]) for i in range(n - 9)) if n >= 10 else False
+ hi = next((i for i, v in enumerate(val) if v >= 0.9), None)
+ trans = None
+ if hi is not None:
+ lo_cands = [i for i in range(hi) if val[i] <= 0.2]
+ if lo_cands:
+ trans = hi - max(lo_cands)
+ if saturated and hi is not None and trans is not None and trans <= 5:
+ code = "O1"
+ elif max(train) >= 0.95 and hi is not None and (trans is None or trans > 5):
+ code = "O3"
+ elif max(train) >= 0.95 and hi is None:
+ code = "O2"
+ else:
+ code = "O4"
+ return {
+ "code": code, "train_saturated": saturated, "val_hi_eval_idx": hi,
+ "transition_width_evals": trans, "n_evals": n,
+ "final_train_em": train[-1], "final_val_em": val[-1],
+ }
+
+
+@torch.no_grad()
+def halting_report(model, cfg: Config) -> dict:
+ """RNN halting structure: mean steps + correlation with gap-to-next-prime (H1-H4)."""
+ primes = sieve_primes(300)
+ gaps, steps = [], []
+ for n in range(cfg.range_start, cfg.range_end + 1):
+ x = torch.tensor(encode_int(n, cfg), dtype=torch.long).unsqueeze(0)
+ h = model._initial_state(x)
+ _, s = model._run_cell(h)
+ gaps.append(next_prime(n, primes) - n)
+ steps.append(float(s.mean()))
+ mean = float(np.mean(steps))
+ rho = float(np.corrcoef(gaps, steps)[0, 1]) if len(set(gaps)) > 1 else 0.0
+ if mean <= cfg.min_steps + 0.5:
+ code = "H1"
+ elif mean >= cfg.max_steps - 0.5:
+ code = "H2"
+ elif abs(rho) >= 0.3:
+ code = "H3"
+ else:
+ code = "H4"
+ return {"code": code, "mean_steps": mean, "corr_gap": rho, "min_steps": cfg.min_steps, "max_steps": cfg.max_steps}
+
+
+def main() -> None:
+ ap = argparse.ArgumentParser(description="prime-grokking eval")
+ ap.add_argument("model")
+ ap.add_argument("seed")
+ ap.add_argument("--ckpt", default="best.pt")
+ a = ap.parse_args()
+ out_dir = os.path.join("runs", a.model, f"seed{a.seed}")
+ cfg = Config.load(os.path.join(out_dir, "config.json"))
+ model = build_model(cfg)
+ model.load_state_dict(torch.load(os.path.join(out_dir, a.ckpt), map_location="cpu"))
+ model.eval()
+
+ _, val_in = get_splits(cfg)
+ val_ex = build_examples(val_in, cfg)
+ tok, em, per = evaluate(model, val_ex, cfg)
+
+ probe = probe_report(model, cfg)
+ sig = grokking_signature(os.path.join(out_dir, "metrics.csv"))
+ hlt = halting_report(model, cfg) if cfg.model == "rnn" else None
+
+ results = {
+ "model": cfg.model, "seed": cfg.seed, "ckpt": a.ckpt,
+ "params": model.param_count(),
+ "val_token_acc": tok, "val_exact_match": em,
+ "per_example": [{"n": "".join(map(str, x)), "target": t, "pred": p, "ok": ok}
+ for x, t, p, ok in per],
+ "probe": probe,
+ "signature": sig,
+ "halting": hlt,
+ }
+ with open(os.path.join(out_dir, "results.json"), "w") as fh:
+ json.dump(results, fh, indent=2)
+ print(json.dumps({"val_token_acc": round(tok, 4), "val_em": round(em, 4),
+ "probe": probe["code"], "probe_acc": round(probe["acc"], 3),
+ "flagged_errors": [e["n"] for e in probe["flagged_errors"]],
+ "signature": sig["code"], "halting": hlt["code"] if hlt else None,
+ "results": os.path.join(out_dir, "results.json")}, indent=2))
+
+
+if __name__ == "__main__":
+ main()
diff --git a/src/model_api.py b/src/model_api.py
new file mode 100644
index 0000000..66c2ace
--- /dev/null
+++ b/src/model_api.py
@@ -0,0 +1,38 @@
+"""Model interface contract + build dispatch + shared greedy decode."""
+import torch
+import torch.nn as nn
+
+from src.config import Config
+
+
+class PrimeModel(nn.Module):
+ """Contract: forward(x, y_in) -> {"logits": (B,T_out,vocab), "halt_steps": (B,) or None}."""
+
+ def forward(self, x, y_in):
+ raise NotImplementedError
+
+ def param_count(self) -> int:
+ return sum(p.numel() for p in self.parameters())
+
+
+def build_model(cfg: Config) -> PrimeModel:
+ if cfg.model == "rnn":
+ from src.models.rnn import TiedRNN
+ return TiedRNN(cfg)
+ if cfg.model == "transformer":
+ from src.models.transformer import TransformerBaseline
+ return TransformerBaseline(cfg)
+ raise ValueError(f"unknown model: {cfg.model}")
+
+
+@torch.no_grad()
+def greedy_decode(model: PrimeModel, x: torch.Tensor, cfg: Config, max_len: int | None = None) -> torch.Tensor:
+ """Autoregressive greedy decode of output digits. Returns (B, max_len) tokens (BOS stripped)."""
+ max_len = max_len or cfg.max_out_len
+ B = x.shape[0]
+ y_in = torch.full((B, 1), cfg.eos_id, dtype=torch.long, device=x.device)
+ for _ in range(max_len):
+ out = model(x, y_in)
+ nxt = out["logits"][:, -1].argmax(-1)
+ y_in = torch.cat([y_in, nxt[:, None]], dim=1)
+ return y_in[:, 1:]
diff --git a/src/models/__init__.py b/src/models/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/src/models/__init__.py
diff --git a/src/models/rnn.py b/src/models/rnn.py
new file mode 100644
index 0000000..9947c66
--- /dev/null
+++ b/src/models/rnn.py
@@ -0,0 +1,98 @@
+"""Weight-tied RNN: one 2-layer cell applied K times, ACT learned halting, GRU digit decoder.
+
+Spec pseudocode (design/experiment-spec.md):
+ state = embed(input_number)
+ for step in range(max_steps):
+ state = step_module(state) # same weights every iteration
+ if halt_condition(state): break
+ output = project(state)
+
+Initial state: masked mean-pool of (digit embedding + sinusoidal positional encoding)
+passed through a small MLP, so digit ORDER reaches the tied cell.
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from src.config import Config
+from src.model_api import PrimeModel
+
+
+class TiedRNN(PrimeModel):
+ def __init__(self, cfg: Config):
+ super().__init__()
+ self.cfg = cfg
+ d = cfg.d_model
+ self.embed = nn.Embedding(cfg.vocab + 1, d) # +1 row = pad
+ self.in_proj = nn.Sequential(nn.Linear(d, d), nn.GELU(), nn.Linear(d, d))
+ self.ln1 = nn.LayerNorm(d)
+ self.cell_ln = nn.LayerNorm(d)
+ self.cell_w1 = nn.Linear(d, d)
+ self.cell_w2 = nn.Linear(d, d)
+ self.halt_head = nn.Linear(d, 1)
+ self.decoder = nn.GRUCell(d, d)
+ self.out_head = nn.Linear(d, cfg.vocab)
+
+ @staticmethod
+ def _sinusoidal(T: int, d: int) -> torch.Tensor:
+ pe = torch.zeros(T, d)
+ pos = torch.arange(T).float().unsqueeze(1)
+ i = torch.arange(d).float().unsqueeze(0)
+ pe[:, 0::2] = torch.sin(pos / 10000 ** (2 * i[:, 0::2] / d))
+ pe[:, 1::2] = torch.cos(pos / 10000 ** (2 * i[:, 1::2] / d))
+ return pe
+
+ def _initial_state(self, x: torch.Tensor) -> torch.Tensor:
+ B, T = x.shape
+ mask = (x != self.cfg.pad_id).float().unsqueeze(-1) # (B,T,1)
+ pos = self._sinusoidal(T, self.cfg.d_model).to(x.device) # (T,d)
+ e = self.embed(x) + pos.unsqueeze(0) # (B,T,d)
+ h = (e * mask).sum(1) / mask.sum(1).clamp(min=1) # (B,d)
+ return self.in_proj(h)
+
+ def _run_cell(self, h0: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ """Tied cell x K steps. Returns (final_state (B,d), mean_steps (B,))."""
+ cfg = self.cfg
+ B = h0.shape[0]
+ device = h0.device
+ if not cfg.halting:
+ h = h0
+ for _ in range(cfg.max_steps):
+ h = h + self.cell_w2(F.gelu(self.cell_w1(self.cell_ln(h))))
+ steps = h0.new_full((B,), float(cfg.max_steps))
+ return h, steps
+ # ACT: run all K steps, accumulate weighted average (K=20 -> no early break needed)
+ h_list, p_list = [], []
+ h = h0
+ for t in range(cfg.max_steps):
+ h = h + self.cell_w2(F.gelu(self.cell_w1(self.cell_ln(h))))
+ p = torch.sigmoid(self.halt_head(self.ln1(h))).squeeze(-1) # (B,)
+ if t < cfg.min_steps:
+ p = p * 0.0
+ h_list.append(h)
+ p_list.append(p)
+ final = torch.zeros_like(h0)
+ steps = torch.zeros(B, device=device)
+ remaining = torch.ones(B, device=device)
+ for t in range(cfg.max_steps):
+ p = p_list[t]
+ w = remaining * p
+ final = final + w.unsqueeze(-1) * h_list[t]
+ steps = steps + (t + 1) * w
+ remaining = remaining * (1 - p)
+ final = final + remaining.unsqueeze(-1) * h_list[-1]
+ steps = steps + remaining * cfg.max_steps
+ return final, steps
+
+ def forward(self, x: torch.Tensor, y_in: torch.Tensor) -> dict:
+ cfg = self.cfg
+ h = self._initial_state(x)
+ h, steps = self._run_cell(h)
+ # autoregressive digit decoder, teacher-forced during training
+ e = self.embed(y_in) # (B,T_out,d)
+ outs = []
+ for t in range(y_in.shape[1]):
+ h = self.decoder(e[:, t], h)
+ outs.append(self.out_head(h))
+ logits = torch.stack(outs, dim=1) # (B,T_out,vocab)
+ return {"logits": logits, "halt_steps": steps}
diff --git a/src/models/transformer.py b/src/models/transformer.py
new file mode 100644
index 0000000..32e158b
--- /dev/null
+++ b/src/models/transformer.py
@@ -0,0 +1,61 @@
+"""Transformer baseline: GPT-style causal decoder over [input digits | output digits].
+
+Fixed d_model (=128) matching the RNN arm. Parameter counts are logged per run but
+NOT gated to parity (design/preregistration.md: weight sharing is the studied variable).
+"""
+import torch
+import torch.nn as nn
+
+from src.config import Config
+from src.model_api import PrimeModel
+
+
+class CausalBlock(nn.Module):
+ def __init__(self, d: int, heads: int):
+ super().__init__()
+ self.ln1 = nn.LayerNorm(d)
+ self.ln2 = nn.LayerNorm(d)
+ self.attn = nn.MultiheadAttention(d, heads, batch_first=True)
+ self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))
+
+ def forward(self, x: torch.Tensor, attn_mask: torch.Tensor, key_pad: torch.Tensor) -> torch.Tensor:
+ a, _ = self.attn(self.ln1(x), self.ln1(x), self.ln1(x),
+ attn_mask=attn_mask, key_padding_mask=key_pad, need_weights=False)
+ x = x + a
+ x = x + self.mlp(self.ln2(x))
+ return x
+
+
+class TransformerBaseline(PrimeModel):
+ def __init__(self, cfg: Config):
+ super().__init__()
+ self.cfg = cfg
+ d = cfg.d_model
+ self.embed = nn.Embedding(cfg.vocab + 1, d) # +1 row = pad
+ self.pos = nn.Embedding(64, d) # 3 in + 6 out max, generous
+ self.blocks = nn.ModuleList([CausalBlock(d, cfg.n_heads) for _ in range(cfg.n_layers)])
+ self.ln_f = nn.LayerNorm(d)
+ self.head = nn.Linear(d, cfg.vocab)
+
+ def forward(self, x: torch.Tensor, y_in: torch.Tensor) -> dict:
+ cfg = self.cfg
+ B, T_in = x.shape
+ T_out = y_in.shape[1]
+ seq = torch.cat([x, y_in], dim=1) # (B, T_in+T_out)
+ T = seq.shape[1]
+ device = seq.device
+ e = self.embed(seq) + self.pos(torch.arange(T, device=device)).unsqueeze(0)
+ # causal mask: input positions (j < T_in) fully visible; output positions causal
+ # (True = blocked, per torch.nn.MultiheadAttention bool convention)
+ blocked = torch.zeros(T, T, dtype=torch.bool, device=device)
+ for i in range(T):
+ for j in range(T):
+ if j >= T_in and j > i:
+ blocked[i, j] = True
+ key_pad = seq == cfg.pad_id # (B,T) True = ignore
+ h = e
+ for blk in self.blocks:
+ h = blk(h, blocked, key_pad)
+ # logits at position p predict token p+1 -> positions [T_in-1, T_in+T_out-2] predict y
+ logits = self.head(self.ln_f(h))[:, T_in - 1: T_in - 1 + T_out]
+ return {"logits": logits, "halt_steps": None}
diff --git a/src/train.py b/src/train.py
new file mode 100644
index 0000000..f0c63e9
--- /dev/null
+++ b/src/train.py
@@ -0,0 +1,154 @@
+"""Training loop. Usage: python -m src.train [model] [seed] [--flag ...]"""
+import csv
+import math
+import os
+import random
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from src.config import Config, parse_args
+from src.data import build_examples, decode_tokens, get_splits, make_batch
+from src.model_api import build_model, greedy_decode
+
+
+def set_seed(s: int) -> None:
+ random.seed(s)
+ np.random.seed(s)
+ torch.manual_seed(s)
+
+
+@torch.no_grad()
+def evaluate(model, examples, cfg: Config):
+ """Token accuracy (teacher-forced) + exact-match accuracy (greedy) + per-example detail."""
+ model.eval()
+ token_correct = 0
+ token_total = 0
+ em_correct = 0
+ per_example = []
+ for x, target in examples:
+ batch = make_batch([(x, target)], cfg)
+ out = model(batch["x"], batch["y_in"])
+ logits = out["logits"][0]
+ y = batch["y"][0]
+ mask = batch["y_mask"][0]
+ pred_tok = logits.argmax(-1)
+ token_correct += int((pred_tok[mask] == y[mask]).sum())
+ token_total += int(mask.sum())
+ xi = torch.tensor(x, dtype=torch.long).unsqueeze(0)
+ gen = greedy_decode(model, xi, cfg)[0].tolist()
+ pred = decode_tokens(gen, cfg)
+ target_n = decode_tokens(target, cfg)
+ ok = pred == target_n
+ em_correct += int(ok)
+ per_example.append((tuple(x), target_n, pred, ok))
+ model.train()
+ return token_correct / max(1, token_total), em_correct / len(examples), per_example
+
+
+def log_example_rows(log_examples, val_per_example):
+ """Format the fixed 10 logged val inputs as '42:ok' strings, in fixed order."""
+ by_x = {x: (target, pred, ok) for x, target, pred, ok in val_per_example}
+ rows = []
+ for x, _target in log_examples:
+ t, p, ok = by_x[tuple(x)]
+ label = "".join(map(str, x)) # "42" in both vocab modes
+ rows.append(f"{label}:{'ok' if ok else f'{p}~{t}'}")
+ return ";".join(rows)
+
+
+def main() -> None:
+ cfg = parse_args()
+ set_seed(cfg.seed)
+ out_dir = os.path.join(cfg.out_dir, cfg.model, f"seed{cfg.seed}")
+ os.makedirs(out_dir, exist_ok=True)
+ cfg.save(os.path.join(out_dir, "config.json"))
+
+ train_in, val_in = get_splits(cfg)
+ train_ex = build_examples(train_in, cfg)
+ val_ex = build_examples(val_in, cfg)
+ model = build_model(cfg)
+ print(f"model={cfg.model} params={model.param_count()} train={len(train_ex)} val={len(val_ex)} "
+ f"vocab={cfg.vocab} eos={cfg.eos_id} pad={cfg.pad_id}")
+
+ opt = torch.optim.AdamW(model.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay)
+ ce = nn.CrossEntropyLoss(reduction="none")
+
+ log_examples = sorted(val_ex, key=lambda x: (len(str(x)), x))[: cfg.log_n_examples]
+
+ csv_path = os.path.join(out_dir, "metrics.csv")
+ fieldnames = ["step", "train_loss", "train_token_acc", "train_em", "val_token_acc",
+ "val_em", "mean_halt_steps", "log_examples", "param_count"]
+ best_val_em = -1.0
+ patience_left = cfg.early_stop_patience
+ best_path = os.path.join(out_dir, "best.pt")
+
+ n_batches = math.ceil(len(train_ex) / cfg.batch_size)
+ rng = random.Random(cfg.seed)
+ step = 0
+ done = False
+ first_row = True
+
+ def _eval_pass(cur_loss, halt_steps):
+ nonlocal best_val_em, patience_left, done, first_row
+ train_tok, train_em, _ = evaluate(model, train_ex, cfg)
+ val_tok, val_em, val_per = evaluate(model, val_ex, cfg)
+ mean_halt = float(halt_steps.mean()) if halt_steps is not None else float("nan")
+ row = {
+ "step": step, "train_loss": float(cur_loss), "train_token_acc": train_tok,
+ "train_em": train_em, "val_token_acc": val_tok, "val_em": val_em,
+ "mean_halt_steps": mean_halt,
+ "log_examples": log_example_rows(log_examples, val_per),
+ "param_count": model.param_count(),
+ }
+ with open(csv_path, "a", newline="") as fh:
+ w = csv.DictWriter(fh, fieldnames=fieldnames)
+ if first_row:
+ w.writeheader()
+ first_row = False
+ w.writerow(row)
+ if val_em > best_val_em:
+ best_val_em = val_em
+ torch.save(model.state_dict(), best_path)
+ if val_em >= cfg.early_stop_em:
+ patience_left -= 1
+ else:
+ patience_left = cfg.early_stop_patience
+ if patience_left <= 0:
+ done = True
+ print(f"step {step} loss {float(cur_loss):.4f} train_em {train_em:.3f} "
+ f"val_em {val_em:.3f} val_tok {val_tok:.3f} halt {mean_halt:.2f}")
+
+ while step < cfg.max_train_steps and not done:
+ rng.shuffle(train_ex)
+ for bi in range(n_batches):
+ sl = train_ex[bi * cfg.batch_size: (bi + 1) * cfg.batch_size]
+ if not sl:
+ continue
+ batch = make_batch(sl, cfg)
+ out = model(batch["x"], batch["y_in"])
+ logits = out["logits"]
+ y_safe = batch["y"].clamp(max=cfg.vocab - 1) # CE index guard for pad positions
+ loss_tokens = ce(logits.reshape(-1, cfg.vocab), y_safe.reshape(-1)).reshape(
+ logits.shape[0], -1) * batch["y_mask"].float()
+ loss_tokens = loss_tokens.sum() / batch["y_mask"].sum().clamp(min=1)
+ lam = cfg.halt_penalty if (cfg.halting and step >= cfg.halt_warmup_steps) else 0.0
+ halt = out["halt_steps"]
+ penalty = halt.float().mean() * lam if halt is not None and lam > 0 else 0.0
+ loss = loss_tokens + penalty
+ opt.zero_grad()
+ loss.backward()
+ opt.step()
+ step += 1
+ if step % cfg.eval_every == 0 or step >= cfg.max_train_steps:
+ _eval_pass(loss.detach(), halt.detach() if halt is not None else None)
+ if done:
+ break
+
+ torch.save(model.state_dict(), os.path.join(out_dir, "last.pt"))
+ print(f"DONE steps={step} best_val_em={best_val_em:.4f}")
+
+
+if __name__ == "__main__":
+ main()
generated by cgit v1.2.3 (git 2.46.0) at 2026-08-31 17:38:38 +0000