| |
| """Decontaminate ingested statement sets and the existing train pool against the evaluation sets |
| (miniF2F test+valid, MOBench, ProofNet# valid+test, optionally FormalMATH). Uses align.data.decontam. |
| |
| .venv/bin/python scripts/p2/supply/decontaminate.py --supply-dir .tmp/p2/supply |
| |
| Writes, per input set `<name>`: `<name>.clean.jsonl` (kept rows), `<name>.flagged.jsonl` (id, reason, hits), |
| and `decontam_report.json` (hit counts per stage per source, independent and first-hit, plus each stage's |
| unique contribution). Nothing is deleted from the inputs. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import importlib.util |
| import json |
| import sys |
| from collections import Counter |
| from pathlib import Path |
|
|
| REPO = Path(__file__).resolve().parents[3] |
| sys.path.insert(0, str(REPO)) |
|
|
| from align.data import hub |
| from align.data.decontam import STAGES, ContaminationLedger, Decontaminator, EvalItem |
|
|
| spec = importlib.util.spec_from_file_location("p2_common", REPO / "scripts" / "p2" / "p2_common.py") |
| C = sys.modules.get("p2_common") or importlib.util.module_from_spec(spec) |
| if "p2_common" not in sys.modules: |
| sys.modules["p2_common"] = C |
| spec.loader.exec_module(C) |
|
|
| import os |
|
|
| |
| P2_TMP = Path(os.environ.get("ALIGN_P2_TMP", "/apps/development/rlhf/align/.tmp/p2")) |
| PROBLEMS = P2_TMP / "problems" |
| GOEDEL_MOBENCH = Path("/apps/development/rlhf/reference-repos/Goedel-Prover-V2/dataset/MOBench.jsonl") |
| PROOFNET = "PAug/ProofNetSharp" |
|
|
|
|
| def fetch_proofnet(out_dir: Path) -> list[EvalItem]: |
| """ProofNet# (Lean 4 port, valid 185 + test 186) through the datasets-server JSON API; cached as JSONL.""" |
| items = [] |
| with urllib_open(f"https://huggingface.co/api/datasets/{PROOFNET}") as r: |
| rev = json.loads(r.read())["sha"] |
| for split in ("valid", "test"): |
| cache = out_dir / f"proofnet_{split}.jsonl" |
| if not cache.exists(): |
| rows = list(hub.iter_rows(PROOFNET, "default", split)) |
| with open(cache, "w", encoding="utf-8") as f: |
| for i, row in enumerate(rows): |
| f.write(json.dumps({"row_index": i, "revision": rev, **row}, ensure_ascii=False) + "\n") |
| for row in C.iter_jsonl(cache): |
| stmt = (row.get("lean4_src_header") or "").rstrip() + "\n\n" + (row.get("lean4_formalization") or "") |
| items.append(EvalItem(f"proofnet-{split}:{row.get('id')}", f"proofnet_{split}", stmt, row.get("nl_statement"))) |
| return items |
|
|
|
|
| def urllib_open(url): |
| import urllib.request |
| return urllib.request.urlopen(url, timeout=60) |
|
|
|
|
| def eval_items(supply_dir: Path, formalmath_as_eval: bool) -> tuple[list[EvalItem], dict]: |
| items: list[EvalItem] = [] |
| for name, path in (("minif2f_test", PROBLEMS / "minif2f.jsonl"), ("minif2f_valid", PROBLEMS / "minif2f_valid.jsonl")): |
| for r in C.iter_jsonl(path): |
| items.append(EvalItem(f"{name}:{r['problem_id']}", name, r["formal_statement"], r.get("informal_prefix"))) |
| |
| for r in C.iter_jsonl(GOEDEL_MOBENCH): |
| items.append(EvalItem(f"mobench:{r['problem_id']}", "mobench", r.get("lean4_code") or r["formal_statement"], r.get("informal_prefix"))) |
| items += fetch_proofnet(supply_dir / "evalsets") |
| if formalmath_as_eval: |
| for r in C.iter_jsonl(supply_dir / "formalmath-all.jsonl"): |
| items.append(EvalItem(f"formalmath:{r['problem_id']}", "formalmath", r["formal_statement"], r.get("informal_prefix"))) |
| return items, dict(Counter(i.eval_set for i in items)) |
|
|
|
|
| def run_set(name: str, path: Path, dec: Decontaminator, out_dir: Path, elab: dict[str, str]) -> dict: |
| ledger = ContaminationLedger() |
| rows_meta: list[tuple[str, dict]] = [] |
| n = 0 |
| per_stage: Counter = Counter() |
| first: Counter = Counter() |
| by_eval: Counter = Counter() |
| unique: Counter = Counter() |
| flagged_f = open(out_dir / f"{name}.flagged.jsonl", "w", encoding="utf-8") |
| clean_f = open(out_dir / f"{name}.clean.jsonl", "w", encoding="utf-8") |
| inherited = 0 |
| |
| |
| pending: list[tuple[str, str]] = [] |
| for r in C.iter_jsonl(path): |
| n += 1 |
| hits = dec.check(r["formal_statement"], r.get("informal_prefix"), elab.get(r["problem_id"])) |
| ledger.add(r["problem_id"], hits, r.get("ancestry")) |
| pending.append((r["problem_id"], json.dumps(r, ensure_ascii=False))) |
| if hits: |
| stages = {h.stage for h in hits} |
| for s in stages: |
| per_stage[s] += 1 |
| first[min(stages, key=STAGES.index)] += 1 |
| if len(stages) == 1: |
| unique[next(iter(stages))] += 1 |
| for es in {dec.items[h.eval_id].eval_set for h in hits}: |
| by_eval[es] += 1 |
| flagged = 0 |
| for pid, line in pending: |
| bad, why = ledger.resolve(pid) |
| if bad: |
| flagged += 1 |
| inherited += why.startswith("inherited") |
| h = ledger.direct.get(pid, []) |
| flagged_f.write(json.dumps({"problem_id": pid, "reason": why, "hits": [ |
| {"stage": x.stage, "eval_id": x.eval_id, "eval_set": dec.items[x.eval_id].eval_set, "score": round(x.score, 3)} |
| for x in h[:8]]}, ensure_ascii=False) + "\n") |
| else: |
| clean_f.write(line + "\n") |
| flagged_f.close(); clean_f.close() |
| return {"rows": n, "flagged": flagged, "clean": n - flagged, "inherited": inherited, |
| "flagged_by_stage_independent": {s: per_stage[s] for s in STAGES if per_stage[s]}, |
| "first_hit_stage": {s: first[s] for s in STAGES if first[s]}, |
| "only_this_stage_fired": {s: unique[s] for s in STAGES if unique[s]}, |
| "flagged_rows_by_eval_set": dict(by_eval)} |
|
|
|
|
| def main(argv=None) -> int: |
| ap = argparse.ArgumentParser(description=__doc__.split("\n")[0]) |
| ap.add_argument("--supply-dir", default=str(P2_TMP / "supply")) |
| ap.add_argument("--pool", default=str(PROBLEMS / "train_pool.jsonl")) |
| ap.add_argument("--formalmath-as-eval", action="store_true") |
| ap.add_argument("--elab-hashes", help="JSONL {problem_id, elab_sha} from elab_hash.py (eval + candidate sides)") |
| a = ap.parse_args(argv) |
| sd = Path(a.supply_dir) |
| (sd / "evalsets").mkdir(parents=True, exist_ok=True) |
| out = sd / "decontam" |
| out.mkdir(exist_ok=True) |
| items, eval_counts = eval_items(sd, a.formalmath_as_eval) |
| elab: dict[str, str] = {} |
| if a.elab_hashes: |
| for r in C.iter_jsonl(a.elab_hashes): |
| elab[r["problem_id"]] = r["elab_sha"] |
| for it in items: |
| if it.eval_id in elab: |
| it.elab_sha = elab[it.eval_id] |
| dec = Decontaminator(items) |
| report = {"eval_sets": eval_counts, "eval_index_skipped": dict(dec.skipped), "sets": {}} |
| inputs = {"train_pool": Path(a.pool)} |
| for p in sorted(sd.glob("*.jsonl")): |
| if p.name.endswith((".clean.jsonl", ".flagged.jsonl")): |
| continue |
| if a.formalmath_as_eval and p.name == "formalmath-all.jsonl": |
| continue |
| inputs[p.stem] = p |
| for name, path in inputs.items(): |
| report["sets"][name] = run_set(name, path, dec, out, elab) |
| print(name, json.dumps(report["sets"][name]), flush=True) |
| (out / "decontam_report.json").write_text(json.dumps(report, indent=1)) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
