Zayden-Jung
d3555d9a52
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.
331 lines
12 KiB
Python
331 lines
12 KiB
Python
"""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())
|