| |
| """Stream candidate Lean statement sets into the P2 Problem JSONL schema (plan v3, Phase 0a item 4). |
| |
| Sources (docs/workbook/data/Datasets.json): FormalMATH-All, NuminaMath-LEAN-Sol, |
| Nemotron-Math-Proofs-v1, RL-Goedel-Pset (levels 3-5 / 4-5), Lean4-FormalConjectures. |
| |
| Properties |
| * Streaming and memory-bounded: JSONL files are read line by line over HTTP (Range-resumable, |
| `align.data.hub.iter_lines`), parquet-only datasets through the datasets-server JSON /rows API |
| (no pyarrow). Nothing is ever held beyond one row plus the per-source dedup set of statement shas. |
| * Every row is pinned to the dataset revision sha (HF API, resolved once at start) and carries |
| `provenance = {source_id, revision, row_index, license, ...}` plus an `ancestry` record |
| (`root` = own problem_id, `parents` = []), the field `align.data.decontam` propagates to mutation |
| children. |
| * Output rows are PROBLEM_FIELDS (+ `provenance`, `ancestry`): `formal_statement` is |
| `GOEDEL_HEADER + ... := by sorry` (the header is REWRITTEN when the source's differs, recorded in |
| `provenance.header_rewritten` / `original_header`; whether the rewrite still elaborates is the |
| elaboration check's job, `elab_hash.py` / `build_problems.py --phase check`). |
| * Nothing is dropped silently: every drop is counted by reason in `ingest_manifest.json` and the first |
| few examples per reason are kept in `<source>.drops.sample.json`. |
| |
| HF_HOME=/apps/development/rlhf/hf_cache HF_HUB_DISABLE_XET=1 \ |
| .venv/bin/python scripts/p2/supply/ingest.py --out-dir .tmp/p2/supply --workers 2 |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import importlib.util |
| import json |
| import logging |
| import os |
| import re |
| import sys |
| import time |
| import urllib.request |
| from collections import Counter |
| from concurrent.futures import ThreadPoolExecutor |
| from pathlib import Path |
| from typing import Any, Callable, Iterator |
|
|
| REPO = Path(__file__).resolve().parents[3] |
| sys.path.insert(0, str(REPO)) |
| os.environ.setdefault("HF_HOME", "/apps/development/rlhf/hf_cache") |
| os.environ.setdefault("HF_HUB_DISABLE_XET", "1") |
|
|
|
|
| def _load_common(): |
| if "p2_common" in sys.modules: |
| return sys.modules["p2_common"] |
| spec = importlib.util.spec_from_file_location("p2_common", REPO / "scripts" / "p2" / "p2_common.py") |
| mod = importlib.util.module_from_spec(spec) |
| sys.modules["p2_common"] = mod |
| spec.loader.exec_module(mod) |
| return mod |
|
|
|
|
| C = _load_common() |
| from align.data import hub |
|
|
| log = logging.getLogger("p2.supply.ingest") |
| MAX_STATEMENT_CHARS = 4000 |
| DEFAULT_OUT = Path(os.environ.get("ALIGN_P2_TMP", "/apps/development/rlhf/align/.tmp/p2")) / "supply" |
| _BY_SORRY_TAIL = re.compile(r":=\s*sorry\s*$") |
| _BY_TAIL = re.compile(r":=\s*by\s*$") |
|
|
|
|
| |
|
|
|
|
| class Source: |
| def __init__(self, key, repo_id, license_, iterate, split="train", problem_source=None, note=""): |
| self.key, self.repo_id, self.license, self.iterate = key, repo_id, license_, iterate |
| self.split, self.problem_source, self.note = split, problem_source or key, note |
|
|
|
|
| def _it_formalmath(src: Source, rev: str, args) -> Iterator[dict]: |
| """FormalMATH_all.json is one 13 MB JSON array (the repo has no JSONL); read once, yield rows.""" |
| req = urllib.request.Request(hub.resolve_url(src.repo_id, "FormalMATH_all.json", rev)) |
| with urllib.request.urlopen(req, timeout=120) as r: |
| data = json.loads(r.read()) |
| for i, row in enumerate(data): |
| yield {"row_index": i, "statement": row.get("autoformalization"), "header": None, |
| "informal": row.get("refined_statement"), "name": row.get("theorem_names"), |
| "extra": {"domain": row.get("domain"), "orig_source": row.get("source")}} |
|
|
|
|
| def _it_numina(src: Source, rev: str, args) -> Iterator[dict]: |
| for i, raw in enumerate(hub.iter_lines(hub.resolve_url(src.repo_id, "train.jsonl", rev))): |
| row = json.loads(raw) if raw.strip() else None |
| if row is None: |
| continue |
| yield {"row_index": i, "statement": row.get("formal_statement"), "header": None, |
| "informal": row.get("problem"), "name": row.get("uuid"), |
| "english": str(row.get("is_english")).lower() != "false", |
| "extra": {"uuid": row.get("uuid"), "orig_source": row.get("source"), |
| "statement_source": row.get("statement_source")}} |
|
|
|
|
| def _it_nemotron(src: Source, rev: str, args) -> Iterator[dict]: |
| """29.5 GB file (whole trajectories): never downloaded in full. `--nemotron-stripes` evenly spaced |
| Range windows of `--nemotron-rows/stripes` lines each, so the sample is not biased to the file head. |
| Row locator = byte offset of the line (stable for a pinned revision).""" |
| url = hub.resolve_url(src.repo_id, "data/lean.jsonl", rev) |
| req = urllib.request.Request(url, method="HEAD") |
| with hub._urlopen(req) as r: |
| size = int(r.headers["Content-Length"]) |
| stripes = max(1, args.nemotron_stripes) |
| per = max(1, args.nemotron_rows // stripes) |
| for s in range(stripes): |
| start = (size // stripes) * s |
| pos = start |
| n = 0 |
| first = start > 0 |
| for raw in hub.iter_lines(url, start_byte=start): |
| here = pos |
| pos += len(raw) + 1 |
| if first: |
| first = False |
| continue |
| if not raw.strip(): |
| continue |
| try: |
| row = json.loads(raw) |
| except json.JSONDecodeError: |
| continue |
| yield {"row_index": here, "statement": row.get("formal_statement"), "header": row.get("lean_header"), |
| "informal": row.get("problem"), "name": row.get("uuid"), |
| "extra": {"uuid": row.get("uuid"), "orig_source": row.get("source"), "byte_offset": here, |
| "stripe": s, "row_license": row.get("license")}} |
| n += 1 |
| if n >= per: |
| break |
|
|
|
|
| def _parquet_or_server_rows(repo_id: str, rev: str, max_rows: int | None) -> Iterator[dict]: |
| """Rows of the single `train` split. With pyarrow (run under .venv-rollout; the import lives in |
| align.data.ingest_parquet, the one whitelisted module) the pinned parquet shard(s) are fetched once to |
| HF_HOME and streamed in row batches. Without it: the datasets-server JSON /rows API, paced at 1 page/s |
| with a long backoff on HTTP 429 (the server rate-limits sequential paging).""" |
| from align.data import ingest_parquet as ip |
|
|
| try: |
| ip._require_pyarrow() |
| have_pa = True |
| except ip.IngestUnavailable: |
| have_pa = False |
| if have_pa: |
| dest = Path(os.environ["HF_HOME"]) / "align_streams" / "supply" / repo_id.replace("/", "__") / rev |
| dest.mkdir(parents=True, exist_ok=True) |
| files = [e for e in hub.list_repo_files(repo_id, rev, "data") if e.get("type") == "file" and e["path"].endswith(".parquet")] |
| emitted = 0 |
| for e in sorted(files, key=lambda e: e["path"]): |
| target = dest / Path(e["path"]).name |
| want = (e.get("lfs") or {}).get("oid") |
| if not (target.exists() and target.stat().st_size == e["size"]): |
| tmp = target.with_name(target.name + f".tmp{os.getpid()}") |
| try: |
| with urllib.request.urlopen(hub.resolve_url(repo_id, e["path"], rev), timeout=300) as r, open(tmp, "wb") as f: |
| while chunk := r.read(4 << 20): |
| f.write(chunk) |
| if want and C.sha256_file(tmp) != want: |
| raise RuntimeError(f"{e['path']}: sha256 differs from the hub's LFS oid") |
| os.replace(tmp, target) |
| finally: |
| if tmp.exists(): |
| tmp.unlink() |
| for row in ip.iter_parquet_rows(target, prefer_ingested=False): |
| if max_rows and emitted >= max_rows: |
| return |
| emitted += 1 |
| yield row |
| return |
| offset, emitted = 0, 0 |
| while not max_rows or emitted < max_rows: |
| for attempt in range(8): |
| try: |
| page = hub.fetch_rows_page(repo_id, "default", "train", offset, 100, use_cache=False) |
| break |
| except hub.HubError as err: |
| if "429" not in str(err) or attempt == 7: |
| raise |
| time.sleep(60) |
| rows = page["rows"] |
| if not rows: |
| return |
| for r in rows: |
| if r.get("truncated_cells"): |
| raise hub.HubError(f"truncated cells at row {r.get('row_idx')}") |
| if max_rows and emitted >= max_rows: |
| return |
| emitted += 1 |
| yield r["row"] |
| offset += len(rows) |
| if offset >= page.get("num_rows_total", offset + 1): |
| return |
| time.sleep(1.0) |
|
|
|
|
| def _it_rows(src: Source, rev: str, args) -> Iterator[dict]: |
| for i, row in enumerate(_parquet_or_server_rows(src.repo_id, rev, args.max_rows)): |
| yield {"row_index": i, "statement": row.get("formal_statement"), "header": None, |
| "informal": row.get("informal_statement"), "name": row.get("problem_id"), |
| "extra": {"orig_problem_id": row.get("problem_id"), "orig_source": row.get("source")}} |
|
|
|
|
| def _it_conjectures(src: Source, rev: str, args) -> Iterator[dict]: |
| """Rows are declarations: statement = everything after the name (binders + type).""" |
| for i, row in enumerate(_parquet_or_server_rows(src.repo_id, rev, args.max_rows)): |
| kind = (row.get("type") or "").strip() |
| imports = row.get("imports") or "[]" |
| yield {"row_index": i, "statement": None, "header": None, "informal": row.get("docstring") or None, |
| "name": row.get("symbolic_name"), "decl_kind": kind, "decl_tail": row.get("statement"), |
| "imports": imports, |
| "extra": {"library": row.get("library"), "filename": row.get("filename"), |
| "source_url": row.get("source_url"), "upstream_commit": row.get("commit")}} |
|
|
|
|
| SOURCES: dict[str, Source] = {s.key: s for s in [ |
| Source("formalmath-all", "SphereLab/FormalMATH-All", "mit", _it_formalmath), |
| Source("numinamath-lean-sol", "iiis-lean/NuminaMath-LEAN-Sol", "apache-2.0", _it_numina), |
| Source("nemotron-math-proofs-v1", "nvidia/Nemotron-Math-Proofs-v1", "cc-by-sa-4.0", _it_nemotron, |
| note="sampled in Range stripes; the full file is 29.5 GB of trajectories"), |
| Source("rl-goedel-pset-l45", "HangHor/RL-Goedel-Pset-Level-4-5", "mit", _it_rows, |
| note="levels 4-5 proper (46,290 rows)"), |
| Source("rl-goedel-pset-l35", "HangHor/RL-Goedel-Pset-Level-3-5", "mit", _it_rows, |
| note="levels 3-5 superset of l45; not in the default set"), |
| Source("lean4-formalconjectures", "phanerozoic/Lean4-FormalConjectures", "apache-2.0", _it_conjectures), |
| ]} |
| DEFAULT_SOURCES = ["formalmath-all", "numinamath-lean-sol", "nemotron-math-proofs-v1", "rl-goedel-pset-l45", |
| "lean4-formalconjectures"] |
|
|
|
|
| def repo_revision(repo_id: str) -> str: |
| with urllib.request.urlopen(f"https://huggingface.co/api/datasets/{repo_id}", timeout=60) as r: |
| return json.loads(r.read())["sha"] |
|
|
|
|
| |
|
|
|
|
| def assemble_statement(rec: dict) -> tuple[str, dict]: |
| """-> (candidate full source text, extra provenance). Raises ContractError('drop:<reason>').""" |
| extra: dict[str, Any] = {} |
| if rec.get("decl_kind") is not None: |
| kind = rec["decl_kind"] |
| if kind not in ("theorem", "lemma"): |
| raise C.ContractError("drop:not_theorem_or_lemma") |
| imps = str(rec.get("imports", "")).lower() |
| if "formalconjectures" in imps or "formal_conjectures" in imps: |
| raise C.ContractError("drop:needs_nonmathlib_imports") |
| if not rec.get("name") or not rec.get("decl_tail"): |
| raise C.ContractError("drop:empty_statement") |
| raw = f"{kind} {rec['name']} {rec['decl_tail'].strip()} := by sorry" |
| return C.GOEDEL_HEADER + raw, {"constructed_from_declaration": True} |
| raw = rec.get("statement") |
| if not isinstance(raw, str) or not raw.strip(): |
| raise C.ContractError("drop:empty_statement") |
| raw = raw.strip("\n") |
| if rec.get("header"): |
| raw = rec["header"].rstrip() + "\n\n" + raw.lstrip() |
| if _BY_SORRY_TAIL.search(raw): |
| raw = _BY_SORRY_TAIL.sub(":= by sorry", raw) |
| extra["sorry_term_rewritten"] = True |
| elif _BY_TAIL.search(raw): |
| raw = raw.rstrip() + " sorry" |
| extra["sorry_appended"] = True |
| elif ":= by" not in raw and re.search(r"\bsorry\b", raw) is None: |
| raw = raw.rstrip() + " := by sorry" |
| extra["body_appended"] = True |
| return raw, extra |
|
|
|
|
| def _open_names(header: str) -> set[str]: |
| out: set[str] = set() |
| for ln in header.splitlines(): |
| ln = ln.strip() |
| if ln.startswith("open "): |
| out.update(ln[5:].split()) |
| return out |
|
|
|
|
| def to_problem(rec: dict, source: Source, revision: str, pid: str) -> tuple[dict, dict]: |
| """Full conversion of one raw row -> (Problem record, side info). Raises ContractError('drop:<reason>').""" |
| raw, extra = assemble_statement(rec) |
| if len(raw) > MAX_STATEMENT_CHARS: |
| raise C.ContractError("drop:too_long") |
| try: |
| stmt = C.normalize_formal_statement(raw) |
| except C.ContractError: |
| raise C.ContractError("drop:not_single_by_sorry") |
| try: |
| name, header, _sig = C.statement_key(stmt) |
| except C.ContractError as e: |
| why = re.sub(r"`[^`]*`", "`X`", str(e).replace("unsupported statement: ", "")) |
| raise C.ContractError(f"drop:unsupported_shape: {why}"[:120]) |
| prov_header = {} |
| if header.rstrip() != C.GOEDEL_HEADER.rstrip(): |
| parsed = C._strict().parse_statement(stmt) |
| body = stmt[len(parsed.header_text):].lstrip("\n") |
| stmt = C.GOEDEL_HEADER + body |
| prov_header = {"header_rewritten": True, "original_header": header.strip()[:600]} |
| extra_open = sorted(_open_names(header) - _open_names(C.GOEDEL_HEADER)) |
| if extra_open: |
| prov_header["header_extra_opens"] = extra_open |
| try: |
| C.statement_key(stmt) |
| except C.ContractError: |
| raise C.ContractError("drop:unsupported_after_header_rewrite") |
| prob = { |
| "problem_id": pid, "source": source.problem_source, "split": "train", "formal_statement": stmt, |
| "informal_prefix": rec.get("informal") if isinstance(rec.get("informal"), str) and rec["informal"].strip() else None, |
| "statement_sha": C.statement_sha(stmt), |
| "provenance": {"source_id": source.repo_id, "revision": revision, "row_index": rec["row_index"], |
| "license": source.license, "theorem_name": name, **extra, **prov_header, **rec.get("extra", {})}, |
| "ancestry": {"root": pid, "parents": []}, |
| } |
| C.validate_problem(prob) |
| return prob, {"header_rewritten": bool(prov_header)} |
|
|
|
|
| def ingest_source(src: Source, revision: str, out_dir: Path, args) -> dict: |
| counts: Counter = Counter() |
| drop_examples: dict[str, list] = {} |
| seen: set[str] = set() |
| t0 = time.time() |
| out_path = out_dir / f"{src.key}.jsonl" |
| tmp = out_path.with_name(out_path.name + f".tmp{os.getpid()}") |
| try: |
| with open(tmp, "w", encoding="utf-8") as f: |
| for rec in src.iterate(src, revision, args): |
| counts["rows_seen"] += 1 |
| if rec.get("english") is False: |
| counts["drop:not_english"] += 1 |
| continue |
| pid = f"{src.key}-{rec['row_index']}" |
| try: |
| prob, side = to_problem(rec, src, revision, pid) |
| except C.ContractError as e: |
| why = str(e) |
| if not why.startswith("drop:"): |
| raise |
| counts[why] += 1 |
| ex = drop_examples.setdefault(why, []) |
| if len(ex) < 3: |
| ex.append({"row_index": rec["row_index"], "name": rec.get("name")}) |
| continue |
| if prob["statement_sha"] in seen: |
| counts["drop:duplicate_statement"] += 1 |
| continue |
| seen.add(prob["statement_sha"]) |
| f.write(json.dumps(prob, ensure_ascii=False) + "\n") |
| counts["kept"] += 1 |
| counts["header_rewritten"] += int(side["header_rewritten"]) |
| counts["header_extra_opens"] += int("header_extra_opens" in prob["provenance"]) |
| if counts["rows_seen"] % 20000 == 0: |
| log.info("%s: %d rows, %d kept", src.key, counts["rows_seen"], counts["kept"]) |
| f.flush() |
| os.fsync(f.fileno()) |
| os.replace(tmp, out_path) |
| finally: |
| if tmp.exists(): |
| tmp.unlink() |
| (out_dir / f"{src.key}.drops.sample.json").write_text(json.dumps(drop_examples, indent=1, ensure_ascii=False)) |
| return {"source_id": src.repo_id, "revision": revision, "license": src.license, "note": src.note, |
| "counts": dict(counts), "seconds": round(time.time() - t0, 1), "output": str(out_path), |
| "output_sha256": C.sha256_file(out_path)} |
|
|
|
|
| def main(argv=None) -> int: |
| ap = argparse.ArgumentParser(description=__doc__.split("\n")[0]) |
| ap.add_argument("--out-dir", default=str(DEFAULT_OUT)) |
| ap.add_argument("--sources", nargs="+", default=DEFAULT_SOURCES, choices=sorted(SOURCES)) |
| ap.add_argument("--workers", type=int, default=2, help="parallel sources; hard-capped at 2") |
| ap.add_argument("--max-rows", type=int, default=0, help="cap rows per datasets-server source (0 = all)") |
| ap.add_argument("--nemotron-rows", type=int, default=60000) |
| ap.add_argument("--nemotron-stripes", type=int, default=8) |
| a = ap.parse_args(argv) |
| logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(message)s") |
| out_dir = Path(a.out_dir) |
| out_dir.mkdir(parents=True, exist_ok=True) |
| revs = {k: repo_revision(SOURCES[k].repo_id) for k in a.sources} |
| log.info("revisions: %s", revs) |
| workers = max(1, min(2, a.workers)) |
| results: dict[str, Any] = {} |
| with ThreadPoolExecutor(max_workers=workers) as ex: |
| futs = {k: ex.submit(ingest_source, SOURCES[k], revs[k], out_dir, a) for k in a.sources} |
| for k, fu in futs.items(): |
| try: |
| results[k] = fu.result() |
| log.info("%s: %s", k, results[k]["counts"]) |
| except Exception as e: |
| log.error("%s FAILED: %r", k, e) |
| results[k] = {"source_id": SOURCES[k].repo_id, "revision": revs[k], "error": repr(e)} |
| man = out_dir / "ingest_manifest.json" |
| prev = json.loads(man.read_text()) if man.exists() else {} |
| prev.update(results) |
| man.write_text(json.dumps(prev, indent=1, ensure_ascii=False)) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
