Zayden-Jung/lean-bench
2
1
Fork 0
Code Issues 2 Pull Requests Actions Packages Projects Releases Wiki Activity

Rebuild lean-bench: pages in, verified indexed Lean out
ci / test (push) Failing after 12s
Details
ci / lean-project (push) Successful in 5s
Details
docker / image (push) Failing after 5s
Details

Browse Source 
Replaces the previous tree entirely. The pipeline is four separable pieces, each
testable on its own:

1. leanbench/llm.py      provider-neutral OpenAI-compatible connector, stdlib only,
                        with retries, JSON coercion and a scripted client so the
                        rest of the code is testable without a network.
2. leanbench/detect.py   per-page detection of formalizable mathematics. Case 1 is
                        mathematics asserted as fact (definition/theorem/lemma/...,
                        deliberately not sub-classified); case 2 is a problem left
                        open to the reader. Hits come back as spans -- start page
                        and index, end page and index -- so several per page is
                        normal. A hit cut off by the page end makes the detector
                        re-read that page together with the next, repeatedly, so an
                        item spread over three pages is found whole.
                        Offsets come from verbatim quotes resolved by
                        leanbench/anchors.py (exact -> whitespace-insensitive ->
                        fuzzy), because models quote well and count badly.
3. leanbench/formalize.py + leanbench/verify.py
                        each span becomes one Lean 4 declaration, judged by the real
                        compiler: full elaboration for case 1, statement legality
                        only for case 2, which has no proof to check. Lean's own
                        diagnostics feed back for a repair attempt. "verified" (Lean
                        accepted it) and "proved" (accepted with no sorry) are kept
                        strictly apart.
4. leanbench/db.py       indexed SQLite plus JSONL export. Every item is indexed by
                        the types of its inputs, or by its own type when it is not a
                        proof, and records the source book (sha256) and the page it
                        starts on -- so `drop --source H --pages A-B` followed by a
                        re-ingest redoes exactly those pages and nothing else.

Also: lean/ Lake project pinned to Mathlib, docker/ image carrying Lean 4 and a
prebuilt Mathlib (the Mathlib revision and the toolchain are resolved against each
other at build time rather than both hardcoded), and .gitea/workflows for tests,
image build and push to the Gitea registry, and a manual "submit a book" run.

58 tests; they need neither network nor Lean.
This commit is contained in:
Zayden-Jung
2026-09-19 07:05:52 +00:00
commit d3555d9a52
37 changed files with 4251 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
leanbench/__init__.py
+16
View File
@@ -0,0 +1,16 @@
"""lean-bench: turn mathematical source documents into verified Lean 4 items.
Pipeline
--------
1. ``leanbench.llm`` provider-neutral LLM connector (the base need).
2. ``leanbench.detect`` find formalizable math on each page, with spans that may
run across several pages.
3. ``leanbench.formalize`` turn each detected span into Lean 4 code;
``leanbench.verify`` check it with the real Lean compiler.
4. ``leanbench.db`` store everything in an indexed database, keyed by the
types an item consumes and produces, plus book hash and page.
"""
__version__ = "0.2.0"
__all__ = ["__version__"]
leanbench/anchors.py
+154
View File
@@ -0,0 +1,154 @@
"""Turning quoted snippets back into exact character offsets.
Language models are bad at counting characters and good at quoting. So the detector
asks for short verbatim quotes at the start and end of each item, and this module
converts those quotes into hard indices into the page text:
1. exact substring match;
2. whitespace-insensitive match (PDF text is full of stray line breaks);
3. fuzzy match with difflib, above a similarity floor.
Anything that fails all three is reported as unresolved rather than guessed at.
"""
from __future__ import annotations
import difflib
import re
from dataclasses import dataclass
_WS = re.compile(r"\s+")
MIN_RATIO = 0.72
@dataclass(frozen=True)
class Match:
start: int
end: int
method: str
score: float
def _collapse(text: str) -> tuple[str, list[int]]:
"""Whitespace-collapsed text plus a map from collapsed index -> original index."""
out: list[str] = []
mapping: list[int] = []
previous_space = True # strip leading whitespace
for index, char in enumerate(text):
if char.isspace():
if previous_space:
continue
out.append(" ")
mapping.append(index)
previous_space = True
else:
out.append(char)
mapping.append(index)
previous_space = False
return "".join(out), mapping
def find(text: str, needle: str, start_at: int = 0) -> Match | None:
"""Locate ``needle`` in ``text`` at or after ``start_at``."""
needle = (needle or "").strip()
if not needle:
return None
exact = text.find(needle, start_at)
if exact != -1:
return Match(exact, exact + len(needle), "exact", 1.0)
haystack, mapping = _collapse(text)
if not haystack:
return None
probe, _ = _collapse(needle)
if not probe:
return None
# where in the collapsed text does start_at land?
floor = 0
for collapsed_index, original_index in enumerate(mapping):
if original_index >= start_at:
floor = collapsed_index
break
else:
floor = len(mapping)
loose = haystack.find(probe, floor)
if loose == -1 and floor: # allow wrapping back if nothing ahead
loose = haystack.find(probe)
if loose != -1:
start = mapping[loose]
last = min(loose + len(probe) - 1, len(mapping) - 1)
return Match(start, mapping[last] + 1, "whitespace", 1.0)
best = _fuzzy(haystack, probe, floor)
if best is None:
return None
lo, hi, score = best
last = min(hi - 1, len(mapping) - 1)
return Match(mapping[lo], mapping[last] + 1, "fuzzy", score)
def _fuzzy(haystack: str, probe: str, floor: int) -> tuple[int, int, float] | None:
"""Slide a window of the probe's length and keep the best-scoring position."""
size = len(probe)
if size < 8 or len(haystack) < size:
return None
step = max(1, size // 4)
matcher = difflib.SequenceMatcher(autojunk=False, b=probe)
best: tuple[int, int, float] | None = None
for start in range(floor, len(haystack) - size + 1, step):
window = haystack[start : start + size]
matcher.set_seq1(window)
if matcher.real_quick_ratio() < MIN_RATIO or matcher.quick_ratio() < MIN_RATIO:
continue
score = matcher.ratio()
if score >= MIN_RATIO and (best is None or score > best[2]):
best = (start, start + size, score)
if score > 0.97:
break
if best is None:
return None
# refine around the coarse hit
lo, hi, score = best
for start in range(max(floor, lo - step), min(len(haystack) - size, lo + step) + 1):
window = haystack[start : start + size]
matcher.set_seq1(window)
candidate = matcher.ratio()
if candidate > score:
lo, hi, score = start, start + size, candidate
return lo, hi, score
@dataclass(frozen=True)
class Span:
start: int
end: int
method: str
score: float
@property
def resolved(self) -> bool:
return self.end > self.start
def resolve_span(
text: str,
start_quote: str,
end_quote: str,
*,
search_from: int = 0,
fallback_length: int = 1200,
) -> Span | None:
"""Resolve a (start quote, end quote) pair into one span of ``text``."""
head = find(text, start_quote, search_from)
if head is None:
return None
tail = find(text, end_quote, head.start)
if tail is None or tail.end <= head.start:
# The tail quote is the fragile one (models paraphrase endings). Rather than
# dropping a real item, fall back to a bounded chunk from the start quote.
end = min(len(text), head.start + fallback_length)
return Span(head.start, end, head.method + "+truncated", head.score * 0.5)
method = head.method if head.method == tail.method else f"{head.method}/{tail.method}"
return Span(head.start, tail.end, method, min(head.score, tail.score))
leanbench/cli.py
+330
View File
@@ -0,0 +1,330 @@
"""Command line: ``lean-bench <command>``.
lean-bench doctor is everything wired up?
lean-bench ingest book.pdf --pages 30-45
lean-bench detect book.pdf --pages 30-31 --out detections.jsonl
lean-bench search --type Int --role input --kind problem
lean-bench export datasets/items.jsonl
lean-bench drop --source <hash> --pages 30-45 redo part of a book
"""
from __future__ import annotations
import argparse
import json
import sys
from .config import Settings
from .db import Database
from .detect import detect_document, write_jsonl
from .document import Document
from .llm import LLMError, build_client
from .pipeline import build_checker, ingest
def _pages(value: str | None) -> tuple[int | None, int | None]:
if not value:
return None, None
text = value.strip()
if "-" in text:
first, _, last = text.partition("-")
return (int(first) if first.strip() else None,
int(last) if last.strip() else None)
page = int(text)
return page, page
def _say(message: str) -> None:
print(message, flush=True)
def _add_common(parser: argparse.ArgumentParser) -> None:
parser.add_argument("--db", help="SQLite path (default: $LEANBENCH_DATABASE)")
parser.add_argument("--env", default=".env", help="dotenv file to load (default: .env)")
# -- commands --------------------------------------------------------------
def cmd_doctor(args: argparse.Namespace) -> int:
settings = Settings.from_env(args.env)
print(f"database {settings.database}")
print(f"llm model {settings.llm.model}")
print(f"llm endpoint {settings.llm.base_url}")
print(f"llm key env {settings.llm.api_key_env}")
try:
import pymupdf # noqa: F401
print("pdf backend pymupdf")
except ImportError:
try:
import pypdf # noqa: F401
print("pdf backend pypdf (install pymupdf for better text)")
except ImportError:
print("pdf backend MISSING - pip install pymupdf")
checker = build_checker(settings)
available = getattr(checker, "available", lambda: False)()
print(f"lean {'ready: ' + ' '.join(settings.lean.command) if available else 'not available (verification will be skipped)'}")
if settings.lean.project:
print(f"lean project {settings.lean.project}")
if args.skip_llm:
return 0
try:
client = build_client(settings.llm)
reply = client.complete(
[{"role": "user", "content": "Reply with the single word: ready"}],
max_tokens=16,
)
print(f"llm ready ({reply.strip()[:40]})")
except LLMError as error:
print(f"llm FAILED: {error}")
return 1
return 0
def cmd_ingest(args: argparse.Namespace) -> int:
settings = Settings.from_env(args.env)
if args.db:
settings.database = args.db
if args.window:
settings.max_window_pages = args.window
first, last = _pages(args.pages)
database = Database(settings.database)
try:
result = ingest(
args.source,
settings=settings,
database=database,
checker=build_checker(settings, enabled=not args.no_verify),
first_page=first,
last_page=last,
title=args.title or "",
replace=not args.keep_existing,
store_results=not args.no_store,
progress=None if args.quiet else _say,
)
summary = result.summary()
print(json.dumps(summary, indent=2))
if args.detections:
write_jsonl(args.detections, result.report.detections)
print(f"detections -> {args.detections}")
if args.export:
count = database.export_jsonl(args.export, source_hash=result.document.sha256)
print(f"exported {count} item(s) -> {args.export}")
finally:
database.close()
return 0
def cmd_detect(args: argparse.Namespace) -> int:
settings = Settings.from_env(args.env)
first, last = _pages(args.pages)
document = Document.load(args.source)
report = detect_document(
build_client(settings.llm),
document,
first_page=first,
last_page=last,
max_window_pages=args.window or settings.max_window_pages,
progress=None if args.quiet else _say,
)
print(report.summary())
if args.out:
write_jsonl(args.out, report.detections)
print(f"-> {args.out}")
else:
for detection in report.detections:
print(json.dumps(detection.as_dict(), ensure_ascii=False))
return 0
def cmd_search(args: argparse.Namespace) -> int:
settings = Settings.from_env(args.env)
database = Database(args.db or settings.database)
try:
rows = database.search(
types=args.type or (),
roles=args.role or (),
kind=args.kind,
source_hash=args.source,
page=args.page,
verified=args.verified,
proved=args.proved,
text=args.text,
match_all_types=args.all_types,
limit=args.limit,
)
if args.json:
print(json.dumps(rows, indent=2, ensure_ascii=False))
return 0
if not rows:
print("no matches")
return 0
for row in rows:
types = ", ".join(f"{t['name']}({t['role']})" for t in row["types"]) or "-"
flag = "proved" if row["proved"] else "verified" if row["verified"] else "unverified"
print(f"{row['id']} {row['kind']:<9} p{row['start_page']:<4} {flag:<10} "
f"{row['label'] or row['lean_name']}")
print(f" types: {types}")
if args.code:
for line in (row["lean_code"] or "").splitlines():
print(" " + line)
finally:
database.close()
return 0
def cmd_sources(args: argparse.Namespace) -> int:
settings = Settings.from_env(args.env)
database = Database(args.db or settings.database)
try:
rows = database.sources()
if args.json:
print(json.dumps(rows, indent=2, ensure_ascii=False))
return 0
for row in rows:
print(f"{row['hash'][:16]} {row['items']:>5} items p{row['pages']:<5} "
f"{row['title']} ({row['extractor']}, {row['ingested_at']})")
finally:
database.close()
return 0
def cmd_stats(args: argparse.Namespace) -> int:
settings = Settings.from_env(args.env)
database = Database(args.db or settings.database)
try:
stats = database.stats()
stats["top_types"] = database.type_counts(limit=args.top)
print(json.dumps(stats, indent=2))
finally:
database.close()
return 0
def cmd_export(args: argparse.Namespace) -> int:
settings = Settings.from_env(args.env)
database = Database(args.db or settings.database)
try:
count = database.export_jsonl(
args.out, source_hash=args.source, include_source_text=args.with_text
)
print(f"exported {count} item(s) -> {args.out}")
finally:
database.close()
return 0
def cmd_drop(args: argparse.Namespace) -> int:
settings = Settings.from_env(args.env)
database = Database(args.db or settings.database)
try:
first, last = _pages(args.pages)
if first or last:
removed = database.delete_pages(args.source, first or 1, last or 10**9)
else:
removed = database.delete_source(args.source)
print(f"removed {removed} item(s)")
finally:
database.close()
return 0
# -- wiring ----------------------------------------------------------------
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(prog="lean-bench", description=__doc__,
formatter_class=argparse.RawDescriptionHelpFormatter)
subparsers = parser.add_subparsers(dest="command", required=True)
doctor = subparsers.add_parser("doctor", help="check LLM, Lean and PDF wiring")
_add_common(doctor)
doctor.add_argument("--skip-llm", action="store_true", help="do not call the model")
doctor.set_defaults(func=cmd_doctor)
ingest_parser = subparsers.add_parser("ingest", help="run the whole pipeline on a source")
_add_common(ingest_parser)
ingest_parser.add_argument("source", help="PDF (or .txt with form-feed page breaks)")
ingest_parser.add_argument("--pages", help="page or range, e.g. 30-45")
ingest_parser.add_argument("--title", help="human title for this source")
ingest_parser.add_argument("--window", type=int, help="max pages per detection window")
ingest_parser.add_argument("--no-verify", action="store_true", help="skip Lean")
ingest_parser.add_argument("--no-store", action="store_true", help="do not write the database")
ingest_parser.add_argument("--keep-existing", action="store_true",
help="do not remove this source's previous items first")
ingest_parser.add_argument("--detections", help="also write raw detections to this JSONL")
ingest_parser.add_argument("--export", help="export this source's items to JSONL afterwards")
ingest_parser.add_argument("--quiet", action="store_true")
ingest_parser.set_defaults(func=cmd_ingest)
detect_parser = subparsers.add_parser("detect", help="detection only, no Lean, no database")
_add_common(detect_parser)
detect_parser.add_argument("source")
detect_parser.add_argument("--pages")
detect_parser.add_argument("--window", type=int)
detect_parser.add_argument("--out", help="write detections to this JSONL")
detect_parser.add_argument("--quiet", action="store_true")
detect_parser.set_defaults(func=cmd_detect)
search = subparsers.add_parser("search", help="find items by the types they are about")
_add_common(search)
search.add_argument("--type", action="append", help="type name; repeatable")
search.add_argument("--role", action="append", choices=["input", "goal", "own"])
search.add_argument("--kind", choices=["statement", "problem"])
search.add_argument("--source", help="source hash")
search.add_argument("--page", type=int)
search.add_argument("--text", help="full-text query")
search.add_argument("--all-types", action="store_true", help="require every --type")
search.add_argument("--verified", action="store_true", default=None)
search.add_argument("--proved", action="store_true", default=None)
search.add_argument("--code", action="store_true", help="print the Lean too")
search.add_argument("--json", action="store_true")
search.add_argument("--limit", type=int, default=25)
search.set_defaults(func=cmd_search)
sources = subparsers.add_parser("sources", help="list ingested books")
_add_common(sources)
sources.add_argument("--json", action="store_true")
sources.set_defaults(func=cmd_sources)
stats = subparsers.add_parser("stats", help="counts, and the most common types")
_add_common(stats)
stats.add_argument("--top", type=int, default=20)
stats.set_defaults(func=cmd_stats)
export = subparsers.add_parser("export", help="dump items as JSONL")
_add_common(export)
export.add_argument("out")
export.add_argument("--source", help="only this source hash")
export.add_argument("--with-text", action="store_true", help="include source passages")
export.set_defaults(func=cmd_export)
drop = subparsers.add_parser("drop", help="remove a source, or some of its pages")
_add_common(drop)
drop.add_argument("--source", required=True, help="source hash")
drop.add_argument("--pages", help="page or range; omit for the whole source")
drop.set_defaults(func=cmd_drop)
return parser
def main(argv: list[str] | None = None) -> int:
args = build_parser().parse_args(argv)
try:
return args.func(args)
except FileNotFoundError as error:
print(f"not found: {error}", file=sys.stderr)
return 2
except LLMError as error:
print(f"llm error: {error}", file=sys.stderr)
return 3
except KeyboardInterrupt:
print("interrupted", file=sys.stderr)
return 130
if __name__ == "__main__":
raise SystemExit(main())
leanbench/config.py
+116
View File
@@ -0,0 +1,116 @@
"""Runtime configuration, read from the environment (and an optional .env file).
Everything the pipeline needs to talk to an LLM or to Lean is settable without
touching code, so the same image runs against any OpenAI-compatible endpoint and
any Lean project.
"""
from __future__ import annotations
import json
import os
from dataclasses import dataclass, field
from pathlib import Path
DEFAULT_BASE_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1"
DEFAULT_MODEL = "qwen-max"
DEFAULT_API_KEY_ENV = "DASHSCOPE_API_KEY"
def load_dotenv(path: Path | str = ".env") -> None:
"""Load ``KEY=VALUE`` lines into os.environ without overwriting what is set."""
path = Path(path)
if not path.is_file():
return
for line in path.read_text(encoding="utf-8").splitlines():
line = line.strip()
if not line or line.startswith("#") or "=" not in line:
continue
key, _, value = line.partition("=")
os.environ.setdefault(key.strip(), value.strip().strip('"').strip("'"))
def _env(name: str, default: str) -> str:
value = os.environ.get(name)
return default if value is None or value == "" else value
def _env_int(name: str, default: int) -> int:
try:
return int(_env(name, str(default)))
except ValueError:
return default
def _env_float(name: str, default: float) -> float:
try:
return float(_env(name, str(default)))
except ValueError:
return default
@dataclass
class LLMSettings:
model: str = DEFAULT_MODEL
base_url: str = DEFAULT_BASE_URL
api_key_env: str = DEFAULT_API_KEY_ENV
timeout: float = 300.0
max_retries: int = 4
temperature: float = 0.0
@classmethod
def from_env(cls) -> LLMSettings:
return cls(
model=_env("LEANBENCH_LLM_MODEL", DEFAULT_MODEL),
base_url=_env("LEANBENCH_LLM_BASE_URL", DEFAULT_BASE_URL),
api_key_env=_env("LEANBENCH_LLM_API_KEY_ENV", DEFAULT_API_KEY_ENV),
timeout=_env_float("LEANBENCH_LLM_TIMEOUT", 300.0),
max_retries=_env_int("LEANBENCH_LLM_MAX_RETRIES", 4),
temperature=_env_float("LEANBENCH_LLM_TEMPERATURE", 0.0),
)
@dataclass
class LeanSettings:
"""How to invoke Lean. ``command`` is an argv prefix that is given a .lean path."""
command: list[str] = field(default_factory=lambda: ["lake", "env", "lean"])
project: str | None = None
timeout: float = 120.0
header: str = "import Mathlib\n"
@classmethod
def from_env(cls) -> LeanSettings:
raw = _env("LEANBENCH_LEAN_COMMAND", '["lake","env","lean"]')
try:
command = json.loads(raw)
if not isinstance(command, list) or not all(isinstance(x, str) for x in command):
raise ValueError
except ValueError:
command = raw.split()
return cls(
command=command,
project=os.environ.get("LEANBENCH_LEAN_PROJECT") or None,
timeout=_env_float("LEANBENCH_LEAN_TIMEOUT", 120.0),
header=_env("LEANBENCH_LEAN_HEADER", "import Mathlib\n"),
)
@dataclass
class Settings:
llm: LLMSettings = field(default_factory=LLMSettings)
lean: LeanSettings = field(default_factory=LeanSettings)
database: str = "database/leanbench.sqlite"
max_window_pages: int = 3
repair_attempts: int = 1
@classmethod
def from_env(cls, dotenv: Path | str | None = ".env") -> Settings:
if dotenv is not None:
load_dotenv(dotenv)
return cls(
llm=LLMSettings.from_env(),
lean=LeanSettings.from_env(),
database=_env("LEANBENCH_DATABASE", "database/leanbench.sqlite"),
max_window_pages=_env_int("LEANBENCH_MAX_WINDOW_PAGES", 3),
repair_attempts=_env_int("LEANBENCH_REPAIR_ATTEMPTS", 1),
)
leanbench/db.py
+431
View File
@@ -0,0 +1,431 @@
"""Step 4: the indexed store.
SQLite, one file, real indexes. Every item records
* the Lean code and what Lean said about it,
* the types it is indexed under -- its inputs, or its own type when it is not a
proof -- which is what makes search by relevance possible,
* which book it came from (by content hash) and which page it starts on,
so re-running any part of any source replaces exactly the old rows and nothing else:
``delete_source`` for a whole book, ``delete_pages`` for a page range.
``export_jsonl`` writes the same rows out as text, which is what belongs in git.
"""
from __future__ import annotations
import hashlib
import json
import sqlite3
from collections.abc import Iterable, Iterator, Sequence
from contextlib import contextmanager
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
SCHEMA_VERSION = 1
SCHEMA = """
PRAGMA journal_mode = WAL;
PRAGMA foreign_keys = ON;
CREATE TABLE IF NOT EXISTS meta (
key TEXT PRIMARY KEY,
value TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS sources (
hash TEXT PRIMARY KEY, -- sha256 of the source file
title TEXT NOT NULL DEFAULT '',
path TEXT NOT NULL DEFAULT '',
pages INTEGER NOT NULL DEFAULT 0,
extractor TEXT NOT NULL DEFAULT '',
ingested_at TEXT NOT NULL,
meta TEXT NOT NULL DEFAULT '{}'
);
CREATE TABLE IF NOT EXISTS items (
id TEXT PRIMARY KEY,
source_hash TEXT NOT NULL REFERENCES sources(hash) ON DELETE CASCADE,
kind TEXT NOT NULL, -- 'statement' (case 1) | 'problem' (case 2)
label TEXT NOT NULL DEFAULT '',
summary TEXT NOT NULL DEFAULT '',
informal TEXT NOT NULL DEFAULT '',
start_page INTEGER NOT NULL, -- the page an item is filed under
start_index INTEGER NOT NULL,
end_page INTEGER NOT NULL,
end_index INTEGER NOT NULL,
page_checksum TEXT NOT NULL DEFAULT '',
source_text TEXT NOT NULL DEFAULT '',
lean_name TEXT NOT NULL DEFAULT '',
lean_code TEXT NOT NULL DEFAULT '',
imports TEXT NOT NULL DEFAULT '[]',
has_proof INTEGER NOT NULL DEFAULT 0,
verified INTEGER NOT NULL DEFAULT 0,
proved INTEGER NOT NULL DEFAULT 0,
verify_mode TEXT NOT NULL DEFAULT '',
diagnostics TEXT NOT NULL DEFAULT '',
confidence REAL NOT NULL DEFAULT 0,
match_method TEXT NOT NULL DEFAULT '',
match_score REAL NOT NULL DEFAULT 0,
truncated INTEGER NOT NULL DEFAULT 0,
created_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS item_types (
item_id TEXT NOT NULL REFERENCES items(id) ON DELETE CASCADE,
name TEXT NOT NULL, -- 'Nat', 'Polynomial', 'Group', ...
role TEXT NOT NULL, -- 'input' | 'goal' | 'own'
PRIMARY KEY (item_id, name, role)
);
CREATE INDEX IF NOT EXISTS idx_items_source_page ON items(source_hash, start_page);
CREATE INDEX IF NOT EXISTS idx_items_kind ON items(kind);
CREATE INDEX IF NOT EXISTS idx_items_verified ON items(verified, proved);
CREATE INDEX IF NOT EXISTS idx_types_name ON item_types(name);
CREATE INDEX IF NOT EXISTS idx_types_name_role ON item_types(name, role);
"""
# A standalone (not external-content) FTS index: it costs a second copy of the text
# and buys ordinary INSERT/DELETE, which external-content tables do not give you.
FTS_SCHEMA = """
CREATE VIRTUAL TABLE IF NOT EXISTS items_fts USING fts5(
id UNINDEXED, label, summary, informal, source_text, lean_code
);
"""
def now() -> str:
return datetime.now(timezone.utc).isoformat(timespec="seconds")
def item_id(source_hash: str, start_page: int, start_index: int,
end_page: int, end_index: int) -> str:
"""Stable identity: same span of the same book is the same item, always."""
key = f"{source_hash}:{start_page}:{start_index}:{end_page}:{end_index}"
return hashlib.sha1(key.encode("utf-8")).hexdigest()[:24]
@dataclass
class ItemRow:
id: str
source_hash: str
kind: str
label: str
summary: str
informal: str
start_page: int
start_index: int
end_page: int
end_index: int
page_checksum: str
source_text: str
lean_name: str
lean_code: str
imports: list[str]
has_proof: bool
verified: bool
proved: bool
verify_mode: str
diagnostics: str
confidence: float
match_method: str
match_score: float
truncated: bool
types: list[tuple[str, str]]
def as_dict(self) -> dict[str, Any]:
data = {k: v for k, v in self.__dict__.items() if k != "types"}
data["types"] = [{"name": n, "role": r} for n, r in self.types]
return data
class Database:
"""Thin, explicit wrapper. No ORM, no magic, no surprises at 3am."""
def __init__(self, path: str | Path = "database/leanbench.sqlite") -> None:
self.path = Path(path)
self.path.parent.mkdir(parents=True, exist_ok=True)
self.connection = sqlite3.connect(str(self.path))
self.connection.row_factory = sqlite3.Row
self.connection.executescript(SCHEMA)
self.fts = self._try_fts()
self.connection.execute(
"INSERT INTO meta(key, value) VALUES('schema_version', ?) "
"ON CONFLICT(key) DO UPDATE SET value = excluded.value",
(str(SCHEMA_VERSION),),
)
self.connection.commit()
def _try_fts(self) -> bool:
try:
self.connection.executescript(FTS_SCHEMA)
return True
except sqlite3.OperationalError:
return False # SQLite built without FTS5; everything else still works
def close(self) -> None:
self.connection.close()
def __enter__(self) -> Database:
return self
def __exit__(self, *exc: object) -> None:
self.close()
@contextmanager
def transaction(self) -> Iterator[sqlite3.Connection]:
try:
yield self.connection
self.connection.commit()
except Exception:
self.connection.rollback()
raise
# -- sources -----------------------------------------------------------
def upsert_source(
self,
hash: str,
*,
title: str = "",
path: str = "",
pages: int = 0,
extractor: str = "",
meta: dict | None = None,
) -> None:
with self.transaction() as connection:
connection.execute(
"""
INSERT INTO sources(hash, title, path, pages, extractor, ingested_at, meta)
VALUES(?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(hash) DO UPDATE SET
title = excluded.title, path = excluded.path, pages = excluded.pages,
extractor = excluded.extractor, ingested_at = excluded.ingested_at,
meta = excluded.meta
""",
(hash, title, path, pages, extractor, now(),
json.dumps(meta or {}, ensure_ascii=False)),
)
def sources(self) -> list[dict[str, Any]]:
rows = self.connection.execute(
"""
SELECT s.*, (SELECT COUNT(*) FROM items i WHERE i.source_hash = s.hash) AS items
FROM sources s ORDER BY s.ingested_at DESC
"""
).fetchall()
return [dict(row) for row in rows]
# -- writing items -----------------------------------------------------
def upsert_item(self, row: ItemRow) -> str:
with self.transaction() as connection:
connection.execute(
"""
INSERT INTO items(
id, source_hash, kind, label, summary, informal,
start_page, start_index, end_page, end_index, page_checksum,
source_text, lean_name, lean_code, imports, has_proof, verified,
proved, verify_mode, diagnostics, confidence, match_method,
match_score, truncated, created_at)
VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)
ON CONFLICT(id) DO UPDATE SET
kind=excluded.kind, label=excluded.label, summary=excluded.summary,
informal=excluded.informal, page_checksum=excluded.page_checksum,
source_text=excluded.source_text, lean_name=excluded.lean_name,
lean_code=excluded.lean_code, imports=excluded.imports,
has_proof=excluded.has_proof, verified=excluded.verified,
proved=excluded.proved, verify_mode=excluded.verify_mode,
diagnostics=excluded.diagnostics, confidence=excluded.confidence,
match_method=excluded.match_method, match_score=excluded.match_score,
truncated=excluded.truncated, created_at=excluded.created_at
""",
(
row.id, row.source_hash, row.kind, row.label, row.summary, row.informal,
row.start_page, row.start_index, row.end_page, row.end_index,
row.page_checksum, row.source_text, row.lean_name, row.lean_code,
json.dumps(row.imports, ensure_ascii=False),
int(row.has_proof), int(row.verified), int(row.proved), row.verify_mode,
row.diagnostics[-4000:], float(row.confidence), row.match_method,
float(row.match_score), int(row.truncated), now(),
),
)
connection.execute("DELETE FROM item_types WHERE item_id = ?", (row.id,))
connection.executemany(
"INSERT OR IGNORE INTO item_types(item_id, name, role) VALUES(?,?,?)",
[(row.id, name, role) for name, role in row.types],
)
self._reindex(row.id)
return row.id
def _reindex(self, item_id_: str) -> None:
if not self.fts:
return
with self.transaction() as connection:
connection.execute("DELETE FROM items_fts WHERE id = ?", (item_id_,))
connection.execute(
"""
INSERT INTO items_fts(id, label, summary, informal, source_text, lean_code)
SELECT id, label, summary, informal, source_text, lean_code
FROM items WHERE id = ?
""",
(item_id_,),
)
# -- removing, so a source can be redone -------------------------------
def delete_source(self, source_hash: str) -> int:
with self.transaction() as connection:
cursor = connection.execute("DELETE FROM items WHERE source_hash = ?", (source_hash,))
connection.execute("DELETE FROM sources WHERE hash = ?", (source_hash,))
return cursor.rowcount
def delete_pages(self, source_hash: str, first_page: int, last_page: int) -> int:
"""Drop items *filed under* pages in this range, so those pages can be redone."""
with self.transaction() as connection:
cursor = connection.execute(
"DELETE FROM items WHERE source_hash = ? AND start_page BETWEEN ? AND ?",
(source_hash, first_page, last_page),
)
return cursor.rowcount
# -- reading -----------------------------------------------------------
def _types_for(self, ids: Sequence[str]) -> dict[str, list[tuple[str, str]]]:
if not ids:
return {}
marks = ",".join("?" * len(ids))
rows = self.connection.execute(
f"SELECT item_id, name, role FROM item_types WHERE item_id IN ({marks})", tuple(ids)
).fetchall()
out: dict[str, list[tuple[str, str]]] = {}
for row in rows:
out.setdefault(row["item_id"], []).append((row["name"], row["role"]))
return out
def _hydrate(self, rows: Iterable[sqlite3.Row]) -> list[dict[str, Any]]:
rows = list(rows)
types = self._types_for([row["id"] for row in rows])
items = []
for row in rows:
item = dict(row)
item["imports"] = json.loads(item.get("imports") or "[]")
item["types"] = [{"name": n, "role": r} for n, r in types.get(row["id"], [])]
items.append(item)
return items
def get(self, id_: str) -> dict[str, Any] | None:
row = self.connection.execute("SELECT * FROM items WHERE id = ?", (id_,)).fetchone()
return self._hydrate([row])[0] if row else None
def search(
self,
*,
types: Sequence[str] = (),
roles: Sequence[str] = (),
kind: str | None = None,
source_hash: str | None = None,
page: int | None = None,
verified: bool | None = None,
proved: bool | None = None,
text: str | None = None,
match_all_types: bool = False,
limit: int = 50,
) -> list[dict[str, Any]]:
"""The point of the index: find items by the types they are about."""
where: list[str] = []
params: list[Any] = []
if types:
marks = ",".join("?" * len(types))
clause = (
f"SELECT item_id FROM item_types WHERE name IN ({marks})"
)
params.extend(types)
if roles:
clause += " AND role IN (" + ",".join("?" * len(roles)) + ")"
params.extend(roles)
if match_all_types:
clause += " GROUP BY item_id HAVING COUNT(DISTINCT name) = ?"
params.append(len(set(types)))
where.append(f"i.id IN ({clause})")
if kind:
where.append("i.kind = ?")
params.append(kind)
if source_hash:
where.append("i.source_hash = ?")
params.append(source_hash)
if page is not None:
where.append("? BETWEEN i.start_page AND i.end_page")
params.append(page)
if verified is not None:
where.append("i.verified = ?")
params.append(int(verified))
if proved is not None:
where.append("i.proved = ?")
params.append(int(proved))
if text:
if self.fts:
where.append("i.id IN (SELECT id FROM items_fts WHERE items_fts MATCH ?)")
params.append(text)
else:
where.append("(i.source_text LIKE ? OR i.lean_code LIKE ? OR i.label LIKE ?)")
params.extend([f"%{text}%"] * 3)
sql = "SELECT i.* FROM items i"
if where:
sql += " WHERE " + " AND ".join(where)
sql += " ORDER BY i.proved DESC, i.verified DESC, i.start_page, i.start_index LIMIT ?"
params.append(limit)
return self._hydrate(self.connection.execute(sql, params).fetchall())
def type_counts(self, *, role: str | None = None, limit: int = 40) -> list[dict[str, Any]]:
sql = "SELECT name, role, COUNT(*) AS items FROM item_types"
params: list[Any] = []
if role:
sql += " WHERE role = ?"
params.append(role)
sql += " GROUP BY name, role ORDER BY items DESC LIMIT ?"
params.append(limit)
return [dict(row) for row in self.connection.execute(sql, params).fetchall()]
def stats(self) -> dict[str, Any]:
one = self.connection.execute(
"""
SELECT COUNT(*) AS items,
SUM(kind = 'statement') AS statements,
SUM(kind = 'problem') AS problems,
SUM(verified) AS verified,
SUM(proved) AS proved
FROM items
"""
).fetchone()
stats = {k: (one[k] or 0) for k in one.keys()}
stats["sources"] = self.connection.execute("SELECT COUNT(*) FROM sources").fetchone()[0]
stats["types"] = self.connection.execute(
"SELECT COUNT(DISTINCT name) FROM item_types"
).fetchone()[0]
return stats
# -- export ------------------------------------------------------------
def export_jsonl(self, path: str | Path, *, source_hash: str | None = None,
include_source_text: bool = False) -> int:
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
sql = "SELECT * FROM items"
params: tuple = ()
if source_hash:
sql += " WHERE source_hash = ?"
params = (source_hash,)
sql += " ORDER BY source_hash, start_page, start_index"
rows = self._hydrate(self.connection.execute(sql, params).fetchall())
with path.open("w", encoding="utf-8") as stream:
for row in rows:
if not include_source_text:
row.pop("source_text", None)
row.pop("diagnostics", None)
stream.write(json.dumps(row, ensure_ascii=False) + "\n")
return len(rows)
leanbench/detect.py
+293
View File
@@ -0,0 +1,293 @@
"""Step 2: decide what on a page is worth turning into Lean, and exactly where it is.
Two kinds of hit:
``statement`` mathematics asserted as fact -- definition, theorem, proposition,
lemma, corollary, a precise claim inside an example. We deliberately
do *not* record which of those it is: all that matters is that it can
become useful Lean 4.
``problem`` a question posed and left open to the reader -- an exercise, an
unsolved problem. Also formalizable, but as a statement with no proof.
Each hit comes back as a span: (start page, start index) -> (end page, end index),
indices being offsets into that page's text. Several hits per page are normal.
If a hit runs off the bottom of the window, the model says so and the window grows
by one page and the page is re-read, repeatedly, so an item spread over three pages
is found whole rather than in pieces.
"""
from __future__ import annotations
import json
from collections.abc import Callable, Sequence
from dataclasses import asdict, dataclass, field
from typing import Any
from .anchors import resolve_span
from .document import Document, Page, PageWindow
KIND_STATEMENT = "statement"
KIND_PROBLEM = "problem"
KINDS = (KIND_STATEMENT, KIND_PROBLEM)
MAX_WINDOW_CHARS = 28000
SYSTEM_PROMPT = """\
You are an exacting extraction engine. You are given the raw text of one or more
consecutive pages of a mathematics book, exactly as extracted from the PDF (line
breaks, hyphenation and all).
Find every passage that could be turned into useful Lean 4 code. There are two kinds:
- "statement": mathematics asserted as fact -- a definition, theorem, proposition,
lemma, corollary, or a precise claim made inside a remark or worked example. Do NOT
tell me which of those it is; the only question is whether it can become useful
Lean 4.
- "problem": a precise mathematical question left open to the reader -- an exercise,
a problem, an unsolved question -- stated without its solution or proof here.
Ignore: running heads and page numbers, table of contents, index, bibliography,
motivational prose with no precise claim, figure captions, and instructions with no
mathematical content ("read section 2", "discuss with a classmate").
Rules:
- Quote verbatim. "start_quote" must be copied character-for-character from the text
given to you and must begin at the very first character of the passage (including
its printed label, e.g. "Theorem 1.6.3."). "end_quote" must likewise be copied
verbatim and end at the very last character of the passage.
- start_quote: 40-90 characters. end_quote: 30-60 characters. Never paraphrase,
never normalise whitespace, never fix hyphenation.
- List passages in order of appearance. Passages must not overlap.
- "continues_past_end": true only when the passage is cut off by the end of the text
you were given -- that is, the book continues it on the following page.
- If nothing qualifies, return an empty list.
Reply with JSON only, in exactly this shape:
{"items": [{"kind": "statement" | "problem",
"label": "Theorem 1.6.3" or "",
"summary": "one short English line",
"start_quote": "...",
"end_quote": "...",
"continues_past_end": false}]}
"""
@dataclass
class Detection:
"""One located item, ready for step 3."""
kind: str
label: str
summary: str
start_page: int
start_index: int
end_page: int
end_index: int
text: str
truncated: bool = False
match_method: str = "exact"
match_score: float = 1.0
pages_scanned: list[int] = field(default_factory=list)
@property
def spans_pages(self) -> int:
return self.end_page - self.start_page + 1
@property
def locator(self) -> str:
return f"p{self.start_page}:{self.start_index}-p{self.end_page}:{self.end_index}"
def as_dict(self) -> dict[str, Any]:
data = asdict(self)
data["spans_pages"] = self.spans_pages
data["locator"] = self.locator
return data
def _window_text(window: PageWindow) -> str:
text = window.text
if len(text) <= MAX_WINDOW_CHARS:
return text
return text[:MAX_WINDOW_CHARS]
def detect_window(client: Any, window: PageWindow, *, temperature: float = 0.0) -> list[dict]:
"""Ask the model what is in this window. Returns raw (unresolved) item dicts."""
pages = window.pages_covered()
label = f"page {pages[0]}" if len(pages) == 1 else f"pages {pages[0]}-{pages[-1]}"
user = (
f"Text of {label} of the book follows between the markers.\n"
f"<<<BEGIN {label.upper()}>>>\n{_window_text(window)}\n<<<END>>>"
)
messages = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": user},
]
if hasattr(client, "complete_json"):
payload = client.complete_json(messages, temperature=temperature)
else: # any bare ChatClient
from .llm import parse_json
payload = parse_json(client.complete(messages, temperature=temperature))
items = payload.get("items") if isinstance(payload, dict) else payload
if not isinstance(items, list):
return []
clean: list[dict] = []
for item in items:
if not isinstance(item, dict):
continue
kind = str(item.get("kind", "")).strip().lower()
if kind not in KINDS:
kind = KIND_PROBLEM if "problem" in kind or "exercise" in kind else KIND_STATEMENT
clean.append(
{
"kind": kind,
"label": str(item.get("label") or "").strip(),
"summary": str(item.get("summary") or "").strip(),
"start_quote": str(item.get("start_quote") or ""),
"end_quote": str(item.get("end_quote") or ""),
"continues_past_end": bool(item.get("continues_past_end")),
}
)
return clean
def resolve_items(
window: PageWindow,
items: Sequence[dict],
*,
still_truncated: bool = False,
) -> tuple[list[Detection], list[dict]]:
"""Turn quoted items into page/index spans. Returns (resolved, unresolved)."""
resolved: list[Detection] = []
unresolved: list[dict] = []
text = window.text
cursor = 0
for item in items:
span = resolve_span(text, item["start_quote"], item["end_quote"], search_from=cursor)
if span is None or not span.resolved:
unresolved.append({**item, "reason": "quotes not found in page text"})
continue
start_page, start_index = window.locate(span.start)
end_page, end_index = window.locate(span.end)
if (end_page, end_index) <= (start_page, start_index):
unresolved.append({**item, "reason": "empty span"})
continue
cursor = max(cursor, span.start + 1)
resolved.append(
Detection(
kind=item["kind"],
label=item["label"],
summary=item["summary"],
start_page=start_page,
start_index=start_index,
end_page=end_page,
end_index=end_index,
text=text[span.start : span.end],
truncated=bool(item.get("continues_past_end")) and still_truncated,
match_method=span.method,
match_score=round(span.score, 3),
pages_scanned=window.pages_covered(),
)
)
return resolved, unresolved
@dataclass
class DetectionReport:
detections: list[Detection] = field(default_factory=list)
unresolved: list[dict] = field(default_factory=list)
windows_used: int = 0
pages_scanned: int = 0
llm_calls: int = 0
def summary(self) -> str:
statements = sum(1 for d in self.detections if d.kind == KIND_STATEMENT)
problems = sum(1 for d in self.detections if d.kind == KIND_PROBLEM)
multipage = sum(1 for d in self.detections if d.spans_pages > 1)
return (
f"{len(self.detections)} items "
f"({statements} statements, {problems} problems, {multipage} spanning pages) "
f"from {self.pages_scanned} pages in {self.llm_calls} model calls; "
f"{len(self.unresolved)} unresolved"
)
def detect_document(
client: Any,
document: Document,
*,
first_page: int | None = None,
last_page: int | None = None,
max_window_pages: int = 3,
progress: Callable[[str], None] | None = None,
on_window_error: str = "skip",
) -> DetectionReport:
"""Scan a document page by page, growing the window when an item is cut off."""
pages: list[Page] = document.page_range(first_page, last_page)
report = DetectionReport(pages_scanned=len(pages))
position = 0
while position < len(pages):
size = 1
window = PageWindow(pages[position : position + 1])
raw: list[dict] = []
needs_more = False
while True:
window = PageWindow(pages[position : position + size])
try:
raw = detect_window(client, window)
except Exception as error: # noqa: BLE001 - one bad page must not kill a book
if on_window_error == "raise":
raise
report.unresolved.append(
{"pages": window.pages_covered(), "reason": f"detector failed: {error}"}
)
raw = []
report.llm_calls += 1
report.windows_used += 1
needs_more = any(item.get("continues_past_end") for item in raw)
room_left = position + size < len(pages)
if needs_more and size < max_window_pages and room_left:
size += 1
if progress:
progress(
f" page {window.first_page}: item runs on, re-reading with "
f"{size} pages"
)
continue
break
resolved, unresolved = resolve_items(window, raw, still_truncated=needs_more)
report.detections.extend(resolved)
report.unresolved.extend(unresolved)
if progress:
progress(
f" pages {window.first_page}-{window.last_page}: "
f"{len(resolved)} item(s)" + (f", {len(unresolved)} unresolved" if unresolved else "")
)
position += size
report.detections = deduplicate(report.detections)
return report
def deduplicate(detections: Sequence[Detection]) -> list[Detection]:
"""Keep the widest span when two hits start at the same place."""
best: dict[tuple[int, int], Detection] = {}
for detection in detections:
key = (detection.start_page, detection.start_index)
current = best.get(key)
if current is None or (detection.end_page, detection.end_index) > (
current.end_page,
current.end_index,
):
best[key] = detection
return sorted(best.values(), key=lambda d: (d.start_page, d.start_index))
def write_jsonl(path: str, detections: Sequence[Detection]) -> None:
with open(path, "w", encoding="utf-8") as stream:
for detection in detections:
stream.write(json.dumps(detection.as_dict(), ensure_ascii=False) + "\n")
leanbench/document.py
+219
View File
@@ -0,0 +1,219 @@
"""Source documents, pages, and the page windows detection runs over.
Every index this project ever reports is an offset into ``Page.text`` as produced
here. That makes the extractor part of the contract, so the extractor name and a
per-page checksum are recorded with the source: if you change extractors, stored
offsets from the old run no longer mean anything, and ``verify_offsets`` says so.
"""
from __future__ import annotations
import hashlib
import re
from collections.abc import Iterator, Sequence
from dataclasses import dataclass, field
from pathlib import Path
PAGE_SEPARATOR = "\n"
def sha256_file(path: Path | str, chunk: int = 1 << 20) -> str:
digest = hashlib.sha256()
with Path(path).open("rb") as stream:
while block := stream.read(chunk):
digest.update(block)
return digest.hexdigest()
def sha1_text(text: str) -> str:
return hashlib.sha1(text.encode("utf-8")).hexdigest()
_LIGATURES = {
"ff": "ff", "fi": "fi", "fl": "fl", "ffi": "ffi", "ffl": "ffl",
" ": " ", "’": "'", "‘": "'", "“": '"', "”": '"',
"–": "-", "—": "--", "−": "-",
}
_TRAILING_SPACE = re.compile(r"[ \t]+\n")
def normalise_page_text(text: str) -> str:
"""Light, *stable* cleanup. Deliberately not aggressive: offsets depend on it."""
for bad, good in _LIGATURES.items():
text = text.replace(bad, good)
text = text.replace("\r\n", "\n").replace("\r", "\n")
text = _TRAILING_SPACE.sub("\n", text)
return text
@dataclass(frozen=True)
class Page:
number: int # 1-based, physical page order in the file
text: str
@property
def checksum(self) -> str:
return sha1_text(self.text)
def __len__(self) -> int:
return len(self.text)
# --------------------------------------------------------------------------
# extraction
def extract_pdf_pages(path: Path | str) -> tuple[list[str], str]:
"""Return (page texts, extractor id). PyMuPDF if present, else pypdf."""
path = Path(path)
backend = None
try: # the modern name; `fitz` is the old one and warns on import
import pymupdf as backend # type: ignore
except ImportError:
try:
import fitz as backend # type: ignore
except ImportError:
backend = None
if backend is not None:
with backend.open(path) as document:
pages = [page.get_text("text") for page in document]
version = getattr(backend, "__version__", "") or getattr(backend, "VersionBind", "")
return pages, f"pymupdf:{version}" if version else "pymupdf"
try:
from pypdf import PdfReader # type: ignore
except ImportError as error: # pragma: no cover - environment problem
raise RuntimeError(
"No PDF backend available: pip install pymupdf (preferred) or pypdf"
) from error
reader = PdfReader(str(path))
return [page.extract_text() or "" for page in reader.pages], "pypdf"
def extract_text_pages(path: Path | str, separator: str = "\f") -> tuple[list[str], str]:
"""Plain-text sources: form feed marks a page break, like ``pdftotext`` emits."""
text = Path(path).read_text(encoding="utf-8", errors="replace")
return text.split(separator), "text"
@dataclass
class Document:
"""A source book: its pages, and the hash the database keys everything by."""
path: Path
sha256: str
pages: list[Page]
extractor: str
title: str = ""
meta: dict = field(default_factory=dict)
@classmethod
def load(cls, path: Path | str, title: str = "") -> Document:
path = Path(path)
if not path.is_file():
raise FileNotFoundError(path)
if path.suffix.lower() == ".pdf":
raw_pages, extractor = extract_pdf_pages(path)
else:
raw_pages, extractor = extract_text_pages(path)
pages = [Page(i + 1, normalise_page_text(t)) for i, t in enumerate(raw_pages)]
return cls(
path=path,
sha256=sha256_file(path),
pages=pages,
extractor=extractor,
title=title or path.stem,
)
def __len__(self) -> int:
return len(self.pages)
def page(self, number: int) -> Page:
if not 1 <= number <= len(self.pages):
raise IndexError(f"page {number} out of range 1..{len(self.pages)}")
return self.pages[number - 1]
def page_range(self, first: int | None = None, last: int | None = None) -> list[Page]:
first = max(1, first or 1)
last = min(len(self.pages), last or len(self.pages))
return [p for p in self.pages if first <= p.number <= last]
def slice(self, start_page: int, start_index: int, end_page: int, end_index: int) -> str:
"""Re-read a stored span straight out of the document."""
if end_page < start_page or (end_page == start_page and end_index < start_index):
raise ValueError("span ends before it starts")
chunks: list[str] = []
for number in range(start_page, end_page + 1):
text = self.page(number).text
begin = start_index if number == start_page else 0
finish = end_index if number == end_page else len(text)
chunks.append(text[begin:finish])
return PAGE_SEPARATOR.join(chunks)
def checksums(self) -> dict[int, str]:
return {page.number: page.checksum for page in self.pages}
# --------------------------------------------------------------------------
# windows
@dataclass
class PageWindow:
"""One or more consecutive pages presented to the model as a single text.
``locate`` maps any offset in that joined text back to (page number, offset in
page), which is what the detector reports and the database stores.
"""
pages: Sequence[Page]
def __post_init__(self) -> None:
if not self.pages:
raise ValueError("a window needs at least one page")
parts: list[str] = []
spans: list[tuple[int, int, int]] = [] # (global start, global end, page number)
cursor = 0
for position, page in enumerate(self.pages):
if position:
cursor += len(PAGE_SEPARATOR)
spans.append((cursor, cursor + len(page.text), page.number))
parts.append(page.text)
cursor += len(page.text)
self.text = PAGE_SEPARATOR.join(parts)
self._spans = spans
@property
def first_page(self) -> int:
return self.pages[0].number
@property
def last_page(self) -> int:
return self.pages[-1].number
def __len__(self) -> int:
return len(self.text)
def locate(self, index: int) -> tuple[int, int]:
"""Global offset -> (page number, offset within that page). Clamps politely."""
index = max(0, min(index, len(self.text)))
for start, end, number in self._spans:
if index < start: # landed inside a separator: belongs to the page ahead
return number, 0
if index <= end:
return number, index - start
start, end, number = self._spans[-1]
return number, end - start
def globalise(self, page_number: int, offset: int) -> int:
for start, end, number in self._spans:
if number == page_number:
return start + max(0, min(offset, end - start))
raise ValueError(f"page {page_number} is not in this window")
def pages_covered(self) -> list[int]:
return [page.number for page in self.pages]
def windows(pages: Sequence[Page], size: int) -> Iterator[PageWindow]:
"""Non-overlapping windows of at most ``size`` pages (used by tests/tools)."""
for start in range(0, len(pages), size):
yield PageWindow(pages[start : start + size])
leanbench/formalize.py
+284
View File
@@ -0,0 +1,284 @@
"""Step 3: every detected span becomes Lean 4 code, then Lean gets the last word.
A ``statement`` item becomes whatever declaration fits -- ``def``, ``structure``,
``theorem`` -- with a proof only if the source actually gave one. A ``problem`` item
is an open question, so it becomes a statement with ``sorry`` for the proof and only
the statement's legality is checked.
When Lean rejects the code, the diagnostics go back to the model and it gets another
go: a compiler error is the cheapest, most honest feedback available.
"""
from __future__ import annotations
import re
from collections.abc import Callable
from dataclasses import asdict, dataclass, field
from typing import Any
from .detect import KIND_PROBLEM, Detection
from .llm import extract_code
from .typeindex import index_types
from .verify import MODE_FULL, MODE_STATEMENT, VerifyResult, contains_sorry
SYSTEM_PROMPT = """\
You translate mathematics from books into Lean 4 (Mathlib). You write Lean 4 with
Mathlib naming and syntax -- never Lean 3, never `begin ... end`.
You are given one passage, extracted from a PDF, so expect broken line breaks,
hyphenation and mangled symbols; reconstruct what the mathematics must have been.
Produce one self-contained top-level declaration that captures the passage:
- a definition in the passage -> `def`, `abbrev`, `structure`, `class` or `instance`;
- a theorem / proposition / lemma / corollary -> `theorem`;
- an exercise or open problem -> `theorem` whose statement is the thing to be shown.
Rules:
- Prefer existing Mathlib definitions over re-defining them. Never invent Mathlib
lemma or definition names: if you are unsure a name exists, restate the maths with
primitives you are sure of.
- State the mathematics faithfully. A statement that is trivially true, or that
quantifies over the wrong thing, is a failure even if it compiles.
- Give the declaration a descriptive snake_case name.
- Include a proof ONLY when the passage itself contains the proof and you are
confident in it; otherwise end the declaration with `sorry`.
- If the passage is not formalizable as a single declaration, say so with
"skip": true and a short reason.
Reply with JSON only:
{"skip": false,
"reason": "",
"imports": ["Mathlib"],
"name": "snake_case_name",
"lean_code": "theorem snake_case_name (n : \\u2115) : ... := by\\n sorry",
"has_proof": false,
"informal": "one-line English restatement",
"index_types": ["Int", "Nat"],
"confidence": 0.0}
"index_types" lists the Lean types this item is about: for a theorem, the types of
its inputs and hypotheses; for a definition, the type of the thing being defined.
"""
REPAIR_PROMPT = """\
Lean rejected that code. Fix it and reply with the same JSON shape.
Do not paper over the error with `sorry` in place of a statement, and do not change
what the passage says in order to make it compile. If a Mathlib name does not exist,
restate the mathematics with primitives that do.
Lean diagnostics:
{diagnostics}
"""
@dataclass
class Formalization:
"""The Lean rendering of one detected item, plus what Lean thought of it."""
detection: Detection
lean_code: str = ""
imports: list[str] = field(default_factory=lambda: ["Mathlib"])
name: str = ""
informal: str = ""
has_proof: bool = False
confidence: float = 0.0
skipped: bool = False
skip_reason: str = ""
attempts: int = 0
verify: VerifyResult | None = None
index_types: list[tuple[str, str]] = field(default_factory=list)
@property
def verified(self) -> bool:
return bool(self.verify and self.verify.ok)
@property
def proved(self) -> bool:
return bool(self.verify and self.verify.proved)
def as_dict(self) -> dict[str, Any]:
return {
"detection": self.detection.as_dict(),
"lean_code": self.lean_code,
"imports": self.imports,
"name": self.name,
"informal": self.informal,
"has_proof": self.has_proof,
"confidence": self.confidence,
"skipped": self.skipped,
"skip_reason": self.skip_reason,
"attempts": self.attempts,
"verify": self.verify.as_dict() if self.verify else None,
"index_types": [{"name": n, "role": r} for n, r in self.index_types],
}
def verification_mode(detection: Detection) -> str:
"""Open problems carry no proof, so only their statement can be checked."""
return MODE_STATEMENT if detection.kind == KIND_PROBLEM else MODE_FULL
def _user_prompt(detection: Detection, document_title: str = "") -> str:
where = (
f"pages {detection.start_page}-{detection.end_page}"
if detection.end_page != detection.start_page
else f"page {detection.start_page}"
)
kind_note = (
"This is an exercise / open problem: state it, do not solve it. End with `sorry`."
if detection.kind == KIND_PROBLEM
else "This is asserted as mathematics in the book."
)
header = f"Source: {document_title or 'book'}, {where}"
if detection.label:
header += f", labelled {detection.label!r}"
return (
f"{header}\n{kind_note}\n\nPassage:\n<<<BEGIN>>>\n{detection.text}\n<<<END>>>"
)
def _clean_code(payload: dict) -> str:
code = payload.get("lean_code") or payload.get("code") or ""
if not isinstance(code, str):
return ""
code = extract_code(code, "lean")
# models occasionally keep the import lines inside the declaration body
lines = [line for line in code.splitlines() if not re.match(r"^\s*import\s", line)]
return "\n".join(lines).strip()
def _imports(payload: dict) -> list[str]:
raw = payload.get("imports")
if isinstance(raw, str):
raw = [raw]
if not isinstance(raw, list):
raw = []
modules = [str(m).strip() for m in raw if str(m).strip()]
return modules or ["Mathlib"]
def formalize(
client: Any,
detection: Detection,
*,
checker: Any = None,
document_title: str = "",
repair_attempts: int = 1,
) -> Formalization:
"""One detected span -> one Lean declaration, verified as far as it can be."""
messages = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": _user_prompt(detection, document_title)},
]
result = Formalization(detection=detection)
mode = verification_mode(detection)
for attempt in range(repair_attempts + 1):
payload = client.complete_json(messages)
if not isinstance(payload, dict):
payload = {}
result.attempts = attempt + 1
if payload.get("skip"):
result.skipped = True
result.skip_reason = str(payload.get("reason") or "model declined")
return result
code = _clean_code(payload)
if not code:
result.skipped = True
result.skip_reason = "model returned no Lean code"
return result
result.lean_code = code
result.imports = _imports(payload)
result.name = str(payload.get("name") or "").strip()
result.informal = str(payload.get("informal") or "").strip()
result.has_proof = bool(payload.get("has_proof")) and not contains_sorry(code)
try:
result.confidence = float(payload.get("confidence") or 0.0)
except (TypeError, ValueError):
result.confidence = 0.0
extra = payload.get("index_types")
result.index_types = index_types(code, extra if isinstance(extra, list) else [])
if checker is None:
return result
result.verify = checker.check(code, mode=mode, imports=result.imports)
if result.verify.ok or result.verify.skipped or attempt >= repair_attempts:
return result
messages = messages + [
{"role": "assistant", "content": _echo(payload, code)},
{
"role": "user",
"content": REPAIR_PROMPT.format(
diagnostics=(result.verify.diagnostics or "")[-3000:]
),
},
]
return result
def _echo(payload: dict, code: str) -> str:
import json
return json.dumps({**payload, "lean_code": code}, ensure_ascii=False)[:6000]
def formalize_all(
client: Any,
detections: list[Detection],
*,
checker: Any = None,
document_title: str = "",
repair_attempts: int = 1,
progress: Callable[[str], None] | None = None,
on_error: str = "skip",
) -> list[Formalization]:
results: list[Formalization] = []
for position, detection in enumerate(detections, 1):
try:
formalization = formalize(
client,
detection,
checker=checker,
document_title=document_title,
repair_attempts=repair_attempts,
)
except Exception as error: # noqa: BLE001 - one bad item must not kill a book
if on_error == "raise":
raise
formalization = Formalization(
detection=detection, skipped=True, skip_reason=f"formalizer failed: {error}"
)
results.append(formalization)
if progress:
status = (
"skipped"
if formalization.skipped
else "proved"
if formalization.proved
else "verified"
if formalization.verified
else "unverified"
)
progress(
f" [{position}/{len(detections)}] {detection.label or detection.kind} "
f"@ {detection.locator}: {status}"
)
return results
def as_dicts(formalizations: list[Formalization]) -> list[dict]:
return [f.as_dict() for f in formalizations]
__all__ = [
"Formalization",
"formalize",
"formalize_all",
"verification_mode",
"as_dicts",
"asdict",
]
leanbench/llm.py
+259
View File
@@ -0,0 +1,259 @@
"""The LLM connector: step 1 of the pipeline, and the thing everything else sits on.
Deliberately dependency-free (stdlib ``urllib`` only) and provider-neutral: any
OpenAI-compatible ``/chat/completions`` endpoint works -- DashScope/Qwen by default,
but also vLLM, Ollama, OpenRouter, DeepSeek, or OpenAI itself. Swap it out entirely
by passing any object with a ``complete(messages, **kwargs) -> str`` method.
"""
from __future__ import annotations
import json
import os
import random
import re
import time
import urllib.error
import urllib.request
from collections.abc import Iterable
from dataclasses import dataclass, field
from typing import Any, Protocol
from .config import LLMSettings
RETRY_STATUS = {408, 409, 425, 429, 500, 502, 503, 504}
_FENCE = re.compile(r"```(?:json|lean4|lean)?\s*(.*?)```", re.DOTALL)
class LLMError(RuntimeError):
"""Any failure to obtain a usable answer from the provider."""
class ChatClient(Protocol):
"""The whole contract the rest of the pipeline depends on."""
def complete(self, messages: list[dict[str, str]], **kwargs: Any) -> str:
...
@dataclass
class Usage:
calls: int = 0
prompt_tokens: int = 0
completion_tokens: int = 0
def add(self, data: dict[str, Any]) -> None:
self.calls += 1
usage = data.get("usage") or {}
self.prompt_tokens += int(usage.get("prompt_tokens") or 0)
self.completion_tokens += int(usage.get("completion_tokens") or 0)
def as_dict(self) -> dict[str, int]:
return {
"calls": self.calls,
"prompt_tokens": self.prompt_tokens,
"completion_tokens": self.completion_tokens,
}
@dataclass
class OpenAICompatibleClient:
"""POST to ``{base_url}/chat/completions``, with retries and JSON helpers."""
model: str
base_url: str
api_key_env: str = "DASHSCOPE_API_KEY"
timeout: float = 300.0
max_retries: int = 4
temperature: float = 0.0
usage: Usage = field(default_factory=Usage)
@classmethod
def from_settings(cls, settings: LLMSettings) -> OpenAICompatibleClient:
return cls(
model=settings.model,
base_url=settings.base_url,
api_key_env=settings.api_key_env,
timeout=settings.timeout,
max_retries=settings.max_retries,
temperature=settings.temperature,
)
# -- transport ---------------------------------------------------------
def _api_key(self) -> str:
key = os.environ.get(self.api_key_env, "").strip()
if not key:
raise LLMError(
f"No API key: set {self.api_key_env} in the environment or in .env"
)
return key
def _post(self, payload: dict[str, Any]) -> dict[str, Any]:
url = self.base_url.rstrip("/") + "/chat/completions"
body = json.dumps(payload).encode("utf-8")
headers = {
"Content-Type": "application/json",
"Authorization": "Bearer " + self._api_key(),
}
last: Exception | None = None
for attempt in range(self.max_retries + 1):
request = urllib.request.Request(url, data=body, headers=headers, method="POST")
try:
with urllib.request.urlopen(request, timeout=self.timeout) as response:
return json.loads(response.read().decode("utf-8"))
except urllib.error.HTTPError as error:
detail = error.read().decode("utf-8", "replace")[:600]
last = LLMError(f"HTTP {error.code} from {url}: {detail}")
if error.code not in RETRY_STATUS:
raise last from error
except (urllib.error.URLError, TimeoutError, OSError) as error:
last = LLMError(f"Network failure talking to {url}: {error}")
if attempt < self.max_retries:
time.sleep(min(2 ** attempt, 30) + random.random())
raise last or LLMError("Request failed for an unknown reason")
# -- api ---------------------------------------------------------------
def complete(
self,
messages: list[dict[str, str]],
*,
temperature: float | None = None,
json_object: bool = False,
max_tokens: int | None = None,
**extra: Any,
) -> str:
payload: dict[str, Any] = {
"model": self.model,
"messages": messages,
"temperature": self.temperature if temperature is None else temperature,
}
if max_tokens:
payload["max_tokens"] = max_tokens
if json_object:
payload["response_format"] = {"type": "json_object"}
payload.update(extra)
try:
data = self._post(payload)
except LLMError as error:
# Not every OpenAI-compatible server implements response_format.
if json_object and "response_format" in str(error):
payload.pop("response_format", None)
data = self._post(payload)
else:
raise
self.usage.add(data)
try:
content = data["choices"][0]["message"]["content"]
except (KeyError, IndexError, TypeError) as error:
raise LLMError(f"Unexpected response shape: {json.dumps(data)[:600]}") from error
if isinstance(content, list): # some servers return content parts
content = "".join(part.get("text", "") for part in content if isinstance(part, dict))
if not isinstance(content, str) or not content.strip():
raise LLMError("Provider returned empty content")
return content
def complete_json(
self,
messages: list[dict[str, str]],
*,
attempts: int = 3,
temperature: float | None = None,
**extra: Any,
) -> Any:
"""Ask for JSON and keep asking until it parses (or give up loudly)."""
conversation = list(messages)
error_text = ""
for attempt in range(attempts):
raw = self.complete(
conversation,
temperature=temperature,
json_object=attempt == 0,
**extra,
)
try:
return parse_json(raw)
except ValueError as error:
error_text = str(error)
conversation = list(messages) + [
{"role": "assistant", "content": raw[:4000]},
{
"role": "user",
"content": (
f"That was not valid JSON ({error_text}). Reply with the JSON "
"object only: no prose, no markdown fences."
),
},
]
raise LLMError(f"Model never produced valid JSON: {error_text}")
def parse_json(raw: str) -> Any:
"""Parse JSON out of a model reply, tolerating fences and surrounding prose."""
text = raw.strip()
try:
return json.loads(text)
except json.JSONDecodeError:
pass
fenced = _FENCE.search(text)
if fenced:
try:
return json.loads(fenced.group(1).strip())
except json.JSONDecodeError:
text = fenced.group(1).strip()
candidates = [i for i in (text.find("{"), text.find("[")) if i != -1]
if candidates:
start = min(candidates)
end = max(text.rfind("}"), text.rfind("]"))
if end > start:
try:
return json.loads(text[start : end + 1])
except json.JSONDecodeError as error:
raise ValueError(str(error)) from error
raise ValueError("no JSON found in reply")
def extract_code(raw: str, language: str = "lean") -> str:
"""Pull a fenced code block out of a reply, or return the reply unchanged."""
fenced = re.search(r"```(?:" + language + r"4?)?[ \t]*\n(.*?)```", raw, re.DOTALL)
return (fenced.group(1) if fenced else raw).strip()
class RecordingClient:
"""Wraps a client and keeps every exchange -- handy for debugging a bad run."""
def __init__(self, inner: ChatClient) -> None:
self.inner = inner
self.log: list[dict[str, Any]] = []
def complete(self, messages: list[dict[str, str]], **kwargs: Any) -> str:
reply = self.inner.complete(messages, **kwargs)
self.log.append({"messages": messages, "reply": reply})
return reply
def complete_json(self, messages: list[dict[str, str]], **kwargs: Any) -> Any:
result = self.inner.complete_json(messages, **kwargs) # type: ignore[attr-defined]
self.log.append({"messages": messages, "reply": result})
return result
class ScriptedClient:
"""A canned-reply client for tests; replies are consumed in order."""
def __init__(self, replies: Iterable[str]) -> None:
self.replies = list(replies)
self.calls: list[list[dict[str, str]]] = []
def complete(self, messages: list[dict[str, str]], **_: Any) -> str:
self.calls.append(messages)
if not self.replies:
raise LLMError("ScriptedClient ran out of replies")
return self.replies.pop(0)
def complete_json(self, messages: list[dict[str, str]], **_: Any) -> Any:
return parse_json(self.complete(messages))
def build_client(settings: LLMSettings | None = None) -> OpenAICompatibleClient:
return OpenAICompatibleClient.from_settings(settings or LLMSettings.from_env())
leanbench/pipeline.py
+181
View File
@@ -0,0 +1,181 @@
"""The whole run, in order: load -> detect -> formalize -> verify -> store.
``ingest`` is what the CLI, the Docker entrypoint and CI all call. Re-ingesting a
book (or just some of its pages) deletes exactly the rows those pages produced last
time and writes the new ones, so a source can be redone without collateral damage.
"""
from __future__ import annotations
from collections.abc import Callable, Sequence
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any
from .config import Settings
from .db import Database, ItemRow, item_id
from .detect import DetectionReport, detect_document
from .document import Document
from .formalize import Formalization, formalize_all, verification_mode
from .llm import build_client
from .verify import LeanChecker, NullChecker
@dataclass
class IngestResult:
document: Document
report: DetectionReport
formalizations: list[Formalization] = field(default_factory=list)
stored: list[str] = field(default_factory=list)
removed: int = 0
def summary(self) -> dict[str, Any]:
verified = sum(1 for f in self.formalizations if f.verified)
proved = sum(1 for f in self.formalizations if f.proved)
skipped = sum(1 for f in self.formalizations if f.skipped)
return {
"source": self.document.title,
"source_hash": self.document.sha256,
"pages": len(self.document),
"detected": len(self.report.detections),
"statements": sum(1 for d in self.report.detections if d.kind == "statement"),
"problems": sum(1 for d in self.report.detections if d.kind == "problem"),
"multi_page": sum(1 for d in self.report.detections if d.spans_pages > 1),
"unresolved": len(self.report.unresolved),
"formalized": len(self.formalizations) - skipped,
"skipped": skipped,
"verified": verified,
"proved": proved,
"stored": len(self.stored),
"removed": self.removed,
}
def build_checker(settings: Settings, *, enabled: bool = True):
if not enabled:
return NullChecker()
checker = LeanChecker.from_settings(settings.lean)
return checker if checker.available() else NullChecker()
def to_row(document: Document, formalization: Formalization) -> ItemRow:
detection = formalization.detection
identifier = item_id(
document.sha256,
detection.start_page,
detection.start_index,
detection.end_page,
detection.end_index,
)
verify = formalization.verify
return ItemRow(
id=identifier,
source_hash=document.sha256,
kind=detection.kind,
label=detection.label,
summary=detection.summary,
informal=formalization.informal,
start_page=detection.start_page,
start_index=detection.start_index,
end_page=detection.end_page,
end_index=detection.end_index,
page_checksum=document.page(detection.start_page).checksum,
source_text=detection.text,
lean_name=formalization.name,
lean_code=formalization.lean_code,
imports=formalization.imports,
has_proof=formalization.has_proof,
verified=bool(verify and verify.ok),
proved=bool(verify and verify.proved),
verify_mode=verify.mode if verify else verification_mode(detection),
diagnostics=(verify.diagnostics if verify else ""),
confidence=formalization.confidence,
match_method=detection.match_method,
match_score=detection.match_score,
truncated=detection.truncated,
types=formalization.index_types,
)
def store(database: Database, document: Document,
formalizations: Sequence[Formalization]) -> list[str]:
database.upsert_source(
document.sha256,
title=document.title,
path=str(document.path),
pages=len(document),
extractor=document.extractor,
meta=document.meta,
)
stored: list[str] = []
for formalization in formalizations:
if formalization.skipped or not formalization.lean_code:
continue
stored.append(database.upsert_item(to_row(document, formalization)))
return stored
def ingest(
source: str | Path,
*,
settings: Settings | None = None,
client: Any = None,
database: Database | None = None,
checker: Any = None,
first_page: int | None = None,
last_page: int | None = None,
title: str = "",
replace: bool = True,
store_results: bool = True,
progress: Callable[[str], None] | None = None,
) -> IngestResult:
settings = settings or Settings.from_env()
client = client or build_client(settings.llm)
checker = checker if checker is not None else build_checker(settings)
owns_database = database is None
database = database or Database(settings.database)
say = progress or (lambda _message: None)
document = Document.load(source, title=title)
say(f"{document.title}: {len(document)} pages, sha256 {document.sha256[:12]}, "
f"extractor {document.extractor}")
result = IngestResult(document=document, report=DetectionReport())
try:
if store_results and replace:
if first_page or last_page:
result.removed = database.delete_pages(
document.sha256, first_page or 1, last_page or len(document)
)
else:
result.removed = database.delete_source(document.sha256)
if result.removed:
say(f"removed {result.removed} item(s) from the previous run")
say("detecting formalizable passages...")
result.report = detect_document(
client,
document,
first_page=first_page,
last_page=last_page,
max_window_pages=settings.max_window_pages,
progress=progress,
)
say(result.report.summary())
say("formalizing...")
result.formalizations = formalize_all(
client,
result.report.detections,
checker=checker,
document_title=document.title,
repair_attempts=settings.repair_attempts,
progress=progress,
)
if store_results:
result.stored = store(database, document, result.formalizations)
say(f"stored {len(result.stored)} item(s) in {database.path}")
finally:
if owns_database:
database.close()
return result
leanbench/typeindex.py
+204
View File
@@ -0,0 +1,204 @@
"""What an item is *about*, in terms the database can index and search on.
An item is indexed by the types it takes in -- or, when it is not a proof, by the
type it is itself. Those types are what make one item relevant to another, so the
index is the thing retrieval later hangs off.
Types are read straight out of the Lean signature (so they reflect the code that
actually elaborated), and merged with whatever the model volunteered.
"""
from __future__ import annotations
import re
from dataclasses import dataclass
PROOF_KEYWORDS = {"theorem", "lemma", "example"}
DECL_KEYWORDS = PROOF_KEYWORDS | {
"def", "abbrev", "instance", "structure", "class", "inductive", "opaque", "axiom",
}
ROLE_INPUT = "input" # a type the item consumes (binders, hypotheses)
ROLE_GOAL = "goal" # the proposition a proof concludes
ROLE_OWN = "own" # the type of a definition, i.e. what it *is*
_DECL = re.compile(
r"^[ \t]*(?:@\[[^\]]*\][ \t\n]*)*"
r"(?:(?:private|protected|noncomputable|partial|unsafe|scoped|local)[ \t]+)*"
r"(" + "|".join(sorted(DECL_KEYWORDS, key=len, reverse=True)) + r")[ \t]+"
r"([^\s:(\[{⦃]+)",
re.MULTILINE,
)
_IDENT = re.compile(r"[A-Za-z_][A-Za-z0-9_'!?₀-₉]*(?:\.[A-Za-z_][A-Za-z0-9_'!?₀-₉]*)*")
UNICODE_TYPES = {
"ℕ": "Nat", "ℤ": "Int", "ℚ": "Rat", "ℝ": "Real", "ℂ": "Complex",
"𝔽": "Field", "𝕜": "Field", "Ω": "Omega", "ℍ": "Quaternion",
}
NOISE = {
"fun", "by", "at", "with", "then", "else", "if", "do", "let", "have", "show",
"from", "this", "sorry", "in", "forall", "exists", "and", "or", "not", "iff",
"true", "false", "Type", "Sort", "Prop", "_root_", "match", "deriving", "where",
}
def _strip_comments(code: str) -> str:
code = re.sub(r"/-.*?-/", " ", code, flags=re.DOTALL)
return re.sub(r"--[^\n]*", " ", code)
def _split_top_level(text: str) -> list[tuple[str, str]]:
"""Split a signature into binder groups: [(opener, inner text), ...]."""
groups: list[tuple[str, str]] = []
pairs = {"(": ")", "{": "}", "[": "]", "⦃": "⦄", "⟨": "⟩"}
depth = 0
opener = ""
start = 0
for index, char in enumerate(text):
if char in pairs:
if depth == 0:
opener, start = char, index + 1
depth += 1
elif char in pairs.values():
depth -= 1
if depth == 0:
groups.append((opener, text[start:index]))
depth = max(depth, 0)
return groups
def _top_level_result(signature: str) -> str:
"""Whatever follows the final colon that is not inside a binder group."""
depth = 0
last_colon = -1
pairs = {"(": ")", "{": "}", "[": "]", "⦃": "⦄", "⟨": "⟩"}
for index, char in enumerate(signature):
if char in pairs:
depth += 1
elif char in pairs.values():
depth = max(0, depth - 1)
elif char == ":" and depth == 0:
if signature[index : index + 2] == ":=":
continue
last_colon = index
return signature[last_colon + 1 :] if last_colon != -1 else ""
SORTS = {"Prop", "Type", "Sort"}
def type_names(fragment: str, *, allow_sorts: bool = False) -> list[str]:
"""Candidate type names inside a type expression.
``allow_sorts`` keeps ``Prop``/``Type``: noise in a hypothesis, but the whole
answer for ``def divides (a b : Int) : Prop``, whose own type *is* Prop.
"""
found: list[str] = []
for symbol, name in UNICODE_TYPES.items():
if symbol in fragment:
found.append(name)
for match in _IDENT.finditer(fragment):
token = match.group(0)
if token in NOISE and not (allow_sorts and token in SORTS):
continue
head = token.split(".")[0]
if not head[:1].isupper():
continue
found.append(token)
ordered: list[str] = []
for name in found:
if name not in ordered:
ordered.append(name)
return ordered
@dataclass
class Signature:
keyword: str
name: str
is_proof: bool
inputs: list[str]
goal: list[str]
own: list[str]
def as_pairs(self) -> list[tuple[str, str]]:
pairs = [(name, ROLE_INPUT) for name in self.inputs]
pairs += [(name, ROLE_GOAL) for name in self.goal]
pairs += [(name, ROLE_OWN) for name in self.own]
seen: set[tuple[str, str]] = set()
out: list[tuple[str, str]] = []
for pair in pairs:
if pair not in seen:
seen.add(pair)
out.append(pair)
return out
def parse_signature(code: str) -> Signature | None:
"""Parse the first declaration in ``code`` into its indexable types."""
clean = _strip_comments(code)
match = _DECL.search(clean)
if not match:
return None
keyword, name = match.group(1), match.group(2)
rest = clean[match.end() :]
depth = 0
cut = len(rest)
pairs = {"(": ")", "{": "}", "[": "]", "⦃": "⦄", "⟨": "⟩"}
for index, char in enumerate(rest):
if char in pairs:
depth += 1
elif char in pairs.values():
depth = max(0, depth - 1)
elif depth == 0:
if rest[index : index + 2] == ":=":
cut = index
break
if rest[index : index + 6] == "\nwhere":
cut = index
break
signature = rest[:cut]
inputs: list[str] = []
for opener, inner in _split_top_level(signature):
body = inner.split(":", 1)[1] if ":" in inner and opener != "[" else inner
inputs.extend(type_names(body))
is_proof = keyword in PROOF_KEYWORDS
result_text = _top_level_result(signature)
result = type_names(result_text, allow_sorts=not is_proof)
return Signature(
keyword=keyword,
name=name,
is_proof=is_proof,
inputs=_unique(inputs),
goal=_unique(result) if is_proof else [],
own=[] if is_proof else _unique(result),
)
def _unique(names: list[str]) -> list[str]:
out: list[str] = []
for name in names:
if name not in out:
out.append(name)
return out
def index_types(code: str, extra: list[str] | None = None) -> list[tuple[str, str]]:
"""The (type name, role) pairs an item should be indexed under."""
signature = parse_signature(code)
pairs = signature.as_pairs() if signature else []
known = {name for name, _ in pairs}
role = ROLE_INPUT if (signature is None or signature.is_proof) else ROLE_OWN
for name in extra or []:
cleaned = str(name).strip()
if not cleaned:
continue
for candidate in type_names(cleaned) or [cleaned]:
if candidate not in known and candidate not in NOISE:
known.add(candidate)
pairs.append((candidate, role))
return pairs
leanbench/verify.py
+148
View File
@@ -0,0 +1,148 @@
"""Step 3, second half: hand the generated Lean to the real compiler and believe it.
Two verification modes, matching the two kinds of item:
``full`` elaborate the declaration as written. If the model supplied a proof,
this is a genuine proof check.
``statement`` the item is an open problem, so there is no proof to check: we only
ask whether the *statement* is legal Lean -- it elaborates with
``sorry`` standing in for the proof.
``verified`` never means "proved". ``proved`` does, and it is only true when Lean
accepted the declaration with no ``sorry`` anywhere in it.
"""
from __future__ import annotations
import os
import re
import shutil
import subprocess
import tempfile
from dataclasses import dataclass
from pathlib import Path
MODE_FULL = "full"
MODE_STATEMENT = "statement"
_ERROR = re.compile(r"^.*?:\d+:\d+: error:", re.MULTILINE)
_SORRY_WARNING = re.compile(r"declaration uses 'sorry'")
_SORRY_TOKEN = re.compile(r"(?<![A-Za-z0-9_'])sorry(?![A-Za-z0-9_'])")
@dataclass
class VerifyResult:
ok: bool
mode: str
used_sorry: bool
diagnostics: str
skipped: bool = False
@property
def proved(self) -> bool:
return self.ok and not self.used_sorry and self.mode == MODE_FULL
def as_dict(self) -> dict:
return {
"ok": self.ok,
"mode": self.mode,
"used_sorry": self.used_sorry,
"proved": self.proved,
"skipped": self.skipped,
"diagnostics": self.diagnostics[-4000:],
}
def contains_sorry(code: str) -> bool:
return bool(_SORRY_TOKEN.search(code))
@dataclass
class LeanChecker:
"""Runs ``lake env lean`` (or whatever you configure) on a scratch file."""
command: list[str]
project: str | None = None
timeout: float = 120.0
header: str = "import Mathlib\n"
@classmethod
def from_settings(cls, settings) -> LeanChecker: # noqa: ANN001 - LeanSettings
return cls(
command=list(settings.command),
project=settings.project,
timeout=settings.timeout,
header=settings.header,
)
def available(self) -> bool:
if not self.command:
return False
if shutil.which(self.command[0]) is None:
return False
if self.project and not Path(self.project).is_dir():
return False
return True
def source_for(self, code: str, imports: list[str] | None = None) -> str:
header = self.header
for module in imports or []:
module = module.strip()
if not module:
continue
line = module if module.startswith("import ") else f"import {module}"
if line not in header:
header += line + "\n"
return header.rstrip("\n") + "\n\n" + code.strip() + "\n"
def check(
self,
code: str,
*,
mode: str = MODE_FULL,
imports: list[str] | None = None,
) -> VerifyResult:
if not self.available():
return VerifyResult(
ok=False,
mode=mode,
used_sorry=contains_sorry(code),
diagnostics=(
f"Lean not available: command {self.command!r}"
+ (f", project {self.project!r}" if self.project else "")
),
skipped=True,
)
source = self.source_for(code, imports)
workdir = self.project or os.getcwd()
with tempfile.TemporaryDirectory(prefix="leanbench-", dir=workdir) as directory:
path = Path(directory) / "LeanBenchCheck.lean"
path.write_text(source, encoding="utf-8")
try:
completed = subprocess.run( # noqa: S603 - command is operator-supplied
[*self.command, str(path)],
cwd=workdir,
capture_output=True,
text=True,
timeout=self.timeout,
)
except subprocess.TimeoutExpired:
return VerifyResult(False, mode, contains_sorry(code),
f"Lean timed out after {self.timeout}s")
except FileNotFoundError as error:
return VerifyResult(False, mode, contains_sorry(code),
f"Lean not runnable: {error}", skipped=True)
diagnostics = (completed.stdout or "") + (completed.stderr or "")
ok = completed.returncode == 0 and not _ERROR.search(diagnostics)
used_sorry = bool(_SORRY_WARNING.search(diagnostics)) or contains_sorry(code)
return VerifyResult(ok, mode, used_sorry, diagnostics[-8000:])
class NullChecker:
"""Stands in when Lean is unavailable (unit tests, a machine without elan)."""
def available(self) -> bool:
return False
def check(self, code: str, *, mode: str = MODE_FULL, imports=None) -> VerifyResult:
return VerifyResult(False, mode, contains_sorry(code), "Lean checking disabled",
skipped=True)