File size: 7,734 Bytes
63f41a6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 | #!/usr/bin/env python
"""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 # noqa: E402
from align.data.decontam import STAGES, ContaminationLedger, Decontaminator, EvalItem # noqa: E402
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 # noqa: E402
#: .tmp/ is gitignored and lives in the main checkout; worktrees share it (override with ALIGN_P2_TMP).
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")))
# all 360 MOBench problems (the verifier supports only 135 of them, contamination does not care)
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
# single pass: the ledger needs ancestry, and mutation children may precede their parent in a file, so
# rows are held only as (id, line) -- the statement text is dropped once checked
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())
|