From b9e37952e6681951e5fa8e23c87c2f073900a026 Mon Sep 17 00:00:00 2001 From: claude-bot <10+claude-bot@noreply.localhost> Date: Sat, 19 Sep 2026 07:16:37 +0000 Subject: [PATCH] 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. 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 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 Claude-Session: https://claude.ai/code/session_01Ci5T3gWDhhmqDec3e1XtYZ --- .env.example | 22 ++ .gitea/workflows/ci.yml | 57 +++++ .gitea/workflows/docker.yml | 221 ++++++++++++++++++ .gitea/workflows/ingest.yml | 129 +++++++++++ .gitignore | 39 ++++ README.md | 199 +++++++++++++++++ docker/Dockerfile | 72 ++++++ docker/docker-compose.yml | 25 +++ docker/entrypoint.sh | 32 +++ docker/resolve-mathlib.sh | 83 +++++++ docker/smoke.py | 134 +++++++++++ examples/sample-book.txt | 21 ++ lean/LeanBench.lean | 18 ++ lean/lakefile.toml | 12 + lean/lean-toolchain | 1 + leanbench/__init__.py | 16 ++ leanbench/anchors.py | 154 +++++++++++++ leanbench/cli.py | 330 +++++++++++++++++++++++++++ leanbench/config.py | 116 ++++++++++ leanbench/db.py | 431 ++++++++++++++++++++++++++++++++++++ leanbench/detect.py | 293 ++++++++++++++++++++++++ leanbench/document.py | 219 ++++++++++++++++++ leanbench/formalize.py | 284 ++++++++++++++++++++++++ leanbench/llm.py | 259 ++++++++++++++++++++++ leanbench/pipeline.py | 181 +++++++++++++++ leanbench/typeindex.py | 204 +++++++++++++++++ leanbench/verify.py | 148 +++++++++++++ pyproject.toml | 36 +++ tests/__init__.py | 0 tests/conftest.py | 79 +++++++ tests/test_anchors.py | 54 +++++ tests/test_db.py | 133 +++++++++++ tests/test_detect.py | 133 +++++++++++ tests/test_document.py | 46 ++++ tests/test_formalize.py | 90 ++++++++ tests/test_llm.py | 38 ++++ tests/test_pipeline.py | 158 +++++++++++++ tests/test_typeindex.py | 81 +++++++ 38 files changed, 4548 insertions(+) create mode 100644 .env.example create mode 100644 .gitea/workflows/ci.yml create mode 100644 .gitea/workflows/docker.yml create mode 100644 .gitea/workflows/ingest.yml create mode 100644 .gitignore create mode 100644 README.md create mode 100644 docker/Dockerfile create mode 100644 docker/docker-compose.yml create mode 100755 docker/entrypoint.sh create mode 100755 docker/resolve-mathlib.sh create mode 100644 docker/smoke.py create mode 100644 examples/sample-book.txt create mode 100644 lean/LeanBench.lean create mode 100644 lean/lakefile.toml create mode 100644 lean/lean-toolchain create mode 100644 leanbench/__init__.py create mode 100644 leanbench/anchors.py create mode 100644 leanbench/cli.py create mode 100644 leanbench/config.py create mode 100644 leanbench/db.py create mode 100644 leanbench/detect.py create mode 100644 leanbench/document.py create mode 100644 leanbench/formalize.py create mode 100644 leanbench/llm.py create mode 100644 leanbench/pipeline.py create mode 100644 leanbench/typeindex.py create mode 100644 leanbench/verify.py create mode 100644 pyproject.toml create mode 100644 tests/__init__.py create mode 100644 tests/conftest.py create mode 100644 tests/test_anchors.py create mode 100644 tests/test_db.py create mode 100644 tests/test_detect.py create mode 100644 tests/test_document.py create mode 100644 tests/test_formalize.py create mode 100644 tests/test_llm.py create mode 100644 tests/test_pipeline.py create mode 100644 tests/test_typeindex.py diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..c6ba27f --- /dev/null +++ b/.env.example @@ -0,0 +1,22 @@ +# Copy to .env and fill in. .env is gitignored. + +# --- LLM connector ------------------------------------------------------- +# Any OpenAI-compatible /chat/completions endpoint. +DASHSCOPE_API_KEY=sk-your-key-here +LEANBENCH_LLM_MODEL=qwen-max +LEANBENCH_LLM_BASE_URL=https://dashscope.aliyuncs.com/compatible-mode/v1 +# Name of the variable the key lives in, if not DASHSCOPE_API_KEY: +# LEANBENCH_LLM_API_KEY_ENV=OPENAI_API_KEY + +# --- pipeline ------------------------------------------------------------ +LEANBENCH_DATABASE=database/leanbench.sqlite +# How many consecutive pages one item may be reassembled from: +LEANBENCH_MAX_WINDOW_PAGES=3 +# Extra tries after Lean rejects the generated code: +LEANBENCH_REPAIR_ATTEMPTS=1 + +# --- Lean ---------------------------------------------------------------- +# A Lake project with Mathlib already built. Leave unset to skip verification. +# LEANBENCH_LEAN_PROJECT=/opt/lean-bench/lean +# LEANBENCH_LEAN_COMMAND=["lake","env","lean"] +# LEANBENCH_LEAN_TIMEOUT=120 diff --git a/.gitea/workflows/ci.yml b/.gitea/workflows/ci.yml new file mode 100644 index 0000000..5dba081 --- /dev/null +++ b/.gitea/workflows/ci.yml @@ -0,0 +1,57 @@ +name: ci + +on: + push: + branches: [main] + pull_request: + workflow_dispatch: + +jobs: + test: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + # The runner image ships a Debian-managed Python (PEP 668), so installing + # into it is refused and `pip install --upgrade pip` cannot replace the + # distro pip. A venv sidesteps both, and costs a couple of seconds. + - name: Set up a virtualenv + run: | + set -eu + python3 --version + python3 -m venv .venv || { + apt-get update -qq + apt-get install -y -qq python3-venv + python3 -m venv .venv + } + .venv/bin/python -m pip install --quiet --upgrade pip + + - name: Install + run: .venv/bin/python -m pip install --quiet -e ".[dev]" + + - name: Lint + run: .venv/bin/python -m ruff check . + + - name: Test + run: .venv/bin/python -m pytest -q + + - name: CLI is wired up + run: | + set -eu + .venv/bin/lean-bench --help > /dev/null + .venv/bin/lean-bench doctor --skip-llm --env /dev/null + + lean-project: + # Cheap guard: the Lean project's pins must at least resolve. + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - name: Resolve the Mathlib / toolchain pair + run: | + set -eu + command -v curl >/dev/null || (apt-get update -qq && apt-get install -y -qq curl) + command -v git >/dev/null || (apt-get update -qq && apt-get install -y -qq git) + chmod +x docker/resolve-mathlib.sh + docker/resolve-mathlib.sh lean "${MATHLIB_REV:-auto}" + echo "toolchain: $(cat lean/lean-toolchain)" + grep -E '^rev' lean/lakefile.toml diff --git a/.gitea/workflows/docker.yml b/.gitea/workflows/docker.yml new file mode 100644 index 0000000..47b68e9 --- /dev/null +++ b/.gitea/workflows/docker.yml @@ -0,0 +1,221 @@ +name: docker + +# Builds the Lean 4 + Mathlib image and LEAVES IT IN PLACE on the runner host. +# +# This is not a throwaway test build. When the job finishes, the runner still has: +# * the image, tagged lean-bench:latest and lean-bench:sha- +# * the Docker layer cache, so the next build reuses Mathlib instead of +# re-downloading it (nothing here ever calls `docker system prune`) +# * a long-lived container `lean-bench`, idle and ready for `docker exec` +# * named volumes `lean-bench-data` (the item database) and `lean-bench-books` +# +# so you can submit a book straight afterwards: +# docker exec -e DASHSCOPE_API_KEY=sk-... lean-bench \ +# lean-bench ingest /books/your-book.pdf --pages 30-60 +# +# Pushing to the registry is optional and never fails the build: the built-in +# Actions token is not accepted by Gitea's container registry, so add a personal +# access token with write:package as the REGISTRY_TOKEN secret to enable it. + +on: + push: + branches: [main] + tags: ["v*"] + workflow_dispatch: + inputs: + mathlib_rev: + description: "Mathlib revision (tag, branch, commit, or 'auto')" + required: false + default: "auto" + push_image: + description: "Push to the registry if a token is configured" + required: false + default: "true" + refresh_service: + description: "Recreate the long-lived lean-bench container even if one is running" + required: false + default: "false" + no_cache: + description: "Build without the layer cache (slow: refetches Mathlib)" + required: false + default: "false" + +env: + REGISTRY: gitea.entroia.com + IMAGE_NAME: lean-bench + LOCAL_IMAGE: lean-bench + SERVICE_CONTAINER: lean-bench + DATA_VOLUME: lean-bench-data + BOOKS_VOLUME: lean-bench-books + +jobs: + image: + runs-on: ubuntu-latest + timeout-minutes: 240 + steps: + - uses: actions/checkout@v4 + + - name: Check the runner can build images + run: | + set -eu + if ! command -v docker >/dev/null; then + echo "::error::No docker CLI on this runner. Give the act_runner the host" + echo "::error::docker socket (-v /var/run/docker.sock:/var/run/docker.sock)" + echo "::error::or run it with the docker-in-docker image." + exit 1 + fi + docker version --format 'client {{.Client.Version}} / server {{.Server.Version}}' + echo "disk on the docker host:" + docker system df || true + + - name: Work out the tags + id: meta + run: | + set -eu + owner="$(echo "${{ github.repository_owner }}" | tr '[:upper:]' '[:lower:]')" + remote="${REGISTRY}/${owner}/${IMAGE_NAME}" + ref="${GITHUB_REF_NAME:-main}" + sha="$(echo "${GITHUB_SHA:-unknown}" | cut -c1-12)" + case "${GITHUB_REF:-}" in + refs/tags/*) version="${ref}" ;; + *) version="${ref//\//-}" ;; + esac + # Local tags are what stays on the runner; remote tags are for the registry. + local_tags="${LOCAL_IMAGE}:${version} ${LOCAL_IMAGE}:sha-${sha}" + remote_tags="${remote}:${version} ${remote}:sha-${sha}" + if [ "${ref}" = "main" ]; then + local_tags="${local_tags} ${LOCAL_IMAGE}:latest" + remote_tags="${remote_tags} ${remote}:latest" + fi + { + echo "remote=${remote}" + echo "primary=${LOCAL_IMAGE}:${version}" + echo "local_tags=${local_tags}" + echo "remote_tags=${remote_tags}" + } >> "$GITHUB_OUTPUT" + echo "local: ${local_tags}" + echo "remote: ${remote_tags}" + + - name: Build + run: | + set -eu + args="" + for tag in ${{ steps.meta.outputs.local_tags }} ${{ steps.meta.outputs.remote_tags }}; do + args="$args -t $tag" + done + # Reuse the previous build's layers explicitly as well as implicitly, so a + # rebuild does not refetch Mathlib. + cache="" + if [ '${{ inputs.no_cache }}' = 'true' ]; then + cache="--no-cache" + elif docker image inspect "${LOCAL_IMAGE}:latest" >/dev/null 2>&1; then + cache="--cache-from ${LOCAL_IMAGE}:latest" + fi + echo "building with: $cache" + docker build \ + -f docker/Dockerfile \ + --build-arg "MATHLIB_REV=${{ inputs.mathlib_rev || 'auto' }}" \ + --build-arg BUILDKIT_INLINE_CACHE=1 \ + $cache $args . + + - name: Check the image actually works + run: | + set -eu + image='${{ steps.meta.outputs.primary }}' + docker run --rm "$image" doctor --skip-llm + docker run --rm --entrypoint bash "$image" -lc \ + 'cd "$LEANBENCH_LEAN_PROJECT" && printf "import Mathlib\ntheorem ci_smoke (n : Nat) : n + 0 = n := by simp\n" > /tmp/CI.lean && lake env lean /tmp/CI.lean && echo "lean + mathlib ok"' + # Whole pipeline, scripted LLM, real compiler. Piped in over stdin so the + # job container needs no bind mount into the host daemon. + docker run --rm -i --entrypoint python3 "$image" - < docker/smoke.py + + - name: Log in to the registry + id: login + if: ${{ github.event_name != 'workflow_dispatch' || inputs.push_image == 'true' }} + continue-on-error: true + run: | + set -eu + token='${{ secrets.REGISTRY_TOKEN }}' + [ -n "$token" ] || token='${{ secrets.GITEA_TOKEN }}' + [ -n "$token" ] || token='${{ secrets.GITHUB_TOKEN }}' + user='${{ vars.REGISTRY_USER }}' + [ -n "$user" ] || user='${{ github.actor }}' + if [ -z "$token" ]; then + echo "::warning::No registry token; keeping the image on the runner only." + exit 1 + fi + echo "$token" | docker login "${REGISTRY}" -u "$user" --password-stdin + + - name: Push + id: push + if: ${{ steps.login.outcome == 'success' }} + continue-on-error: true + run: | + set -eu + for tag in ${{ steps.meta.outputs.remote_tags }}; do + echo "pushing $tag" + docker push "$tag" + done + + - name: Leave the environment ready for use + run: | + set -eu + image='${{ steps.meta.outputs.primary }}' + docker volume create "${DATA_VOLUME}" >/dev/null + docker volume create "${BOOKS_VOLUME}" >/dev/null + + running="" + if docker ps --format '{{.Names}}' | grep -qx "${SERVICE_CONTAINER}"; then + running="yes" + fi + if [ -n "$running" ] && [ '${{ inputs.refresh_service }}' != 'true' ]; then + echo "${SERVICE_CONTAINER} is already running; leaving it alone so nothing" + echo "in flight is interrupted. Re-run with refresh_service=true to move it" + echo "onto the image just built." + else + docker rm -f "${SERVICE_CONTAINER}" >/dev/null 2>&1 || true + docker run -d \ + --name "${SERVICE_CONTAINER}" \ + --restart unless-stopped \ + -v "${DATA_VOLUME}:/data" \ + -v "${BOOKS_VOLUME}:/books" \ + -e DASHSCOPE_API_KEY \ + --entrypoint sleep \ + "$image" infinity + echo "${SERVICE_CONTAINER} is up on $image" + fi + docker exec "${SERVICE_CONTAINER}" lean-bench doctor --skip-llm || true + + echo "images kept on this runner:" + docker images "${LOCAL_IMAGE}" --format ' {{.Repository}}:{{.Tag}} {{.Size}} {{.CreatedSince}}' + env: + DASHSCOPE_API_KEY: ${{ secrets.DASHSCOPE_API_KEY }} + + - name: Summary + if: always() + run: | + { + echo "### lean-bench image" + echo "" + echo "Kept on the runner (not cleaned up):" + for tag in ${{ steps.meta.outputs.local_tags }}; do echo "- \`$tag\`"; done + echo "" + if [ '${{ steps.push.outcome }}' = 'success' ]; then + echo "Pushed to the registry:" + for tag in ${{ steps.meta.outputs.remote_tags }}; do echo "- \`$tag\`"; done + else + echo "**Not pushed.** Gitea's container registry rejects the built-in" + echo "Actions token. Create a personal access token with \`write:package\`" + echo "and add it as the \`REGISTRY_TOKEN\` secret (and \`REGISTRY_USER\`" + echo "as a variable if it is not \`${{ github.actor }}\`). The image is on" + echo "the runner either way." + fi + echo "" + echo "Submit a book to the container that is already running:" + echo '```' + echo "docker cp your-book.pdf lean-bench:/books/" + echo "docker exec -e DASHSCOPE_API_KEY=sk-... lean-bench \\" + echo " lean-bench ingest /books/your-book.pdf --pages 30-60" + echo "docker exec lean-bench lean-bench stats" + echo '```' + } >> "$GITHUB_STEP_SUMMARY" diff --git a/.gitea/workflows/ingest.yml b/.gitea/workflows/ingest.yml new file mode 100644 index 0000000..fcb640d --- /dev/null +++ b/.gitea/workflows/ingest.yml @@ -0,0 +1,129 @@ +name: ingest + +# Submit a book and run the whole pipeline, inside the long-lived `lean-bench` +# container the docker workflow leaves running on this runner. The book is copied +# into the `lean-bench-books` volume and the items land in `lean-bench-data`, so +# everything survives the job and is there for the next one -- `--pages` of the +# same book can be redone later without re-ingesting the rest. +# +# The source can be a path inside the repository or a URL the runner can reach. +# Set the DASHSCOPE_API_KEY secret (or whichever key your provider uses) first. + +on: + workflow_dispatch: + inputs: + source: + description: "Path in the repo, or a URL, of the book to ingest" + required: true + pages: + description: "Page range, e.g. 30-60 (blank = the whole book)" + required: false + default: "" + title: + description: "Human title for this source" + required: false + default: "" + verify: + description: "Run Lean verification (slow, but the point)" + required: false + default: "true" + image: + description: "Image to use if the lean-bench container is not running" + required: false + default: "lean-bench:latest" + +env: + SERVICE_CONTAINER: lean-bench + DATA_VOLUME: lean-bench-data + BOOKS_VOLUME: lean-bench-books + +jobs: + ingest: + runs-on: ubuntu-latest + timeout-minutes: 720 + steps: + - uses: actions/checkout@v4 + + - name: Fetch the source + id: source + run: | + set -eu + mkdir -p books out + src='${{ inputs.source }}' + case "$src" in + http://*|https://*) + name="$(basename "${src%%\?*}")" + curl -fsSL --retry 3 -o "books/$name" "$src" + ;; + *) + [ -f "$src" ] || { echo "::error::$src is not in the repository"; exit 1; } + name="$(basename "$src")" + cp "$src" "books/$name" + ;; + esac + ls -l "books/$name" + echo "name=$name" >> "$GITHUB_OUTPUT" + + - name: Make sure the lean-bench container is up + run: | + set -eu + docker volume create "${DATA_VOLUME}" >/dev/null + docker volume create "${BOOKS_VOLUME}" >/dev/null + if ! docker ps --format '{{.Names}}' | grep -qx "${SERVICE_CONTAINER}"; then + image='${{ inputs.image }}' + if ! docker image inspect "$image" >/dev/null 2>&1; then + echo "::error::No $image on this runner. Run the docker workflow first." + exit 1 + fi + docker rm -f "${SERVICE_CONTAINER}" >/dev/null 2>&1 || true + docker run -d --name "${SERVICE_CONTAINER}" --restart unless-stopped \ + -v "${DATA_VOLUME}:/data" -v "${BOOKS_VOLUME}:/books" \ + --entrypoint sleep "$image" infinity + fi + docker exec "${SERVICE_CONTAINER}" lean-bench doctor --skip-llm + + - name: Ingest + env: + DASHSCOPE_API_KEY: ${{ secrets.DASHSCOPE_API_KEY }} + LEANBENCH_LLM_MODEL: ${{ vars.LEANBENCH_LLM_MODEL }} + LEANBENCH_LLM_BASE_URL: ${{ vars.LEANBENCH_LLM_BASE_URL }} + run: | + set -eu + [ -n "${DASHSCOPE_API_KEY:-}" ] || { echo "::error::Set the DASHSCOPE_API_KEY secret"; exit 1; } + name='${{ steps.source.outputs.name }}' + docker cp "books/$name" "${SERVICE_CONTAINER}:/books/$name" + + set -- ingest "/books/$name" \ + --export "/data/${name%.*}.items.jsonl" \ + --detections "/data/${name%.*}.detections.jsonl" + # `cond && set -- ...` would abort the step under `set -e` whenever the + # condition is false, so these are spelled out. + if [ -n '${{ inputs.pages }}' ]; then set -- "$@" --pages '${{ inputs.pages }}'; fi + if [ -n '${{ inputs.title }}' ]; then set -- "$@" --title '${{ inputs.title }}'; fi + if [ '${{ inputs.verify }}' != 'true' ]; then set -- "$@" --no-verify; fi + + docker exec \ + -e DASHSCOPE_API_KEY \ + -e LEANBENCH_LLM_MODEL \ + -e LEANBENCH_LLM_BASE_URL \ + "${SERVICE_CONTAINER}" lean-bench "$@" + docker exec "${SERVICE_CONTAINER}" lean-bench stats + + - name: Collect a copy of the results + if: always() + run: | + set -eu + # The originals stay in the lean-bench-data volume; these are just copies + # for the artifact. + docker exec "${SERVICE_CONTAINER}" sh -c 'ls -1 /data' || true + for f in $(docker exec "${SERVICE_CONTAINER}" sh -c 'ls -1 /data' 2>/dev/null || true); do + docker cp "${SERVICE_CONTAINER}:/data/$f" "out/$f" || true + done + ls -l out || true + + - uses: actions/upload-artifact@v3 + if: always() + with: + name: leanbench-${{ steps.source.outputs.name }} + path: out/ + retention-days: 30 diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..33e3fe9 --- /dev/null +++ b/.gitignore @@ -0,0 +1,39 @@ +# Source books and anything derived from their full text: copyrighted material, +# never committed. Only code and the short Lean datasets belong in git. +*.pdf +*.epub +*.djvu +books/ +textbook/*.txt +*_raw.txt + +# Local state +.env +data/ +database/ +results/ +out/ +*.sqlite +*.sqlite-wal +*.sqlite-shm + +# Lean build output +lean/.lake/ +lean/lake-packages/ +lean/build/ + +# Python +__pycache__/ +*.py[cod] +.venv/ +venv/ +*.egg-info/ +.pytest_cache/ +.ruff_cache/ +.coverage +htmlcov/ + +# Editors and OS +.DS_Store +.idea/ +.vscode/ diff --git a/README.md b/README.md new file mode 100644 index 0000000..35f777e --- /dev/null +++ b/README.md @@ -0,0 +1,199 @@ +# lean-bench + +Point it at a mathematics book. It finds the mathematics worth formalizing, turns +each piece into Lean 4, asks the Lean compiler whether that Lean is legal, and files +the result in a database indexed by the types each item is about. + +``` +book.pdf ──▶ detect ──▶ formalize ──▶ verify (Lean) ──▶ indexed database + │ │ │ │ + spans across Lean 4 code statements for searchable by the + pages, two per item open problems, types an item takes + kinds of hit proofs when given in, or is +``` + +## The four pieces + +**1. The LLM connector** — `leanbench/llm.py`. Stdlib only, no SDK. Any +OpenAI-compatible `/chat/completions` endpoint: DashScope/Qwen by default, equally +happy with vLLM, Ollama, DeepSeek or OpenAI. Retries, JSON coercion, and a +`ScriptedClient` so the rest of the code is testable without a network. + +**2. Detection** — `leanbench/detect.py`. For every page, two questions: + +- *Case 1, `statement`* — is there mathematics asserted as fact here? A definition, + theorem, proposition, lemma, corollary, a precise claim inside an example. The + kind is deliberately **not** recorded: all that matters is that it can become + useful Lean 4. +- *Case 2, `problem`* — is there a question left open to the reader? An exercise, an + unsolved problem. Formalizable too, as a statement with no proof. + +Each hit is returned as a span: + +```json +{"kind": "statement", "label": "Theorem 1.6.3", + "start_page": 41, "start_index": 1832, + "end_page": 42, "end_index": 96, + "spans_pages": 2, "truncated": false} +``` + +Indices are character offsets into that page's text, so several hits per page are +normal and each is addressable. **If a passage runs off the bottom of the page**, the +model says so, the page is re-read together with the next one, and that repeats — +so an item spread over three pages is found whole, not in fragments. The limit is +`--window` (default 3); an item still cut off at the limit is stored with +`truncated: true` rather than silently clipped. + +Models cannot count characters, so they are never asked to. They quote the first and +last few words verbatim, and `leanbench/anchors.py` turns those quotes into exact +offsets — exact match, then whitespace-insensitive match (PDF line breaks), then +fuzzy match. A quote that matches none of the three is reported as unresolved, never +guessed at. + +**3. Formalization and verification** — `leanbench/formalize.py`, `leanbench/verify.py`. +Every detected span becomes one self-contained Lean 4 declaration, which Lean then +judges: + +| item | what is checked | +|---|---| +| `statement` with a proof in the book | the whole declaration, proof included | +| `statement` without one | the statement elaborates; the proof is `sorry` | +| `problem` (case 2) | **statement legality only** — there is no proof to check | + +When Lean rejects the code, its diagnostics go straight back to the model for one +more attempt (`LEANBENCH_REPAIR_ATTEMPTS`). Two flags are kept apart on purpose: +`verified` means Lean accepted it, `proved` means Lean accepted it **with no `sorry` +anywhere**. + +**4. The database** — `leanbench/db.py`. SQLite, indexed, plus a JSONL export for +git. Each item records: + +- its Lean code, imports, and what Lean said; +- **the types it is indexed under** — for a proof, the types of its inputs and + hypotheses; for a definition, the type of the thing being defined. That is what + makes one item relevant to another, and it is what search runs on; +- **which book** (sha256 of the file) and **which page it starts on**. + +So any part of any source can be redone without touching anything else: + +```bash +lean-bench drop --source --pages 30-45 # forget those pages +lean-bench ingest book.pdf --pages 30-45 # and do them again +``` + +Re-ingesting does this for you: a run over pages 30–45 deletes exactly the items +filed under those pages first. Item identity is `sha1(book hash + span)`, so the same +passage of the same book is always the same row. + +## Quick start + +```bash +pip install -e ".[dev]" +cp .env.example .env # put your API key in it +lean-bench doctor # is the LLM reachable? is Lean installed? + +lean-bench ingest book.pdf --pages 30-60 --export datasets/items.jsonl +lean-bench stats +lean-bench search --type Int --role input --kind problem --code +``` + +Lean verification needs a Lean project with Mathlib built. Without one, everything +still runs and nothing claims to be verified. With Docker you get one for free: + +```bash +docker run --rm -e DASHSCOPE_API_KEY \ + -v "$PWD/books:/books" -v "$PWD/data:/data" \ + gitea.entroia.com/zayden-jung/lean-bench:latest \ + ingest /books/book.pdf --pages 30-60 +``` + +## Commands + +| | | +|---|---| +| `lean-bench doctor` | check the LLM, Lean and PDF wiring | +| `lean-bench ingest SRC` | the whole pipeline; `--pages`, `--no-verify`, `--export` | +| `lean-bench detect SRC` | detection only — no Lean, no database | +| `lean-bench search` | by `--type`, `--role`, `--kind`, `--page`, `--text` | +| `lean-bench sources` / `stats` | what is in the database | +| `lean-bench export OUT` | JSONL dump | +| `lean-bench drop --source H` | forget a source, or `--pages` of it | + +Sources may be PDFs, or `.txt` files whose pages are separated by form feeds +(what `pdftotext` emits). + +## Configuration + +All environment, all optional, all overridable in `.env`: + +| variable | default | +|---|---| +| `LEANBENCH_LLM_MODEL` | `qwen-max` | +| `LEANBENCH_LLM_BASE_URL` | DashScope's OpenAI-compatible endpoint | +| `LEANBENCH_LLM_API_KEY_ENV` | `DASHSCOPE_API_KEY` | +| `LEANBENCH_DATABASE` | `database/leanbench.sqlite` | +| `LEANBENCH_MAX_WINDOW_PAGES` | `3` | +| `LEANBENCH_REPAIR_ATTEMPTS` | `1` | +| `LEANBENCH_LEAN_COMMAND` | `["lake","env","lean"]` | +| `LEANBENCH_LEAN_PROJECT` | unset (a Lake project with Mathlib built) | + +## CI/CD + +Three Gitea Actions workflows, for the `ubuntu-latest` / `ubuntu-24.04` / +`ubuntu-22.04` runner: + +- **`ci`** — lint and the full test suite on every push. No network, no Lean: the + LLM is scripted and the compiler is faked, so it is fast and deterministic. It + installs into a venv, because the runner image's Python is Debian-managed and + refuses installs into itself (PEP 668). +- **`docker`** — builds the Lean 4 + Mathlib image and checks it honestly: Mathlib + really elaborates inside it, and a scripted end-to-end run proves the pipeline + writes a verified item to the database with the real compiler attached. The + Mathlib revision and the Lean toolchain are resolved *against each other* at + build time (`docker/resolve-mathlib.sh`), because hardcoding both is how image + builds break two months later. +- **`ingest`** — run by hand: submit a book (repo path or URL), pick a page range. + It runs inside the long-lived container and the results stay in the volume; the + artifact is just a copy. + +### The build is not a throwaway + +When the `docker` workflow finishes, the runner host still has everything it built: + +| what | name | +|---|---| +| image | `lean-bench:latest`, `lean-bench:sha-` | +| layer cache | kept — nothing ever runs `docker system prune`, so a rebuild reuses Mathlib instead of refetching it | +| idle container | `lean-bench` (`--restart unless-stopped`) | +| volumes | `lean-bench-data` (the database), `lean-bench-books` | + +So you can use it directly, with no rebuild and no `docker run`: + +```bash +docker cp your-book.pdf lean-bench:/books/ +docker exec -e DASHSCOPE_API_KEY=sk-... lean-bench \ + lean-bench ingest /books/your-book.pdf --pages 30-60 +docker exec lean-bench lean-bench search --type Int --role input --code +``` + +A rebuild leaves a running container alone rather than interrupting whatever is in +it; re-run the workflow with `refresh_service=true` to move it onto the new image. + +### Secrets + +| name | kind | why | +|---|---|---| +| `DASHSCOPE_API_KEY` | secret | the LLM, for `ingest` | +| `REGISTRY_TOKEN` | secret | a personal access token with `write:package`. **The built-in Actions token is not accepted by Gitea's container registry**, so without this the image is not pushed — it is still built and kept on the runner, and the build does not fail | +| `REGISTRY_USER` | variable | only if the pushing account is not the actor | + +## What this does not claim + +Lean says whether the code is *legal*, and — when there is a proof — whether it is +*correct*. It cannot say the Lean faithfully renders what the book meant. A +well-formed statement of the wrong theorem passes every check here. `confidence`, +`match_method` and the stored `source_text` are kept so that judgement stays +possible; spot-check before trusting a dataset. + +Source books are not committed. `.gitignore` keeps PDFs and extracted text out of +the repository; only code and the short Lean datasets belong in git. diff --git a/docker/Dockerfile b/docker/Dockerfile new file mode 100644 index 0000000..5f2f8cb --- /dev/null +++ b/docker/Dockerfile @@ -0,0 +1,72 @@ +# lean-bench: Lean 4 + Mathlib + the ingestion pipeline, in one image. +# +# Build from the repository root: +# docker build -f docker/Dockerfile -t lean-bench:latest . +# +# The Mathlib revision and the Lean toolchain are resolved against each other at +# build time (see docker/resolve-mathlib.sh); override with --build-arg MATHLIB_REV=v4.33.0. +FROM ubuntu:24.04 + +LABEL org.opencontainers.image.title="lean-bench" \ + org.opencontainers.image.description="Detect formalizable mathematics in books and turn it into verified, indexed Lean 4" \ + org.opencontainers.image.source="https://gitea.entroia.com/Zayden-Jung/lean-bench" + +ENV DEBIAN_FRONTEND=noninteractive \ + ELAN_HOME=/opt/elan \ + PATH=/opt/elan/bin:$PATH \ + PYTHONUNBUFFERED=1 + +RUN apt-get update && apt-get install -y --no-install-recommends \ + build-essential ca-certificates curl git jq sqlite3 \ + libgmp-dev python3 python3-pip \ + && rm -rf /var/lib/apt/lists/* + +# -- elan (the Lean toolchain manager) -------------------------------------- +RUN curl -fsSL https://raw.githubusercontent.com/leanprover/elan/master/elan-init.sh -o /tmp/elan-init.sh \ + && sh /tmp/elan-init.sh -y --no-modify-path --default-toolchain none \ + && rm /tmp/elan-init.sh \ + && elan --version + +WORKDIR /opt/lean-bench + +# -- Lean project: toolchain, Mathlib, and its prebuilt cache --------------- +# Its own layer, because it is the expensive one and it changes rarely. +ARG MATHLIB_REV=auto +COPY lean/ lean/ +COPY docker/resolve-mathlib.sh docker/resolve-mathlib.sh +RUN chmod +x docker/resolve-mathlib.sh && docker/resolve-mathlib.sh lean "$MATHLIB_REV" +RUN elan toolchain install "$(cat lean/lean-toolchain)" \ + && elan default "$(cat lean/lean-toolchain)" +RUN cd lean \ + && lake update \ + && (lake exe cache get || echo "no Mathlib cache available; building from source") \ + && lake build \ + && lake env lean --version + +# -- the pipeline ----------------------------------------------------------- +COPY pyproject.toml README.md ./ +COPY leanbench/ leanbench/ +RUN pip install --break-system-packages --no-cache-dir . \ + && lean-bench --help > /dev/null + +COPY docker/entrypoint.sh /usr/local/bin/leanbench-entrypoint +RUN chmod +x /usr/local/bin/leanbench-entrypoint + +ENV LEANBENCH_LEAN_PROJECT=/opt/lean-bench/lean \ + LEANBENCH_LEAN_COMMAND='["lake","env","lean"]' \ + LEANBENCH_DATABASE=/data/leanbench.sqlite \ + LEANBENCH_LLM_MODEL=qwen-max \ + LEANBENCH_LLM_BASE_URL=https://dashscope.aliyuncs.com/compatible-mode/v1 \ + LEANBENCH_LLM_API_KEY_ENV=DASHSCOPE_API_KEY + +# /books is where you put the source you are submitting; /data holds the database. +VOLUME ["/data", "/books"] +WORKDIR /work + +# Prove the environment works while the image is still being built. +RUN mkdir -p /work && cd /opt/lean-bench/lean \ + && printf 'import Mathlib\ntheorem image_smoke (n : Nat) : n + 0 = n := by simp\n' > /tmp/Smoke.lean \ + && lake env lean /tmp/Smoke.lean && rm /tmp/Smoke.lean + +ENTRYPOINT ["leanbench-entrypoint"] +CMD ["doctor"] diff --git a/docker/docker-compose.yml b/docker/docker-compose.yml new file mode 100644 index 0000000..1cff794 --- /dev/null +++ b/docker/docker-compose.yml @@ -0,0 +1,25 @@ +# Run the pipeline against a book on your machine: +# +# export DASHSCOPE_API_KEY=sk-... +# mkdir -p books data && cp ~/my-book.pdf books/ +# docker compose -f docker/docker-compose.yml run --rm leanbench \ +# ingest /books/my-book.pdf --pages 30-60 --export /data/items.jsonl +# +# The database persists in ./data, so re-running a page range replaces just those items. +services: + leanbench: + image: ${LEANBENCH_IMAGE:-gitea.entroia.com/zayden-jung/lean-bench:latest} + build: + context: .. + dockerfile: docker/Dockerfile + args: + MATHLIB_REV: ${MATHLIB_REV:-auto} + environment: + DASHSCOPE_API_KEY: ${DASHSCOPE_API_KEY:?set DASHSCOPE_API_KEY (or your provider key)} + LEANBENCH_LLM_MODEL: ${LEANBENCH_LLM_MODEL:-qwen-max} + LEANBENCH_LLM_BASE_URL: ${LEANBENCH_LLM_BASE_URL:-https://dashscope.aliyuncs.com/compatible-mode/v1} + LEANBENCH_DATABASE: /data/leanbench.sqlite + volumes: + - ../books:/books + - ../data:/data + command: ["doctor"] diff --git a/docker/entrypoint.sh b/docker/entrypoint.sh new file mode 100755 index 0000000..b6f2dd7 --- /dev/null +++ b/docker/entrypoint.sh @@ -0,0 +1,32 @@ +#!/usr/bin/env bash +# Thin front door for the image. `docker run ...` +# works, `submit` is an alias for `ingest`, and anything else runs verbatim. +set -euo pipefail + +if [ -f /run/secrets/llm_api_key ] && [ -z "${DASHSCOPE_API_KEY:-}" ]; then + DASHSCOPE_API_KEY="$(cat /run/secrets/llm_api_key)" + export DASHSCOPE_API_KEY +fi + +mkdir -p "$(dirname "${LEANBENCH_DATABASE:-/data/leanbench.sqlite}")" 2>/dev/null || true + +case "${1:-doctor}" in + bash | sh) + shift + exec bash "$@" + ;; + submit) + shift + exec lean-bench ingest "$@" + ;; + doctor | ingest | detect | search | export | sources | stats | drop) + exec lean-bench "$@" + ;; + lean-bench) + shift + exec lean-bench "$@" + ;; + *) + exec "$@" + ;; +esac diff --git a/docker/resolve-mathlib.sh b/docker/resolve-mathlib.sh new file mode 100755 index 0000000..b1098f0 --- /dev/null +++ b/docker/resolve-mathlib.sh @@ -0,0 +1,83 @@ +#!/usr/bin/env bash +# Pin Mathlib and the Lean toolchain to each other. +# +# Mathlib tags do not exist for every Lean point release, so hardcoding both is how +# image builds break six weeks later. Instead: choose a Mathlib revision, then take +# *its* lean-toolchain as the truth and write both into the Lean project. +# +# resolve-mathlib.sh [revision] +# +# revision defaults to "auto": the newest Mathlib tag whose toolchain matches the +# toolchain already in the project, else the newest tag, else master. +set -euo pipefail + +PROJECT="${1:?usage: resolve-mathlib.sh [revision]}" +REQUESTED="${2:-auto}" +REPO="https://github.com/leanprover-community/mathlib4.git" +RAW="https://raw.githubusercontent.com/leanprover-community/mathlib4" +TOOLCHAIN_FILE="$PROJECT/lean-toolchain" +LAKEFILE="$PROJECT/lakefile.toml" + +toolchain_of() { # toolchain_of -> prints e.g. leanprover/lean4:v4.33.0 + curl -fsSL --retry 3 --max-time 30 "$RAW/$1/lean-toolchain" 2>/dev/null | tr -d '[:space:]' +} + +candidates() { + git ls-remote --tags --refs "$REPO" 'v*' 2>/dev/null \ + | awk '{print $2}' | sed 's#refs/tags/##' \ + | grep -E '^v[0-9]+\.[0-9]+\.[0-9]+$' \ + | sort -V -r +} + +wanted="$(tr -d '[:space:]' < "$TOOLCHAIN_FILE" 2>/dev/null || true)" +rev="" + +if [ "$REQUESTED" != "auto" ]; then + rev="$REQUESTED" + echo "resolve-mathlib: using requested revision $rev" +else + echo "resolve-mathlib: looking for a Mathlib tag built on ${wanted:-}" + all="$(candidates || true)" + if [ -n "$wanted" ]; then + # only the tags sharing a major.minor with what we asked for are plausible + series="$(printf '%s' "$wanted" | sed -E 's#.*:v([0-9]+\.[0-9]+).*#v\1#')" + for tag in $(printf '%s\n' "$all" | grep -E "^${series}\." || true); do + if [ "$(toolchain_of "$tag")" = "$wanted" ]; then + rev="$tag" + echo "resolve-mathlib: $tag matches $wanted exactly" + break + fi + done + if [ -z "$rev" ]; then + rev="$(printf '%s\n' "$all" | grep -E "^${series}\." | head -1 || true)" + [ -n "$rev" ] && echo "resolve-mathlib: no exact match; newest in $series series is $rev" + fi + fi + if [ -z "$rev" ]; then + rev="$(printf '%s\n' "$all" | head -1 || true)" + echo "resolve-mathlib: falling back to newest tag ${rev:-master}" + fi + rev="${rev:-master}" +fi + +actual="$(toolchain_of "$rev" || true)" +if [ -z "$actual" ]; then + echo "resolve-mathlib: could not read the toolchain for $rev; keeping ${wanted:-project default}" >&2 + actual="$wanted" +fi + +printf '%s\n' "$actual" > "$TOOLCHAIN_FILE" +# rewrite only the rev line inside the [[require]] block for mathlib +python3 - "$LAKEFILE" "$rev" <<'PY' +import re +import sys + +path, rev = sys.argv[1], sys.argv[2] +with open(path, encoding="utf-8") as stream: + text = stream.read() +text = re.sub(r'(?m)^rev\s*=\s*".*"$', f'rev = "{rev}"', text, count=1) +with open(path, "w", encoding="utf-8") as stream: + stream.write(text) +PY + +echo "resolve-mathlib: mathlib $rev on $actual" diff --git a/docker/smoke.py b/docker/smoke.py new file mode 100644 index 0000000..99de0f1 --- /dev/null +++ b/docker/smoke.py @@ -0,0 +1,134 @@ +"""End-to-end smoke test for the image, with the real Lean compiler attached. + +The LLM is scripted (the image has no API key and CI should not need one), but +everything else is real: pages are extracted, spans are resolved to page/index, +Lean elaborates the generated code, and the items land in an indexed database. + +Run it against a built image without mounting anything: + + docker run --rm -i --entrypoint python3 lean-bench:latest - < docker/smoke.py +""" +from __future__ import annotations + +import json +import os +import sys +import tempfile + +from leanbench.config import Settings +from leanbench.db import Database +from leanbench.llm import ScriptedClient +from leanbench.pipeline import build_checker, ingest + +BOOK_PAGE_ONE = "Theorem 1. For every natural number n, n + 0 = n.\n" +BOOK_PAGE_TWO = "Exercise 2. Show that 2 + 2 = 4.\n" + + +def detection(items: list[dict]) -> str: + return json.dumps({"items": items}) + + +def main() -> int: + workdir = tempfile.mkdtemp(prefix="leanbench-smoke-") + book = os.path.join(workdir, "book.txt") + with open(book, "w", encoding="utf-8") as stream: + stream.write(BOOK_PAGE_ONE + "\f" + BOOK_PAGE_TWO) + + client = ScriptedClient( + [ + detection( + [ + { + "kind": "statement", + "label": "Theorem 1", + "summary": "adding zero changes nothing", + "start_quote": "Theorem 1. For every natural number n", + "end_quote": "n + 0 = n.", + "continues_past_end": False, + } + ] + ), + detection( + [ + { + "kind": "problem", + "label": "Exercise 2", + "summary": "two plus two", + "start_quote": "Exercise 2. Show that", + "end_quote": "2 + 2 = 4.", + "continues_past_end": False, + } + ] + ), + json.dumps( + { + "skip": False, + "imports": ["Mathlib"], + "name": "add_zero_nat", + "lean_code": "theorem add_zero_nat (n : Nat) : n + 0 = n := by simp", + "has_proof": True, + "informal": "n + 0 = n", + "index_types": ["Nat"], + "confidence": 1.0, + } + ), + json.dumps( + { + "skip": False, + "imports": ["Mathlib"], + "name": "two_add_two", + "lean_code": "theorem two_add_two : (2 : Nat) + 2 = 4 := by sorry", + "has_proof": False, + "informal": "2 + 2 = 4", + "index_types": ["Nat"], + "confidence": 1.0, + } + ), + ] + ) + + settings = Settings.from_env(dotenv=None) + settings.database = os.path.join(workdir, "smoke.sqlite") + settings.repair_attempts = 0 + + checker = build_checker(settings) + if not getattr(checker, "available", lambda: False)(): + print("FAIL: this image has no working Lean; verification would be skipped") + return 1 + + database = Database(settings.database) + result = ingest( + book, + settings=settings, + client=client, + database=database, + checker=checker, + title="smoke", + ) + summary = result.summary() + print(json.dumps(summary, indent=2)) + + failures = [] + if summary["detected"] != 2: + failures.append(f"expected 2 detections, got {summary['detected']}") + if summary["verified"] != 2: + failures.append(f"Lean should have accepted both, got {summary['verified']}") + if summary["proved"] != 1: + failures.append(f"exactly one item carries a real proof, got {summary['proved']}") + if not database.search(types=["Nat"]): + failures.append("the type index is empty") + problems = database.search(kind="problem") + if len(problems) != 1 or problems[0]["verify_mode"] != "statement": + failures.append("the open problem should be checked as a statement only") + database.close() + + for line in failures: + print("FAIL:", line) + if failures: + return 1 + print("pipeline + lean end to end: ok") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/examples/sample-book.txt b/examples/sample-book.txt new file mode 100644 index 0000000..76a3184 --- /dev/null +++ b/examples/sample-book.txt @@ -0,0 +1,21 @@ +1.6 Divisibility in the Integers + +Definition 1.6.1. Let a and b be integers. We say that a divides b, written a | b, +if there exists an integer k such that b = a k. + +Theorem 1.6.2. Let a, b and c be integers. If a divides b and a divides c, then a +divides b + c. + +Proof. Write b = a k and c = a m for integers k and m. Then b + c = a (k + m), so a +divides b + c. + +Theorem 1.6.3. Let a and b be integers, not both zero. Then the set of common +divisors of a and b has a greatest element, written gcd(a, b), and there exist + integers s and t with gcd(a, b) = s a + t b. + +Exercise 1.6.4. Show that if a divides b and b divides a, then a = b or a = -b. + +Exercise 1.6.5. Prove that the sum of two odd integers is even. + +Exercise 1.6.6. Determine whether every integer greater than 2 is the sum of two +primes. diff --git a/lean/LeanBench.lean b/lean/LeanBench.lean new file mode 100644 index 0000000..2dfe4b1 --- /dev/null +++ b/lean/LeanBench.lean @@ -0,0 +1,18 @@ +/- +The verification environment. + +Candidate declarations are never added to this file: each one is elaborated in its +own scratch file against this project's imports, so a broken candidate can never +break the environment the next candidate is checked in. + +Anything you want every candidate to see -- a shared convention, a local notation -- +belongs here. +-/ +import Mathlib + +namespace LeanBench + +/-- Sanity check: if this fails, the toolchain and Mathlib do not agree. -/ +theorem environment_ok : (2 : ℕ) + 2 = 4 := by norm_num + +end LeanBench diff --git a/lean/lakefile.toml b/lean/lakefile.toml new file mode 100644 index 0000000..de43134 --- /dev/null +++ b/lean/lakefile.toml @@ -0,0 +1,12 @@ +name = "leanbench" +defaultTargets = ["LeanBench"] + +# The revision is rewritten by docker/resolve-mathlib.sh at image build time, so the +# Mathlib revision and the Lean toolchain always agree. Edit MATHLIB_REV, not this. +[[require]] +name = "mathlib" +git = "https://github.com/leanprover-community/mathlib4.git" +rev = "v4.33.1" + +[[lean_lib]] +name = "LeanBench" diff --git a/lean/lean-toolchain b/lean/lean-toolchain new file mode 100644 index 0000000..a8afa7d --- /dev/null +++ b/lean/lean-toolchain @@ -0,0 +1 @@ +leanprover/lean4:v4.33.1 diff --git a/leanbench/__init__.py b/leanbench/__init__.py new file mode 100644 index 0000000..895ce20 --- /dev/null +++ b/leanbench/__init__.py @@ -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__"] diff --git a/leanbench/anchors.py b/leanbench/anchors.py new file mode 100644 index 0000000..60162aa --- /dev/null +++ b/leanbench/anchors.py @@ -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)) diff --git a/leanbench/cli.py b/leanbench/cli.py new file mode 100644 index 0000000..e63785f --- /dev/null +++ b/leanbench/cli.py @@ -0,0 +1,330 @@ +"""Command line: ``lean-bench ``. + + lean-bench doctor is everything wired up? + lean-bench ingest book.pdf --pages 30-45 + lean-bench detect book.pdf --pages 30-31 --out detections.jsonl + lean-bench search --type Int --role input --kind problem + lean-bench export datasets/items.jsonl + lean-bench drop --source --pages 30-45 redo part of a book +""" +from __future__ import annotations + +import argparse +import json +import sys + +from .config import Settings +from .db import Database +from .detect import detect_document, write_jsonl +from .document import Document +from .llm import LLMError, build_client +from .pipeline import build_checker, ingest + + +def _pages(value: str | None) -> tuple[int | None, int | None]: + if not value: + return None, None + text = value.strip() + if "-" in text: + first, _, last = text.partition("-") + return (int(first) if first.strip() else None, + int(last) if last.strip() else None) + page = int(text) + return page, page + + +def _say(message: str) -> None: + print(message, flush=True) + + +def _add_common(parser: argparse.ArgumentParser) -> None: + parser.add_argument("--db", help="SQLite path (default: $LEANBENCH_DATABASE)") + parser.add_argument("--env", default=".env", help="dotenv file to load (default: .env)") + + +# -- commands -------------------------------------------------------------- + + +def cmd_doctor(args: argparse.Namespace) -> int: + settings = Settings.from_env(args.env) + print(f"database {settings.database}") + print(f"llm model {settings.llm.model}") + print(f"llm endpoint {settings.llm.base_url}") + print(f"llm key env {settings.llm.api_key_env}") + + try: + import pymupdf # noqa: F401 + print("pdf backend pymupdf") + except ImportError: + try: + import pypdf # noqa: F401 + print("pdf backend pypdf (install pymupdf for better text)") + except ImportError: + print("pdf backend MISSING - pip install pymupdf") + + checker = build_checker(settings) + available = getattr(checker, "available", lambda: False)() + print(f"lean {'ready: ' + ' '.join(settings.lean.command) if available else 'not available (verification will be skipped)'}") + if settings.lean.project: + print(f"lean project {settings.lean.project}") + + if args.skip_llm: + return 0 + try: + client = build_client(settings.llm) + reply = client.complete( + [{"role": "user", "content": "Reply with the single word: ready"}], + max_tokens=16, + ) + print(f"llm ready ({reply.strip()[:40]})") + except LLMError as error: + print(f"llm FAILED: {error}") + return 1 + return 0 + + +def cmd_ingest(args: argparse.Namespace) -> int: + settings = Settings.from_env(args.env) + if args.db: + settings.database = args.db + if args.window: + settings.max_window_pages = args.window + first, last = _pages(args.pages) + + database = Database(settings.database) + try: + result = ingest( + args.source, + settings=settings, + database=database, + checker=build_checker(settings, enabled=not args.no_verify), + first_page=first, + last_page=last, + title=args.title or "", + replace=not args.keep_existing, + store_results=not args.no_store, + progress=None if args.quiet else _say, + ) + summary = result.summary() + print(json.dumps(summary, indent=2)) + if args.detections: + write_jsonl(args.detections, result.report.detections) + print(f"detections -> {args.detections}") + if args.export: + count = database.export_jsonl(args.export, source_hash=result.document.sha256) + print(f"exported {count} item(s) -> {args.export}") + finally: + database.close() + return 0 + + +def cmd_detect(args: argparse.Namespace) -> int: + settings = Settings.from_env(args.env) + first, last = _pages(args.pages) + document = Document.load(args.source) + report = detect_document( + build_client(settings.llm), + document, + first_page=first, + last_page=last, + max_window_pages=args.window or settings.max_window_pages, + progress=None if args.quiet else _say, + ) + print(report.summary()) + if args.out: + write_jsonl(args.out, report.detections) + print(f"-> {args.out}") + else: + for detection in report.detections: + print(json.dumps(detection.as_dict(), ensure_ascii=False)) + return 0 + + +def cmd_search(args: argparse.Namespace) -> int: + settings = Settings.from_env(args.env) + database = Database(args.db or settings.database) + try: + rows = database.search( + types=args.type or (), + roles=args.role or (), + kind=args.kind, + source_hash=args.source, + page=args.page, + verified=args.verified, + proved=args.proved, + text=args.text, + match_all_types=args.all_types, + limit=args.limit, + ) + if args.json: + print(json.dumps(rows, indent=2, ensure_ascii=False)) + return 0 + if not rows: + print("no matches") + return 0 + for row in rows: + types = ", ".join(f"{t['name']}({t['role']})" for t in row["types"]) or "-" + flag = "proved" if row["proved"] else "verified" if row["verified"] else "unverified" + print(f"{row['id']} {row['kind']:<9} p{row['start_page']:<4} {flag:<10} " + f"{row['label'] or row['lean_name']}") + print(f" types: {types}") + if args.code: + for line in (row["lean_code"] or "").splitlines(): + print(" " + line) + finally: + database.close() + return 0 + + +def cmd_sources(args: argparse.Namespace) -> int: + settings = Settings.from_env(args.env) + database = Database(args.db or settings.database) + try: + rows = database.sources() + if args.json: + print(json.dumps(rows, indent=2, ensure_ascii=False)) + return 0 + for row in rows: + print(f"{row['hash'][:16]} {row['items']:>5} items p{row['pages']:<5} " + f"{row['title']} ({row['extractor']}, {row['ingested_at']})") + finally: + database.close() + return 0 + + +def cmd_stats(args: argparse.Namespace) -> int: + settings = Settings.from_env(args.env) + database = Database(args.db or settings.database) + try: + stats = database.stats() + stats["top_types"] = database.type_counts(limit=args.top) + print(json.dumps(stats, indent=2)) + finally: + database.close() + return 0 + + +def cmd_export(args: argparse.Namespace) -> int: + settings = Settings.from_env(args.env) + database = Database(args.db or settings.database) + try: + count = database.export_jsonl( + args.out, source_hash=args.source, include_source_text=args.with_text + ) + print(f"exported {count} item(s) -> {args.out}") + finally: + database.close() + return 0 + + +def cmd_drop(args: argparse.Namespace) -> int: + settings = Settings.from_env(args.env) + database = Database(args.db or settings.database) + try: + first, last = _pages(args.pages) + if first or last: + removed = database.delete_pages(args.source, first or 1, last or 10**9) + else: + removed = database.delete_source(args.source) + print(f"removed {removed} item(s)") + finally: + database.close() + return 0 + + +# -- wiring ---------------------------------------------------------------- + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(prog="lean-bench", description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + subparsers = parser.add_subparsers(dest="command", required=True) + + doctor = subparsers.add_parser("doctor", help="check LLM, Lean and PDF wiring") + _add_common(doctor) + doctor.add_argument("--skip-llm", action="store_true", help="do not call the model") + doctor.set_defaults(func=cmd_doctor) + + ingest_parser = subparsers.add_parser("ingest", help="run the whole pipeline on a source") + _add_common(ingest_parser) + ingest_parser.add_argument("source", help="PDF (or .txt with form-feed page breaks)") + ingest_parser.add_argument("--pages", help="page or range, e.g. 30-45") + ingest_parser.add_argument("--title", help="human title for this source") + ingest_parser.add_argument("--window", type=int, help="max pages per detection window") + ingest_parser.add_argument("--no-verify", action="store_true", help="skip Lean") + ingest_parser.add_argument("--no-store", action="store_true", help="do not write the database") + ingest_parser.add_argument("--keep-existing", action="store_true", + help="do not remove this source's previous items first") + ingest_parser.add_argument("--detections", help="also write raw detections to this JSONL") + ingest_parser.add_argument("--export", help="export this source's items to JSONL afterwards") + ingest_parser.add_argument("--quiet", action="store_true") + ingest_parser.set_defaults(func=cmd_ingest) + + detect_parser = subparsers.add_parser("detect", help="detection only, no Lean, no database") + _add_common(detect_parser) + detect_parser.add_argument("source") + detect_parser.add_argument("--pages") + detect_parser.add_argument("--window", type=int) + detect_parser.add_argument("--out", help="write detections to this JSONL") + detect_parser.add_argument("--quiet", action="store_true") + detect_parser.set_defaults(func=cmd_detect) + + search = subparsers.add_parser("search", help="find items by the types they are about") + _add_common(search) + search.add_argument("--type", action="append", help="type name; repeatable") + search.add_argument("--role", action="append", choices=["input", "goal", "own"]) + search.add_argument("--kind", choices=["statement", "problem"]) + search.add_argument("--source", help="source hash") + search.add_argument("--page", type=int) + search.add_argument("--text", help="full-text query") + search.add_argument("--all-types", action="store_true", help="require every --type") + search.add_argument("--verified", action="store_true", default=None) + search.add_argument("--proved", action="store_true", default=None) + search.add_argument("--code", action="store_true", help="print the Lean too") + search.add_argument("--json", action="store_true") + search.add_argument("--limit", type=int, default=25) + search.set_defaults(func=cmd_search) + + sources = subparsers.add_parser("sources", help="list ingested books") + _add_common(sources) + sources.add_argument("--json", action="store_true") + sources.set_defaults(func=cmd_sources) + + stats = subparsers.add_parser("stats", help="counts, and the most common types") + _add_common(stats) + stats.add_argument("--top", type=int, default=20) + stats.set_defaults(func=cmd_stats) + + export = subparsers.add_parser("export", help="dump items as JSONL") + _add_common(export) + export.add_argument("out") + export.add_argument("--source", help="only this source hash") + export.add_argument("--with-text", action="store_true", help="include source passages") + export.set_defaults(func=cmd_export) + + drop = subparsers.add_parser("drop", help="remove a source, or some of its pages") + _add_common(drop) + drop.add_argument("--source", required=True, help="source hash") + drop.add_argument("--pages", help="page or range; omit for the whole source") + drop.set_defaults(func=cmd_drop) + + return parser + + +def main(argv: list[str] | None = None) -> int: + args = build_parser().parse_args(argv) + try: + return args.func(args) + except FileNotFoundError as error: + print(f"not found: {error}", file=sys.stderr) + return 2 + except LLMError as error: + print(f"llm error: {error}", file=sys.stderr) + return 3 + except KeyboardInterrupt: + print("interrupted", file=sys.stderr) + return 130 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/leanbench/config.py b/leanbench/config.py new file mode 100644 index 0000000..cd7922c --- /dev/null +++ b/leanbench/config.py @@ -0,0 +1,116 @@ +"""Runtime configuration, read from the environment (and an optional .env file). + +Everything the pipeline needs to talk to an LLM or to Lean is settable without +touching code, so the same image runs against any OpenAI-compatible endpoint and +any Lean project. +""" +from __future__ import annotations + +import json +import os +from dataclasses import dataclass, field +from pathlib import Path + +DEFAULT_BASE_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1" +DEFAULT_MODEL = "qwen-max" +DEFAULT_API_KEY_ENV = "DASHSCOPE_API_KEY" + + +def load_dotenv(path: Path | str = ".env") -> None: + """Load ``KEY=VALUE`` lines into os.environ without overwriting what is set.""" + path = Path(path) + if not path.is_file(): + return + for line in path.read_text(encoding="utf-8").splitlines(): + line = line.strip() + if not line or line.startswith("#") or "=" not in line: + continue + key, _, value = line.partition("=") + os.environ.setdefault(key.strip(), value.strip().strip('"').strip("'")) + + +def _env(name: str, default: str) -> str: + value = os.environ.get(name) + return default if value is None or value == "" else value + + +def _env_int(name: str, default: int) -> int: + try: + return int(_env(name, str(default))) + except ValueError: + return default + + +def _env_float(name: str, default: float) -> float: + try: + return float(_env(name, str(default))) + except ValueError: + return default + + +@dataclass +class LLMSettings: + model: str = DEFAULT_MODEL + base_url: str = DEFAULT_BASE_URL + api_key_env: str = DEFAULT_API_KEY_ENV + timeout: float = 300.0 + max_retries: int = 4 + temperature: float = 0.0 + + @classmethod + def from_env(cls) -> LLMSettings: + return cls( + model=_env("LEANBENCH_LLM_MODEL", DEFAULT_MODEL), + base_url=_env("LEANBENCH_LLM_BASE_URL", DEFAULT_BASE_URL), + api_key_env=_env("LEANBENCH_LLM_API_KEY_ENV", DEFAULT_API_KEY_ENV), + timeout=_env_float("LEANBENCH_LLM_TIMEOUT", 300.0), + max_retries=_env_int("LEANBENCH_LLM_MAX_RETRIES", 4), + temperature=_env_float("LEANBENCH_LLM_TEMPERATURE", 0.0), + ) + + +@dataclass +class LeanSettings: + """How to invoke Lean. ``command`` is an argv prefix that is given a .lean path.""" + + command: list[str] = field(default_factory=lambda: ["lake", "env", "lean"]) + project: str | None = None + timeout: float = 120.0 + header: str = "import Mathlib\n" + + @classmethod + def from_env(cls) -> LeanSettings: + raw = _env("LEANBENCH_LEAN_COMMAND", '["lake","env","lean"]') + try: + command = json.loads(raw) + if not isinstance(command, list) or not all(isinstance(x, str) for x in command): + raise ValueError + except ValueError: + command = raw.split() + return cls( + command=command, + project=os.environ.get("LEANBENCH_LEAN_PROJECT") or None, + timeout=_env_float("LEANBENCH_LEAN_TIMEOUT", 120.0), + header=_env("LEANBENCH_LEAN_HEADER", "import Mathlib\n"), + ) + + +@dataclass +class Settings: + llm: LLMSettings = field(default_factory=LLMSettings) + lean: LeanSettings = field(default_factory=LeanSettings) + database: str = "database/leanbench.sqlite" + max_window_pages: int = 3 + repair_attempts: int = 1 + + @classmethod + def from_env(cls, dotenv: Path | str | None = ".env") -> Settings: + if dotenv is not None: + load_dotenv(dotenv) + return cls( + llm=LLMSettings.from_env(), + lean=LeanSettings.from_env(), + database=_env("LEANBENCH_DATABASE", "database/leanbench.sqlite"), + max_window_pages=_env_int("LEANBENCH_MAX_WINDOW_PAGES", 3), + repair_attempts=_env_int("LEANBENCH_REPAIR_ATTEMPTS", 1), + ) diff --git a/leanbench/db.py b/leanbench/db.py new file mode 100644 index 0000000..3e830b1 --- /dev/null +++ b/leanbench/db.py @@ -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) diff --git a/leanbench/detect.py b/leanbench/detect.py new file mode 100644 index 0000000..3a288df --- /dev/null +++ b/leanbench/detect.py @@ -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"<<>>\n{_window_text(window)}\n<<>>" + ) + 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") diff --git a/leanbench/document.py b/leanbench/document.py new file mode 100644 index 0000000..09a49fe --- /dev/null +++ b/leanbench/document.py @@ -0,0 +1,219 @@ +"""Source documents, pages, and the page windows detection runs over. + +Every index this project ever reports is an offset into ``Page.text`` as produced +here. That makes the extractor part of the contract, so the extractor name and a +per-page checksum are recorded with the source: if you change extractors, stored +offsets from the old run no longer mean anything, and ``verify_offsets`` says so. +""" +from __future__ import annotations + +import hashlib +import re +from collections.abc import Iterator, Sequence +from dataclasses import dataclass, field +from pathlib import Path + +PAGE_SEPARATOR = "\n" + + +def sha256_file(path: Path | str, chunk: int = 1 << 20) -> str: + digest = hashlib.sha256() + with Path(path).open("rb") as stream: + while block := stream.read(chunk): + digest.update(block) + return digest.hexdigest() + + +def sha1_text(text: str) -> str: + return hashlib.sha1(text.encode("utf-8")).hexdigest() + + +_LIGATURES = { + "ff": "ff", "fi": "fi", "fl": "fl", "ffi": "ffi", "ffl": "ffl", + " ": " ", "’": "'", "‘": "'", "“": '"', "”": '"', + "–": "-", "—": "--", "−": "-", +} +_TRAILING_SPACE = re.compile(r"[ \t]+\n") + + +def normalise_page_text(text: str) -> str: + """Light, *stable* cleanup. Deliberately not aggressive: offsets depend on it.""" + for bad, good in _LIGATURES.items(): + text = text.replace(bad, good) + text = text.replace("\r\n", "\n").replace("\r", "\n") + text = _TRAILING_SPACE.sub("\n", text) + return text + + +@dataclass(frozen=True) +class Page: + number: int # 1-based, physical page order in the file + text: str + + @property + def checksum(self) -> str: + return sha1_text(self.text) + + def __len__(self) -> int: + return len(self.text) + + +# -------------------------------------------------------------------------- +# extraction + + +def extract_pdf_pages(path: Path | str) -> tuple[list[str], str]: + """Return (page texts, extractor id). PyMuPDF if present, else pypdf.""" + path = Path(path) + backend = None + try: # the modern name; `fitz` is the old one and warns on import + import pymupdf as backend # type: ignore + except ImportError: + try: + import fitz as backend # type: ignore + except ImportError: + backend = None + if backend is not None: + with backend.open(path) as document: + pages = [page.get_text("text") for page in document] + version = getattr(backend, "__version__", "") or getattr(backend, "VersionBind", "") + return pages, f"pymupdf:{version}" if version else "pymupdf" + try: + from pypdf import PdfReader # type: ignore + except ImportError as error: # pragma: no cover - environment problem + raise RuntimeError( + "No PDF backend available: pip install pymupdf (preferred) or pypdf" + ) from error + reader = PdfReader(str(path)) + return [page.extract_text() or "" for page in reader.pages], "pypdf" + + +def extract_text_pages(path: Path | str, separator: str = "\f") -> tuple[list[str], str]: + """Plain-text sources: form feed marks a page break, like ``pdftotext`` emits.""" + text = Path(path).read_text(encoding="utf-8", errors="replace") + return text.split(separator), "text" + + +@dataclass +class Document: + """A source book: its pages, and the hash the database keys everything by.""" + + path: Path + sha256: str + pages: list[Page] + extractor: str + title: str = "" + meta: dict = field(default_factory=dict) + + @classmethod + def load(cls, path: Path | str, title: str = "") -> Document: + path = Path(path) + if not path.is_file(): + raise FileNotFoundError(path) + if path.suffix.lower() == ".pdf": + raw_pages, extractor = extract_pdf_pages(path) + else: + raw_pages, extractor = extract_text_pages(path) + pages = [Page(i + 1, normalise_page_text(t)) for i, t in enumerate(raw_pages)] + return cls( + path=path, + sha256=sha256_file(path), + pages=pages, + extractor=extractor, + title=title or path.stem, + ) + + def __len__(self) -> int: + return len(self.pages) + + def page(self, number: int) -> Page: + if not 1 <= number <= len(self.pages): + raise IndexError(f"page {number} out of range 1..{len(self.pages)}") + return self.pages[number - 1] + + def page_range(self, first: int | None = None, last: int | None = None) -> list[Page]: + first = max(1, first or 1) + last = min(len(self.pages), last or len(self.pages)) + return [p for p in self.pages if first <= p.number <= last] + + def slice(self, start_page: int, start_index: int, end_page: int, end_index: int) -> str: + """Re-read a stored span straight out of the document.""" + if end_page < start_page or (end_page == start_page and end_index < start_index): + raise ValueError("span ends before it starts") + chunks: list[str] = [] + for number in range(start_page, end_page + 1): + text = self.page(number).text + begin = start_index if number == start_page else 0 + finish = end_index if number == end_page else len(text) + chunks.append(text[begin:finish]) + return PAGE_SEPARATOR.join(chunks) + + def checksums(self) -> dict[int, str]: + return {page.number: page.checksum for page in self.pages} + + +# -------------------------------------------------------------------------- +# windows + + +@dataclass +class PageWindow: + """One or more consecutive pages presented to the model as a single text. + + ``locate`` maps any offset in that joined text back to (page number, offset in + page), which is what the detector reports and the database stores. + """ + + pages: Sequence[Page] + + def __post_init__(self) -> None: + if not self.pages: + raise ValueError("a window needs at least one page") + parts: list[str] = [] + spans: list[tuple[int, int, int]] = [] # (global start, global end, page number) + cursor = 0 + for position, page in enumerate(self.pages): + if position: + cursor += len(PAGE_SEPARATOR) + spans.append((cursor, cursor + len(page.text), page.number)) + parts.append(page.text) + cursor += len(page.text) + self.text = PAGE_SEPARATOR.join(parts) + self._spans = spans + + @property + def first_page(self) -> int: + return self.pages[0].number + + @property + def last_page(self) -> int: + return self.pages[-1].number + + def __len__(self) -> int: + return len(self.text) + + def locate(self, index: int) -> tuple[int, int]: + """Global offset -> (page number, offset within that page). Clamps politely.""" + index = max(0, min(index, len(self.text))) + for start, end, number in self._spans: + if index < start: # landed inside a separator: belongs to the page ahead + return number, 0 + if index <= end: + return number, index - start + start, end, number = self._spans[-1] + return number, end - start + + def globalise(self, page_number: int, offset: int) -> int: + for start, end, number in self._spans: + if number == page_number: + return start + max(0, min(offset, end - start)) + raise ValueError(f"page {page_number} is not in this window") + + def pages_covered(self) -> list[int]: + return [page.number for page in self.pages] + + +def windows(pages: Sequence[Page], size: int) -> Iterator[PageWindow]: + """Non-overlapping windows of at most ``size`` pages (used by tests/tools).""" + for start in range(0, len(pages), size): + yield PageWindow(pages[start : start + size]) diff --git a/leanbench/formalize.py b/leanbench/formalize.py new file mode 100644 index 0000000..ccffdcf --- /dev/null +++ b/leanbench/formalize.py @@ -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<<>>\n{detection.text}\n<<>>" + ) + + +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", +] diff --git a/leanbench/llm.py b/leanbench/llm.py new file mode 100644 index 0000000..b6e5c16 --- /dev/null +++ b/leanbench/llm.py @@ -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()) diff --git a/leanbench/pipeline.py b/leanbench/pipeline.py new file mode 100644 index 0000000..4b201d1 --- /dev/null +++ b/leanbench/pipeline.py @@ -0,0 +1,181 @@ +"""The whole run, in order: load -> detect -> formalize -> verify -> store. + +``ingest`` is what the CLI, the Docker entrypoint and CI all call. Re-ingesting a +book (or just some of its pages) deletes exactly the rows those pages produced last +time and writes the new ones, so a source can be redone without collateral damage. +""" +from __future__ import annotations + +from collections.abc import Callable, Sequence +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from .config import Settings +from .db import Database, ItemRow, item_id +from .detect import DetectionReport, detect_document +from .document import Document +from .formalize import Formalization, formalize_all, verification_mode +from .llm import build_client +from .verify import LeanChecker, NullChecker + + +@dataclass +class IngestResult: + document: Document + report: DetectionReport + formalizations: list[Formalization] = field(default_factory=list) + stored: list[str] = field(default_factory=list) + removed: int = 0 + + def summary(self) -> dict[str, Any]: + verified = sum(1 for f in self.formalizations if f.verified) + proved = sum(1 for f in self.formalizations if f.proved) + skipped = sum(1 for f in self.formalizations if f.skipped) + return { + "source": self.document.title, + "source_hash": self.document.sha256, + "pages": len(self.document), + "detected": len(self.report.detections), + "statements": sum(1 for d in self.report.detections if d.kind == "statement"), + "problems": sum(1 for d in self.report.detections if d.kind == "problem"), + "multi_page": sum(1 for d in self.report.detections if d.spans_pages > 1), + "unresolved": len(self.report.unresolved), + "formalized": len(self.formalizations) - skipped, + "skipped": skipped, + "verified": verified, + "proved": proved, + "stored": len(self.stored), + "removed": self.removed, + } + + +def build_checker(settings: Settings, *, enabled: bool = True): + if not enabled: + return NullChecker() + checker = LeanChecker.from_settings(settings.lean) + return checker if checker.available() else NullChecker() + + +def to_row(document: Document, formalization: Formalization) -> ItemRow: + detection = formalization.detection + identifier = item_id( + document.sha256, + detection.start_page, + detection.start_index, + detection.end_page, + detection.end_index, + ) + verify = formalization.verify + return ItemRow( + id=identifier, + source_hash=document.sha256, + kind=detection.kind, + label=detection.label, + summary=detection.summary, + informal=formalization.informal, + start_page=detection.start_page, + start_index=detection.start_index, + end_page=detection.end_page, + end_index=detection.end_index, + page_checksum=document.page(detection.start_page).checksum, + source_text=detection.text, + lean_name=formalization.name, + lean_code=formalization.lean_code, + imports=formalization.imports, + has_proof=formalization.has_proof, + verified=bool(verify and verify.ok), + proved=bool(verify and verify.proved), + verify_mode=verify.mode if verify else verification_mode(detection), + diagnostics=(verify.diagnostics if verify else ""), + confidence=formalization.confidence, + match_method=detection.match_method, + match_score=detection.match_score, + truncated=detection.truncated, + types=formalization.index_types, + ) + + +def store(database: Database, document: Document, + formalizations: Sequence[Formalization]) -> list[str]: + database.upsert_source( + document.sha256, + title=document.title, + path=str(document.path), + pages=len(document), + extractor=document.extractor, + meta=document.meta, + ) + stored: list[str] = [] + for formalization in formalizations: + if formalization.skipped or not formalization.lean_code: + continue + stored.append(database.upsert_item(to_row(document, formalization))) + return stored + + +def ingest( + source: str | Path, + *, + settings: Settings | None = None, + client: Any = None, + database: Database | None = None, + checker: Any = None, + first_page: int | None = None, + last_page: int | None = None, + title: str = "", + replace: bool = True, + store_results: bool = True, + progress: Callable[[str], None] | None = None, +) -> IngestResult: + settings = settings or Settings.from_env() + client = client or build_client(settings.llm) + checker = checker if checker is not None else build_checker(settings) + owns_database = database is None + database = database or Database(settings.database) + + say = progress or (lambda _message: None) + document = Document.load(source, title=title) + say(f"{document.title}: {len(document)} pages, sha256 {document.sha256[:12]}, " + f"extractor {document.extractor}") + + result = IngestResult(document=document, report=DetectionReport()) + try: + if store_results and replace: + if first_page or last_page: + result.removed = database.delete_pages( + document.sha256, first_page or 1, last_page or len(document) + ) + else: + result.removed = database.delete_source(document.sha256) + if result.removed: + say(f"removed {result.removed} item(s) from the previous run") + + say("detecting formalizable passages...") + result.report = detect_document( + client, + document, + first_page=first_page, + last_page=last_page, + max_window_pages=settings.max_window_pages, + progress=progress, + ) + say(result.report.summary()) + + say("formalizing...") + result.formalizations = formalize_all( + client, + result.report.detections, + checker=checker, + document_title=document.title, + repair_attempts=settings.repair_attempts, + progress=progress, + ) + + if store_results: + result.stored = store(database, document, result.formalizations) + say(f"stored {len(result.stored)} item(s) in {database.path}") + finally: + if owns_database: + database.close() + return result diff --git a/leanbench/typeindex.py b/leanbench/typeindex.py new file mode 100644 index 0000000..a77f92c --- /dev/null +++ b/leanbench/typeindex.py @@ -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 diff --git a/leanbench/verify.py b/leanbench/verify.py new file mode 100644 index 0000000..1d28aa6 --- /dev/null +++ b/leanbench/verify.py @@ -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"(? 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) diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..501d4bc --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,36 @@ +[build-system] +requires = ["setuptools>=68"] +build-backend = "setuptools.build_meta" + +[project] +name = "lean-bench" +version = "0.2.0" +description = "Detect formalizable mathematics in source documents and turn it into verified, indexed Lean 4" +readme = "README.md" +requires-python = ">=3.10" +license = { text = "MIT" } +dependencies = [ + "pymupdf>=1.24", +] + +[project.optional-dependencies] +pypdf = ["pypdf>=4.0"] +dev = ["pytest>=8.0", "ruff>=0.5"] + +[project.scripts] +lean-bench = "leanbench.cli:main" + +[tool.setuptools.packages.find] +include = ["leanbench*"] + +[tool.pytest.ini_options] +testpaths = ["tests"] +addopts = "-q" + +[tool.ruff] +line-length = 100 +target-version = "py310" + +[tool.ruff.lint] +select = ["E", "F", "W", "I", "UP", "B"] +ignore = ["E501", "B008"] diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..cd76556 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,79 @@ +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from leanbench.document import Document +from leanbench.verify import MODE_FULL, VerifyResult + +PAGE_ONE = """Chapter 1. Divisibility + +1.6 Divisibility in the Integers + +Definition 1.6.1. Let a and b be integers. We say a divides b, written a | b, +if there exists an integer k such that b = a k. + +Theorem 1.6.2. Let a, b and c be integers. If a divides b and a divides c, +then a divides b + c. + +Proof. Write b = a k and c = a m. Then b + c = a (k + m). QED + +Theorem 1.6.3. Let a and b be nonzero integers. Then the set of common +""" + +PAGE_TWO = """divisors of a and b has a greatest element, the greatest common divisor. + +Exercise 1.6.4. Show that if a divides b and b divides a, then a = b or a = -b. + +Exercise 1.6.5. Prove that the sum of two odd integers is even. +""" + + +@pytest.fixture() +def two_page_book(tmp_path: Path) -> Document: + path = tmp_path / "book.txt" + path.write_text(PAGE_ONE + "\f" + PAGE_TWO, encoding="utf-8") + return Document.load(path, title="Test Book") + + +class FakeChecker: + """A Lean stand-in: fails anything containing ``BAD``, otherwise accepts.""" + + def __init__(self) -> None: + self.calls: list[str] = [] + + def available(self) -> bool: + return True + + def check(self, code: str, *, mode: str = MODE_FULL, imports=None) -> VerifyResult: + self.calls.append(code) + if "BAD" in code: + return VerifyResult(False, mode, False, "LeanBenchCheck.lean:3:0: error: bad") + return VerifyResult(True, mode, "sorry" in code, "") + + +@pytest.fixture() +def fake_checker() -> FakeChecker: + return FakeChecker() + + +def detection_reply(items: list[dict]) -> str: + return json.dumps({"items": items}) + + +def formalize_reply(**overrides) -> str: + payload = { + "skip": False, + "reason": "", + "imports": ["Mathlib"], + "name": "test_decl", + "lean_code": "theorem test_decl (a b : Int) (h : a ∣ b) : a ∣ b := by\n sorry", + "has_proof": False, + "informal": "a divides b", + "index_types": ["Int"], + "confidence": 0.8, + } + payload.update(overrides) + return json.dumps(payload) diff --git a/tests/test_anchors.py b/tests/test_anchors.py new file mode 100644 index 0000000..0608f87 --- /dev/null +++ b/tests/test_anchors.py @@ -0,0 +1,54 @@ +from __future__ import annotations + +from leanbench.anchors import find, resolve_span + +TEXT = ( + "Theorem 1.6.2. Let a, b and c be integers. If a divides b and a divides c,\n" + "then a divides b + c.\n\nProof. Write b = a k." +) + + +def test_exact_match() -> None: + match = find(TEXT, "Theorem 1.6.2.") + assert match is not None and match.method == "exact" + assert TEXT[match.start : match.end] == "Theorem 1.6.2." + + +def test_whitespace_insensitive_match_across_a_line_break() -> None: + match = find(TEXT, "a divides c, then a divides b + c.") + assert match is not None + assert match.method == "whitespace" + assert TEXT[match.start : match.end].replace("\n", " ") == "a divides c, then a divides b + c." + + +def test_fuzzy_match_survives_a_typo() -> None: + match = find(TEXT, "Let a, b and c be integres. If a divides b") + assert match is not None + assert match.method == "fuzzy" + assert "integers" in TEXT[match.start : match.end] + + +def test_missing_anchor_is_reported_not_guessed() -> None: + assert find(TEXT, "nothing like this appears anywhere in the page") is None + + +def test_resolve_span_covers_start_to_end() -> None: + span = resolve_span(TEXT, "Theorem 1.6.2.", "a divides b + c.") + assert span is not None and span.resolved + assert TEXT[span.start : span.end].startswith("Theorem 1.6.2.") + assert TEXT[span.start : span.end].endswith("a divides b + c.") + + +def test_resolve_span_falls_back_when_the_end_quote_is_wrong() -> None: + span = resolve_span(TEXT, "Theorem 1.6.2.", "this ending was invented", fallback_length=40) + assert span is not None + assert "truncated" in span.method + assert span.end - span.start == 40 + + +def test_resolve_span_searches_forward_from_the_cursor() -> None: + text = "alpha beta alpha gamma" + first = resolve_span(text, "alpha", "beta") + assert first is not None and first.start == 0 + second = resolve_span(text, "alpha", "gamma", search_from=first.end) + assert second is not None and second.start == text.rindex("alpha") diff --git a/tests/test_db.py b/tests/test_db.py new file mode 100644 index 0000000..5af1b89 --- /dev/null +++ b/tests/test_db.py @@ -0,0 +1,133 @@ +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from leanbench.db import Database, ItemRow, item_id + + +def row(**overrides) -> ItemRow: + data = dict( + id="", + source_hash="book-hash-a", + kind="statement", + label="Theorem 1.6.2", + summary="divisibility is additive", + informal="a | b and a | c implies a | b + c", + start_page=12, + start_index=100, + end_page=12, + end_index=260, + page_checksum="abc", + source_text="Theorem 1.6.2 ...", + lean_name="dvd_add'", + lean_code="theorem dvd_add' (a b c : Int) : True := trivial", + imports=["Mathlib"], + has_proof=True, + verified=True, + proved=True, + verify_mode="full", + diagnostics="", + confidence=0.9, + match_method="exact", + match_score=1.0, + truncated=False, + types=[("Int", "input")], + ) + data.update(overrides) + if not data["id"]: + data["id"] = item_id(data["source_hash"], data["start_page"], data["start_index"], + data["end_page"], data["end_index"]) + return ItemRow(**data) + + +@pytest.fixture() +def database(tmp_path: Path) -> Database: + db = Database(tmp_path / "test.sqlite") + db.upsert_source("book-hash-a", title="Algebra", path="a.pdf", pages=400) + yield db + db.close() + + +def test_item_id_is_stable_for_the_same_span() -> None: + first = item_id("h", 1, 2, 3, 4) + assert first == item_id("h", 1, 2, 3, 4) + assert first != item_id("h", 1, 2, 3, 5) + + +def test_store_and_read_back(database: Database) -> None: + identifier = database.upsert_item(row()) + stored = database.get(identifier) + assert stored is not None + assert stored["imports"] == ["Mathlib"] + assert stored["types"] == [{"name": "Int", "role": "input"}] + + +def test_reingesting_the_same_span_updates_rather_than_duplicates(database: Database) -> None: + database.upsert_item(row()) + database.upsert_item(row(lean_name="renamed", types=[("Nat", "input")])) + assert database.stats()["items"] == 1 + stored = database.get(item_id("book-hash-a", 12, 100, 12, 260)) + assert stored is not None and stored["lean_name"] == "renamed" + assert stored["types"] == [{"name": "Nat", "role": "input"}], "stale types are dropped" + + +def test_search_by_type_is_the_point(database: Database) -> None: + database.upsert_item(row(types=[("Int", "input")])) + database.upsert_item(row(start_index=500, end_index=700, kind="problem", + types=[("Polynomial", "input"), ("Int", "input")])) + assert len(database.search(types=["Int"])) == 2 + assert len(database.search(types=["Polynomial"])) == 1 + assert len(database.search(types=["Int", "Polynomial"], match_all_types=True)) == 1 + assert len(database.search(types=["Int"], kind="problem")) == 1 + + +def test_search_by_role_separates_definitions_from_theorems(database: Database) -> None: + database.upsert_item(row(types=[("Ideal", "own")])) + database.upsert_item(row(start_index=900, end_index=950, types=[("Ideal", "input")])) + assert len(database.search(types=["Ideal"], roles=["own"])) == 1 + assert len(database.search(types=["Ideal"], roles=["input"])) == 1 + + +def test_search_by_page_finds_items_spanning_it(database: Database) -> None: + database.upsert_item(row(start_page=12, end_page=14)) + assert len(database.search(page=13)) == 1 + assert len(database.search(page=20)) == 0 + + +def test_dropping_a_page_range_lets_you_redo_part_of_a_book(database: Database) -> None: + database.upsert_item(row(start_page=12)) + database.upsert_item(row(start_page=40, start_index=1, end_page=40, end_index=50)) + removed = database.delete_pages("book-hash-a", 10, 20) + assert removed == 1 + assert database.stats()["items"] == 1 + assert database.search(page=40) + + +def test_dropping_a_source_takes_its_items_with_it(database: Database) -> None: + database.upsert_item(row()) + database.delete_source("book-hash-a") + assert database.stats()["items"] == 0 + assert database.sources() == [] + + +def test_export_jsonl_leaves_out_bulk_by_default(database: Database, tmp_path: Path) -> None: + database.upsert_item(row()) + out = tmp_path / "items.jsonl" + assert database.export_jsonl(out) == 1 + record = json.loads(out.read_text(encoding="utf-8").splitlines()[0]) + assert "source_text" not in record and "diagnostics" not in record + assert record["types"] == [{"name": "Int", "role": "input"}] + assert database.export_jsonl(out, include_source_text=True) == 1 + assert "source_text" in json.loads(out.read_text(encoding="utf-8").splitlines()[0]) + + +def test_stats_counts_the_things_that_matter(database: Database) -> None: + database.upsert_item(row()) + database.upsert_item(row(start_index=2, kind="problem", proved=False, verified=True)) + stats = database.stats() + assert stats["items"] == 2 and stats["statements"] == 1 and stats["problems"] == 1 + assert stats["verified"] == 2 and stats["proved"] == 1 + assert database.type_counts()[0]["name"] == "Int" diff --git a/tests/test_detect.py b/tests/test_detect.py new file mode 100644 index 0000000..b1b483a --- /dev/null +++ b/tests/test_detect.py @@ -0,0 +1,133 @@ +from __future__ import annotations + +from leanbench.detect import KIND_PROBLEM, KIND_STATEMENT, detect_document +from leanbench.document import Document +from leanbench.llm import ScriptedClient + +from .conftest import detection_reply + + +def test_single_page_detection_reports_page_and_index(two_page_book: Document) -> None: + client = ScriptedClient( + [ + detection_reply( + [ + { + "kind": "statement", + "label": "Definition 1.6.1", + "summary": "divisibility", + "start_quote": "Definition 1.6.1. Let a and b be integers.", + "end_quote": "b = a k.", + "continues_past_end": False, + } + ] + ), + detection_reply([]), + ] + ) + report = detect_document(client, two_page_book, max_window_pages=3) + assert len(report.detections) == 1 + found = report.detections[0] + assert found.kind == KIND_STATEMENT + assert found.start_page == 1 and found.end_page == 1 + page = two_page_book.page(1).text + assert page[found.start_index : found.end_index].startswith("Definition 1.6.1") + assert page[found.start_index : found.end_index].endswith("b = a k.") + + +def test_item_running_off_the_page_triggers_a_second_page(two_page_book: Document) -> None: + truncated = detection_reply( + [ + { + "kind": "statement", + "label": "Theorem 1.6.3", + "summary": "gcd exists", + "start_quote": "Theorem 1.6.3. Let a and b be nonzero integers.", + "end_quote": "the set of common", + "continues_past_end": True, + } + ] + ) + whole = detection_reply( + [ + { + "kind": "statement", + "label": "Theorem 1.6.3", + "summary": "gcd exists", + "start_quote": "Theorem 1.6.3. Let a and b be nonzero integers.", + "end_quote": "the greatest common divisor.", + "continues_past_end": False, + }, + { + "kind": "problem", + "label": "Exercise 1.6.5", + "summary": "odd + odd is even", + "start_quote": "Exercise 1.6.5. Prove that the sum of two odd", + "end_quote": "odd integers is even.", + "continues_past_end": False, + }, + ] + ) + client = ScriptedClient([truncated, whole]) + report = detect_document(client, two_page_book, max_window_pages=3) + + assert len(client.calls) == 2, "the page should be re-read with its successor" + theorem = report.detections[0] + assert theorem.start_page == 1 + assert theorem.end_page == 2, "the span must cross the page boundary" + assert theorem.spans_pages == 2 + assert not theorem.truncated + assert two_page_book.page(2).text[: theorem.end_index].endswith("greatest common divisor.") + + exercise = report.detections[1] + assert exercise.kind == KIND_PROBLEM + assert exercise.start_page == 2 + + +def test_window_stops_growing_at_the_limit_and_says_so(two_page_book: Document) -> None: + always_truncated = detection_reply( + [ + { + "kind": "statement", + "label": "", + "summary": "", + "start_quote": "Definition 1.6.1. Let a and b be integers.", + "end_quote": "b = a k.", + "continues_past_end": True, + } + ] + ) + client = ScriptedClient([always_truncated] * 4) + report = detect_document(client, two_page_book, max_window_pages=2) + assert len(client.calls) == 2 + assert report.detections[0].truncated is True + + +def test_unresolvable_quotes_are_recorded_not_invented(two_page_book: Document) -> None: + client = ScriptedClient( + [ + detection_reply( + [ + { + "kind": "statement", + "label": "Hallucination 9", + "summary": "not in the book", + "start_quote": "Theorem 9.9.9 about topological quantum widgets", + "end_quote": "which is absurd.", + "continues_past_end": False, + } + ] + ), + detection_reply([]), + ] + ) + report = detect_document(client, two_page_book) + assert report.detections == [] + assert len(report.unresolved) == 1 + + +def test_a_failing_page_does_not_kill_the_run(two_page_book: Document) -> None: + client = ScriptedClient(["not json at all", detection_reply([])]) + report = detect_document(client, two_page_book) + assert report.pages_scanned == 2 + assert any("detector failed" in u.get("reason", "") for u in report.unresolved) diff --git a/tests/test_document.py b/tests/test_document.py new file mode 100644 index 0000000..9d48dfa --- /dev/null +++ b/tests/test_document.py @@ -0,0 +1,46 @@ +from __future__ import annotations + +from leanbench.document import Document, Page, PageWindow, normalise_page_text + + +def test_pages_are_numbered_from_one(two_page_book: Document) -> None: + assert len(two_page_book) == 2 + assert [p.number for p in two_page_book.pages] == [1, 2] + assert two_page_book.page(1).text.startswith("Chapter 1") + + +def test_window_locate_maps_back_to_the_right_page() -> None: + window = PageWindow([Page(7, "abcdef"), Page(8, "ghij")]) + assert window.text == "abcdef\nghij" + assert window.locate(0) == (7, 0) + assert window.locate(5) == (7, 5) + assert window.locate(6) == (7, 6) # end of page 7 + assert window.locate(7) == (8, 0) # first char of page 8 + assert window.locate(10) == (8, 3) + assert window.locate(999) == (8, 4) # clamped + + +def test_window_globalise_is_the_inverse() -> None: + window = PageWindow([Page(1, "hello"), Page(2, "world")]) + for page, offset in [(1, 0), (1, 4), (2, 0), (2, 4)]: + assert window.locate(window.globalise(page, offset)) == (page, offset) + + +def test_document_slice_reconstructs_a_cross_page_span(two_page_book: Document) -> None: + page_one = two_page_book.page(1).text + start = page_one.index("Theorem 1.6.3") + end = two_page_book.page(2).text.index("has a greatest") + text = two_page_book.slice(1, start, 2, end) + assert text.startswith("Theorem 1.6.3") + assert text.rstrip().endswith("divisors of a and b") + + +def test_normalisation_is_stable_and_gentle() -> None: + text = "first line \nsecond" + assert normalise_page_text(text) == "first line\nsecond" + assert normalise_page_text(normalise_page_text(text)) == normalise_page_text(text) + + +def test_page_checksum_changes_with_content() -> None: + assert Page(1, "a").checksum != Page(1, "b").checksum + assert Page(1, "a").checksum == Page(9, "a").checksum diff --git a/tests/test_formalize.py b/tests/test_formalize.py new file mode 100644 index 0000000..69f3352 --- /dev/null +++ b/tests/test_formalize.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +from leanbench.detect import KIND_PROBLEM, KIND_STATEMENT, Detection +from leanbench.formalize import formalize, formalize_all, verification_mode +from leanbench.llm import ScriptedClient +from leanbench.verify import MODE_FULL, MODE_STATEMENT + +from .conftest import formalize_reply + + +def make_detection(kind: str = KIND_STATEMENT) -> Detection: + return Detection( + kind=kind, + label="Theorem 1.6.2", + summary="divisibility is additive", + start_page=12, + start_index=100, + end_page=12, + end_index=260, + text="Theorem 1.6.2. If a divides b and a divides c then a divides b + c.", + ) + + +def test_open_problems_are_only_checked_as_statements() -> None: + assert verification_mode(make_detection(KIND_PROBLEM)) == MODE_STATEMENT + assert verification_mode(make_detection(KIND_STATEMENT)) == MODE_FULL + + +def test_formalize_produces_verified_code_and_types(fake_checker) -> None: + client = ScriptedClient([formalize_reply()]) + result = formalize(client, make_detection(), checker=fake_checker) + assert result.lean_code.startswith("theorem test_decl") + assert result.verified + assert not result.proved, "a sorry in the code means it is not proved" + assert ("Int", "input") in result.index_types + + +def test_a_real_proof_counts_as_proved(fake_checker) -> None: + client = ScriptedClient( + [formalize_reply(lean_code="theorem t (n : Nat) : n = n := rfl", has_proof=True)] + ) + result = formalize(client, make_detection(), checker=fake_checker) + assert result.has_proof and result.proved + + +def test_lean_diagnostics_drive_a_repair_attempt(fake_checker) -> None: + client = ScriptedClient( + [ + formalize_reply(lean_code="theorem t : BAD := by sorry"), + formalize_reply(lean_code="theorem t (n : Nat) : n = n := rfl", has_proof=True), + ] + ) + result = formalize(client, make_detection(), checker=fake_checker, repair_attempts=1) + assert result.attempts == 2 + assert result.verified + assert len(client.calls) == 2 + repair_prompt = client.calls[1][-1]["content"] + assert "error: bad" in repair_prompt, "the compiler's own words go back to the model" + + +def test_repair_gives_up_after_the_budget(fake_checker) -> None: + client = ScriptedClient([formalize_reply(lean_code="theorem t : BAD := by sorry")] * 3) + result = formalize(client, make_detection(), checker=fake_checker, repair_attempts=1) + assert result.attempts == 2 + assert not result.verified + + +def test_model_may_decline_an_unformalizable_passage(fake_checker) -> None: + client = ScriptedClient(['{"skip": true, "reason": "prose only"}']) + result = formalize(client, make_detection(), checker=fake_checker) + assert result.skipped and result.skip_reason == "prose only" + assert fake_checker.calls == [] + + +def test_imports_inside_the_body_are_stripped(fake_checker) -> None: + client = ScriptedClient( + [formalize_reply(lean_code="import Mathlib\ntheorem t (n : Nat) : n = n := rfl")] + ) + result = formalize(client, make_detection(), checker=fake_checker) + assert "import" not in result.lean_code + + +def test_one_bad_item_does_not_stop_the_batch(fake_checker) -> None: + client = ScriptedClient(["garbage, not json", formalize_reply()]) + results = formalize_all( + client, [make_detection(), make_detection(KIND_PROBLEM)], checker=fake_checker + ) + assert len(results) == 2 + assert results[0].skipped and "formalizer failed" in results[0].skip_reason + assert results[1].verified diff --git a/tests/test_llm.py b/tests/test_llm.py new file mode 100644 index 0000000..89c1c26 --- /dev/null +++ b/tests/test_llm.py @@ -0,0 +1,38 @@ +from __future__ import annotations + +import pytest + +from leanbench.llm import LLMError, ScriptedClient, extract_code, parse_json + + +def test_parse_json_plain() -> None: + assert parse_json('{"a": 1}') == {"a": 1} + + +def test_parse_json_through_a_fence() -> None: + assert parse_json('```json\n{"a": [1, 2]}\n```') == {"a": [1, 2]} + + +def test_parse_json_with_chatter_around_it() -> None: + assert parse_json('Sure!\n{"ok": true}\nHope that helps.') == {"ok": True} + + +def test_parse_json_raises_when_there_is_none() -> None: + with pytest.raises(ValueError): + parse_json("no json here") + + +def test_extract_code_pulls_the_lean_block() -> None: + raw = "Here you go:\n```lean4\ntheorem t : True := trivial\n```\nDone." + assert extract_code(raw) == "theorem t : True := trivial" + + +def test_extract_code_passes_bare_code_through() -> None: + assert extract_code("theorem t : True := trivial") == "theorem t : True := trivial" + + +def test_scripted_client_runs_out_loudly() -> None: + client = ScriptedClient(["one"]) + assert client.complete([]) == "one" + with pytest.raises(LLMError): + client.complete([]) diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py new file mode 100644 index 0000000..6ecea49 --- /dev/null +++ b/tests/test_pipeline.py @@ -0,0 +1,158 @@ +from __future__ import annotations + +from pathlib import Path + +from leanbench.config import Settings +from leanbench.db import Database +from leanbench.document import Document +from leanbench.llm import ScriptedClient +from leanbench.pipeline import ingest + +from .conftest import detection_reply, formalize_reply + +DETECT_PAGE_ONE = detection_reply( + [ + { + "kind": "statement", + "label": "Definition 1.6.1", + "summary": "divisibility", + "start_quote": "Definition 1.6.1. Let a and b be integers.", + "end_quote": "b = a k.", + "continues_past_end": False, + }, + { + "kind": "statement", + "label": "Theorem 1.6.3", + "summary": "gcd exists", + "start_quote": "Theorem 1.6.3. Let a and b be nonzero integers.", + "end_quote": "the set of common", + "continues_past_end": True, + }, + ] +) + +DETECT_BOTH_PAGES = detection_reply( + [ + { + "kind": "statement", + "label": "Definition 1.6.1", + "summary": "divisibility", + "start_quote": "Definition 1.6.1. Let a and b be integers.", + "end_quote": "b = a k.", + "continues_past_end": False, + }, + { + "kind": "statement", + "label": "Theorem 1.6.3", + "summary": "gcd exists", + "start_quote": "Theorem 1.6.3. Let a and b be nonzero integers.", + "end_quote": "the greatest common divisor.", + "continues_past_end": False, + }, + { + "kind": "problem", + "label": "Exercise 1.6.4", + "summary": "mutual divisibility", + "start_quote": "Exercise 1.6.4. Show that if a divides b and b divides a", + "end_quote": "a = b or a = -b.", + "continues_past_end": False, + }, + ] +) + + +def _settings(tmp_path: Path) -> Settings: + settings = Settings.from_env(dotenv=None) + settings.database = str(tmp_path / "leanbench.sqlite") + settings.repair_attempts = 0 + return settings + + +def test_end_to_end_run_stores_indexed_items(tmp_path: Path, two_page_book, fake_checker) -> None: + client = ScriptedClient( + [DETECT_PAGE_ONE, DETECT_BOTH_PAGES] + + [ + formalize_reply(name="divides_def", + lean_code="def divides (a b : Int) : Prop := ∃ k : Int, b = a * k"), + formalize_reply(name="gcd_exists", + lean_code="theorem gcd_exists (a b : Int) : True := by\n sorry"), + formalize_reply(name="dvd_antisymm'", + lean_code="theorem dvd_antisymm' (a b : Int) : True := by\n sorry"), + ] + ) + settings = _settings(tmp_path) + database = Database(settings.database) + result = ingest( + two_page_book.path, + settings=settings, + client=client, + database=database, + checker=fake_checker, + title="Test Book", + ) + summary = result.summary() + assert summary["detected"] == 3 + assert summary["problems"] == 1 + assert summary["multi_page"] == 1 + assert summary["stored"] == 3 + assert summary["verified"] == 3 + + assert database.stats()["items"] == 3 + by_type = database.search(types=["Int"]) + assert len(by_type) == 3, "every item is reachable through its types" + problems = database.search(kind="problem") + assert len(problems) == 1 and problems[0]["start_page"] == 2 + database.close() + + +def test_reingesting_a_page_range_replaces_only_that_range( + tmp_path: Path, two_page_book: Document, fake_checker +) -> None: + settings = _settings(tmp_path) + database = Database(settings.database) + + first = ScriptedClient([DETECT_PAGE_ONE, DETECT_BOTH_PAGES] + [formalize_reply()] * 3) + ingest(two_page_book.path, settings=settings, client=first, database=database, + checker=fake_checker) + assert database.stats()["items"] == 3 + + # redo page 2 only: the page-1 item survives, page-2 items are rebuilt + second = ScriptedClient( + [ + detection_reply( + [ + { + "kind": "problem", + "label": "Exercise 1.6.5", + "summary": "odd plus odd", + "start_quote": "Exercise 1.6.5. Prove that the sum of two odd", + "end_quote": "odd integers is even.", + "continues_past_end": False, + } + ] + ), + formalize_reply(name="odd_add_odd"), + ] + ) + ingest(two_page_book.path, settings=settings, client=second, database=database, + checker=fake_checker, first_page=2, last_page=2) + + pages = sorted(item["start_page"] for item in database.search(limit=100)) + assert pages == [1, 1, 2], "page 1 untouched, page 2 replaced by the new run" + database.close() + + +def test_a_run_without_lean_stores_items_but_claims_nothing( + tmp_path: Path, two_page_book: Document +) -> None: + from leanbench.verify import NullChecker + + client = ScriptedClient([DETECT_PAGE_ONE, DETECT_BOTH_PAGES] + [formalize_reply()] * 3) + settings = _settings(tmp_path) + database = Database(settings.database) + result = ingest(two_page_book.path, settings=settings, client=client, database=database, + checker=NullChecker()) + assert result.summary()["stored"] == 3 + assert result.summary()["verified"] == 0 + assert database.search(verified=True) == [] + database.close() diff --git a/tests/test_typeindex.py b/tests/test_typeindex.py new file mode 100644 index 0000000..cf8ac61 --- /dev/null +++ b/tests/test_typeindex.py @@ -0,0 +1,81 @@ +from __future__ import annotations + +from leanbench.typeindex import ROLE_GOAL, ROLE_INPUT, ROLE_OWN, index_types, parse_signature + + +def test_theorem_is_indexed_by_the_types_of_its_inputs() -> None: + code = "theorem add_dvd (a b c : Int) (h : a ∣ b) (k : a ∣ c) : a ∣ b + c := by\n sorry" + signature = parse_signature(code) + assert signature is not None + assert signature.is_proof + assert signature.name == "add_dvd" + assert "Int" in signature.inputs + pairs = index_types(code) + assert ("Int", ROLE_INPUT) in pairs + + +def test_definition_is_indexed_by_its_own_type() -> None: + code = "def divides (a b : Int) : Prop := ∃ k : Int, b = a * k" + signature = parse_signature(code) + assert signature is not None + assert not signature.is_proof + pairs = index_types(code) + assert ("Int", ROLE_INPUT) in pairs + assert ("Prop", ROLE_OWN) in pairs, "a predicate's own type is the point of it" + assert all(role != ROLE_GOAL for _, role in pairs) + + +def test_a_proof_is_not_indexed_by_Prop() -> None: + pairs = index_types("theorem t (n : Nat) : n = n := rfl") + assert all(name != "Prop" for name, _ in pairs) + + +def test_definition_of_a_concrete_type_records_it() -> None: + pairs = index_types("def double (n : Nat) : Nat := 2 * n") + assert ("Nat", ROLE_OWN) in pairs + + +def test_instance_binders_count_as_inputs() -> None: + code = "theorem foo {G : Type*} [Group G] (x : G) : x * 1 = x := by simp" + signature = parse_signature(code) + assert signature is not None + assert "Group" in signature.inputs + + +def test_unicode_number_types_are_normalised() -> None: + code = "theorem sq_nonneg' (n : ℤ) : 0 ≤ n ^ 2 := by positivity" + pairs = dict(index_types(code)) + assert pairs.get("Int") == ROLE_INPUT + + +def test_structure_records_its_own_type() -> None: + code = "structure Monoid' (M : Type) : Type where\n one : M" + signature = parse_signature(code) + assert signature is not None + assert signature.keyword == "structure" + assert signature.own or signature.inputs + + +def test_model_supplied_types_are_merged_without_duplicates() -> None: + code = "theorem t (n : Nat) : n = n := rfl" + pairs = index_types(code, ["Nat", "Polynomial", " "]) + names = [name for name, _ in pairs] + assert names.count("Nat") == 1 + assert "Polynomial" in names + + +def test_comments_do_not_leak_into_the_index() -> None: + code = "-- about Matrix stuff\n/- Group -/\ntheorem t (n : Nat) : n = n := rfl" + names = [name for name, _ in index_types(code)] + assert "Matrix" not in names and "Group" not in names + + +def test_declaration_keyword_detected_after_attributes() -> None: + code = "@[simp]\nprivate theorem t (n : Nat) : n = n := rfl" + signature = parse_signature(code) + assert signature is not None and signature.keyword == "theorem" + + +def test_unparsable_code_yields_nothing_rather_than_nonsense() -> None: + assert parse_signature("this is not lean") is None + assert index_types("this is not lean") == []