Download code/scripts/p2/build_problems.py from viai957/align: direct link, hf CLI and curl.
- Browser
- Download file 24.9 kB
-
https://huggingface.co/viai957/align/resolve/main/code/scripts/p2/build_problems.py
- Command line
-
hf download hf://viai957/align/code/scripts/p2/build_problems.py
-
curl -L -o build_problems.py https://huggingface.co/viai957/align/resolve/main/code/scripts/p2/build_problems.py
24.9 kB
| #!/usr/bin/env python | |
| """P2 problem sets: evaluation sets (miniF2F, MOBench) and a decontaminated training pool. | |
| Everything is written as Problem JSONL (docs/P2.md, "Contract"): | |
| {"problem_id", "source", "split", "formal_statement", "informal_prefix", "statement_sha"} | |
| Three phases, run in this order (each is idempotent and resumable by re-running): | |
| eval Goedel's dataset/minif2f.jsonl (244) and dataset/MOBench.jsonl (360) | |
| -> minif2f.jsonl, mobench.jsonl (the verifier-supported subset) and | |
| mobench_unsupported.json (listed, never silently dropped). | |
| pool-sample Stream the PINNED revision of Goedel-LM/Goedel-Pset-v1 (MIT), apply cheap | |
| filters + decontamination against both eval sets, rank by a seeded hash and | |
| keep the top `--candidates` -> pool_candidates.jsonl. | |
| Needs pyarrow, which is NOT an align dependency: run this phase under | |
| .venv-rollout (`align.data.ingest_parquet` raises IngestUnavailable otherwise). | |
| check Run the strict verifier's own reference elaboration (`theorem _align_ref <sig> | |
| := by sorry` under the statement's header, exactly what StrictVerifier does | |
| before it looks at a candidate) on the candidates in rank order and keep the | |
| first N that elaborate -> train_pool.jsonl + problems_manifest.json. | |
| Needs the Lean toolchain; takes the heavy-job lock when --workers > 2. | |
| Why each filter exists (every one is counted in problems_manifest.json): | |
| * unparseable / unsupported shape: the strict verifier scores 0 for ANY candidate against a | |
| reference it cannot parse (`unsupported_statement`); such a problem is unwinnable. | |
| * `:= by` count != 1: Goedel's prompt builder does `split(":= by")[0] + ":= by sorry"`, which | |
| silently TRUNCATES a statement with an earlier `:= by`. | |
| * contamination, independent layers against miniF2F+MOBench: (1) signature sha, (2) | |
| alpha-renaming-invariant signature sha, (3) normalised theorem-name / problem-id match, and | |
| on the ranked candidates (4) overlap of the informal text: any 13-word window, or at least | |
| three distinct 8-word windows. miniF2F and MOBench are | |
| drawn from public olympiad problems that NuminaMath (the source of Goedel-Pset) also | |
| contains, so equal informal problems with DIFFERENT formalisations are the realistic leak. | |
| * elaboration: a statement that does not elaborate under the pinned Mathlib is a wasted | |
| sample budget and a guaranteed 0. | |
| .venv-rollout/bin/python scripts/p2/build_problems.py --phase pool-sample --out-dir D | |
| .venv/bin/python scripts/p2/build_problems.py --phase eval --out-dir D | |
| .venv/bin/python scripts/p2/build_problems.py --phase check --out-dir D | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import heapq | |
| import importlib.util | |
| import json | |
| import logging | |
| import os | |
| import re | |
| import sys | |
| import time | |
| import urllib.request | |
| from collections import Counter | |
| from pathlib import Path | |
| 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.build_problems") | |
| GOEDEL_REPO = Path("/apps/development/rlhf/reference-repos/Goedel-Prover-V2") | |
| POOL_REPO = "Goedel-LM/Goedel-Pset-v1" | |
| #: Pinned revision (HF commit sha, last modified 2025-04-18). Licence: the dataset card has NO | |
| #: `license:` metadata; the repo's LICENSE file says "licensed under the MIT License". The | |
| #: informal problems are "from Numina" (README; AI-MO/NuminaMath-CoT and -1.5 are both Apache-2.0 | |
| #: on the hub, the README does not say which one). | |
| POOL_REVISION = "79fc4fbe87b4c694a941bc63ad92dc38d39d41db" | |
| POOL_LICENSE = "MIT (LICENSE file at the pinned revision; no license field in the dataset card metadata); informal problems from Numina (Apache-2.0)" | |
| NGRAM = 13 | |
| #: Second informal layer: LaTeX formatting differences break long windows (an IMO 1965 problem | |
| #: shared 7 distinct 8-grams with a pool problem but no 13-gram), so >= SHORT_HITS distinct | |
| #: 8-gram matches also count. A single 8-gram is common phrasing ("what is the smallest | |
| #: possible value of the") and is not enough. | |
| SHORT_NGRAM = 8 | |
| SHORT_HITS = 3 | |
| MAX_STATEMENT_CHARS = 4000 | |
| DEFAULT_SEED = 20261012 | |
| # =========================================================================== eval sets | |
| def build_eval_problems(minif2f_path, mobench_path): | |
| """-> (minif2f, mobench_supported, mobench_excluded[(id, reason)], counts).""" | |
| minif2f, mob, excluded = [], [], [] | |
| for r in C.iter_jsonl(minif2f_path): | |
| minif2f.append(_eval_problem(r["name"], "minif2f", r["split"], r["formal_statement"], r.get("informal_prefix"))) | |
| seen = set() | |
| for r in C.iter_jsonl(mobench_path): | |
| pid = r["problem_id"] | |
| if pid in seen: | |
| raise C.ContractError(f"duplicate MOBench problem_id {pid!r}") | |
| seen.add(pid) | |
| # `formal_statement` in this file is the bare theorem; `lean4_code` is header + the | |
| # definitions it needs + docstring + theorem, i.e. what the verifier has to be given. | |
| try: | |
| mob.append(_eval_problem(pid, "mobench", r["split"], r["lean4_code"].lstrip(), r.get("informal_prefix"))) | |
| except C.ContractError as e: | |
| excluded.append((pid, str(e))) | |
| ids = [p["problem_id"] for p in minif2f] | |
| if len(set(ids)) != len(ids): | |
| raise C.ContractError("duplicate miniF2F problem_id") | |
| reasons = Counter(re.sub(r"found \d+", "found N", re.sub(r"\d+$", "", why)) for _, why in excluded) | |
| counts = { | |
| "minif2f": len(minif2f), | |
| "mobench_total": len(mob) + len(excluded), | |
| "mobench_supported": len(mob), | |
| "mobench_excluded": len(excluded), | |
| "mobench_excluded_reasons": dict(reasons), | |
| } | |
| return minif2f, mob, excluded, counts | |
| def _eval_problem(pid, source, split, raw_statement, informal): | |
| stmt = C.normalize_formal_statement(raw_statement) | |
| return { | |
| "problem_id": pid, | |
| "source": source, | |
| "split": split, | |
| "formal_statement": stmt, | |
| "informal_prefix": informal, | |
| "statement_sha": C.statement_sha(stmt), # raises ContractError if the verifier can't use it | |
| } | |
| # =========================================================================== decontamination | |
| class Decontaminator: | |
| """Layered contamination test against a set of evaluation problems.""" | |
| def __init__(self, eval_problems, ngram: int = NGRAM, short_ngram: int = SHORT_NGRAM, short_hits: int = SHORT_HITS): | |
| self.ngram, self.short_ngram, self.short_hits = ngram, short_ngram, short_hits | |
| self.short_grams: set[tuple[str, ...]] = set() | |
| self.sha: set[str] = set() | |
| self.alpha: set[str] = set() | |
| self.names: set[str] = set() | |
| self.grams: set[tuple[str, ...]] = set() | |
| for p in eval_problems: | |
| st = p["formal_statement"] | |
| self.sha.add(p["statement_sha"]) | |
| self.alpha.add(C.alpha_sha(st)) | |
| self.names.add(C.norm_name(p["problem_id"])) | |
| self.names.add(C.norm_name(C.statement_key(st)[0])) | |
| if p.get("informal_prefix"): | |
| toks = C.informal_tokens(p["informal_prefix"]) | |
| self.grams |= C.ngrams(toks, ngram) | |
| self.short_grams |= C.ngrams(toks, short_ngram) | |
| # Substring layer for names long enough that a substring hit is not noise. | |
| long_names = sorted((n for n in self.names if len(n) >= 8), key=len, reverse=True) | |
| self._name_re = re.compile("|".join(re.escape(n) for n in long_names)) if long_names else None | |
| def check_statement(self, statement: str, problem_id: str) -> str | None: | |
| """Cheap layers (1-3). Returns the first reason or None.""" | |
| name = C.statement_key(statement)[0] | |
| if C.statement_sha(statement) in self.sha: | |
| return "statement_sha" | |
| if C.alpha_sha(statement) in self.alpha: | |
| return "alpha_sha" | |
| for n in (C.norm_name(name), C.norm_name(problem_id)): | |
| if n in self.names or (self._name_re is not None and self._name_re.search(n)): | |
| return "name" | |
| return None | |
| def check_informal(self, informal: str | None) -> str | None: | |
| """'informal_13gram' / 'informal_8gram3' / None.""" | |
| if not informal or not self.grams: | |
| return None | |
| toks = C.informal_tokens(informal) | |
| if not self.grams.isdisjoint(C.ngrams(toks, self.ngram)): | |
| return "informal_13gram" | |
| if len(self.short_grams & C.ngrams(toks, self.short_ngram)) >= self.short_hits: | |
| return "informal_8gram3" | |
| return None | |
| # =========================================================================== pool sampling | |
| def rank_hash(seed: int, problem_id: str) -> str: | |
| return C.sha256_text(f"{seed}|{problem_id}") | |
| def sample_pool(rows, decon: Decontaminator, n_candidates: int, seed: int): | |
| """Stream `rows` ({problem_id, formal_statement, informal_statement}) -> (ranked candidates, | |
| counts). Memory is O(n_candidates): a bounded max-heap on the seeded rank hash.""" | |
| counts: Counter = Counter() | |
| heap: list = [] # (-rank_int, tiebreak, rec): root is the WORST (largest hash) kept | |
| for r in rows: | |
| counts["rows"] += 1 | |
| pid, raw = r["problem_id"], r["formal_statement"] | |
| if len(raw) > MAX_STATEMENT_CHARS: | |
| counts["drop_too_long"] += 1 | |
| continue | |
| try: | |
| stmt = C.normalize_formal_statement(raw) | |
| except C.ContractError: | |
| counts["drop_not_single_by_sorry"] += 1 | |
| continue | |
| try: | |
| _name, header, _sig = C.statement_key(stmt) | |
| except C.ContractError: | |
| counts["drop_unsupported_shape"] += 1 | |
| continue | |
| if header.rstrip() != C.GOEDEL_HEADER.rstrip(): | |
| counts["info_nonstandard_header"] += 1 # kept: problem_check re-inserts the header | |
| why = decon.check_statement(stmt, pid) | |
| if why is not None: | |
| counts[f"contaminated_{why}"] += 1 | |
| continue | |
| counts["eligible"] += 1 | |
| h = rank_hash(seed, pid) | |
| rec = { | |
| "problem_id": pid, | |
| "source": "goedel-pset-v1", | |
| "split": "train", | |
| "formal_statement": stmt, | |
| "informal_prefix": r.get("informal_statement"), | |
| "statement_sha": C.statement_sha(stmt), | |
| "rank_hash": h, | |
| } | |
| item = (-int(h, 16), pid, rec) | |
| if len(heap) < n_candidates: | |
| heapq.heappush(heap, item) | |
| elif item > heap[0]: | |
| heapq.heapreplace(heap, item) | |
| ranked = [t[2] for t in sorted(heap, key=lambda t: -t[0])] # ascending rank hash | |
| counts["ranked_kept"] = len(ranked) | |
| return ranked, counts | |
| def finalize_candidates(ranked, decon: Decontaminator): | |
| """Rank-order post-filters that need the candidate list: informal-text overlap and | |
| duplicate statements (the same statement under two problem_ids keeps the better rank).""" | |
| counts: Counter = Counter() | |
| out, seen = [], set() | |
| for rec in ranked: | |
| why = decon.check_informal(rec["informal_prefix"]) | |
| if why is not None: | |
| counts[f"contaminated_{why}_among_ranked"] += 1 | |
| continue | |
| if rec["statement_sha"] in seen: | |
| counts["duplicate_statement_among_ranked"] += 1 | |
| continue | |
| seen.add(rec["statement_sha"]) | |
| out.append(rec) | |
| counts["candidates"] = len(out) | |
| return out, counts | |
| def download_pool_parquets(dest: Path) -> list[dict]: | |
| """Fetch the pinned parquet shards (about 830 MB, once) and verify each against the LFS oid | |
| the hub reports for that exact revision. Returns [{path, sha256, size}].""" | |
| from align.data import hub | |
| dest.mkdir(parents=True, exist_ok=True) | |
| entries = [ | |
| e for e in hub.list_repo_files(POOL_REPO, POOL_REVISION, "data") if e.get("type") == "file" and e["path"].endswith(".parquet") | |
| ] | |
| if not entries: | |
| raise RuntimeError(f"no parquet files under {POOL_REPO}@{POOL_REVISION}/data") | |
| out = [] | |
| for e in sorted(entries, key=lambda e: e["path"]): | |
| want = (e.get("lfs") or {}).get("oid") | |
| target = dest / Path(e["path"]).name | |
| if not (target.exists() and target.stat().st_size == e["size"] and (want is None or C.sha256_file(target) == want)): | |
| tmp = target.with_name(target.name + f".tmp{os.getpid()}") | |
| try: | |
| url = hub.resolve_url(POOL_REPO, e["path"], POOL_REVISION) | |
| log.info("downloading %s", url) | |
| with urllib.request.urlopen(url, timeout=120) as resp, open(tmp, "wb") as f: | |
| while True: | |
| b = resp.read(4 << 20) | |
| if not b: | |
| break | |
| f.write(b) | |
| if want is not None and C.sha256_file(tmp) != want: | |
| raise RuntimeError(f"{e['path']}: sha256 differs from the hub's LFS oid {want}") | |
| os.replace(tmp, target) | |
| finally: | |
| if tmp.exists(): | |
| tmp.unlink() | |
| out.append({"path": str(target), "sha256": C.sha256_file(target), "size": target.stat().st_size, "lfs_oid": want}) | |
| return out | |
| # =========================================================================== elaboration check | |
| def check_elaborates(worker, statement: str, timeout_s: float) -> tuple[bool, str]: | |
| """Does `statement` give a valid reference for StrictVerifier? Mirrors `_run`: header env | |
| (imports, audit prelude, header opens), then `theorem _align_ref <sig> := by sorry`.""" | |
| from align.verify.lean.repl_pool import ReplError, ReplProtocolError | |
| from align.verify.lean.strict import AUDIT_PRELUDE, REF_NAME, parse_statement | |
| ref = parse_statement(statement) | |
| if not ref.ok: | |
| return False, f"unsupported_statement: {ref.problem}" | |
| try: | |
| env = worker.header_env(ref.header_text, AUDIT_PRELUDE) | |
| except ReplProtocolError as e: # the header itself (imports / opens) does not elaborate | |
| return False, "header_error: " + " ".join(str(e).split())[:200] | |
| except ReplError as e: | |
| return False, e.code | |
| try: | |
| resp = worker.command(f"theorem {REF_NAME} {ref.sig_text} := by\n sorry", env, timeout_s) | |
| except ReplError as e: | |
| return False, e.code | |
| msgs = resp.get("messages", []) if isinstance(resp, dict) else [] | |
| errs = [m for m in msgs if isinstance(m, dict) and m.get("severity") == "error"] | |
| if errs or "env" not in resp: | |
| return False, "elab_error: " + " ".join(str(errs[0].get("data", resp))[:200].split()) if errs else "no_env" | |
| return True, "ok" | |
| SELF_TEST_GOOD = C.GOEDEL_HEADER + "theorem selftest_good (x : ℕ) (h : x = 2) : x + x = 4 := by sorry" | |
| SELF_TEST_BAD = C.GOEDEL_HEADER + 'theorem selftest_bad (x : ℕ) : x + "a" = 4 := by sorry' | |
| def run_check(candidates, n_target: int, pool, timeout_s: float, chunk: int = 64): | |
| """Elaborate candidates in rank order; stop once `n_target` passed. Returns | |
| (selected, log_rows, counts). Selection depends only on rank order and verdicts, never on | |
| chunking: the failures counted are those ranked BEFORE the n_target-th pass.""" | |
| ok_pass, ok_fail = pool.map(lambda w, s: check_elaborates(w, s, timeout_s), [SELF_TEST_GOOD, SELF_TEST_BAD]) | |
| for what, res, want in (("known-good", ok_pass, True), ("known-bad", ok_fail, False)): | |
| if isinstance(res, BaseException): | |
| raise res | |
| if res[0] is not want: | |
| raise RuntimeError(f"elaboration instrument failed its ceiling test: {what} statement gave {res}") | |
| log_rows, selected = [], [] | |
| i = 0 | |
| while i < len(candidates) and len(selected) < n_target: | |
| batch = candidates[i : i + chunk] | |
| t0 = time.monotonic() | |
| res = pool.map(lambda w, rec: check_elaborates(w, rec["formal_statement"], timeout_s), batch) | |
| for rec, r in zip(batch, res): | |
| if isinstance(r, BaseException): | |
| raise r # a programming error or a dead pool must stop the build, not become a "failure" | |
| log_rows.append({"problem_id": rec["problem_id"], "ok": r[0], "why": r[1]}) | |
| i += len(batch) | |
| log.info("checked %d/%d candidates, %d ok (%.0fs/chunk)", i, len(candidates), sum(1 for x in log_rows if x["ok"]), time.monotonic() - t0) | |
| selected = [] | |
| for rec, row in zip(candidates, log_rows): | |
| if row["ok"]: | |
| selected.append(rec) | |
| if len(selected) == n_target: | |
| break | |
| if len(selected) < n_target: | |
| raise RuntimeError( | |
| f"only {len(selected)} of {n_target} requested statements elaborate among {len(candidates)} candidates; " | |
| "re-run pool-sample with a larger --candidates" | |
| ) | |
| cutoff = candidates.index(selected[-1]) + 1 | |
| considered = log_rows[:cutoff] | |
| counts = { | |
| "checked_in_rank_order": cutoff, | |
| "elaborated": sum(1 for x in considered if x["ok"]), | |
| "failed_to_elaborate": sum(1 for x in considered if not x["ok"]), | |
| "failure_reasons": dict(Counter(re.sub(r":.*", "", x["why"]) for x in considered if not x["ok"])), | |
| } | |
| return selected, log_rows[:cutoff], counts | |
| # =========================================================================== CLI | |
| def _paths(out_dir: Path): | |
| return { | |
| "minif2f": out_dir / "minif2f.jsonl", | |
| "mobench": out_dir / "mobench.jsonl", | |
| "mob_excl": out_dir / "mobench_unsupported.json", | |
| "cand": out_dir / "pool_candidates.jsonl", | |
| "cand_stats": out_dir / "pool_sample_stats.json", | |
| "pool": out_dir / "train_pool.jsonl", | |
| "elab_log": out_dir / "elaboration_log.jsonl", | |
| "eval_elab": out_dir / "eval_elaboration.json", | |
| "manifest": out_dir / "problems_manifest.json", | |
| } | |
| def _read_json(p: Path): | |
| return json.loads(p.read_text()) if p.exists() else {} | |
| def phase_eval(a, P): | |
| ds = GOEDEL_REPO / "dataset" | |
| minif2f, mob, excl, counts = build_eval_problems(ds / "minif2f.jsonl", ds / "MOBench.jsonl") | |
| C.write_jsonl_atomic(P["minif2f"], minif2f) | |
| C.write_jsonl_atomic(P["mobench"], mob) | |
| P["mob_excl"].write_text(json.dumps([{"problem_id": i, "reason": r} for i, r in excl], indent=1, ensure_ascii=False)) | |
| log.info("eval: %s", counts) | |
| return {"counts": counts, "files": {k: C.sha256_file(P[k]) for k in ("minif2f", "mobench", "mob_excl")}, | |
| "inputs": {f: C.sha256_file(ds / f) for f in ("minif2f.jsonl", "MOBench.jsonl")}} | |
| def phase_pool_sample(a, P): | |
| from align.data.ingest_parquet import iter_parquet_rows | |
| for k in ("minif2f", "mobench"): | |
| if not P[k].exists(): | |
| raise SystemExit(f"{P[k]} missing: run --phase eval first (decontamination needs both eval sets)") | |
| evalp = list(C.iter_jsonl(P["minif2f"])) + list(C.iter_jsonl(P["mobench"])) | |
| decon = Decontaminator(evalp) | |
| files = download_pool_parquets(Path(os.environ.get("HF_HOME", "/apps/development/rlhf/hf_cache")) / "align_streams" / "p2" / POOL_REPO.replace("/", "--") / POOL_REVISION) | |
| def rows(): | |
| for f in files: | |
| log.info("streaming %s", f["path"]) | |
| yield from iter_parquet_rows(f["path"], columns=["problem_id", "formal_statement", "informal_statement"], prefer_ingested=False) | |
| ranked, counts = sample_pool(rows(), decon, a.candidates, a.seed) | |
| cands, c2 = finalize_candidates(ranked, decon) | |
| counts.update(c2) | |
| C.write_jsonl_atomic(P["cand"], cands) | |
| stats = {"dataset": POOL_REPO, "revision": POOL_REVISION, "license": POOL_LICENSE, "seed": a.seed, "n_target_candidates": a.candidates, | |
| "ngram": NGRAM, "short_ngram": SHORT_NGRAM, "short_hits": SHORT_HITS, "parquet_files": files, "counts": dict(counts), "candidates_sha256": C.sha256_file(P["cand"]), | |
| "eval_decontam_sets": {"statement_sha": len(decon.sha), "alpha_sha": len(decon.alpha), "names": len(decon.names), "informal_13grams": len(decon.grams), "informal_8grams": len(decon.short_grams)}} | |
| P["cand_stats"].write_text(json.dumps(stats, indent=1)) | |
| log.info("pool-sample: %s", dict(counts)) | |
| return stats | |
| def phase_check(a, P): | |
| from align.utils import guard | |
| from align.verify.lean import ReplPool, find_toolchain | |
| from align.verify.lean.repl_pool import ReplLimits | |
| if not 1 <= a.workers <= 8: | |
| raise SystemExit("--workers must be in [1, 8]") | |
| for k in ("cand", "minif2f", "mobench"): | |
| if not P[k].exists(): | |
| raise SystemExit(f"{P[k]} missing: run the earlier phases first") | |
| if a.workers > 2: | |
| log.info("waiting for the heavy-job lock (Lean pool with %d workers)", a.workers) | |
| guard.acquire_heavy_job_lock("p2-build-problems", blocking=True) | |
| guard.preflight_host_check(strict=True) | |
| tc = find_toolchain() | |
| pool = ReplPool.from_toolchain(tc, n_workers=a.workers, limits=ReplLimits(), preflight=True, | |
| scratch_dir=str(Path(os.environ.get("TMPDIR", str(REPO / ".tmp"))) / "p2_build_pool")) | |
| cands = list(C.iter_jsonl(P["cand"])) | |
| selected, log_rows, counts = run_check(cands, a.n, pool, a.elab_timeout) | |
| # Eval sets: elaboration is REPORTED, not filtered (dropping eval problems would change the benchmark). | |
| ev = {} | |
| for k in ("minif2f", "mobench"): | |
| probs = list(C.iter_jsonl(P[k])) | |
| res = pool.map(lambda w, rec: check_elaborates(w, rec["formal_statement"], a.elab_timeout), probs) | |
| bad = [] | |
| for rec, r in zip(probs, res): | |
| if isinstance(r, BaseException): | |
| raise r | |
| if not r[0]: | |
| bad.append({"problem_id": rec["problem_id"], "why": r[1]}) | |
| ev[k] = {"n": len(probs), "fail_to_elaborate": len(bad), "failures": bad} | |
| P["eval_elab"].write_text(json.dumps(ev, indent=1, ensure_ascii=False)) | |
| C.write_jsonl_atomic(P["elab_log"], log_rows) | |
| final = [{k: r[k] for k in C.PROBLEM_FIELDS} for r in selected] | |
| for r in final: | |
| C.validate_problem(r) | |
| C.write_jsonl_atomic(P["pool"], final) | |
| stats = _read_json(P["cand_stats"]) | |
| manifest = { | |
| "phase_eval": _read_json(P["manifest"]).get("phase_eval"), | |
| "pool": { | |
| "dataset": POOL_REPO, "revision": POOL_REVISION, "license": POOL_LICENSE, "seed": stats.get("seed"), | |
| "requested_n": a.n, "selected": len(final), | |
| "source_rows": stats["counts"].get("rows"), "sampling_counts": stats["counts"], "elaboration": counts, | |
| "elaboration_timeout_s": a.elab_timeout, "decontamination_sets": stats["eval_decontam_sets"], | |
| "train_pool_sha256": C.sha256_file(P["pool"]), "candidates_sha256": stats["candidates_sha256"], | |
| }, | |
| "eval": {"minif2f_sha256": C.sha256_file(P["minif2f"]), "mobench_sha256": C.sha256_file(P["mobench"]), "elaboration": {k: {"n": v["n"], "fail_to_elaborate": v["fail_to_elaborate"]} for k, v in ev.items()}}, | |
| "verifier_pins": tc.pins, | |
| } | |
| P["manifest"].write_text(json.dumps(manifest, indent=1, ensure_ascii=False)) | |
| log.info("pool: selected %d; manifest %s", len(final), P["manifest"]) | |
| return manifest | |
| def main(argv=None) -> int: | |
| ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) | |
| ap.add_argument("--phase", required=True, choices=["eval", "pool-sample", "check"]) | |
| ap.add_argument("--out-dir", required=True) | |
| ap.add_argument("--n", type=int, default=2000, help="training pool size") | |
| ap.add_argument("--seed", type=int, default=DEFAULT_SEED) | |
| ap.add_argument("--candidates", type=int, default=5000, help="ranked candidates kept for the elaboration check") | |
| ap.add_argument("--workers", type=int, default=8) | |
| ap.add_argument("--elab-timeout", type=float, default=60.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 (never / or /tmp)") | |
| out.mkdir(parents=True, exist_ok=True) | |
| P = _paths(out) | |
| if a.phase == "eval": | |
| res = phase_eval(a, P) | |
| m = _read_json(P["manifest"]) | |
| m["phase_eval"] = res | |
| P["manifest"].write_text(json.dumps(m, indent=1, ensure_ascii=False)) | |
| elif a.phase == "pool-sample": | |
| phase_pool_sample(a, P) | |
| else: | |
| phase_check(a, P) | |
| return 0 | |
| if __name__ == "__main__": | |
| raise SystemExit(main()) | |