lean-bench/experiments/P0-pilot/run_pilot.py

"""Run the pilot detectors over the pilot items. Resumable; one JSON line per item.

    python experiments/P0-pilot/run_pilot.py --items experiments/P0-pilot/data/items.jsonl \
        --out experiments/P0-pilot/data/results.jsonl --dry-run          # no API calls: shows the plan
    python experiments/P0-pilot/run_pilot.py --items ... --out ... --max-calls 400

Detectors (all see the passage and/or the Lean code, never the human labels):
  lint    static patterns on the projected simplex          (no LLM)
  role    source-side simplex (LLM sees only the passage) aligned to the Lean simplex
  judge   one LLM call that sees passage + Lean and says faithful / not
  compile the stored Lean verdict (``verified``), read from the items file

Hard stops: --max-calls bounds API calls; a rerun skips items already in --out.
"""
from __future__ import annotations

import argparse
import fcntl
import hashlib
import inspect
import json
import os
import sys
from pathlib import Path

from leanbench import srcsimplex
from leanbench.align import align
from leanbench.config import Settings
from leanbench.lint import lint
from leanbench.llm import OpenAICompatibleClient, Usage
from leanbench.project import project

JUDGE_PROMPT_VERSION = "judge-v1"
JUDGE_SYSTEM = (
    "You check a Lean 4 (Mathlib) formalization against a textbook passage. Decide whether the "
    "Lean statement says what the passage says: same objects, same hypotheses, same conclusion, "
    "same quantifiers. Ignore whether it compiles. Reply with JSON only: "
    '{"faithful": true|false, "confidence": 0..1, "reason": "one sentence"}'
)


def judge(client, passage: str, lean: str) -> dict:
    reply = client.complete_json([
        {"role": "system", "content": JUDGE_SYSTEM},
        {"role": "user", "content": f"Passage:\n{passage}\n\nLean:\n{lean}"},
    ])
    if not isinstance(reply, dict) or type(reply.get("faithful")) is not bool:
        raise ValueError(f"unusable judge reply: {str(reply)[:200]}")
    return {"faithful": reply["faithful"], "confidence": reply.get("confidence"),
            "reason": str(reply.get("reason", ""))[:300]}


def run_item(client, item: dict) -> dict:
    if not item.get("source_text", "").strip():
        raise ValueError("Missing source_text; use a private source-inclusive export")
    code = item.get("lean_code") or ""
    row: dict = {"id": item["id"], "compiled": bool(item.get("verified")), "errors": {}}
    proj = project(code)
    row["projection"] = proj.as_dict()
    row["lint"] = lint(proj)
    try:
        src = srcsimplex.extract_source_simplex(client, item.get("source_text", ""),
                                                item.get("label", ""))
        al = align(src, proj)
        row["source"] = [r.__dict__ for r in src.roles]
        row["role"] = al.as_dict()
    except CallBudgetExceeded:
        raise
    except Exception as error:  # a failed extraction is data, not a crash
        row["errors"]["role"] = str(error)[:300]
    try:
        row["judge"] = judge(client, item.get("source_text", ""), code)
    except CallBudgetExceeded:
        raise
    except Exception as error:
        row["errors"]["judge"] = str(error)[:300]
    return row


class CallBudgetExceeded(RuntimeError):
    """Stop before issuing a request that would exceed the cumulative budget."""


def digest(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def atomic_json(path: Path, value: dict) -> None:
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(json.dumps(value, indent=2, allow_nan=False) + "\n", encoding="utf-8")
    os.replace(temporary, path)


class AuditedClient(OpenAICompatibleClient):
    """Persist reservations before HTTP requests, including transport and JSON retries."""

    def complete(self, *args, **kwargs):
        try:
            return super().complete(*args, **kwargs)
        finally:
            self.checkpoint()


def execute(args, items, settings) -> int:
    rows = []
    if args.out.exists():
        rows = [json.loads(x) for x in args.out.read_text(encoding="utf-8").splitlines() if x.strip()]
    ids = {item["id"] for item in items}
    done = {row["id"] for row in rows}
    if len(done) != len(rows) or not done <= ids:
        raise ValueError("Results have duplicate or unknown item IDs")
    hashes = {
        "source": hashlib.sha256((srcsimplex.SYSTEM_PROMPT + inspect.getsource(srcsimplex.passage_prompt)).encode()).hexdigest(),
        "judge": hashlib.sha256(JUDGE_SYSTEM.encode()).hexdigest(),
    }
    identity = {"model": settings.llm.model, "base_url": settings.llm.base_url,
                "temperature": settings.llm.temperature, "items_sha256": digest(args.items),
                "prompt_hashes": hashes,
                "prompts": [srcsimplex.PROMPT_VERSION, JUDGE_PROMPT_VERSION]}
    labels = getattr(args, "labels", None)
    if labels is not None:
        identity["labels_sha256"] = digest(labels)
    meta_path = args.out.with_suffix(".meta.json")
    meta = {**identity, "request_attempts": 0, "usage": Usage().as_dict(), "status": "ready"}
    if meta_path.exists():
        meta = json.loads(meta_path.read_text(encoding="utf-8"))
        if any(meta.get(k) != v for k, v in identity.items()):
            raise ValueError("Run identity changed; use a new output path instead of mixing runs")
    elif rows:
        raise ValueError("Cannot resume results without their audit metadata")
    todo = [item for item in items if item["id"] not in done]
    print(f"items={len(items)} done={len(done)} todo={len(todo)} model={settings.llm.model} "
          f"nominal_calls={len(todo) * 2} request_attempts={meta['request_attempts']} cap={args.max_calls}")
    print("prompt_hashes:", json.dumps(hashes))
    if args.dry_run:
        return 0
    if settings.llm.model != "qwen-max" or settings.llm.temperature != 0:
        raise ValueError("P0 is preregistered for qwen-max at temperature 0")
    client = AuditedClient.from_settings(settings.llm)
    client.usage = Usage(**meta["usage"])

    def checkpoint():
        meta["usage"] = client.usage.as_dict()
        meta["completed_items"] = len(done)
        atomic_json(meta_path, meta)

    def reserve():
        if meta["request_attempts"] >= args.max_calls:
            raise CallBudgetExceeded("Cumulative request cap reached; no further HTTP request was sent")
        meta["request_attempts"] += 1
        checkpoint()

    client.checkpoint = checkpoint
    client.request_guard = reserve
    meta["max_calls"] = args.max_calls
    meta["status"] = "running"
    checkpoint()
    try:
        with args.out.open("a", encoding="utf-8") as stream:
            for item in todo:
                row = run_item(client, item)
                stream.write(json.dumps(row, ensure_ascii=False) + "\n")
                stream.flush()
                os.fsync(stream.fileno())
                done.add(item["id"])
                checkpoint()
                print(f"completed={len(done)}/{len(items)} requests={meta['request_attempts']}", flush=True)
        meta["status"] = "complete"
    except CallBudgetExceeded as error:
        meta["status"] = "budget_exhausted"
        print(str(error), file=sys.stderr)
    except BaseException:
        meta["status"] = "interrupted"
        raise
    finally:
        if args.out.exists():
            meta["results_sha256"] = digest(args.out)
        checkpoint()
    print(json.dumps(meta))
    return 0 if meta["status"] == "complete" else 2


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--items", required=True, type=Path)
    ap.add_argument("--out", required=True, type=Path)
    ap.add_argument("--max-calls", type=int, default=400)
    ap.add_argument("--dry-run", action="store_true")
    ap.add_argument("--labels", type=Path, help="Frozen labels; their hash is bound to the run")
    args = ap.parse_args()
    if not 1 <= args.max_calls <= 400:
        ap.error("P0 requires a request cap between 1 and 400")
    items = [json.loads(x) for x in args.items.read_text(encoding="utf-8").splitlines() if x.strip()]
    if not items or len({item["id"] for item in items}) != len(items):
        ap.error("Input must contain unique item IDs")
    if any(not item.get("source_text", "").strip() for item in items):
        ap.error("Every item needs nonempty source_text; use a private source-inclusive export")
    if not args.dry_run and args.labels is None:
        ap.error("Freeze all blind labels first, then pass --labels")
    if args.labels is not None:
        annotation = json.loads(args.labels.read_text(encoding="utf-8"))
        key_path = args.labels.with_name("key.json")
        key = json.loads(key_path.read_text(encoding="utf-8"))["key"]
        verdicts = {"faithful", "minor_drift", "wrong", "not_formalizable"}
        if (set(key.values()) != {item["id"] for item in items}
                or set(annotation["labels"]) != set(key)
                or any(row.get("verdict") not in verdicts for row in annotation["labels"].values())
                or annotation.get("items_sha256") != digest(args.items)
                or annotation.get("key_sha256") != digest(key_path)):
            ap.error("Frozen labels, key and source-inclusive items must be complete and match")
    args.out.parent.mkdir(parents=True, exist_ok=True)
    with args.out.with_suffix(args.out.suffix + ".lock").open("w") as lock:
        fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
        return execute(args, items, Settings.from_env())


if __name__ == "__main__":
    raise SystemExit(main())