| """Shared helpers for the P2 data / verification / evaluation scripts (stdlib + numpy). |
| |
| Loaded BY PATH from the sibling scripts (`scripts/` is not a package and the dependency |
| firewall rejects an import it cannot place in stdlib / allowed / align), the same way |
| `scripts/lean_ceiling_test.py` loads its case tables. |
| |
| Contents: the three JSONL record schemas of the P2 contract with strict validators, an |
| append-only JSONL writer that survives a SIGKILL mid-line, and the statement / proof |
| normalisations used for decontamination and for "distinct proof" selection. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import hashlib |
| import importlib.util |
| import json |
| import os |
| import re |
| import threading |
| from pathlib import Path |
| from typing import Any, Iterable, Iterator, Optional |
|
|
| REPO = Path(__file__).resolve().parents[2] |
|
|
| |
| |
| |
| |
| GOEDEL_HEADER = "import Mathlib\nimport Aesop\n\nset_option maxHeartbeats 0\n\nopen BigOperators Real Nat Topology Rat\n\n" |
|
|
| PROBLEM_FIELDS = ("problem_id", "source", "split", "formal_statement", "informal_prefix", "statement_sha") |
| GENERATION_FIELDS = ( |
| "problem_id", "sample_idx", "round", "cap_tokens", "prompt_ids", "completion_ids", "prompt_text", |
| "completion_text", "finish_reason", "lean_code", "messages", "model_path", "model_rev", "sampling", |
| "parent", |
| ) |
| KEY_FIELDS = ("problem_id", "sample_idx", "round", "cap_tokens") |
|
|
|
|
| class ContractError(ValueError): |
| """A record violates the P2 shared contract. Always loud: a lenient reader is how a |
| cap-mislabelled or stale-verdict row ends up in a table.""" |
|
|
|
|
| |
|
|
|
|
| def sha256_text(s: str) -> str: |
| return hashlib.sha256(s.encode("utf-8")).hexdigest() |
|
|
|
|
| def sha256_file(path: str | os.PathLike, chunk: int = 4 << 20) -> str: |
| h = hashlib.sha256() |
| with open(path, "rb") as f: |
| while True: |
| b = f.read(chunk) |
| if not b: |
| break |
| h.update(b) |
| return h.hexdigest() |
|
|
|
|
| |
|
|
|
|
| def iter_jsonl(path: str | os.PathLike) -> Iterator[dict]: |
| """Strict JSONL reader: a malformed line raises with its line number (no skipping).""" |
| with open(path, encoding="utf-8") as f: |
| for i, line in enumerate(f, 1): |
| if not line.strip(): |
| continue |
| try: |
| obj = json.loads(line) |
| except json.JSONDecodeError as e: |
| raise ContractError(f"{path}:{i}: malformed JSON ({e.msg})") from e |
| if not isinstance(obj, dict): |
| raise ContractError(f"{path}:{i}: expected a JSON object, got {type(obj).__name__}") |
| yield obj |
|
|
|
|
| def write_jsonl_atomic(path: str | os.PathLike, rows: Iterable[dict]) -> int: |
| """tmp + fsync + rename. A half-written problem file looks complete; never leave one.""" |
| p = Path(path) |
| p.parent.mkdir(parents=True, exist_ok=True) |
| tmp = p.with_name(p.name + f".tmp{os.getpid()}") |
| n = 0 |
| try: |
| with open(tmp, "w", encoding="utf-8") as f: |
| for r in rows: |
| f.write(json.dumps(r, ensure_ascii=False) + "\n") |
| n += 1 |
| f.flush() |
| os.fsync(f.fileno()) |
| os.replace(tmp, p) |
| finally: |
| if tmp.exists(): |
| tmp.unlink() |
| return n |
|
|
|
|
| def repair_torn_tail(path: str | os.PathLike) -> int: |
| """Make an append-only JSONL safe to resume: if the last line was cut off by a kill, drop |
| it. Returns the number of bytes removed. A MIDDLE line that does not parse is not a torn |
| write (appends only ever damage the tail), so that raises when the file is read.""" |
| p = Path(path) |
| if not p.exists() or p.stat().st_size == 0: |
| return 0 |
| with open(p, "rb") as f: |
| data = f.read() |
| if data.endswith(b"\n"): |
| return 0 |
| cut = data.rfind(b"\n") + 1 |
| tail = data[cut:] |
| try: |
| json.loads(tail.decode("utf-8")) |
| with open(p, "ab") as f: |
| f.write(b"\n") |
| return 0 |
| except (json.JSONDecodeError, UnicodeDecodeError): |
| with open(p, "r+b") as f: |
| f.truncate(cut) |
| f.flush() |
| os.fsync(f.fileno()) |
| return len(tail) |
|
|
|
|
| class AppendJsonl: |
| """Thread-safe append-only JSONL writer; every `append` is flushed and fsynced so a |
| SIGKILL loses at most the record being written (which `repair_torn_tail` discards).""" |
|
|
| def __init__(self, path: str | os.PathLike) -> None: |
| self.path = Path(path) |
| self.path.parent.mkdir(parents=True, exist_ok=True) |
| self.torn_bytes_removed = repair_torn_tail(self.path) |
| self._f = open(self.path, "ab") |
| self._lock = threading.Lock() |
|
|
| def append(self, rec: dict) -> None: |
| line = (json.dumps(rec, ensure_ascii=False) + "\n").encode("utf-8") |
| with self._lock: |
| self._f.write(line) |
| self._f.flush() |
| os.fsync(self._f.fileno()) |
|
|
| def close(self) -> None: |
| self._f.close() |
|
|
| def __enter__(self) -> "AppendJsonl": |
| return self |
|
|
| def __exit__(self, *exc) -> None: |
| self.close() |
|
|
|
|
| |
|
|
|
|
| def gen_key(rec: dict) -> tuple: |
| return (rec["problem_id"], int(rec["sample_idx"]), int(rec["round"]), int(rec["cap_tokens"])) |
|
|
|
|
| def validate_problem(rec: dict) -> dict: |
| missing = [k for k in PROBLEM_FIELDS if k not in rec] |
| if missing: |
| raise ContractError(f"problem {rec.get('problem_id')!r} lacks fields {missing}") |
| for k in ("problem_id", "source", "split", "formal_statement", "statement_sha"): |
| if not isinstance(rec[k], str) or not rec[k].strip(): |
| raise ContractError(f"problem field {k!r} must be a non-empty string, got {rec[k]!r}") |
| if rec["informal_prefix"] is not None and not isinstance(rec["informal_prefix"], str): |
| raise ContractError("informal_prefix must be a string or null") |
| return rec |
|
|
|
|
| def effective_max_tokens(rec: dict) -> int: |
| """What the sampler was actually allowed to emit. `cap_tokens` is the NOMINAL cap every |
| number is labelled with; `sampling.max_tokens` is smaller only when the context window, |
| not the cap, ended the sample (`context_clipped`): at cap 40960 with a 40960-token window |
| every sample loses its prompt length (~250-600 tokens).""" |
| mt = (rec.get("sampling") or {}).get("max_tokens") |
| return int(mt) if mt is not None else int(rec["cap_tokens"]) |
|
|
|
|
| def validate_generation(rec: dict) -> dict: |
| """Shape checks on one Generation record. Cross-record checks live in the callers.""" |
| missing = [k for k in GENERATION_FIELDS if k not in rec] |
| pid = rec.get("problem_id") |
| if missing: |
| raise ContractError(f"generation {pid!r} lacks fields {missing}") |
| for k in ("sample_idx", "round", "cap_tokens"): |
| if not isinstance(rec[k], int) or isinstance(rec[k], bool) or rec[k] < 0: |
| raise ContractError(f"generation {pid!r}: {k} must be a non-negative int, got {rec[k]!r}") |
| if rec["cap_tokens"] < 1: |
| raise ContractError(f"generation {pid!r}: cap_tokens must be >= 1") |
| if rec["finish_reason"] not in ("stop", "length"): |
| raise ContractError(f"generation {pid!r}: finish_reason {rec['finish_reason']!r} not in stop|length") |
| for k in ("prompt_ids", "completion_ids"): |
| v = rec[k] |
| if not isinstance(v, list) or not v or not all(isinstance(t, int) and not isinstance(t, bool) for t in v): |
| raise ContractError(f"generation {pid!r}: {k} must be a non-empty list of ints") |
| eff = effective_max_tokens(rec) |
| if eff > rec["cap_tokens"]: |
| raise ContractError( |
| f"generation {pid!r}: sampling.max_tokens={eff} exceeds the stated cap_tokens={rec['cap_tokens']} -- the " |
| "cap label understates what the sampler was allowed, and every number must state its true cap" |
| ) |
| if len(rec["completion_ids"]) > eff: |
| raise ContractError( |
| f"generation {pid!r}: {len(rec['completion_ids'])} completion tokens exceed the sampler's max_tokens={eff} " |
| f"(cap_tokens={rec['cap_tokens']}) -- the cap label is wrong, and every number must state its true cap" |
| ) |
| if rec["finish_reason"] == "length" and len(rec["completion_ids"]) != eff: |
| raise ContractError( |
| f"generation {pid!r}: finish_reason=length but {len(rec['completion_ids'])} tokens != max_tokens {eff} " |
| "(the run was cut by a different limit than the one it is labelled with)" |
| ) |
| if rec["lean_code"] is not None and not isinstance(rec["lean_code"], str): |
| raise ContractError(f"generation {pid!r}: lean_code must be a string or null") |
| parent = rec["parent"] |
| if rec["round"] == 0 and parent is not None: |
| raise ContractError(f"generation {pid!r}: round 0 cannot have a parent") |
| if rec["round"] > 0: |
| if not (isinstance(parent, dict) and {"problem_id", "sample_idx", "round"} <= set(parent)): |
| raise ContractError(f"generation {pid!r}: round {rec['round']} needs a parent {{problem_id,sample_idx,round}}") |
| if parent["problem_id"] != pid or parent["round"] != rec["round"] - 1: |
| raise ContractError(f"generation {pid!r}: parent {parent} is not the previous round of the same problem") |
| return rec |
|
|
|
|
| def load_problems(path: str | os.PathLike) -> dict[str, dict]: |
| out: dict[str, dict] = {} |
| for rec in iter_jsonl(path): |
| validate_problem(rec) |
| if rec["problem_id"] in out: |
| raise ContractError(f"{path}: duplicate problem_id {rec['problem_id']!r}") |
| out[rec["problem_id"]] = rec |
| if not out: |
| raise ContractError(f"{path}: no problems") |
| return out |
|
|
|
|
| def code_sha(lean_code: Optional[str]) -> Optional[str]: |
| return None if lean_code is None else sha256_text(lean_code) |
|
|
|
|
| |
|
|
|
|
| _BY = re.compile(r":=\s*by\b") |
|
|
|
|
| def normalize_formal_statement(raw: str) -> str: |
| """Goedel's statement normalisation (`lean4_code.split(":= by")[0] + ":= by sorry"`), made |
| safe. Goedel's split silently TRUNCATES a statement that contains `:= by` earlier (for |
| example a `(h : x = by norm_num)` binder) and so changes the problem; here that raises. A |
| line break between `:=` and `by` (`:=\\n by sorry`) is accepted and canonicalised, which |
| Goedel's literal split cannot handle at all.""" |
| s = raw.strip("\n") |
| hits = list(_BY.finditer(s)) |
| if len(hits) != 1: |
| raise ContractError(f"statement must contain exactly one ':= by', found {len(hits)}") |
| m = hits[0] |
| tail = s[m.end() :] |
| if tail.strip() not in ("", "sorry"): |
| raise ContractError(f"text after ':= by' is not a bare sorry: {tail.strip()[:60]!r}") |
| return s[: m.start()] + ":= by sorry" |
|
|
|
|
| def _strict(): |
| from align.verify.lean import strict |
|
|
| return strict |
|
|
|
|
| def statement_key(statement: str) -> tuple[str, str, tuple[str, ...]]: |
| """(theorem name, header text, signature tokens) of a formal statement, via the SAME parser |
| the verifier uses. Raises ContractError when the verifier could not use the statement.""" |
| p = _strict().parse_statement(statement) |
| if not p.ok: |
| raise ContractError(f"unsupported statement: {p.problem}") |
| return p.name, p.header_text, tuple(p.sig_toks) |
|
|
|
|
| def statement_sha(statement: str) -> str: |
| """sha256 of the whitespace- and comment-normalised signature (lexer tokens joined by one |
| space). The theorem NAME is deliberately excluded: a renamed copy is the same problem.""" |
| _, _, sig = statement_key(statement) |
| return sha256_text(" ".join(sig)) |
|
|
|
|
| _KEEP_LOWER = frozenset( |
| "fun forall exists if then else let have by sorry at with from in do match return where " |
| "sin cos tan exp log sqrt abs max min id fst snd succ pred gcd lcm choose fact floor ceil " |
| "card sum prod range filter map deriv logb arctan sinh cosh cexp nnabs dist norm " |
| "theorem lemma".split() |
| ) |
|
|
|
|
| def alpha_sha(statement: str) -> str: |
| """Like `statement_sha` but invariant to consistent renaming of local (lowercase, undotted) |
| identifiers. Deliberately conservative: it also equates `sin x = 1` with `cos y = 1`; a |
| false positive here costs one pool statement, a false negative costs a contaminated |
| evaluation.""" |
| _, _, sig = statement_key(statement) |
| mp: dict[str, str] = {} |
| out = [] |
| for t in sig: |
| if re.fullmatch(r"[a-z_][A-Za-z0-9_'₀-₉]*", t) and t not in _KEEP_LOWER: |
| out.append(mp.setdefault(t, f"v{len(mp)}")) |
| else: |
| out.append(t) |
| return sha256_text(" ".join(out)) |
|
|
|
|
| _LEAN_PROMPTS = None |
|
|
|
|
| def _lean_prompts(): |
| """`extras/vllm_rollout/lean_prompts.py` (the byte-faithful port of Goedel's extraction), loaded |
| BY PATH: it is stdlib-only at import and is the single source of truth for what the model's |
| code block means, so the SFT selection must not grow a second, slightly different regex.""" |
| global _LEAN_PROMPTS |
| if _LEAN_PROMPTS is None: |
| spec = importlib.util.spec_from_file_location("p2_lean_prompts", REPO / "extras" / "vllm_rollout" / "lean_prompts.py") |
| mod = importlib.util.module_from_spec(spec) |
| spec.loader.exec_module(mod) |
| _LEAN_PROMPTS = mod |
| return _LEAN_PROMPTS |
|
|
|
|
| def displayed_signature(completion_text: str) -> Optional[tuple[str, ...]]: |
| """Signature tokens of the theorem the MODEL WROTE in the code block Goedel's extractor takes |
| (last ```lean4 fence, first `theorem ... := by`), i.e. BEFORE `problem_check` swaps the original |
| statement in. None when there is no such theorem or the verifier's parser cannot read it. |
| |
| Why this exists: `problem_check` replaces whatever statement the model wrote with the original, |
| so a sample whose displayed theorem is `v = 66` but whose tactics happen to prove the original |
| `v = 65` is ACCEPTED, and its completion then teaches the model to state the wrong theorem. |
| Goedel's evaluation protocol credits it (and so does ours, to stay comparable); a TRAINING row |
| must not.""" |
| lp = _lean_prompts() |
| extracted = lp.DeepSeekCoTHandler().extrac_code(completion_text) |
| if extracted == "None": |
| return None |
| proof = lp.remove_comments(extracted) |
| span = lp.return_theorem_to_replace(proof) |
| if span is None: |
| return None |
| shown = re.sub(r":=\s*by$", ":= by sorry", proof[span[0] : span[1]].strip()) |
| parsed = _strict().parse_statement(shown) |
| return tuple(parsed.sig_toks) if parsed.ok else None |
|
|
|
|
| def displayed_statement_matches(formal_statement: str, completion_text: str) -> bool: |
| """True iff the theorem the model displayed has the same signature tokens as `formal_statement` |
| (theorem NAME ignored: `problem_check` renames it anyway and a renamed copy is the same problem).""" |
| shown = displayed_signature(completion_text) |
| return shown is not None and shown == statement_key(formal_statement)[2] |
|
|
|
|
| def norm_name(name: str) -> str: |
| return re.sub(r"[^a-z0-9]", "", name.lower()) |
|
|
|
|
| def informal_tokens(text: str) -> list[str]: |
| return re.findall(r"[a-z0-9]+", text.lower()) |
|
|
|
|
| def ngrams(tokens: list[str], n: int) -> set[tuple[str, ...]]: |
| return {tuple(tokens[i : i + n]) for i in range(len(tokens) - n + 1)} if len(tokens) >= n else set() |
|
|
|
|
| def proof_hash(lean_code: str) -> str: |
| """Hash of a proof invariant to whitespace and comments (lexer tokens), used to prefer |
| DISTINCT proofs. Falls back to whitespace-collapsed text when the lexer gives up.""" |
| lx = _strict().lex(lean_code) |
| if lx.broken: |
| return sha256_text(" ".join(lean_code.split())) |
| return sha256_text(" ".join(t.text for t in lx.toks)) |
|
|
|
|
| def compile_view(full_code: str) -> str: |
| """Goedel's `compile.py:handle`: drop every line whose stripped text starts with import / |
| set_option / open. Goedel compiles THIS text in a REPL environment that already ran the |
| header, so the line numbers of its error dicts index this text, and `get_error_str` must be |
| given this text (not `full_code`). Identical to `extras/vllm_rollout/lean_prompts.compile_view` |
| (a test pins the two together).""" |
| lines = full_code.split("\n") |
| return "\n".join(l for l in lines if not (l.strip().startswith("import") or l.strip().startswith("set_option") or l.strip().startswith("open"))) |
|
|
