Rebuild lean-bench: pages in, verified indexed Lean out
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.
Infrastructure:
- 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, so the build does not rot.
One Docker layer per Lean step, each long one under `timeout`:
the gigabyte cache fetch earns a layer that a later failure never
invalidates, and a stalled step fails with a message instead of
sitting silent for an hour. There is deliberately no `lake build`
of our own library -- verification runs `lake env lean` against
Mathlib's oleans and needs nothing else, and building the library
was pure cost. docker/smoke.py runs the whole pipeline inside the
image against the real compiler, with the LLM scripted.
- .gitea/workflows ci (lint + tests in a venv, since the runner image's Python is
Debian-managed and refuses installs into itself), docker (build,
verify, optionally push), ingest (submit a book by hand). The
docker build runs with --progress=plain, because BuildKit's
default renderer collapses a long step into one rewritten line
and a stall then looks exactly like progress.
The docker workflow is not a throwaway build: it leaves the image, the layer cache,
an idle `lean-bench` container and the `lean-bench-data` / `lean-bench-books`
volumes in place on the runner, and never prunes, so the environment is there to be
used afterwards. Pushing to the registry is optional and cannot fail the build --
Gitea's container registry does not accept the built-in Actions token, so that needs
a REGISTRY_TOKEN secret with write:package.
58 tests; they need neither network nor Lean.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ci5T3gWDhhmqDec3e1XtYZ
This commit is contained in:
@@ -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__"]
|
||||
@@ -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))
|
||||
@@ -0,0 +1,419 @@
|
||||
"""Command line: ``lean-bench <command>``.
|
||||
|
||||
lean-bench doctor is everything wired up?
|
||||
lean-bench ingest book.pdf --pages 30-45
|
||||
lean-bench extract book.pdf --pages 90 --compare which extractor should I use?
|
||||
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, build_ocr, 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)")
|
||||
|
||||
|
||||
def _add_extractor(parser: argparse.ArgumentParser) -> None:
|
||||
parser.add_argument(
|
||||
"--extractor", choices=["text", "auto", "ocr"],
|
||||
help="where page text comes from: the PDF's own text, OCR, or OCR only for "
|
||||
"pages whose embedded text looks broken (auto)")
|
||||
parser.add_argument("--ocr-model", help="vision model for OCR (default qwen-vl-max)")
|
||||
parser.add_argument("--ocr-dpi", type=int, help="render resolution for OCR (default 200)")
|
||||
parser.add_argument("--ocr-backend", choices=["vision", "tesseract"])
|
||||
parser.add_argument("--ocr-cache", help="directory for cached transcriptions")
|
||||
|
||||
|
||||
def _apply_extractor(settings, args) -> None:
|
||||
if getattr(args, "extractor", None):
|
||||
settings.ocr.mode = args.extractor
|
||||
if getattr(args, "ocr_model", None):
|
||||
settings.ocr.model = args.ocr_model
|
||||
if getattr(args, "ocr_dpi", None):
|
||||
settings.ocr.dpi = args.ocr_dpi
|
||||
if getattr(args, "ocr_backend", None):
|
||||
settings.ocr.backend = args.ocr_backend
|
||||
if getattr(args, "ocr_cache", None):
|
||||
settings.ocr.cache_dir = args.ocr_cache
|
||||
|
||||
|
||||
# -- 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)
|
||||
_apply_extractor(settings, args)
|
||||
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)
|
||||
_apply_extractor(settings, args)
|
||||
first, last = _pages(args.pages)
|
||||
client = build_client(settings.llm)
|
||||
document = Document.load(
|
||||
args.source,
|
||||
mode=settings.ocr.mode,
|
||||
ocr=build_ocr(settings, client),
|
||||
threshold=settings.ocr.threshold,
|
||||
first_page=first, last_page=last,
|
||||
progress=None if args.quiet else _say,
|
||||
)
|
||||
report = detect_document(
|
||||
client,
|
||||
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_extract(args: argparse.Namespace) -> int:
|
||||
"""Show what the pipeline would actually read. The decision tool for --extractor."""
|
||||
from .ocr import text_quality
|
||||
|
||||
settings = Settings.from_env(args.env)
|
||||
_apply_extractor(settings, args)
|
||||
first, last = _pages(args.pages)
|
||||
|
||||
plain = Document.load(args.source)
|
||||
pages = plain.page_range(first, last)
|
||||
if not pages:
|
||||
print("no pages in that range")
|
||||
return 1
|
||||
|
||||
if not args.compare and settings.ocr.mode == "text":
|
||||
for page in pages:
|
||||
print(f"----- page {page.number} (quality {text_quality(page.text)}) -----")
|
||||
print(page.text)
|
||||
return 0
|
||||
|
||||
reader = build_ocr(settings if settings.ocr.mode != "text" else _forced_ocr(settings),
|
||||
build_client(settings.llm))
|
||||
for page in pages:
|
||||
embedded = page.text
|
||||
print(f"===== page {page.number} =====")
|
||||
print(f"--- embedded text (quality {text_quality(embedded)}) ---")
|
||||
print(embedded if args.full else embedded[:800])
|
||||
transcribed = reader.page_text(plain.path, plain.sha256, page.number)
|
||||
print(f"--- {reader.backend.name} ---")
|
||||
if transcribed is None:
|
||||
print(f"FAILED: {reader.failures.get(page.number)}")
|
||||
else:
|
||||
print(transcribed if args.full else transcribed[:800])
|
||||
return 0
|
||||
|
||||
|
||||
def _forced_ocr(settings):
|
||||
settings.ocr.mode = "ocr"
|
||||
return settings
|
||||
|
||||
|
||||
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")
|
||||
_add_extractor(ingest_parser)
|
||||
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")
|
||||
_add_extractor(detect_parser)
|
||||
detect_parser.set_defaults(func=cmd_detect)
|
||||
|
||||
extract = subparsers.add_parser(
|
||||
"extract", help="show the page text the pipeline would read (and compare extractors)")
|
||||
_add_common(extract)
|
||||
_add_extractor(extract)
|
||||
extract.add_argument("source")
|
||||
extract.add_argument("--pages")
|
||||
extract.add_argument("--compare", action="store_true",
|
||||
help="show embedded text and OCR side by side")
|
||||
extract.add_argument("--full", action="store_true", help="do not truncate")
|
||||
extract.set_defaults(func=cmd_extract)
|
||||
|
||||
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())
|
||||
@@ -0,0 +1,143 @@
|
||||
"""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 OcrSettings:
|
||||
"""Reading pages as pictures. Off by default; ``auto`` is the useful setting."""
|
||||
|
||||
mode: str = "text" # text | auto | ocr
|
||||
backend: str = "vision" # vision | tesseract
|
||||
model: str = "qwen-vl-max"
|
||||
dpi: int = 200
|
||||
threshold: float = 0.75 # below this quality, `auto` reaches for the camera
|
||||
cache_dir: str = ".leanbench-cache/ocr"
|
||||
language: str = "eng" # tesseract only
|
||||
|
||||
@classmethod
|
||||
def from_env(cls) -> OcrSettings:
|
||||
return cls(
|
||||
mode=_env("LEANBENCH_EXTRACTOR", "text").lower(),
|
||||
backend=_env("LEANBENCH_OCR_BACKEND", "vision").lower(),
|
||||
model=_env("LEANBENCH_OCR_MODEL", "qwen-vl-max"),
|
||||
dpi=_env_int("LEANBENCH_OCR_DPI", 200),
|
||||
threshold=_env_float("LEANBENCH_OCR_THRESHOLD", 0.75),
|
||||
cache_dir=_env("LEANBENCH_OCR_CACHE", ".leanbench-cache/ocr"),
|
||||
language=_env("LEANBENCH_OCR_LANGUAGE", "eng"),
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Settings:
|
||||
llm: LLMSettings = field(default_factory=LLMSettings)
|
||||
lean: LeanSettings = field(default_factory=LeanSettings)
|
||||
ocr: OcrSettings = field(default_factory=OcrSettings)
|
||||
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(),
|
||||
ocr=OcrSettings.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),
|
||||
)
|
||||
+431
@@ -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)
|
||||
@@ -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")
|
||||
@@ -0,0 +1,285 @@
|
||||
"""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
|
||||
from typing import Any
|
||||
|
||||
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 = "",
|
||||
*,
|
||||
mode: str = "text",
|
||||
ocr: Any = None,
|
||||
threshold: float = 0.75,
|
||||
first_page: int | None = None,
|
||||
last_page: int | None = None,
|
||||
progress: Any = None,
|
||||
) -> Document:
|
||||
"""Load a source.
|
||||
|
||||
``mode`` decides where the text comes from:
|
||||
|
||||
``text`` the embedded text, as the PDF stores it. Fast, free, and wrong
|
||||
whenever the document's maths fonts carry a custom encoding.
|
||||
``ocr`` render every page and read the pixels. Slow and costed per page,
|
||||
but it sees what a human sees.
|
||||
``auto`` embedded text, except on pages that fail a quality check -- which
|
||||
is usually exactly the pages with the mathematics on them.
|
||||
|
||||
``first_page``/``last_page`` bound which pages are OCR'd, so a 25-page trial
|
||||
of a 600-page book costs 25 calls rather than 600. Pages outside the range
|
||||
keep their embedded text.
|
||||
"""
|
||||
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(text)) for i, text in enumerate(raw_pages)]
|
||||
|
||||
document = cls(
|
||||
path=path,
|
||||
sha256=sha256_file(path),
|
||||
pages=pages,
|
||||
extractor=extractor,
|
||||
title=title or path.stem,
|
||||
)
|
||||
if mode != "text" and ocr is not None:
|
||||
document._apply_ocr(ocr, mode, threshold, first_page, last_page, progress)
|
||||
return document
|
||||
|
||||
def _apply_ocr(self, ocr, mode, threshold, first_page, last_page, progress) -> None:
|
||||
from .ocr import text_quality
|
||||
|
||||
say = progress or (lambda _message: None)
|
||||
low = max(1, first_page or 1)
|
||||
high = min(len(self.pages), last_page or len(self.pages))
|
||||
replaced: list[int] = []
|
||||
skipped: list[int] = []
|
||||
|
||||
for index, page in enumerate(self.pages):
|
||||
if not low <= page.number <= high:
|
||||
continue
|
||||
quality = text_quality(page.text)
|
||||
if mode == "auto" and quality >= threshold:
|
||||
continue
|
||||
text = ocr.page_text(self.path, self.sha256, page.number)
|
||||
if text is None:
|
||||
skipped.append(page.number)
|
||||
say(f" page {page.number}: OCR failed, keeping embedded text")
|
||||
continue
|
||||
self.pages[index] = Page(page.number, normalise_page_text(text))
|
||||
replaced.append(page.number)
|
||||
say(f" page {page.number}: OCR (embedded text scored {quality})")
|
||||
|
||||
if replaced:
|
||||
self.extractor = f"{self.extractor}+{ocr.backend.name}@{ocr.dpi}dpi"
|
||||
self.meta["ocr"] = {
|
||||
"mode": mode,
|
||||
"backend": ocr.backend.name,
|
||||
"dpi": ocr.dpi,
|
||||
"threshold": threshold,
|
||||
"pages_replaced": replaced,
|
||||
"pages_failed": skipped,
|
||||
}
|
||||
|
||||
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])
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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())
|
||||
@@ -0,0 +1,308 @@
|
||||
"""Reading pages as pictures, for when the embedded text is a lie.
|
||||
|
||||
Many mathematics PDFs ship their formulas in fonts with custom encodings. The
|
||||
glyphs look right on screen and extract as garbage: a real page of this project's
|
||||
test book yields
|
||||
|
||||
Zab Š Za ˚ Zb; defined by Œxab 7! .Œxa; Œxb/.
|
||||
|
||||
for what a reader plainly sees as Z_ab ≅ Z_a ⊕ Z_b, [x]_ab ↦ ([x]_a, [x]_b). No
|
||||
amount of prompting recovers that, because the information is gone before the model
|
||||
sees it. Rendering the page and reading the pixels does recover it.
|
||||
|
||||
The default backend is the vision model on the same OpenAI-compatible endpoint the
|
||||
rest of the pipeline already uses, so this costs no new dependency and no new
|
||||
credential. Tesseract is available as an offline fallback, with the caveat that it
|
||||
is good at prose and poor at mathematics -- which is the opposite of what we need,
|
||||
so it is a last resort rather than an equal option.
|
||||
|
||||
Transcription is cached on disk, keyed by document hash, page, dpi and backend. A
|
||||
600-page book is 600 vision calls; nobody should pay for them twice.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import tempfile
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, Protocol
|
||||
|
||||
DEFAULT_DPI = 200
|
||||
|
||||
TRANSCRIBE_PROMPT = """\
|
||||
Transcribe this page from a mathematics textbook, exactly as printed.
|
||||
|
||||
- Keep every word, in reading order. Do not summarise, correct or comment.
|
||||
- Write mathematics as LaTeX: inline as $...$, displayed as $$...$$.
|
||||
- Keep the printed labels ("Theorem 1.6.3.", "Exercise 4.") exactly as they appear.
|
||||
- Include running heads, page numbers and footnotes as their own lines.
|
||||
- If part of the page is a figure you cannot transcribe, write [figure] on its own line.
|
||||
|
||||
Output only the transcription.
|
||||
"""
|
||||
|
||||
|
||||
class OcrBackend(Protocol):
|
||||
"""Anything that can turn a rendered page into text."""
|
||||
|
||||
name: str
|
||||
|
||||
def transcribe(self, image: bytes) -> str:
|
||||
...
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# rendering
|
||||
|
||||
|
||||
def render_page(path: Path | str, page_number: int, dpi: int = DEFAULT_DPI) -> bytes:
|
||||
"""Render a 1-based page of a PDF to PNG bytes."""
|
||||
try:
|
||||
import pymupdf as backend # type: ignore
|
||||
except ImportError: # pragma: no cover - fitz is the old name
|
||||
import fitz as backend # type: ignore
|
||||
|
||||
with backend.open(str(path)) as document:
|
||||
if not 1 <= page_number <= document.page_count:
|
||||
raise IndexError(f"page {page_number} out of range 1..{document.page_count}")
|
||||
page = document[page_number - 1]
|
||||
return page.get_pixmap(dpi=dpi).tobytes("png")
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# is the embedded text worth trusting?
|
||||
|
||||
# Glyphs that a correctly-encoded mathematics page essentially never contains, but
|
||||
# which mis-mapped maths fonts produce in quantity: ligature slots, spacing
|
||||
# modifiers and stray accents standing in for operators and brackets.
|
||||
_SUSPICIOUS = set("ŒœŠšŽžŸ˚˙˝˘˛˜ˆ¯´`¨¤ÐÞłŁ§¶")
|
||||
_SUSPICIOUS_PAIRS = ("7!", "!7", "¤", "Š", "˚")
|
||||
|
||||
# A second, nastier corruption: some maths fonts map operators onto ordinary ASCII
|
||||
# letters, so the extracted text is made of perfectly innocent characters and the
|
||||
# glyph test above scores it 1.0. On the project's test book this produces lines like
|
||||
#
|
||||
# I C J D fa C b W a 2 I and b 2 J g
|
||||
#
|
||||
# for "I + J = {a + b : a ∈ I and b ∈ J}" -- D is "=", C is "+", 2 is "∈", W is ":",
|
||||
# f and g are the braces. Only the standalone capitals are counted: a bare "2" is
|
||||
# the ∈ of this encoding, but it is also how people write "chapter 2", and including
|
||||
# it produced false positives on ordinary prose. Measured across the test book,
|
||||
# letters alone give clean separation -- 0.00 on the copyright page and on English
|
||||
# sentences full of numerals, 10-11 on corrupted pages, median 5.6.
|
||||
_OPERATOR_LETTERS = re.compile(r"(?<=[ (])[DCW](?=[ )])")
|
||||
_MIN_LENGTH_FOR_OPERATOR_TEST = 200
|
||||
|
||||
|
||||
def text_quality(text: str) -> float:
|
||||
"""Rough confidence that extracted text says what the page says, in 0..1.
|
||||
|
||||
This does not judge OCR accuracy or prose quality. It answers one question:
|
||||
does this look like a font-encoding accident? Two independent signals, because
|
||||
there are two failure modes and the obvious one misses the worse one:
|
||||
|
||||
* exotic glyphs (Œ Š ˚ 7!) standing in for brackets and arrows;
|
||||
* ordinary letters standing in for operators (D for =, C for +, 2 for ∈).
|
||||
|
||||
The second scores a perfect 1.0 on the first test and is the more common of the
|
||||
two, so both are required. Calibrated against this project's test book, where
|
||||
89% of pages carry the second signal and only the copyright page is clean.
|
||||
"""
|
||||
stripped = [c for c in text if not c.isspace()]
|
||||
if not stripped:
|
||||
return 0.0
|
||||
|
||||
# Signal two: operators disguised as letters, measured per 1000 characters.
|
||||
# Short fragments are exempt -- the rate is meaningless over a single sentence.
|
||||
if len(text) >= _MIN_LENGTH_FOR_OPERATOR_TEST:
|
||||
operator_rate = 1000 * len(_OPERATOR_LETTERS.findall(text)) / len(text)
|
||||
operator_score = max(0.0, 1.0 - operator_rate / 4)
|
||||
else:
|
||||
operator_score = 1.0
|
||||
|
||||
suspicious = sum(1 for c in stripped if c in _SUSPICIOUS)
|
||||
if suspicious <= 2:
|
||||
# Ordinary prose has the occasional accent ("Poincare\u00b4", "Erdos"). Only a
|
||||
# pile of these is evidence of a font-encoding accident, and without this
|
||||
# floor a short clean string is condemned by a single character.
|
||||
score = 1.0
|
||||
else:
|
||||
ratio = suspicious / len(stripped)
|
||||
score = max(0.0, 1.0 - ratio * 40) # 2.5% suspicious glyphs -> 0
|
||||
for token in _SUSPICIOUS_PAIRS:
|
||||
if text.count(token) >= 3:
|
||||
score *= 0.5
|
||||
# Either signal alone condemns the page: both mean the maths is unreadable.
|
||||
return round(min(score, operator_score), 3)
|
||||
|
||||
|
||||
def looks_broken(text: str, threshold: float = 0.75) -> bool:
|
||||
return text_quality(text) < threshold
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# backends
|
||||
|
||||
|
||||
@dataclass
|
||||
class VisionLLMBackend:
|
||||
"""Read the page with a vision model over the OpenAI-compatible endpoint."""
|
||||
|
||||
client: Any
|
||||
model: str = "qwen-vl-max"
|
||||
prompt: str = TRANSCRIBE_PROMPT
|
||||
detail: str = "high"
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return f"vision:{self.model}"
|
||||
|
||||
def transcribe(self, image: bytes) -> str:
|
||||
data_url = "data:image/png;base64," + base64.b64encode(image).decode("ascii")
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": self.prompt},
|
||||
{"type": "image_url", "image_url": {"url": data_url}},
|
||||
],
|
||||
}
|
||||
]
|
||||
return self.client.complete(messages, model=self.model, temperature=0.0).strip()
|
||||
|
||||
|
||||
@dataclass
|
||||
class TesseractBackend:
|
||||
"""Offline fallback. Fine on prose, poor on mathematics -- know what you get."""
|
||||
|
||||
language: str = "eng"
|
||||
binary: str = "tesseract"
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return f"tesseract:{self.language}"
|
||||
|
||||
def available(self) -> bool:
|
||||
return shutil.which(self.binary) is not None
|
||||
|
||||
def transcribe(self, image: bytes) -> str:
|
||||
if not self.available():
|
||||
raise RuntimeError(f"{self.binary} is not installed")
|
||||
with tempfile.TemporaryDirectory(prefix="leanbench-ocr-") as directory:
|
||||
source = Path(directory) / "page.png"
|
||||
source.write_bytes(image)
|
||||
result = subprocess.run( # noqa: S603 - operator-chosen binary
|
||||
[self.binary, str(source), "stdout", "-l", self.language],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=300,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
raise RuntimeError(f"tesseract failed: {result.stderr[:400]}")
|
||||
return result.stdout.strip()
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# cache
|
||||
|
||||
|
||||
@dataclass
|
||||
class OcrCache:
|
||||
"""Disk cache of transcriptions. Keyed so a re-run is free, not just faster."""
|
||||
|
||||
directory: Path
|
||||
enabled: bool = True
|
||||
hits: int = 0
|
||||
misses: int = 0
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self.directory = Path(self.directory)
|
||||
|
||||
def _path(self, document_hash: str, page: int, dpi: int, backend: str) -> Path:
|
||||
token = hashlib.sha1(f"{dpi}:{backend}".encode()).hexdigest()[:10]
|
||||
return self.directory / document_hash[:16] / f"p{page:05d}-{token}.txt"
|
||||
|
||||
def get(self, document_hash: str, page: int, dpi: int, backend: str) -> str | None:
|
||||
if not self.enabled:
|
||||
return None
|
||||
path = self._path(document_hash, page, dpi, backend)
|
||||
if path.is_file():
|
||||
self.hits += 1
|
||||
return path.read_text(encoding="utf-8")
|
||||
self.misses += 1
|
||||
return None
|
||||
|
||||
def put(self, document_hash: str, page: int, dpi: int, backend: str, text: str) -> None:
|
||||
if not self.enabled:
|
||||
return
|
||||
path = self._path(document_hash, page, dpi, backend)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(text, encoding="utf-8")
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# the thing the rest of the pipeline calls
|
||||
|
||||
|
||||
@dataclass
|
||||
class OcrReader:
|
||||
backend: OcrBackend
|
||||
cache: OcrCache | None = None
|
||||
dpi: int = DEFAULT_DPI
|
||||
failures: dict[int, str] = field(default_factory=dict)
|
||||
|
||||
def page_text(self, path: Path | str, document_hash: str, page_number: int) -> str | None:
|
||||
"""Transcribe one page. Returns None (and records why) if it cannot."""
|
||||
if self.cache is not None:
|
||||
cached = self.cache.get(document_hash, page_number, self.dpi, self.backend.name)
|
||||
if cached is not None:
|
||||
return cached
|
||||
try:
|
||||
image = render_page(path, page_number, self.dpi)
|
||||
text = self.backend.transcribe(image)
|
||||
except Exception as error: # noqa: BLE001 - one bad page must not kill a book
|
||||
self.failures[page_number] = str(error)
|
||||
return None
|
||||
text = _tidy(text)
|
||||
if self.cache is not None:
|
||||
self.cache.put(document_hash, page_number, self.dpi, self.backend.name, text)
|
||||
return text
|
||||
|
||||
|
||||
_FENCE = re.compile(r"^\s*```[a-z]*\s*\n(.*?)\n```\s*$", re.DOTALL)
|
||||
|
||||
|
||||
def _tidy(text: str) -> str:
|
||||
"""Strip the markdown fence some models wrap transcriptions in."""
|
||||
match = _FENCE.match(text)
|
||||
return (match.group(1) if match else text).strip()
|
||||
|
||||
|
||||
def build_backend(settings: Any, client: Any = None) -> OcrBackend:
|
||||
"""Make the backend named by settings, falling back to tesseract offline."""
|
||||
kind = (getattr(settings, "backend", "vision") or "vision").lower()
|
||||
if kind in ("vision", "llm", "auto"):
|
||||
if client is None:
|
||||
raise ValueError("the vision backend needs an LLM client")
|
||||
return VisionLLMBackend(client=client, model=settings.model)
|
||||
if kind == "tesseract":
|
||||
return TesseractBackend(language=getattr(settings, "language", "eng"))
|
||||
raise ValueError(f"unknown OCR backend {kind!r}")
|
||||
|
||||
|
||||
__all__ = [
|
||||
"DEFAULT_DPI",
|
||||
"OcrBackend",
|
||||
"OcrCache",
|
||||
"OcrReader",
|
||||
"TesseractBackend",
|
||||
"VisionLLMBackend",
|
||||
"build_backend",
|
||||
"looks_broken",
|
||||
"render_page",
|
||||
"text_quality",
|
||||
]
|
||||
@@ -0,0 +1,206 @@
|
||||
"""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_ocr(settings: Settings, client: Any = None):
|
||||
"""An OcrReader for this configuration, or None when OCR is switched off."""
|
||||
if settings.ocr.mode == "text":
|
||||
return None
|
||||
from .ocr import OcrCache, OcrReader, build_backend
|
||||
|
||||
backend = build_backend(settings.ocr, client)
|
||||
cache = OcrCache(settings.ocr.cache_dir) if settings.ocr.cache_dir else None
|
||||
return OcrReader(backend=backend, cache=cache, dpi=settings.ocr.dpi)
|
||||
|
||||
|
||||
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)
|
||||
reader = build_ocr(settings, client) if settings.ocr.mode != "text" else None
|
||||
if reader is not None:
|
||||
say(f"reading pages with {reader.backend.name} ({settings.ocr.mode} mode)")
|
||||
document = Document.load(
|
||||
source,
|
||||
title=title,
|
||||
mode=settings.ocr.mode,
|
||||
ocr=reader,
|
||||
threshold=settings.ocr.threshold,
|
||||
first_page=first_page,
|
||||
last_page=last_page,
|
||||
progress=progress,
|
||||
)
|
||||
say(f"{document.title}: {len(document)} pages, sha256 {document.sha256[:12]}, "
|
||||
f"extractor {document.extractor}")
|
||||
if reader is not None and reader.cache is not None:
|
||||
say(f" OCR cache: {reader.cache.hits} hit(s), {reader.cache.misses} miss(es)")
|
||||
|
||||
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
|
||||
@@ -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
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user