| |
| """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 |
| ROWS_PER_SHARD = 2000 |
| EASY_MIN_SAMPLES = 8 |
|
|
|
|
| class SelectionError(C.ContractError): |
| pass |
|
|
|
|
| def _norm_ws(s: str) -> str: |
| return " ".join(s.split()) |
|
|
|
|
| |
|
|
|
|
| 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]] = {} |
| 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 |
|
|
|
|
| |
|
|
|
|
| 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": |
| |
| 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 |
|
|
| |
| |
| 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" |
| 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 |
|
|
|
|
| |
|
|
|
|
| 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 |
|
|
|
|
| |
|
|
|
|
| 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) |
| 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, |
| "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()) |
|
|
