Blame for code/scripts/p2/build_sft_data.py

#!/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())