| """gridlm: two small models trained on a thegrid fossil record. |
| |
| screen encoder: [CLS] parent [SEP] child -> P(child is born >= 2 times in its epoch) |
| generator causal LM: [BOS] parent [SEP] child [EOS], trained on successful children only |
| |
| Held-out design: rows are split by EPOCH (train = earliest ~80% of rows' epochs, val = next |
| ~10%, test = last ~10%), so the test set is later, unseen runs of the colony. Test metrics are |
| also reported on "novel" children whose genome never appears in the training split. |
| |
| Inputs are the two instruction sequences only. Excluded on purpose: mutation_type (recorded |
| only for successful origins, so it leaks the label) and the parent's epoch birth count |
| (includes births after the child appeared). |
| |
| usage: python train.py data.jsonl.gz outdir [--threads N] [--epochs-screen E] [--epochs-gen E] |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import gzip |
| import json |
| import math |
| import os |
| import random |
| import time |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| MAXG = 64 |
| SPECIALS = ["[PAD]", "[CLS]", "[SEP]", "[BOS]", "[EOS]", "[UNK]"] |
| PAD, CLS, SEP, BOS, EOS, UNK = range(6) |
| DEV = torch.device("cpu") |
| KEY_OPS = ["alloc", "copy", "fork", "harvest", "scan", "move", "ifnotdone", "jmpb", "salvage", "steal"] |
|
|
|
|
| |
| def load(path): |
| rows = [] |
| with gzip.open(path, "rt", encoding="utf-8") as f: |
| for line in f: |
| r = json.loads(line) |
| if 1 <= len(r["child"]) <= MAXG and 1 <= len(r["parent"]) <= MAXG: |
| rows.append(r) |
| return rows |
|
|
|
|
| def split_by_epoch(rows): |
| epochs = sorted({r["epoch"] for r in rows}) |
| per = {e: 0 for e in epochs} |
| for r in rows: |
| per[r["epoch"]] += 1 |
| total, acc, cut1, cut2 = len(rows), 0, None, None |
| for e in epochs: |
| acc += per[e] |
| if cut1 is None and acc >= 0.8 * total: |
| cut1 = e |
| if cut2 is None and acc >= 0.9 * total: |
| cut2 = e |
| tr = [r for r in rows if r["epoch"] <= cut1] |
| va = [r for r in rows if cut1 < r["epoch"] <= cut2] |
| te = [r for r in rows if r["epoch"] > cut2] |
| return tr, va, te, (cut1, cut2) |
|
|
|
|
| def build_vocab(rows): |
| ops = sorted({t for r in rows for t in r["parent"] + r["child"]}) |
| return SPECIALS + ops |
|
|
|
|
| def enc_pair(r, idx): |
| p = [idx.get(t, UNK) for t in r["parent"]] |
| c = [idx.get(t, UNK) for t in r["child"]] |
| return p, c |
|
|
|
|
| def screen_tensor(rows, idx): |
| L = 2 * MAXG + 2 |
| X = np.full((len(rows), L), PAD, dtype=np.int64) |
| y = np.zeros(len(rows), dtype=np.float32) |
| for i, r in enumerate(rows): |
| p, c = enc_pair(r, idx) |
| seq = [CLS] + p + [SEP] + c |
| X[i, :len(seq)] = seq |
| y[i] = r["label"] |
| return torch.from_numpy(X), torch.from_numpy(y) |
|
|
|
|
| def lm_tensor(rows, idx): |
| L = 2 * MAXG + 3 |
| X = np.full((len(rows), L), PAD, dtype=np.int64) |
| M = np.zeros((len(rows), L), dtype=np.float32) |
| for i, r in enumerate(rows): |
| p, c = enc_pair(r, idx) |
| seq = [BOS] + p + [SEP] + c + [EOS] |
| X[i, :len(seq)] = seq |
| start = 1 + len(p) |
| M[i, start:start + len(c) + 1] = 1 |
| return torch.from_numpy(X), torch.from_numpy(M) |
|
|
|
|
| |
| class Encoder(nn.Module): |
| def __init__(self, V, d=128, layers=4, heads=4, ff=512, L=2 * MAXG + 2, causal=False): |
| super().__init__() |
| self.tok = nn.Embedding(V, d, padding_idx=PAD) |
| self.pos = nn.Embedding(L + 1, d) |
| layer = nn.TransformerEncoderLayer(d, heads, ff, dropout=0.1, batch_first=True, norm_first=True) |
| self.body = nn.TransformerEncoder(layer, layers, enable_nested_tensor=False) |
| self.norm = nn.LayerNorm(d) |
| self.causal = causal |
|
|
| def forward(self, x): |
| T = x.shape[1] |
| h = self.tok(x) + self.pos(torch.arange(T, device=x.device))[None] |
| mask = nn.Transformer.generate_square_subsequent_mask(T, device=x.device) if self.causal else None |
| h = self.body(h, mask=mask, src_key_padding_mask=(x == PAD), is_causal=self.causal) |
| return self.norm(h) |
|
|
|
|
| class Screen(nn.Module): |
| def __init__(self, V, **kw): |
| super().__init__() |
| self.enc = Encoder(V, **kw) |
| self.head = nn.Linear(self.enc.tok.embedding_dim, 1) |
|
|
| def forward(self, x): |
| return self.head(self.enc(x)[:, 0]).squeeze(-1) |
|
|
|
|
| class Generator(nn.Module): |
| def __init__(self, V, **kw): |
| super().__init__() |
| self.enc = Encoder(V, L=2 * MAXG + 3, causal=True, **kw) |
| self.head = nn.Linear(self.enc.tok.embedding_dim, V) |
|
|
| def forward(self, x): |
| return self.head(self.enc(x)) |
|
|
|
|
| |
| def auroc(scores, labels): |
| s = np.asarray(scores, dtype=np.float64) |
| y = np.asarray(labels, dtype=np.int64) |
| pos, neg = y.sum(), len(y) - y.sum() |
| if pos == 0 or neg == 0: |
| return float("nan") |
| order = np.argsort(s, kind="mergesort") |
| ranks = np.empty(len(s)) |
| sorted_s = s[order] |
| i = 0 |
| while i < len(s): |
| j = i |
| while j + 1 < len(s) and sorted_s[j + 1] == sorted_s[i]: |
| j += 1 |
| ranks[order[i:j + 1]] = (i + j) / 2 + 1 |
| i = j + 1 |
| return float((ranks[y == 1].sum() - pos * (pos + 1) / 2) / (pos * neg)) |
|
|
|
|
| def avg_precision(scores, labels): |
| order = np.argsort(-np.asarray(scores)) |
| y = np.asarray(labels)[order] |
| tp = np.cumsum(y) |
| prec = tp / np.arange(1, len(y) + 1) |
| return float((prec * y).sum() / max(1, y.sum())) |
|
|
|
|
| def precision_at(scores, labels, frac): |
| k = max(1, int(len(scores) * frac)) |
| order = np.argsort(-np.asarray(scores))[:k] |
| return float(np.asarray(labels)[order].mean()) |
|
|
|
|
| def report(name, scores, labels): |
| return {"model": name, "n": len(labels), "base_rate": float(np.mean(labels)), |
| "auroc": auroc(scores, labels), "avg_precision": avg_precision(scores, labels), |
| "precision_at_top10pct": precision_at(scores, labels, 0.10)} |
|
|
|
|
| |
| def edit_distance(a, b): |
| prev = list(range(len(b) + 1)) |
| for i, x in enumerate(a, 1): |
| cur = [i] |
| for j, yv in enumerate(b, 1): |
| cur.append(min(prev[j] + 1, cur[j - 1] + 1, prev[j - 1] + (x != yv))) |
| prev = cur |
| return prev[-1] |
|
|
|
|
| def features(r): |
| c, p = r["child"], r["parent"] |
| f = [len(c), len(p), len(c) - len(p), edit_distance(p, c)] |
| f += [c.count(op) for op in KEY_OPS] |
| f.append(float(all(op in c for op in ("alloc", "copy", "fork")))) |
| return f |
|
|
|
|
| def logistic(Xtr, ytr, Xte, iters=300): |
| mu, sd = Xtr.mean(0), Xtr.std(0) + 1e-6 |
| A = torch.tensor((Xtr - mu) / sd, dtype=torch.float32) |
| B = torch.tensor((Xte - mu) / sd, dtype=torch.float32) |
| y = torch.tensor(ytr, dtype=torch.float32) |
| w = torch.zeros(A.shape[1], requires_grad=True) |
| b = torch.zeros(1, requires_grad=True) |
| opt = torch.optim.LBFGS([w, b], max_iter=iters) |
|
|
| def closure(): |
| opt.zero_grad() |
| loss = F.binary_cross_entropy_with_logits(A @ w + b, y) + 1e-4 * (w * w).sum() |
| loss.backward() |
| return loss |
| opt.step(closure) |
| with torch.no_grad(): |
| return (B @ w + b).numpy() |
|
|
|
|
| |
| def batches(n, bs, shuffle, rng, lengths=None): |
| """Shuffled batches; with `lengths`, rows are sorted by length inside chunks of 50 |
| batches so each batch can be trimmed to its own longest row.""" |
| idx = list(range(n)) |
| if shuffle: |
| rng.shuffle(idx) |
| if lengths is not None: |
| chunk = bs * 50 |
| idx = [i for c in range(0, n, chunk) for i in sorted(idx[c:c + chunk], key=lambda j: lengths[j])] |
| groups = [idx[i:i + bs] for i in range(0, n, bs)] |
| if shuffle: |
| rng.shuffle(groups) |
| yield from groups |
| return |
| for i in range(0, n, bs): |
| yield idx[i:i + bs] |
|
|
|
|
| def trim(x, *others): |
| """Cut right padding to the longest row in the batch (sequences are left-aligned).""" |
| L = max(1, int((x != PAD).sum(1).max())) |
| return (x[:, :L],) + tuple(o[:, :L] for o in others) |
|
|
|
|
| @torch.no_grad() |
| def screen_scores(model, X, bs=4096): |
| model.eval() |
| return np.concatenate([torch.sigmoid(model(trim(X[i:i + bs])[0].to(DEV))).cpu().numpy() |
| for i in range(0, len(X), bs)]) |
|
|
|
|
| @torch.no_grad() |
| def lm_logprob(model, X, M, bs=2048): |
| """Mean log-prob per target token of each row's child (+EOS) given its parent.""" |
| model.eval() |
| out = [] |
| for i in range(0, len(X), bs): |
| x, m = trim(X[i:i + bs], M[i:i + bs]) |
| x, m = x.to(DEV), m.to(DEV) |
| logp = F.log_softmax(model(x[:, :-1]), -1) |
| tgt = x[:, 1:] |
| lp = logp.gather(-1, tgt.unsqueeze(-1)).squeeze(-1) * m[:, :-1] |
| out.append((lp.sum(1) / m[:, :-1].sum(1).clamp(min=1)).cpu().numpy()) |
| return np.concatenate(out) |
|
|
|
|
| def train_screen(Xtr, ytr, Xva, yva, V, epochs, log, rng): |
| model = Screen(V).to(DEV) |
| opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01) |
| pos_w = ((1 - ytr.mean()) / ytr.mean()).detach().to(DEV) |
| best, best_state = -1, None |
| for ep in range(epochs): |
| model.train() |
| t0, tot = time.time(), 0.0 |
| lengths = (Xtr != PAD).sum(1).tolist() |
| for b in batches(len(Xtr), 512, True, rng, lengths): |
| logits = model(trim(Xtr[b])[0].to(DEV)) |
| loss = F.binary_cross_entropy_with_logits(logits, ytr[b].to(DEV), pos_weight=pos_w) |
| opt.zero_grad() |
| loss.backward() |
| nn.utils.clip_grad_norm_(model.parameters(), 1.0) |
| opt.step() |
| tot += loss.item() * len(b) |
| va = auroc(screen_scores(model, Xva), yva.numpy()) |
| log(f"screen epoch {ep + 1}/{epochs} loss {tot / len(Xtr):.4f} val_auroc {va:.4f} ({time.time() - t0:.0f}s)") |
| if va > best: |
| best, best_state = va, {k: v.clone() for k, v in model.state_dict().items()} |
| model.load_state_dict(best_state) |
| return model, best |
|
|
|
|
| def train_generator(Xtr, Mtr, Xva, Mva, V, epochs, log, rng): |
| model = Generator(V).to(DEV) |
| opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01) |
| best, best_state = -1e9, None |
| for ep in range(epochs): |
| model.train() |
| t0, tot, cnt = time.time(), 0.0, 0.0 |
| lengths = (Xtr != PAD).sum(1).tolist() |
| for b in batches(len(Xtr), 256, True, rng, lengths): |
| x, m = trim(Xtr[b], Mtr[b]) |
| x, m = x.to(DEV), m.to(DEV) |
| logits = model(x[:, :-1]) |
| nll = F.cross_entropy(logits.reshape(-1, V), x[:, 1:].reshape(-1), reduction="none") |
| nll = (nll.view(x.shape[0], -1) * m[:, :-1]).sum() / m[:, :-1].sum() |
| opt.zero_grad() |
| nll.backward() |
| nn.utils.clip_grad_norm_(model.parameters(), 1.0) |
| opt.step() |
| tot += nll.item() |
| cnt += 1 |
| va = float(lm_logprob(model, Xva, Mva).mean()) |
| log(f"generator epoch {ep + 1}/{epochs} nll {tot / cnt:.4f} val_logprob/token {va:.4f} ({time.time() - t0:.0f}s)") |
| if va > best: |
| best, best_state = va, {k: v.clone() for k, v in model.state_dict().items()} |
| model.load_state_dict(best_state) |
| return model, best |
|
|
|
|
| @torch.no_grad() |
| def sample_children(model, parents_ids, vocab, n, temperature=0.9, rng_seed=0): |
| torch.manual_seed(rng_seed) |
| model.eval() |
| out = [] |
| for p in parents_ids[:n]: |
| seq = [BOS] + p + [SEP] |
| child = [] |
| for _ in range(MAXG + 1): |
| logits = model(torch.tensor([seq], device=DEV))[0, -1] / temperature |
| logits[[PAD, CLS, SEP, BOS, UNK]] = -1e9 |
| tok = int(torch.multinomial(F.softmax(logits, -1), 1)) |
| if tok == EOS: |
| break |
| child.append(tok) |
| seq.append(tok) |
| out.append([vocab[t] for t in child]) |
| return out |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("data") |
| ap.add_argument("outdir") |
| ap.add_argument("--threads", type=int, default=os.cpu_count()) |
| ap.add_argument("--epochs-screen", type=int, default=4) |
| ap.add_argument("--epochs-gen", type=int, default=6) |
| ap.add_argument("--max-train", type=int, default=0, help="cap training rows (0 = all), for a smoke run") |
| ap.add_argument("--samples", type=int, default=500) |
| ap.add_argument("--device", default="cpu") |
| a = ap.parse_args() |
| global DEV |
| DEV = torch.device(a.device) |
| torch.set_num_threads(a.threads) |
| os.makedirs(a.outdir, exist_ok=True) |
| logf = open(os.path.join(a.outdir, "train.log"), "a", encoding="utf-8") |
|
|
| def log(msg): |
| line = f"{time.strftime('%H:%M:%S')} {msg}" |
| print(line, flush=True) |
| logf.write(line + "\n") |
| logf.flush() |
|
|
| random.seed(0) |
| np.random.seed(0) |
| torch.manual_seed(0) |
| rng = random.Random(0) |
|
|
| rows = load(a.data) |
| tr, va, te, cuts = split_by_epoch(rows) |
| if a.max_train: |
| rng.shuffle(tr) |
| tr = tr[:a.max_train] |
| vocab = build_vocab(tr) |
| idx = {t: i for i, t in enumerate(vocab)} |
| train_children = {r["child_id"] for r in tr} |
| novel = np.array([r["child_id"] not in train_children for r in te]) |
| log(f"rows {len(rows)} train {len(tr)} val {len(va)} test {len(te)} (novel test {int(novel.sum())}); " |
| f"epoch cuts {cuts}; vocab {len(vocab)}; train base rate {np.mean([r['label'] for r in tr]):.3f}") |
|
|
| |
| yte = np.array([r["label"] for r in te]) |
| ftr = np.array([features(r) for r in tr[:200000]], dtype=np.float64) |
| fte = np.array([features(r) for r in te], dtype=np.float64) |
| results = {"split": {"epoch_cuts": cuts, "train": len(tr), "val": len(va), "test": len(te), |
| "novel_test": int(novel.sum())}, "test": [], "test_novel": []} |
| lr_scores = logistic(ftr, np.array([r["label"] for r in tr[:200000]]), fte) |
| rule = fte[:, -1] |
| for name, s in (("baseline: has alloc+copy+fork", rule), ("baseline: logistic on 15 hand features", lr_scores)): |
| results["test"].append(report(name, s, yte)) |
| results["test_novel"].append(report(name, s[novel], yte[novel])) |
| log(f"{name}: test auroc {results['test'][-1]['auroc']:.4f}") |
|
|
| |
| Xtr, ytr = screen_tensor(tr, idx) |
| Xva, yva = screen_tensor(va, idx) |
| Xte, _ = screen_tensor(te, idx) |
| screen, best = train_screen(Xtr, ytr, Xva, yva, len(vocab), a.epochs_screen, log, rng) |
| s = screen_scores(screen, Xte) |
| results["test"].append(report("gridlm screen", s, yte)) |
| results["test_novel"].append(report("gridlm screen", s[novel], yte[novel])) |
| log(f"screen: test auroc {results['test'][-1]['auroc']:.4f} (novel {results['test_novel'][-1]['auroc']:.4f})") |
| torch.save({k: v.cpu() for k, v in screen.state_dict().items()}, os.path.join(a.outdir, "screen.pt")) |
|
|
| |
| trp = [r for r in tr if r["label"] == 1] |
| vap = [r for r in va if r["label"] == 1] |
| Gtr, Mtr = lm_tensor(trp, idx) |
| Gva, Mva = lm_tensor(vap, idx) |
| gen, _ = train_generator(Gtr, Mtr, Gva, Mva, len(vocab), a.epochs_gen, log, rng) |
| torch.save({k: v.cpu() for k, v in gen.state_dict().items()}, os.path.join(a.outdir, "generator.pt")) |
| Gte, Mte = lm_tensor(te, idx) |
| lp = lm_logprob(gen, Gte, Mte) |
| results["test"].append(report("gridlm generator log p(child|parent)", lp, yte)) |
| results["test_novel"].append(report("gridlm generator log p(child|parent)", lp[novel], yte[novel])) |
| log(f"generator ranking: test auroc {results['test'][-1]['auroc']:.4f}") |
|
|
| |
| test_parents = [enc_pair(r, idx)[0] for r in te] |
| rng.shuffle(test_parents) |
| kids = sample_children(gen, test_parents, vocab, a.samples) |
| parents_txt = [[vocab[t] for t in p] for p in test_parents[:a.samples]] |
| valid = [1 <= len(k) <= MAXG for k in kids] |
| afc = [all(op in k for op in ("alloc", "copy", "fork")) for k in kids] |
| same = [k == p for k, p in zip(kids, parents_txt)] |
| fake = [{"parent": p, "child": k, "label": 0, "child_id": ""} for p, k in zip(parents_txt, kids) if 1 <= len(k) <= MAXG] |
| Xs, _ = screen_tensor(fake, idx) |
| gen_scores = screen_scores(screen, Xs) if len(fake) else np.array([]) |
| real_scores = s[: len(fake)] |
| results["samples"] = { |
| "n": len(kids), "valid_length": float(np.mean(valid)), "has_alloc_copy_fork": float(np.mean(afc)), |
| "identical_to_parent": float(np.mean(same)), |
| "screen_mean_p_generated": float(gen_scores.mean()) if len(gen_scores) else None, |
| "screen_mean_p_real_test_children": float(real_scores.mean()), |
| "note": "screen scores on generated children are one model judging another; only a live colony is ground truth", |
| } |
| with open(os.path.join(a.outdir, "samples.jsonl"), "w", encoding="utf-8") as f: |
| for p, k in zip(parents_txt, kids): |
| f.write(json.dumps({"parent": p, "child": k}) + "\n") |
| results["vocab"] = vocab |
| results["config"] = {"d_model": 128, "layers": 4, "heads": 4, "ff": 512, "max_genome": MAXG, |
| "label": "child genome born >= 2 times in its epoch", |
| "screen_params": sum(p.numel() for p in screen.parameters()), |
| "generator_params": sum(p.numel() for p in gen.parameters())} |
| with open(os.path.join(a.outdir, "results.json"), "w", encoding="utf-8") as f: |
| json.dump(results, f, indent=1) |
| with open(os.path.join(a.outdir, "vocab.json"), "w", encoding="utf-8") as f: |
| json.dump(vocab, f) |
| log("done: " + json.dumps(results["samples"])) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
