spb/wpai0001_uqo
Public
1#!/usr/bin/env python32"""Train the Tier-0 model matrix used by the transplant experiments (documented training stage).34Roles (OPSEQ, modulus 11, scratchpad):5 donor D : operators 8 (incl. Ω = qua), vex excluded6 recipient R : Ω excluded, vex excluded (4 layers, d = 128)7 unrelated U : Ω excluded, vex INCLUDED (competent on an unrelated operator the recipient lacks)8 recipient_w96 : as R with d = 96 (width mismatch)9 recipient_L6 : as R with 6 layers (depth mismatch)10 donor_s2 / recipient_s2 : as D / R with trace_stride = 2 (two operators per emitted result)11Each model directory holds model.pt, configs, vocab, history and the training split hash.12"""13from __future__ import annotations1415import argparse16import json17import sys18import time19from pathlib import Path2021ROOT = Path(__file__).resolve().parents[1]22sys.path.insert(0, str(ROOT / "src"))2324from wp_ai_0001.config import OpSeqTaskConfig, ToyModelConfig, TrainConfig # noqa: E40225from wp_ai_0001.datasets import opseq # noqa: E40226from wp_ai_0001.training.train_toy import train # noqa: E4022728OMEGA = "qua"29MAX_SEQ = 2430BASE = ["zib", "zob", "qua", "lar", "neg", "flip", "dro", "mik"]31ROLES: dict[str, tuple[OpSeqTaskConfig, ToyModelConfig]] = {32 "donor": (OpSeqTaskConfig(operators=BASE), ToyModelConfig(max_seq_len=MAX_SEQ)),33 "recipient": (OpSeqTaskConfig(operators=BASE, excluded_operators=[OMEGA]), ToyModelConfig(max_seq_len=MAX_SEQ)),34 "unrelated": (OpSeqTaskConfig(operators=[o for o in BASE if o != OMEGA] + ["vex"]), ToyModelConfig(max_seq_len=MAX_SEQ)),35 "recipient_w96": (OpSeqTaskConfig(operators=BASE, excluded_operators=[OMEGA]), ToyModelConfig(d_model=96, d_ff=384, max_seq_len=MAX_SEQ)),36 "recipient_L6": (OpSeqTaskConfig(operators=BASE, excluded_operators=[OMEGA]), ToyModelConfig(n_layers=6, max_seq_len=MAX_SEQ)),37 "donor_s2": (OpSeqTaskConfig(operators=BASE, trace_stride=2), ToyModelConfig(max_seq_len=MAX_SEQ)),38 "recipient_s2": (OpSeqTaskConfig(operators=BASE, excluded_operators=[OMEGA], trace_stride=2), ToyModelConfig(max_seq_len=MAX_SEQ)),39}40DEFAULT_SEEDS = {"donor": [0, 1, 2], "recipient": [0, 1, 2], "unrelated": [0, 1, 2], "recipient_w96": [0, 1], "recipient_L6": [0, 1], "donor_s2": [0, 1], "recipient_s2": [0, 1]}41MODELS_DIR = ROOT / "artifacts" / "models" / "matrix_v1"424344def model_dir(role: str, seed: int) -> Path:45 return MODELS_DIR / f"{role}_s{seed}"464748def main() -> None:49 ap = argparse.ArgumentParser()50 ap.add_argument("--roles", nargs="*", default=list(ROLES))51 ap.add_argument("--steps", type=int, default=3000)52 ap.add_argument("--device", default="auto")53 args = ap.parse_args()54 index_path = MODELS_DIR / "index.json"55 index = json.loads(index_path.read_text()) if index_path.exists() else {}56 for role in args.roles:57 task, mcfg = ROLES[role]58 for seed in DEFAULT_SEEDS[role]:59 d = model_dir(role, seed)60 if (d / "model.pt").exists():61 print(f"skip {role} s{seed} (exists)")62 continue63 t0 = time.time()64 # training split seeds are role- and seed-specific; evaluation splits of the experiments use other seeds65 tr = opseq.generate(task, 40_000, 100_000 + 1_000 * list(ROLES).index(role) + seed, f"{role}_train_s{seed}", mode="train")66 dv = opseq.generate(task, 1_000, 200_000 + 1_000 * list(ROLES).index(role) + seed, f"{role}_dev_s{seed}", mode="train", exclude_hashes=tr.item_hashes())67 opseq.validate(tr, task)68 vocab = opseq.vocab_for(task)69 hist: list = []70 train(mcfg, TrainConfig(seed=seed, steps=args.steps, device=args.device), tr, dv, vocab, d, hist)71 tr.save(d / "train_split.json")72 index[f"{role}_s{seed}"] = {"role": role, "seed": seed, "task": task.model_dump(), "model": mcfg.model_dump(), "train_hash": tr.hash, "dev_hash": dv.hash,73 "final_dev_acc": hist[-1]["dev_acc"], "train_seconds": time.time() - t0}74 index_path.write_text(json.dumps(index, indent=2))75 print(f"trained {role} s{seed}: dev_acc={hist[-1]['dev_acc']:.3f} ({time.time()-t0:.0f}s)")767778if __name__ == "__main__":79 main()80