#!/usr/bin/env python """Build the P2 SFT shard set from Generation + Verification JSONL (expert iteration). .venv/bin/python scripts/p2/build_sft_data.py \\ --generations gens_r0.jsonl gens_r1.jsonl --verifications ver_r0.jsonl ver_r1.jsonl \\ --problems train_pool.jsonl --eval-problems minif2f.jsonl mobench.jsonl \\ --cap-tokens 16384 --out-dir DATA_DIR Writes `DATA_DIR/shards/sft/train/sft-train-NNNN.*` through `align.data.store.ShardWriter` (the layout `align.train.stage_smoke.resolve_split_dir` discovers, so `scripts/sft.py` reads it unchanged), `DATA_DIR/manifest.json` (counts by source / round / kind, token totals, hashes) and `DATA_DIR/selection.jsonl` (which sample became which row, and why others did not). What is trained (a) kind "first": accepted round-0 samples; (b) kind "correction": accepted round>0 samples whose prompt, AS GENERATED, already carries the failed attempt and its Lean errors. Rows are `prompt_ids + completion_ids` EXACTLY as sampled: nothing is re-tokenised. Re-encoding text would change the token sequence the policy was sampled from (BPE is not a bijection on text), and the training signal would no longer be the sampled distribution. The loss mask is the store's `prompt_len` extra: `collate_sft` labels positions [prompt_len, L) and nothing before, which is the completion exactly; `readback_check` re-opens the finished shards and asserts it for every row. Defects each rule prevents * eval leakage: `--problems` must be the TRAINING pool; any generation for a problem that is not in it, or whose statement sha is in an `--eval-problems` file, RAISES. * stale verdicts: a verification record whose `code_sha256` differs from the generation's code raises (the generation file was regenerated under the same keys). * cap mixing: only `--cap-tokens` is used (required, never defaulted); samples at another cap are counted and skipped, so a 40960-cap reference run cannot leak into the 16384 training set. * too-easy problems (8/8 at round 0) carry no learning signal and crowd the mix: excluded unless `--keep-easy`. * a correction example whose prompt does NOT contain the failed attempt would teach "write the answer to an unrelated error message": dropped and counted (`correction_prompt_incomplete`); one whose parent was in fact ACCEPTED is dropped (`correction_parent_not_failed`). * length-truncated or EOS-less completions would teach never to stop: dropped; `--allow-no-eos` is the explicit override. * an ACCEPTED sample whose displayed theorem is not the problem's statement (`problem_check` swaps the original in before verification, so tactics that prove the original are credited even when the model wrote a different, possibly false, theorem above them) would teach the model to state the wrong theorem: dropped and counted (`drop_statement_not_faithful`). The evaluation keeps Goedel's crediting rule (comparability); only the TRAINING selection is stricter. * duplicates: distinct proofs (whitespace/comment-insensitive hash) are preferred; duplicates only fill the cap with `--fill-duplicates`. """ from __future__ import annotations import argparse import importlib.util import json import logging import sys from collections import Counter, defaultdict from pathlib import Path from typing import Optional REPO = Path(__file__).resolve().parents[2] sys.path.insert(0, str(REPO)) def _load_common(): spec = importlib.util.spec_from_file_location("p2_common", Path(__file__).with_name("p2_common.py")) mod = importlib.util.module_from_spec(spec) sys.modules["p2_common"] = mod spec.loader.exec_module(mod) return mod C = sys.modules.get("p2_common") or _load_common() log = logging.getLogger("p2.sft_data") BUILD_VERSION = 2 # `assert_build_is_current` floor (cfg.gates.min_build_version); see docs/P2.md ROWS_PER_SHARD = 2000 EASY_MIN_SAMPLES = 8 class SelectionError(C.ContractError): pass def _norm_ws(s: str) -> str: return " ".join(s.split()) # =========================================================================== pass 1: index class Index: """Light per-sample metadata (no token lists): the full records are re-read by byte offset in pass 2, so memory is O(samples), not O(tokens).""" def __init__(self) -> None: self.loc: dict[tuple, tuple[int, int]] = {} # key -> (file index, byte offset) self.meta: dict[tuple, dict] = {} self.files: list[str] = [] def scan_generations(gen_paths, cap_tokens: int, counts: Counter) -> Index: idx = Index() for fi, path in enumerate(gen_paths): idx.files.append(str(path)) with open(path, "rb") as f: off = 0 for lineno, raw in enumerate(f, 1): here, off = off, off + len(raw) if not raw.strip(): continue try: rec = json.loads(raw) except json.JSONDecodeError as e: raise C.ContractError(f"{path}:{lineno}: malformed JSON ({e.msg})") from e C.validate_generation(rec) key = C.gen_key(rec) if key in idx.loc: raise C.ContractError(f"{path}:{lineno}: duplicate generation key {key}") counts["generations_read"] += 1 if rec["cap_tokens"] != cap_tokens: counts["skipped_other_cap"] += 1 continue idx.loc[key] = (fi, here) idx.meta[key] = { "finish_reason": rec["finish_reason"], "code_sha": C.code_sha(rec["lean_code"]), "proof_hash": C.proof_hash(rec["lean_code"]) if rec["lean_code"] and rec["lean_code"].strip() else None, "n_prompt": len(rec["prompt_ids"]), "n_completion": len(rec["completion_ids"]), "last_token": rec["completion_ids"][-1], "parent": rec["parent"], "model_rev": rec["model_rev"], } return idx def read_at(idx: Index, key: tuple) -> dict: fi, off = idx.loc[key] with open(idx.files[fi], "rb") as f: f.seek(off) return json.loads(f.readline()) def load_verifications(paths, idx: Index, cap_tokens: int, counts: Counter) -> dict[tuple, dict]: ver: dict[tuple, dict] = {} for p in paths: for rec in C.iter_jsonl(p): key = C.gen_key(rec) if key[3] != cap_tokens: counts["verifications_other_cap"] += 1 continue if key in ver: raise C.ContractError(f"{p}: duplicate verification key {key} across files") if key not in idx.meta: raise C.ContractError(f"{p}: verification {key} has no generation in the given generation files") m = idx.meta[key] if rec["code_sha256"] != m["code_sha"]: raise C.ContractError( f"{p}: stale verification for {key}: code_sha256 {rec['code_sha256']} != generation's {m['code_sha']}" ) ver[key] = rec missing = [k for k in idx.meta if k not in ver] if missing: raise C.ContractError(f"{len(missing)} generations have no verification record (first: {missing[0]}); verify first") return ver # =========================================================================== selection def eos_ids_for(model_path: str) -> set[int]: """EOS token ids from the model's config.json / generation_config.json.""" out: set[int] = set() for name in ("config.json", "generation_config.json"): p = Path(model_path) / name if p.exists(): v = json.loads(p.read_text()).get("eos_token_id") if isinstance(v, int): out.add(v) elif isinstance(v, list): out |= {int(x) for x in v} if not out: raise SelectionError(f"no eos_token_id found under {model_path}; pass --eos-ids") return out def select(idx: Index, ver: dict[tuple, dict], problems: dict[str, dict], a, counts: Counter, eos: set[int]): """-> (selected [(key, kind)], per_problem stats, dropped log).""" by_problem: dict[str, list[tuple]] = defaultdict(list) for key in idx.meta: by_problem[key[0]].append(key) selected: list[tuple[tuple, str]] = [] pstats: dict[str, dict] = {} dropped: list[dict] = [] for pid in sorted(by_problem): keys = sorted(by_problem[pid], key=lambda k: (k[2], k[1])) r0 = [k for k in keys if k[2] == 0] n0, c0 = len(r0), sum(1 for k in r0 if ver[k]["accepted"]) st = {"n_round0": n0, "accepted_round0": c0, "selected": 0, "easy_excluded": False} pstats[pid] = st if n0 >= EASY_MIN_SAMPLES and c0 == n0 and not a.keep_easy: st["easy_excluded"] = True counts["easy_problems_excluded"] += 1 counts["easy_samples_excluded"] += c0 continue corr, first = [], [] for k in keys: if not ver[k]["accepted"]: continue m = idx.meta[k] why = _reject_sample(k, m, a, eos) if why is None and not C.displayed_statement_matches(problems[pid]["formal_statement"], read_at(idx, k)["completion_text"]): why = "statement_not_faithful" if why is not None: counts[f"drop_{why}"] += 1 dropped.append({"key": list(k), "why": why}) continue if k[2] == 0: first.append(k) else: why = _check_correction(idx, ver, k) if why is not None: counts[f"drop_{why}"] += 1 dropped.append({"key": list(k), "why": why}) continue corr.append(k) if getattr(a, "rank", "order") == "shortest": # expert-iteration keep rule: among accepted samples prefer the shortest completion (ties: key order) first.sort(key=lambda k: (idx.meta[k]["n_completion"], k)) corr.sort(key=lambda k: (idx.meta[k]["n_completion"], k)) chosen: list[tuple[tuple, str]] = [] seen_hash: set[str] = set() def room(kind: str, kind_limit: int) -> bool: return len(chosen) < a.max_per_problem and sum(1 for _, kd in chosen if kd == kind) < kind_limit # Corrections first (rarer, and the only source of the self-correction skill), then # first attempts; within a kind, distinct proofs before any duplicate. for kind, cands, kind_limit in (("correction", corr, a.max_correction_per_problem), ("first", first, a.max_per_problem)): for k in cands: h = idx.meta[k]["proof_hash"] if room(kind, kind_limit) and h not in seen_hash: seen_hash.add(h) chosen.append((k, kind)) if a.fill_duplicates: for k in cands: if room(kind, kind_limit) and (k, kind) not in chosen: chosen.append((k, kind)) n_acc = len(corr) + len(first) counts["accepted_candidates"] += n_acc counts["not_selected_cap_or_duplicate"] += n_acc - len(chosen) st["selected"] = len(chosen) selected += chosen return selected, pstats, dropped def _reject_sample(k: tuple, m: dict, a, eos: set[int]) -> Optional[str]: if m["finish_reason"] != "stop": return "finish_reason_not_stop" # an accepted length-capped sample cannot exist; belt and braces if m["proof_hash"] is None: return "accepted_without_code" if not a.allow_no_eos and m["last_token"] not in eos: return "completion_lacks_eos" if a.max_seq_len is not None and m["n_prompt"] + m["n_completion"] > a.max_seq_len: return "over_max_seq_len" return None def _check_correction(idx: Index, ver: dict[tuple, dict], key: tuple) -> Optional[str]: m = idx.meta[key] par = m["parent"] pkey = (par["problem_id"], int(par["sample_idx"]), int(par["round"]), key[3]) if pkey not in idx.meta: raise SelectionError(f"correction {key} names parent {pkey} which is not among the generations at this cap") if ver[pkey]["accepted"]: return "correction_parent_not_failed" child, parent = read_at(idx, key), read_at(idx, pkey) prompt = _norm_ws(child["prompt_text"]) tail = _norm_ws(parent["completion_text"])[-200:] if not tail or tail not in prompt or "rror" not in prompt: return "correction_prompt_incomplete" return None # =========================================================================== write def write_shards(out_dir: Path, idx: Index, selected, problems: dict[str, dict], held: set[str]): from align.data.store import ShardWriter dirs = {"train": out_dir / "shards" / "sft" / "train", "held": out_dir / "shards" / "sft" / "held"} writers: dict[str, ShardWriter] = {} nrows: Counter = Counter() receipts: dict[str, list] = defaultdict(list) sel_log = [] def writer(split: str) -> ShardWriter: w = writers.get(split) if w is None or w.n_rows >= ROWS_PER_SHARD: if w is not None: receipts[split].append(w.finalize()) w = writers[split] = ShardWriter(dirs[split], f"sft-{split}-{len(receipts[split]):04d}") return w toks = Counter() by: Counter = Counter() for key, kind in selected: rec = read_at(idx, key) pid = key[0] split = "held" if pid in held else "train" ids = rec["prompt_ids"] + rec["completion_ids"] src = problems[pid]["source"] writer(split).add_row( ids, problem_id=pid, subject=src, boundary_positions=[], prompt_len=len(rec["prompt_ids"]), kind=kind, round=key[2], sample_idx=key[1], cap_tokens=key[3], n_completion=len(rec["completion_ids"]), proof_sha=idx.meta[key]["proof_hash"], model_rev=rec["model_rev"], ) nrows[split] += 1 toks[f"{split}_prompt_tokens"] += len(rec["prompt_ids"]) toks[f"{split}_completion_tokens"] += len(rec["completion_ids"]) by[f"{split}/{src}/round{key[2]}/{kind}"] += 1 sel_log.append({"key": list(key), "kind": kind, "split": split}) for split, w in writers.items(): receipts[split].append(w.finalize()) return receipts, nrows, toks, by, sel_log def readback_check(out_dir: Path, idx: Index, selected, held: set[str]) -> int: """Re-open the written shards and assert, for EVERY row, that the stored tokens equal the sampled prompt+completion and that the loss mask (as `collate_sft` builds it) covers exactly the completion positions. Returns the number of rows checked.""" import numpy as np from align.data.loader import collate_sft from align.data.store import ShardDataset n = 0 for split in ("train", "held"): d = out_dir / "shards" / "sft" / split want = [(k, kd) for k, kd in selected if (k[0] in held) == (split == "held")] if not want: continue ds = ShardDataset(d) if len(ds) != len(want): raise SelectionError(f"{split}: {len(ds)} stored rows != {len(want)} selected") for i, (key, _kind) in enumerate(want): rec = read_at(idx, key) row = ds[i] p = len(rec["prompt_ids"]) if not np.array_equal(row.input_ids, np.asarray(rec["prompt_ids"] + rec["completion_ids"], dtype=row.input_ids.dtype)): raise SelectionError(f"{key}: stored tokens differ from the sampled prompt+completion") if int(row.extras["prompt_len"]) != p: raise SelectionError(f"{key}: prompt_len {row.extras['prompt_len']} != {p}") lab = collate_sft([row], pad_token_id=0).labels[0].numpy() if not (np.all(lab[:p] == -100) and np.array_equal(lab[p:], row.input_ids[p:].astype(lab.dtype))): raise SelectionError(f"{key}: loss mask is not exactly the completion span") n += 1 return n # =========================================================================== CLI def main(argv=None) -> int: ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("--generations", nargs="+", required=True) ap.add_argument("--verifications", nargs="+", required=True) ap.add_argument("--problems", nargs="+", required=True, help="TRAINING pool Problem JSONL (split must be train)") ap.add_argument("--eval-problems", nargs="+", required=True, help="evaluation Problem JSONL; none of their statements may appear in --problems") ap.add_argument("--cap-tokens", type=int, required=True, help="only samples generated at this cap are used (no default on purpose)") ap.add_argument("--out-dir", required=True) ap.add_argument("--max-per-problem", type=int, default=4) ap.add_argument("--max-correction-per-problem", type=int, default=1, help="slots reserved for correction examples inside --max-per-problem") ap.add_argument("--keep-easy", action="store_true", help="keep problems solved n/n at round 0 (n >= 8)") ap.add_argument("--rank", choices=("order", "shortest"), default="order", help="which accepted samples win the per-problem cap: sampling order (default) or shortest completion first") ap.add_argument("--fill-duplicates", action="store_true", help="fill the cap with duplicate proofs when distinct ones run out") ap.add_argument("--max-seq-len", type=int, default=None, help="drop (never truncate) rows longer than this") ap.add_argument("--allow-no-eos", action="store_true") ap.add_argument("--eos-ids", type=int, nargs="*", default=None) ap.add_argument("--held-frac", type=float, default=0.0, help="fraction of PROBLEMS (not rows) routed to a held-out split") ap.add_argument("--seed", type=int, default=0) a = ap.parse_args(argv) logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s") out = Path(a.out_dir).resolve() if not str(out).startswith("/apps/"): raise SystemExit(f"--out-dir {out} must be on /apps") if out.exists() and any(out.iterdir()): raise SystemExit(f"{out} is not empty: refusing to mix shards from two builds") if a.max_correction_per_problem > a.max_per_problem or a.max_per_problem < 1: raise SystemExit("need 1 <= max-correction-per-problem <= max-per-problem") problems: dict[str, dict] = {} for p in a.problems: problems.update(C.load_problems(p)) bad = [pid for pid, r in problems.items() if r["split"] != "train"] if bad: raise SelectionError(f"{len(bad)} problems in --problems are not split=train (first: {bad[0]})") eval_sha: set[str] = set() for p in a.eval_problems: eval_sha |= {r["statement_sha"] for r in C.load_problems(p).values()} leaked = [pid for pid, r in problems.items() if r["statement_sha"] in eval_sha] if leaked: raise SelectionError(f"{len(leaked)} training problems share a statement with an eval set (first: {leaked[0]})") counts: Counter = Counter() idx = scan_generations(a.generations, a.cap_tokens, counts) if not idx.meta: raise SelectionError(f"no generations at cap_tokens={a.cap_tokens}") foreign = sorted({k[0] for k in idx.meta} - set(problems)) if foreign: raise SelectionError(f"{len(foreign)} generated problems are not in the training pool (first: {foreign[0]}): refusing to train on them") ver = load_verifications(a.verifications, idx, a.cap_tokens, counts) revs = {m["model_rev"] for m in idx.meta.values()} model_paths = {read_at(idx, next(iter(idx.loc)))["model_path"]} eos = set(a.eos_ids) if a.eos_ids else eos_ids_for(next(iter(model_paths))) selected, pstats, dropped = select(idx, ver, problems, a, counts, eos) if not selected: raise SelectionError("nothing selected: check the verification file (all rejected?) and the filters; counts=" + json.dumps(dict(counts))) import hashlib held: set[str] = set() if a.held_frac > 0: for pid in {k[0] for k, _ in selected}: if int(hashlib.sha256(f"{a.seed}|held|{pid}".encode()).hexdigest(), 16) % 10_000 < a.held_frac * 10_000: held.add(pid) out.mkdir(parents=True, exist_ok=True) try: receipts, nrows, toks, by, sel_log = write_shards(out, idx, selected, problems, held) n_checked = readback_check(out, idx, selected, held) except BaseException: import shutil shutil.rmtree(out / "shards", ignore_errors=True) # a half-written shard set looks usable raise total = sum(nrows.values()) sum_receipt = sum(r.n_rows for rs in receipts.values() for r in rs) if not (total == len(selected) == sum_receipt == n_checked): raise SelectionError(f"row accounting differs: selected {len(selected)}, written {total}, receipts {sum_receipt}, read back {n_checked}") (out / "selection.jsonl").write_text("".join(json.dumps(r) + "\n" for r in sel_log)) (out / "dropped.jsonl").write_text("".join(json.dumps(r) + "\n" for r in dropped)) manifest = { "build_kind": "p2_sft", "build_version": BUILD_VERSION, "pins_verified": True, # the row accounting above reconciled; there are no corpus pins in P2 "pin_violations": [], "cap_tokens": a.cap_tokens, "params": {k: getattr(a, k) for k in ("max_per_problem", "max_correction_per_problem", "keep_easy", "fill_duplicates", "max_seq_len", "allow_no_eos", "held_frac", "seed", "rank")}, "eos_ids": sorted(eos), "model_revs": sorted(revs), "rows": {"total": total, **{s: nrows[s] for s in ("train", "held")}}, "rows_by_split_source_round_kind": dict(sorted(by.items())), "tokens": dict(toks), "problems": {"with_rows": len({k[0] for k, _ in selected}), "in_pool": len(problems), "seen_in_generations": len(pstats), "easy_excluded": counts["easy_problems_excluded"], "held_problems": len(held)}, "counts": dict(counts), "inputs": {"generations": {str(p): C.sha256_file(p) for p in a.generations}, "verifications": {str(p): C.sha256_file(p) for p in a.verifications}, "problems": {str(p): C.sha256_file(p) for p in a.problems}, "eval_problems": {str(p): C.sha256_file(p) for p in a.eval_problems}}, "shards": {s: [r.to_json() for r in rs] for s, rs in receipts.items()}, "readback_rows_checked": n_checked, } (out / "manifest.json").write_text(json.dumps(manifest, indent=1)) print(json.dumps({"rows": manifest["rows"], "tokens": manifest["tokens"], "counts": manifest["counts"]})) return 0 if __name__ == "__main__": raise SystemExit(main())