commit 598f7fd6620d16a9d4c8f12bc2f94cbb981ac638 Author: claude-bot <10+claude-bot@noreply.localhost> Date: Sat Sep 19 16:22:09 2026 +0000 Rebuild lean-bench: pages in, verified indexed Lean out Replaces the previous tree entirely. The pipeline is four separable pieces, each testable on its own: 1. leanbench/llm.py provider-neutral OpenAI-compatible connector, stdlib only, with retries, JSON coercion and a scripted client so the rest of the code is testable without a network. 2. leanbench/detect.py per-page detection of formalizable mathematics. Case 1 is mathematics asserted as fact (definition/theorem/lemma/..., deliberately not sub-classified); case 2 is a problem left open to the reader. Hits come back as spans -- start page and index, end page and index -- so several per page is normal. A hit cut off by the page end makes the detector re-read that page together with the next, repeatedly, so an item spread over three pages is found whole. Offsets come from verbatim quotes resolved by leanbench/anchors.py (exact -> whitespace-insensitive -> fuzzy), because models quote well and count badly. 3. leanbench/formalize.py + leanbench/verify.py each span becomes one Lean 4 declaration, judged by the real compiler: full elaboration for case 1, statement legality only for case 2, which has no proof to check. Lean's own diagnostics feed back for a repair attempt. "verified" (Lean accepted it) and "proved" (accepted with no sorry) are kept strictly apart. 4. leanbench/db.py indexed SQLite plus JSONL export. Every item is indexed by the types of its inputs, or by its own type when it is not a proof, and records the source book (sha256) and the page it starts on -- so `drop --source H --pages A-B` followed by a re-ingest redoes exactly those pages and nothing else. Infrastructure: - lean/ Lake project pinned to Mathlib. - docker/ image carrying Lean 4 and a prebuilt Mathlib. The Mathlib revision and the toolchain are resolved against each other at build time rather than both hardcoded, so the build does not rot. One Docker layer per Lean step, each long one under `timeout`: the gigabyte cache fetch earns a layer that a later failure never invalidates, and a stalled step fails with a message instead of sitting silent for an hour. There is deliberately no `lake build` of our own library -- verification runs `lake env lean` against Mathlib's oleans and needs nothing else, and building the library was pure cost. docker/smoke.py runs the whole pipeline inside the image against the real compiler, with the LLM scripted. - .gitea/workflows ci (lint + tests in a venv, since the runner image's Python is Debian-managed and refuses installs into itself), docker (build, verify, optionally push), ingest (submit a book by hand). The docker build runs with --progress=plain, because BuildKit's default renderer collapses a long step into one rewritten line and a stall then looks exactly like progress. The docker workflow is not a throwaway build: it leaves the image, the layer cache, an idle `lean-bench` container and the `lean-bench-data` / `lean-bench-books` volumes in place on the runner, and never prunes, so the environment is there to be used afterwards. Pushing to the registry is optional and cannot fail the build -- Gitea's container registry does not accept the built-in Actions token, so that needs a REGISTRY_TOKEN secret with write:package. 58 tests; they need neither network nor Lean. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Ci5T3gWDhhmqDec3e1XtYZ 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..b2422c9 --- /dev/null +++ b/.gitea/workflows/docker.yml @@ -0,0 +1,283 @@ +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" + skip_size_check: + description: "Build even if the host looks too small (RAM or disk)" + 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: 150 + 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}}' + # Worth recording: a Lean build wants several GB of RAM, and a silent stall + # in `lake` is almost always memory or disk rather than a real hang. + echo "docker host: $(docker info --format '{{.NCPU}} cpu, {{.MemTotal}} bytes ram, storage {{.Driver}}')" + echo "disk on the docker host:" + docker system df || true + df -h / 2>/dev/null || true + + # Mathlib is the constraint, not the pipeline. Finding that out 40 minutes + # into a build, from silence, is how an afternoon disappears -- so check first. + - name: Check the host is big enough for Mathlib + run: | + set -eu + bytes="$(docker info --format '{{.MemTotal}}')" + gib=$(( bytes / 1073741824 )) + need=4 + echo "docker host memory: ${gib} GiB" + if [ "$gib" -lt "$need" ]; then + echo "::error::This host has ${gib} GiB of RAM. Loading Mathlib's environment" + echo "::error::needs roughly 4-8 GB, and below that Lean does not fail -- it" + echo "::error::thrashes silently until something kills it. Give the runner host" + echo "::error::at least ${need} GB (8 GB is comfortable), or build elsewhere." + echo "::error::Verification needs this memory at RUN time too, not just to build." + if [ '${{ inputs.skip_size_check }}' = 'true' ]; then + echo "::warning::skip_size_check is set; building anyway." + else + exit 1 + fi + elif [ "$gib" -lt 8 ]; then + echo "::warning::${gib} GiB will work but is tight for Mathlib; 8 GB is comfortable." + fi + + # Disk, measured where images actually land. Exporting the image unpacks it + # a second time next to the copy the build already wrote, so the peak is + # roughly twice the image size on top of whatever the layer cache holds. + # An 18 GB margin was not enough: the export died on the last Mathlib olean. + root="$(docker info --format '{{.DockerRootDir}}')" + echo "docker root dir: ${root}" + free_kb="$(docker run --rm -v "${root}":/probe:ro busybox df -P /probe 2>/dev/null \ + | awk 'NR==2{print $4}')" || true + [ -n "${free_kb:-}" ] || free_kb="$(df -P / | awk 'NR==2{print $4}')" + free_gb=$(( free_kb / 1048576 )) + echo "free where images land: ${free_gb} GiB" + if [ "$free_gb" -lt 25 ]; then + echo "::error::Only ${free_gb} GiB free under ${root}. The Lean + Mathlib image" + echo "::error::needs ~25 GB of headroom: it is unpacked again at export time," + echo "::error::and a build that gets all the way there and then dies on ENOSPC" + echo "::error::has wasted the full 25 minutes. Free space, grow the disk, or move" + echo "::error::Docker's storage to a bigger one." + if [ '${{ inputs.skip_size_check }}' = 'true' ]; then + echo "::warning::skip_size_check is set; building anyway." + else + exit 1 + fi + elif [ "$free_gb" -lt 40 ]; then + echo "::warning::${free_gb} GiB free is workable but leaves little room for a second image version." + fi + + - 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" + # --progress=plain: the fancy renderer collapses long-running steps into a + # single rewritten line, which makes a stalled build indistinguishable from a + # quiet one in the Actions log. + docker build \ + --progress=plain \ + -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..63c0f63 --- /dev/null +++ b/README.md @@ -0,0 +1,277 @@ +# 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. + +### Reading the page, when the PDF lies + +Many mathematics PDFs carry their formulas in fonts with custom encodings. They look +right on screen and extract as nonsense. Two different failure modes, both present in +this project's test book: + +| what the page says | what `pymupdf` extracts | +|---|---| +| Z_ab ≅ Z_a ⊕ Z_b, [x]_ab ↦ ([x]_a, [x]_b) | `Zab Š Za ˚ Zb; defined by Œxab 7! .Œxa; Œxb/.` | +| I + J = {a + b : a ∈ I and b ∈ J} | `I C J D fa C b W a 2 I and b 2 J g` | + +The second is the dangerous one: it is made of ordinary ASCII — `D` is `=`, `C` is +`+`, `2` is `∈`, `W` is `:` — so nothing looks wrong until you read it. No prompt +recovers either, because the information is gone before the model sees the text. + +`--extractor ocr` renders each page and reads the pixels with a vision model, over +the same endpoint and key the rest of the pipeline uses. On the page above that +turns the first row's garbage into `$\mathbb{Z}_{ab} \cong \mathbb{Z}_a \oplus +\mathbb{Z}_b$`. Transcriptions are cached on disk by (document hash, page, dpi, +backend), so a re-run costs nothing and an interrupted run resumes. + +```bash +lean-bench extract book.pdf --pages 90 --compare # see both, decide +lean-bench ingest book.pdf --pages 85-110 --extractor ocr +``` + +`--extractor auto` reads embedded text and OCRs only pages that fail a quality +check, which is cheaper on books where the damage is patchy. The check scores both +failure modes above; be aware it is a heuristic over character statistics, not a +judgement about meaning. On a badly encoded book `auto` converges on `ocr` anyway — +in the test book 89% of pages fail it — so if a spot check shows the maths is +mangled, just use `ocr` and skip the guessing. + +`--ocr-backend tesseract` is available offline. It is good at prose and poor at +mathematics, which is the wrong way round for this job; treat it as a fallback. + +**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 extract SRC` | show the page text that would be read; `--compare` puts both extractors side by side | +| `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_EXTRACTOR` | `text` (`text` / `auto` / `ocr`) | +| `LEANBENCH_OCR_MODEL` | `qwen-vl-max` | +| `LEANBENCH_OCR_DPI` | `200` | +| `LEANBENCH_OCR_THRESHOLD` | `0.75` | +| `LEANBENCH_OCR_CACHE` | `.leanbench-cache/ocr` | +| `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 (`docker/smoke.py`) + 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. + + The Lean steps are one Docker layer each, and each long one is wrapped in + `timeout`. That is not tidiness: the Mathlib cache fetch is gigabytes, so it earns + a layer that a later failure never invalidates, and a stalled `lake` step fails + with a diagnosable message instead of sitting silent for hours. The build also + runs with `--progress=plain`, because BuildKit's default renderer collapses a + long step into one rewritten line and a stall then looks exactly like progress. + There is deliberately no `lake build` of our own library: verification runs + `lake env lean `, which needs Mathlib's oleans and nothing else. +- **`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. + +### What the runner needs + +| | | +|---|---| +| RAM | **4 GB minimum, 8 GB comfortable** | +| disk | **~25 GB free where Docker stores images**, 40 GB comfortable | +| docker | the CLI plus the host socket (or a docker-in-docker runner) | + +Memory is the real constraint, and it is not negotiable: loading Mathlib's +environment needs several GB, and the pipeline does exactly that for **every item it +verifies**, not just once at build time. Below the threshold Lean does not fail +cleanly — it thrashes, prints nothing, and eventually something kills it. The +`docker` workflow therefore checks the host's memory before it starts and refuses +with an explanation rather than discovering it forty minutes in. Override with +`skip_size_check` if you know better than the check does. + +Disk is checked the same way, and measured where images actually land +(`docker info --format '{{.DockerRootDir}}'`) rather than wherever the job happens +to be running. The margin looks generous because export unpacks the image a second +time beside the copy the build already wrote: a host with 18 GB free built every +layer of this image and then died on the last Mathlib olean while exporting. + +Note that on Docker's containerd image store, images live under +`/var/lib/containerd`, which `data-root` in `daemon.json` does **not** move — that +setting only relocates `/var/lib/docker`. Moving image storage to another disk means +relocating both. + +### 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..cd706e8 --- /dev/null +++ b/docker/Dockerfile @@ -0,0 +1,93 @@ +# 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 --------------- +# One layer per step on purpose. The cache fetch is the expensive one (gigabytes), +# so it gets a layer of its own: a failure in any later step re-runs that step, never +# the download. Every long step is wrapped in `timeout` so a stall fails loudly with +# a diagnosable message instead of hanging until the CI job's own limit. +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)" + +# Resolve dependencies: clones Mathlib and writes the manifest. +RUN cd lean && timeout 1800 lake update + +# Mathlib's prebuilt artifacts. Slow, large, and cached as its own layer. +RUN cd lean && timeout 5400 lake exe cache get + +# The honest check that this image can do its job, and it is exactly what +# verification does at runtime: elaborate a file against Mathlib with `lake env lean`. +# There is deliberately no `lake build` here -- compiling our own library is something +# nothing in the pipeline ever uses. +# +# Loading Mathlib's environment needs several GB of RAM. On a host that does not have +# it, this does not fail: it thrashes, and the only symptom is silence. So it runs +# under `timeout`, and exit 124 is reported for what it almost always is. +RUN cd lean && timeout 1200 lake env lean LeanBench.lean; status=$?; \ + if [ "$status" -eq 124 ]; then \ + echo "ERROR: elaborating 'import Mathlib' timed out after 20 minutes." >&2; \ + echo "Every olean it needs is already on local disk, so this is not the network." >&2; \ + echo "Loading Mathlib wants ~4-8 GB of RAM. This machine has:" >&2; \ + { free -h || head -3 /proc/meminfo; } >&2 2>/dev/null; \ + echo "Give the build host more memory, or build this image elsewhere." >&2; \ + fi; \ + [ "$status" -eq 0 ] || exit "$status"; \ + 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 + +RUN mkdir -p /work + +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..b51258b --- /dev/null +++ b/leanbench/cli.py @@ -0,0 +1,419 @@ +"""Command line: ``lean-bench ``. + + lean-bench doctor is everything wired up? + lean-bench ingest book.pdf --pages 30-45 + lean-bench extract book.pdf --pages 90 --compare which extractor should I use? + lean-bench detect book.pdf --pages 30-31 --out detections.jsonl + lean-bench search --type Int --role input --kind problem + lean-bench export datasets/items.jsonl + lean-bench drop --source --pages 30-45 redo part of a book +""" +from __future__ import annotations + +import argparse +import json +import sys + +from .config import Settings +from .db import Database +from .detect import detect_document, write_jsonl +from .document import Document +from .llm import LLMError, build_client +from .pipeline import build_checker, build_ocr, ingest + + +def _pages(value: str | None) -> tuple[int | None, int | None]: + if not value: + return None, None + text = value.strip() + if "-" in text: + first, _, last = text.partition("-") + return (int(first) if first.strip() else None, + int(last) if last.strip() else None) + page = int(text) + return page, page + + +def _say(message: str) -> None: + print(message, flush=True) + + +def _add_common(parser: argparse.ArgumentParser) -> None: + parser.add_argument("--db", help="SQLite path (default: $LEANBENCH_DATABASE)") + parser.add_argument("--env", default=".env", help="dotenv file to load (default: .env)") + + +def _add_extractor(parser: argparse.ArgumentParser) -> None: + parser.add_argument( + "--extractor", choices=["text", "auto", "ocr"], + help="where page text comes from: the PDF's own text, OCR, or OCR only for " + "pages whose embedded text looks broken (auto)") + parser.add_argument("--ocr-model", help="vision model for OCR (default qwen-vl-max)") + parser.add_argument("--ocr-dpi", type=int, help="render resolution for OCR (default 200)") + parser.add_argument("--ocr-backend", choices=["vision", "tesseract"]) + parser.add_argument("--ocr-cache", help="directory for cached transcriptions") + + +def _apply_extractor(settings, args) -> None: + if getattr(args, "extractor", None): + settings.ocr.mode = args.extractor + if getattr(args, "ocr_model", None): + settings.ocr.model = args.ocr_model + if getattr(args, "ocr_dpi", None): + settings.ocr.dpi = args.ocr_dpi + if getattr(args, "ocr_backend", None): + settings.ocr.backend = args.ocr_backend + if getattr(args, "ocr_cache", None): + settings.ocr.cache_dir = args.ocr_cache + + +# -- commands -------------------------------------------------------------- + + +def cmd_doctor(args: argparse.Namespace) -> int: + settings = Settings.from_env(args.env) + print(f"database {settings.database}") + print(f"llm model {settings.llm.model}") + print(f"llm endpoint {settings.llm.base_url}") + print(f"llm key env {settings.llm.api_key_env}") + + try: + import pymupdf # noqa: F401 + print("pdf backend pymupdf") + except ImportError: + try: + import pypdf # noqa: F401 + print("pdf backend pypdf (install pymupdf for better text)") + except ImportError: + print("pdf backend MISSING - pip install pymupdf") + + checker = build_checker(settings) + available = getattr(checker, "available", lambda: False)() + print(f"lean {'ready: ' + ' '.join(settings.lean.command) if available else 'not available (verification will be skipped)'}") + if settings.lean.project: + print(f"lean project {settings.lean.project}") + + if args.skip_llm: + return 0 + try: + client = build_client(settings.llm) + reply = client.complete( + [{"role": "user", "content": "Reply with the single word: ready"}], + max_tokens=16, + ) + print(f"llm ready ({reply.strip()[:40]})") + except LLMError as error: + print(f"llm FAILED: {error}") + return 1 + return 0 + + +def cmd_ingest(args: argparse.Namespace) -> int: + settings = Settings.from_env(args.env) + _apply_extractor(settings, args) + if args.db: + settings.database = args.db + if args.window: + settings.max_window_pages = args.window + first, last = _pages(args.pages) + + database = Database(settings.database) + try: + result = ingest( + args.source, + settings=settings, + database=database, + checker=build_checker(settings, enabled=not args.no_verify), + first_page=first, + last_page=last, + title=args.title or "", + replace=not args.keep_existing, + store_results=not args.no_store, + progress=None if args.quiet else _say, + ) + summary = result.summary() + print(json.dumps(summary, indent=2)) + if args.detections: + write_jsonl(args.detections, result.report.detections) + print(f"detections -> {args.detections}") + if args.export: + count = database.export_jsonl(args.export, source_hash=result.document.sha256) + print(f"exported {count} item(s) -> {args.export}") + finally: + database.close() + return 0 + + +def cmd_detect(args: argparse.Namespace) -> int: + settings = Settings.from_env(args.env) + _apply_extractor(settings, args) + first, last = _pages(args.pages) + client = build_client(settings.llm) + document = Document.load( + args.source, + mode=settings.ocr.mode, + ocr=build_ocr(settings, client), + threshold=settings.ocr.threshold, + first_page=first, last_page=last, + progress=None if args.quiet else _say, + ) + report = detect_document( + client, + document, + first_page=first, + last_page=last, + max_window_pages=args.window or settings.max_window_pages, + progress=None if args.quiet else _say, + ) + print(report.summary()) + if args.out: + write_jsonl(args.out, report.detections) + print(f"-> {args.out}") + else: + for detection in report.detections: + print(json.dumps(detection.as_dict(), ensure_ascii=False)) + return 0 + + +def cmd_extract(args: argparse.Namespace) -> int: + """Show what the pipeline would actually read. The decision tool for --extractor.""" + from .ocr import text_quality + + settings = Settings.from_env(args.env) + _apply_extractor(settings, args) + first, last = _pages(args.pages) + + plain = Document.load(args.source) + pages = plain.page_range(first, last) + if not pages: + print("no pages in that range") + return 1 + + if not args.compare and settings.ocr.mode == "text": + for page in pages: + print(f"----- page {page.number} (quality {text_quality(page.text)}) -----") + print(page.text) + return 0 + + reader = build_ocr(settings if settings.ocr.mode != "text" else _forced_ocr(settings), + build_client(settings.llm)) + for page in pages: + embedded = page.text + print(f"===== page {page.number} =====") + print(f"--- embedded text (quality {text_quality(embedded)}) ---") + print(embedded if args.full else embedded[:800]) + transcribed = reader.page_text(plain.path, plain.sha256, page.number) + print(f"--- {reader.backend.name} ---") + if transcribed is None: + print(f"FAILED: {reader.failures.get(page.number)}") + else: + print(transcribed if args.full else transcribed[:800]) + return 0 + + +def _forced_ocr(settings): + settings.ocr.mode = "ocr" + return settings + + +def cmd_search(args: argparse.Namespace) -> int: + settings = Settings.from_env(args.env) + database = Database(args.db or settings.database) + try: + rows = database.search( + types=args.type or (), + roles=args.role or (), + kind=args.kind, + source_hash=args.source, + page=args.page, + verified=args.verified, + proved=args.proved, + text=args.text, + match_all_types=args.all_types, + limit=args.limit, + ) + if args.json: + print(json.dumps(rows, indent=2, ensure_ascii=False)) + return 0 + if not rows: + print("no matches") + return 0 + for row in rows: + types = ", ".join(f"{t['name']}({t['role']})" for t in row["types"]) or "-" + flag = "proved" if row["proved"] else "verified" if row["verified"] else "unverified" + print(f"{row['id']} {row['kind']:<9} p{row['start_page']:<4} {flag:<10} " + f"{row['label'] or row['lean_name']}") + print(f" types: {types}") + if args.code: + for line in (row["lean_code"] or "").splitlines(): + print(" " + line) + finally: + database.close() + return 0 + + +def cmd_sources(args: argparse.Namespace) -> int: + settings = Settings.from_env(args.env) + database = Database(args.db or settings.database) + try: + rows = database.sources() + if args.json: + print(json.dumps(rows, indent=2, ensure_ascii=False)) + return 0 + for row in rows: + print(f"{row['hash'][:16]} {row['items']:>5} items p{row['pages']:<5} " + f"{row['title']} ({row['extractor']}, {row['ingested_at']})") + finally: + database.close() + return 0 + + +def cmd_stats(args: argparse.Namespace) -> int: + settings = Settings.from_env(args.env) + database = Database(args.db or settings.database) + try: + stats = database.stats() + stats["top_types"] = database.type_counts(limit=args.top) + print(json.dumps(stats, indent=2)) + finally: + database.close() + return 0 + + +def cmd_export(args: argparse.Namespace) -> int: + settings = Settings.from_env(args.env) + database = Database(args.db or settings.database) + try: + count = database.export_jsonl( + args.out, source_hash=args.source, include_source_text=args.with_text + ) + print(f"exported {count} item(s) -> {args.out}") + finally: + database.close() + return 0 + + +def cmd_drop(args: argparse.Namespace) -> int: + settings = Settings.from_env(args.env) + database = Database(args.db or settings.database) + try: + first, last = _pages(args.pages) + if first or last: + removed = database.delete_pages(args.source, first or 1, last or 10**9) + else: + removed = database.delete_source(args.source) + print(f"removed {removed} item(s)") + finally: + database.close() + return 0 + + +# -- wiring ---------------------------------------------------------------- + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(prog="lean-bench", description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + subparsers = parser.add_subparsers(dest="command", required=True) + + doctor = subparsers.add_parser("doctor", help="check LLM, Lean and PDF wiring") + _add_common(doctor) + doctor.add_argument("--skip-llm", action="store_true", help="do not call the model") + doctor.set_defaults(func=cmd_doctor) + + ingest_parser = subparsers.add_parser("ingest", help="run the whole pipeline on a source") + _add_common(ingest_parser) + ingest_parser.add_argument("source", help="PDF (or .txt with form-feed page breaks)") + ingest_parser.add_argument("--pages", help="page or range, e.g. 30-45") + ingest_parser.add_argument("--title", help="human title for this source") + ingest_parser.add_argument("--window", type=int, help="max pages per detection window") + ingest_parser.add_argument("--no-verify", action="store_true", help="skip Lean") + ingest_parser.add_argument("--no-store", action="store_true", help="do not write the database") + ingest_parser.add_argument("--keep-existing", action="store_true", + help="do not remove this source's previous items first") + ingest_parser.add_argument("--detections", help="also write raw detections to this JSONL") + ingest_parser.add_argument("--export", help="export this source's items to JSONL afterwards") + ingest_parser.add_argument("--quiet", action="store_true") + _add_extractor(ingest_parser) + ingest_parser.set_defaults(func=cmd_ingest) + + detect_parser = subparsers.add_parser("detect", help="detection only, no Lean, no database") + _add_common(detect_parser) + detect_parser.add_argument("source") + detect_parser.add_argument("--pages") + detect_parser.add_argument("--window", type=int) + detect_parser.add_argument("--out", help="write detections to this JSONL") + detect_parser.add_argument("--quiet", action="store_true") + _add_extractor(detect_parser) + detect_parser.set_defaults(func=cmd_detect) + + extract = subparsers.add_parser( + "extract", help="show the page text the pipeline would read (and compare extractors)") + _add_common(extract) + _add_extractor(extract) + extract.add_argument("source") + extract.add_argument("--pages") + extract.add_argument("--compare", action="store_true", + help="show embedded text and OCR side by side") + extract.add_argument("--full", action="store_true", help="do not truncate") + extract.set_defaults(func=cmd_extract) + + search = subparsers.add_parser("search", help="find items by the types they are about") + _add_common(search) + search.add_argument("--type", action="append", help="type name; repeatable") + search.add_argument("--role", action="append", choices=["input", "goal", "own"]) + search.add_argument("--kind", choices=["statement", "problem"]) + search.add_argument("--source", help="source hash") + search.add_argument("--page", type=int) + search.add_argument("--text", help="full-text query") + search.add_argument("--all-types", action="store_true", help="require every --type") + search.add_argument("--verified", action="store_true", default=None) + search.add_argument("--proved", action="store_true", default=None) + search.add_argument("--code", action="store_true", help="print the Lean too") + search.add_argument("--json", action="store_true") + search.add_argument("--limit", type=int, default=25) + search.set_defaults(func=cmd_search) + + sources = subparsers.add_parser("sources", help="list ingested books") + _add_common(sources) + sources.add_argument("--json", action="store_true") + sources.set_defaults(func=cmd_sources) + + stats = subparsers.add_parser("stats", help="counts, and the most common types") + _add_common(stats) + stats.add_argument("--top", type=int, default=20) + stats.set_defaults(func=cmd_stats) + + export = subparsers.add_parser("export", help="dump items as JSONL") + _add_common(export) + export.add_argument("out") + export.add_argument("--source", help="only this source hash") + export.add_argument("--with-text", action="store_true", help="include source passages") + export.set_defaults(func=cmd_export) + + drop = subparsers.add_parser("drop", help="remove a source, or some of its pages") + _add_common(drop) + drop.add_argument("--source", required=True, help="source hash") + drop.add_argument("--pages", help="page or range; omit for the whole source") + drop.set_defaults(func=cmd_drop) + + return parser + + +def main(argv: list[str] | None = None) -> int: + args = build_parser().parse_args(argv) + try: + return args.func(args) + except FileNotFoundError as error: + print(f"not found: {error}", file=sys.stderr) + return 2 + except LLMError as error: + print(f"llm error: {error}", file=sys.stderr) + return 3 + except KeyboardInterrupt: + print("interrupted", file=sys.stderr) + return 130 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/leanbench/config.py b/leanbench/config.py new file mode 100644 index 0000000..dc9f99b --- /dev/null +++ b/leanbench/config.py @@ -0,0 +1,143 @@ +"""Runtime configuration, read from the environment (and an optional .env file). + +Everything the pipeline needs to talk to an LLM or to Lean is settable without +touching code, so the same image runs against any OpenAI-compatible endpoint and +any Lean project. +""" +from __future__ import annotations + +import json +import os +from dataclasses import dataclass, field +from pathlib import Path + +DEFAULT_BASE_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1" +DEFAULT_MODEL = "qwen-max" +DEFAULT_API_KEY_ENV = "DASHSCOPE_API_KEY" + + +def load_dotenv(path: Path | str = ".env") -> None: + """Load ``KEY=VALUE`` lines into os.environ without overwriting what is set.""" + path = Path(path) + if not path.is_file(): + return + for line in path.read_text(encoding="utf-8").splitlines(): + line = line.strip() + if not line or line.startswith("#") or "=" not in line: + continue + key, _, value = line.partition("=") + os.environ.setdefault(key.strip(), value.strip().strip('"').strip("'")) + + +def _env(name: str, default: str) -> str: + value = os.environ.get(name) + return default if value is None or value == "" else value + + +def _env_int(name: str, default: int) -> int: + try: + return int(_env(name, str(default))) + except ValueError: + return default + + +def _env_float(name: str, default: float) -> float: + try: + return float(_env(name, str(default))) + except ValueError: + return default + + +@dataclass +class LLMSettings: + model: str = DEFAULT_MODEL + base_url: str = DEFAULT_BASE_URL + api_key_env: str = DEFAULT_API_KEY_ENV + timeout: float = 300.0 + max_retries: int = 4 + temperature: float = 0.0 + + @classmethod + def from_env(cls) -> LLMSettings: + return cls( + model=_env("LEANBENCH_LLM_MODEL", DEFAULT_MODEL), + base_url=_env("LEANBENCH_LLM_BASE_URL", DEFAULT_BASE_URL), + api_key_env=_env("LEANBENCH_LLM_API_KEY_ENV", DEFAULT_API_KEY_ENV), + timeout=_env_float("LEANBENCH_LLM_TIMEOUT", 300.0), + max_retries=_env_int("LEANBENCH_LLM_MAX_RETRIES", 4), + temperature=_env_float("LEANBENCH_LLM_TEMPERATURE", 0.0), + ) + + +@dataclass +class LeanSettings: + """How to invoke Lean. ``command`` is an argv prefix that is given a .lean path.""" + + command: list[str] = field(default_factory=lambda: ["lake", "env", "lean"]) + project: str | None = None + timeout: float = 120.0 + header: str = "import Mathlib\n" + + @classmethod + def from_env(cls) -> LeanSettings: + raw = _env("LEANBENCH_LEAN_COMMAND", '["lake","env","lean"]') + try: + command = json.loads(raw) + if not isinstance(command, list) or not all(isinstance(x, str) for x in command): + raise ValueError + except ValueError: + command = raw.split() + return cls( + command=command, + project=os.environ.get("LEANBENCH_LEAN_PROJECT") or None, + timeout=_env_float("LEANBENCH_LEAN_TIMEOUT", 120.0), + header=_env("LEANBENCH_LEAN_HEADER", "import Mathlib\n"), + ) + + +@dataclass +class OcrSettings: + """Reading pages as pictures. Off by default; ``auto`` is the useful setting.""" + + mode: str = "text" # text | auto | ocr + backend: str = "vision" # vision | tesseract + model: str = "qwen-vl-max" + dpi: int = 200 + threshold: float = 0.75 # below this quality, `auto` reaches for the camera + cache_dir: str = ".leanbench-cache/ocr" + language: str = "eng" # tesseract only + + @classmethod + def from_env(cls) -> OcrSettings: + return cls( + mode=_env("LEANBENCH_EXTRACTOR", "text").lower(), + backend=_env("LEANBENCH_OCR_BACKEND", "vision").lower(), + model=_env("LEANBENCH_OCR_MODEL", "qwen-vl-max"), + dpi=_env_int("LEANBENCH_OCR_DPI", 200), + threshold=_env_float("LEANBENCH_OCR_THRESHOLD", 0.75), + cache_dir=_env("LEANBENCH_OCR_CACHE", ".leanbench-cache/ocr"), + language=_env("LEANBENCH_OCR_LANGUAGE", "eng"), + ) + + +@dataclass +class Settings: + llm: LLMSettings = field(default_factory=LLMSettings) + lean: LeanSettings = field(default_factory=LeanSettings) + ocr: OcrSettings = field(default_factory=OcrSettings) + database: str = "database/leanbench.sqlite" + max_window_pages: int = 3 + repair_attempts: int = 1 + + @classmethod + def from_env(cls, dotenv: Path | str | None = ".env") -> Settings: + if dotenv is not None: + load_dotenv(dotenv) + return cls( + llm=LLMSettings.from_env(), + lean=LeanSettings.from_env(), + ocr=OcrSettings.from_env(), + database=_env("LEANBENCH_DATABASE", "database/leanbench.sqlite"), + max_window_pages=_env_int("LEANBENCH_MAX_WINDOW_PAGES", 3), + repair_attempts=_env_int("LEANBENCH_REPAIR_ATTEMPTS", 1), + ) 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..fa388e2 --- /dev/null +++ b/leanbench/document.py @@ -0,0 +1,285 @@ +"""Source documents, pages, and the page windows detection runs over. + +Every index this project ever reports is an offset into ``Page.text`` as produced +here. That makes the extractor part of the contract, so the extractor name and a +per-page checksum are recorded with the source: if you change extractors, stored +offsets from the old run no longer mean anything, and ``verify_offsets`` says so. +""" +from __future__ import annotations + +import hashlib +import re +from collections.abc import Iterator, Sequence +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +PAGE_SEPARATOR = "\n" + + +def sha256_file(path: Path | str, chunk: int = 1 << 20) -> str: + digest = hashlib.sha256() + with Path(path).open("rb") as stream: + while block := stream.read(chunk): + digest.update(block) + return digest.hexdigest() + + +def sha1_text(text: str) -> str: + return hashlib.sha1(text.encode("utf-8")).hexdigest() + + +_LIGATURES = { + "ff": "ff", "fi": "fi", "fl": "fl", "ffi": "ffi", "ffl": "ffl", + " ": " ", "’": "'", "‘": "'", "“": '"', "”": '"', + "–": "-", "—": "--", "−": "-", +} +_TRAILING_SPACE = re.compile(r"[ \t]+\n") + + +def normalise_page_text(text: str) -> str: + """Light, *stable* cleanup. Deliberately not aggressive: offsets depend on it.""" + for bad, good in _LIGATURES.items(): + text = text.replace(bad, good) + text = text.replace("\r\n", "\n").replace("\r", "\n") + text = _TRAILING_SPACE.sub("\n", text) + return text + + +@dataclass(frozen=True) +class Page: + number: int # 1-based, physical page order in the file + text: str + + @property + def checksum(self) -> str: + return sha1_text(self.text) + + def __len__(self) -> int: + return len(self.text) + + +# -------------------------------------------------------------------------- +# extraction + + +def extract_pdf_pages(path: Path | str) -> tuple[list[str], str]: + """Return (page texts, extractor id). PyMuPDF if present, else pypdf.""" + path = Path(path) + backend = None + try: # the modern name; `fitz` is the old one and warns on import + import pymupdf as backend # type: ignore + except ImportError: + try: + import fitz as backend # type: ignore + except ImportError: + backend = None + if backend is not None: + with backend.open(path) as document: + pages = [page.get_text("text") for page in document] + version = getattr(backend, "__version__", "") or getattr(backend, "VersionBind", "") + return pages, f"pymupdf:{version}" if version else "pymupdf" + try: + from pypdf import PdfReader # type: ignore + except ImportError as error: # pragma: no cover - environment problem + raise RuntimeError( + "No PDF backend available: pip install pymupdf (preferred) or pypdf" + ) from error + reader = PdfReader(str(path)) + return [page.extract_text() or "" for page in reader.pages], "pypdf" + + +def extract_text_pages(path: Path | str, separator: str = "\f") -> tuple[list[str], str]: + """Plain-text sources: form feed marks a page break, like ``pdftotext`` emits.""" + text = Path(path).read_text(encoding="utf-8", errors="replace") + return text.split(separator), "text" + + +@dataclass +class Document: + """A source book: its pages, and the hash the database keys everything by.""" + + path: Path + sha256: str + pages: list[Page] + extractor: str + title: str = "" + meta: dict = field(default_factory=dict) + + @classmethod + def load( + cls, + path: Path | str, + title: str = "", + *, + mode: str = "text", + ocr: Any = None, + threshold: float = 0.75, + first_page: int | None = None, + last_page: int | None = None, + progress: Any = None, + ) -> Document: + """Load a source. + + ``mode`` decides where the text comes from: + + ``text`` the embedded text, as the PDF stores it. Fast, free, and wrong + whenever the document's maths fonts carry a custom encoding. + ``ocr`` render every page and read the pixels. Slow and costed per page, + but it sees what a human sees. + ``auto`` embedded text, except on pages that fail a quality check -- which + is usually exactly the pages with the mathematics on them. + + ``first_page``/``last_page`` bound which pages are OCR'd, so a 25-page trial + of a 600-page book costs 25 calls rather than 600. Pages outside the range + keep their embedded text. + """ + path = Path(path) + if not path.is_file(): + raise FileNotFoundError(path) + if path.suffix.lower() == ".pdf": + raw_pages, extractor = extract_pdf_pages(path) + else: + raw_pages, extractor = extract_text_pages(path) + pages = [Page(i + 1, normalise_page_text(text)) for i, text in enumerate(raw_pages)] + + document = cls( + path=path, + sha256=sha256_file(path), + pages=pages, + extractor=extractor, + title=title or path.stem, + ) + if mode != "text" and ocr is not None: + document._apply_ocr(ocr, mode, threshold, first_page, last_page, progress) + return document + + def _apply_ocr(self, ocr, mode, threshold, first_page, last_page, progress) -> None: + from .ocr import text_quality + + say = progress or (lambda _message: None) + low = max(1, first_page or 1) + high = min(len(self.pages), last_page or len(self.pages)) + replaced: list[int] = [] + skipped: list[int] = [] + + for index, page in enumerate(self.pages): + if not low <= page.number <= high: + continue + quality = text_quality(page.text) + if mode == "auto" and quality >= threshold: + continue + text = ocr.page_text(self.path, self.sha256, page.number) + if text is None: + skipped.append(page.number) + say(f" page {page.number}: OCR failed, keeping embedded text") + continue + self.pages[index] = Page(page.number, normalise_page_text(text)) + replaced.append(page.number) + say(f" page {page.number}: OCR (embedded text scored {quality})") + + if replaced: + self.extractor = f"{self.extractor}+{ocr.backend.name}@{ocr.dpi}dpi" + self.meta["ocr"] = { + "mode": mode, + "backend": ocr.backend.name, + "dpi": ocr.dpi, + "threshold": threshold, + "pages_replaced": replaced, + "pages_failed": skipped, + } + + def __len__(self) -> int: + return len(self.pages) + + def page(self, number: int) -> Page: + if not 1 <= number <= len(self.pages): + raise IndexError(f"page {number} out of range 1..{len(self.pages)}") + return self.pages[number - 1] + + def page_range(self, first: int | None = None, last: int | None = None) -> list[Page]: + first = max(1, first or 1) + last = min(len(self.pages), last or len(self.pages)) + return [p for p in self.pages if first <= p.number <= last] + + def slice(self, start_page: int, start_index: int, end_page: int, end_index: int) -> str: + """Re-read a stored span straight out of the document.""" + if end_page < start_page or (end_page == start_page and end_index < start_index): + raise ValueError("span ends before it starts") + chunks: list[str] = [] + for number in range(start_page, end_page + 1): + text = self.page(number).text + begin = start_index if number == start_page else 0 + finish = end_index if number == end_page else len(text) + chunks.append(text[begin:finish]) + return PAGE_SEPARATOR.join(chunks) + + def checksums(self) -> dict[int, str]: + return {page.number: page.checksum for page in self.pages} + + +# -------------------------------------------------------------------------- +# windows + + +@dataclass +class PageWindow: + """One or more consecutive pages presented to the model as a single text. + + ``locate`` maps any offset in that joined text back to (page number, offset in + page), which is what the detector reports and the database stores. + """ + + pages: Sequence[Page] + + def __post_init__(self) -> None: + if not self.pages: + raise ValueError("a window needs at least one page") + parts: list[str] = [] + spans: list[tuple[int, int, int]] = [] # (global start, global end, page number) + cursor = 0 + for position, page in enumerate(self.pages): + if position: + cursor += len(PAGE_SEPARATOR) + spans.append((cursor, cursor + len(page.text), page.number)) + parts.append(page.text) + cursor += len(page.text) + self.text = PAGE_SEPARATOR.join(parts) + self._spans = spans + + @property + def first_page(self) -> int: + return self.pages[0].number + + @property + def last_page(self) -> int: + return self.pages[-1].number + + def __len__(self) -> int: + return len(self.text) + + def locate(self, index: int) -> tuple[int, int]: + """Global offset -> (page number, offset within that page). Clamps politely.""" + index = max(0, min(index, len(self.text))) + for start, end, number in self._spans: + if index < start: # landed inside a separator: belongs to the page ahead + return number, 0 + if index <= end: + return number, index - start + start, end, number = self._spans[-1] + return number, end - start + + def globalise(self, page_number: int, offset: int) -> int: + for start, end, number in self._spans: + if number == page_number: + return start + max(0, min(offset, end - start)) + raise ValueError(f"page {page_number} is not in this window") + + def pages_covered(self) -> list[int]: + return [page.number for page in self.pages] + + +def windows(pages: Sequence[Page], size: int) -> Iterator[PageWindow]: + """Non-overlapping windows of at most ``size`` pages (used by tests/tools).""" + for start in range(0, len(pages), size): + yield PageWindow(pages[start : start + size]) 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/ocr.py b/leanbench/ocr.py new file mode 100644 index 0000000..cefbdef --- /dev/null +++ b/leanbench/ocr.py @@ -0,0 +1,308 @@ +"""Reading pages as pictures, for when the embedded text is a lie. + +Many mathematics PDFs ship their formulas in fonts with custom encodings. The +glyphs look right on screen and extract as garbage: a real page of this project's +test book yields + + Zab Š Za ˚ Zb; defined by Œxab 7! .Œxa; Œxb/. + +for what a reader plainly sees as Z_ab ≅ Z_a ⊕ Z_b, [x]_ab ↦ ([x]_a, [x]_b). No +amount of prompting recovers that, because the information is gone before the model +sees it. Rendering the page and reading the pixels does recover it. + +The default backend is the vision model on the same OpenAI-compatible endpoint the +rest of the pipeline already uses, so this costs no new dependency and no new +credential. Tesseract is available as an offline fallback, with the caveat that it +is good at prose and poor at mathematics -- which is the opposite of what we need, +so it is a last resort rather than an equal option. + +Transcription is cached on disk, keyed by document hash, page, dpi and backend. A +600-page book is 600 vision calls; nobody should pay for them twice. +""" +from __future__ import annotations + +import base64 +import hashlib +import re +import shutil +import subprocess +import tempfile +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Protocol + +DEFAULT_DPI = 200 + +TRANSCRIBE_PROMPT = """\ +Transcribe this page from a mathematics textbook, exactly as printed. + +- Keep every word, in reading order. Do not summarise, correct or comment. +- Write mathematics as LaTeX: inline as $...$, displayed as $$...$$. +- Keep the printed labels ("Theorem 1.6.3.", "Exercise 4.") exactly as they appear. +- Include running heads, page numbers and footnotes as their own lines. +- If part of the page is a figure you cannot transcribe, write [figure] on its own line. + +Output only the transcription. +""" + + +class OcrBackend(Protocol): + """Anything that can turn a rendered page into text.""" + + name: str + + def transcribe(self, image: bytes) -> str: + ... + + +# -------------------------------------------------------------------------- +# rendering + + +def render_page(path: Path | str, page_number: int, dpi: int = DEFAULT_DPI) -> bytes: + """Render a 1-based page of a PDF to PNG bytes.""" + try: + import pymupdf as backend # type: ignore + except ImportError: # pragma: no cover - fitz is the old name + import fitz as backend # type: ignore + + with backend.open(str(path)) as document: + if not 1 <= page_number <= document.page_count: + raise IndexError(f"page {page_number} out of range 1..{document.page_count}") + page = document[page_number - 1] + return page.get_pixmap(dpi=dpi).tobytes("png") + + +# -------------------------------------------------------------------------- +# is the embedded text worth trusting? + +# Glyphs that a correctly-encoded mathematics page essentially never contains, but +# which mis-mapped maths fonts produce in quantity: ligature slots, spacing +# modifiers and stray accents standing in for operators and brackets. +_SUSPICIOUS = set("ŒœŠšŽžŸ˚˙˝˘˛˜ˆ¯´`¨¤ÐÞłŁ§¶") +_SUSPICIOUS_PAIRS = ("7!", "!7", "¤", "Š", "˚") + +# A second, nastier corruption: some maths fonts map operators onto ordinary ASCII +# letters, so the extracted text is made of perfectly innocent characters and the +# glyph test above scores it 1.0. On the project's test book this produces lines like +# +# I C J D fa C b W a 2 I and b 2 J g +# +# for "I + J = {a + b : a ∈ I and b ∈ J}" -- D is "=", C is "+", 2 is "∈", W is ":", +# f and g are the braces. Only the standalone capitals are counted: a bare "2" is +# the ∈ of this encoding, but it is also how people write "chapter 2", and including +# it produced false positives on ordinary prose. Measured across the test book, +# letters alone give clean separation -- 0.00 on the copyright page and on English +# sentences full of numerals, 10-11 on corrupted pages, median 5.6. +_OPERATOR_LETTERS = re.compile(r"(?<=[ (])[DCW](?=[ )])") +_MIN_LENGTH_FOR_OPERATOR_TEST = 200 + + +def text_quality(text: str) -> float: + """Rough confidence that extracted text says what the page says, in 0..1. + + This does not judge OCR accuracy or prose quality. It answers one question: + does this look like a font-encoding accident? Two independent signals, because + there are two failure modes and the obvious one misses the worse one: + + * exotic glyphs (Œ Š ˚ 7!) standing in for brackets and arrows; + * ordinary letters standing in for operators (D for =, C for +, 2 for ∈). + + The second scores a perfect 1.0 on the first test and is the more common of the + two, so both are required. Calibrated against this project's test book, where + 89% of pages carry the second signal and only the copyright page is clean. + """ + stripped = [c for c in text if not c.isspace()] + if not stripped: + return 0.0 + + # Signal two: operators disguised as letters, measured per 1000 characters. + # Short fragments are exempt -- the rate is meaningless over a single sentence. + if len(text) >= _MIN_LENGTH_FOR_OPERATOR_TEST: + operator_rate = 1000 * len(_OPERATOR_LETTERS.findall(text)) / len(text) + operator_score = max(0.0, 1.0 - operator_rate / 4) + else: + operator_score = 1.0 + + suspicious = sum(1 for c in stripped if c in _SUSPICIOUS) + if suspicious <= 2: + # Ordinary prose has the occasional accent ("Poincare\u00b4", "Erdos"). Only a + # pile of these is evidence of a font-encoding accident, and without this + # floor a short clean string is condemned by a single character. + score = 1.0 + else: + ratio = suspicious / len(stripped) + score = max(0.0, 1.0 - ratio * 40) # 2.5% suspicious glyphs -> 0 + for token in _SUSPICIOUS_PAIRS: + if text.count(token) >= 3: + score *= 0.5 + # Either signal alone condemns the page: both mean the maths is unreadable. + return round(min(score, operator_score), 3) + + +def looks_broken(text: str, threshold: float = 0.75) -> bool: + return text_quality(text) < threshold + + +# -------------------------------------------------------------------------- +# backends + + +@dataclass +class VisionLLMBackend: + """Read the page with a vision model over the OpenAI-compatible endpoint.""" + + client: Any + model: str = "qwen-vl-max" + prompt: str = TRANSCRIBE_PROMPT + detail: str = "high" + + @property + def name(self) -> str: + return f"vision:{self.model}" + + def transcribe(self, image: bytes) -> str: + data_url = "data:image/png;base64," + base64.b64encode(image).decode("ascii") + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": self.prompt}, + {"type": "image_url", "image_url": {"url": data_url}}, + ], + } + ] + return self.client.complete(messages, model=self.model, temperature=0.0).strip() + + +@dataclass +class TesseractBackend: + """Offline fallback. Fine on prose, poor on mathematics -- know what you get.""" + + language: str = "eng" + binary: str = "tesseract" + + @property + def name(self) -> str: + return f"tesseract:{self.language}" + + def available(self) -> bool: + return shutil.which(self.binary) is not None + + def transcribe(self, image: bytes) -> str: + if not self.available(): + raise RuntimeError(f"{self.binary} is not installed") + with tempfile.TemporaryDirectory(prefix="leanbench-ocr-") as directory: + source = Path(directory) / "page.png" + source.write_bytes(image) + result = subprocess.run( # noqa: S603 - operator-chosen binary + [self.binary, str(source), "stdout", "-l", self.language], + capture_output=True, + text=True, + timeout=300, + ) + if result.returncode != 0: + raise RuntimeError(f"tesseract failed: {result.stderr[:400]}") + return result.stdout.strip() + + +# -------------------------------------------------------------------------- +# cache + + +@dataclass +class OcrCache: + """Disk cache of transcriptions. Keyed so a re-run is free, not just faster.""" + + directory: Path + enabled: bool = True + hits: int = 0 + misses: int = 0 + + def __post_init__(self) -> None: + self.directory = Path(self.directory) + + def _path(self, document_hash: str, page: int, dpi: int, backend: str) -> Path: + token = hashlib.sha1(f"{dpi}:{backend}".encode()).hexdigest()[:10] + return self.directory / document_hash[:16] / f"p{page:05d}-{token}.txt" + + def get(self, document_hash: str, page: int, dpi: int, backend: str) -> str | None: + if not self.enabled: + return None + path = self._path(document_hash, page, dpi, backend) + if path.is_file(): + self.hits += 1 + return path.read_text(encoding="utf-8") + self.misses += 1 + return None + + def put(self, document_hash: str, page: int, dpi: int, backend: str, text: str) -> None: + if not self.enabled: + return + path = self._path(document_hash, page, dpi, backend) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(text, encoding="utf-8") + + +# -------------------------------------------------------------------------- +# the thing the rest of the pipeline calls + + +@dataclass +class OcrReader: + backend: OcrBackend + cache: OcrCache | None = None + dpi: int = DEFAULT_DPI + failures: dict[int, str] = field(default_factory=dict) + + def page_text(self, path: Path | str, document_hash: str, page_number: int) -> str | None: + """Transcribe one page. Returns None (and records why) if it cannot.""" + if self.cache is not None: + cached = self.cache.get(document_hash, page_number, self.dpi, self.backend.name) + if cached is not None: + return cached + try: + image = render_page(path, page_number, self.dpi) + text = self.backend.transcribe(image) + except Exception as error: # noqa: BLE001 - one bad page must not kill a book + self.failures[page_number] = str(error) + return None + text = _tidy(text) + if self.cache is not None: + self.cache.put(document_hash, page_number, self.dpi, self.backend.name, text) + return text + + +_FENCE = re.compile(r"^\s*```[a-z]*\s*\n(.*?)\n```\s*$", re.DOTALL) + + +def _tidy(text: str) -> str: + """Strip the markdown fence some models wrap transcriptions in.""" + match = _FENCE.match(text) + return (match.group(1) if match else text).strip() + + +def build_backend(settings: Any, client: Any = None) -> OcrBackend: + """Make the backend named by settings, falling back to tesseract offline.""" + kind = (getattr(settings, "backend", "vision") or "vision").lower() + if kind in ("vision", "llm", "auto"): + if client is None: + raise ValueError("the vision backend needs an LLM client") + return VisionLLMBackend(client=client, model=settings.model) + if kind == "tesseract": + return TesseractBackend(language=getattr(settings, "language", "eng")) + raise ValueError(f"unknown OCR backend {kind!r}") + + +__all__ = [ + "DEFAULT_DPI", + "OcrBackend", + "OcrCache", + "OcrReader", + "TesseractBackend", + "VisionLLMBackend", + "build_backend", + "looks_broken", + "render_page", + "text_quality", +] diff --git a/leanbench/pipeline.py b/leanbench/pipeline.py new file mode 100644 index 0000000..256bf5b --- /dev/null +++ b/leanbench/pipeline.py @@ -0,0 +1,206 @@ +"""The whole run, in order: load -> detect -> formalize -> verify -> store. + +``ingest`` is what the CLI, the Docker entrypoint and CI all call. Re-ingesting a +book (or just some of its pages) deletes exactly the rows those pages produced last +time and writes the new ones, so a source can be redone without collateral damage. +""" +from __future__ import annotations + +from collections.abc import Callable, Sequence +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from .config import Settings +from .db import Database, ItemRow, item_id +from .detect import DetectionReport, detect_document +from .document import Document +from .formalize import Formalization, formalize_all, verification_mode +from .llm import build_client +from .verify import LeanChecker, NullChecker + + +@dataclass +class IngestResult: + document: Document + report: DetectionReport + formalizations: list[Formalization] = field(default_factory=list) + stored: list[str] = field(default_factory=list) + removed: int = 0 + + def summary(self) -> dict[str, Any]: + verified = sum(1 for f in self.formalizations if f.verified) + proved = sum(1 for f in self.formalizations if f.proved) + skipped = sum(1 for f in self.formalizations if f.skipped) + return { + "source": self.document.title, + "source_hash": self.document.sha256, + "pages": len(self.document), + "detected": len(self.report.detections), + "statements": sum(1 for d in self.report.detections if d.kind == "statement"), + "problems": sum(1 for d in self.report.detections if d.kind == "problem"), + "multi_page": sum(1 for d in self.report.detections if d.spans_pages > 1), + "unresolved": len(self.report.unresolved), + "formalized": len(self.formalizations) - skipped, + "skipped": skipped, + "verified": verified, + "proved": proved, + "stored": len(self.stored), + "removed": self.removed, + } + + +def build_ocr(settings: Settings, client: Any = None): + """An OcrReader for this configuration, or None when OCR is switched off.""" + if settings.ocr.mode == "text": + return None + from .ocr import OcrCache, OcrReader, build_backend + + backend = build_backend(settings.ocr, client) + cache = OcrCache(settings.ocr.cache_dir) if settings.ocr.cache_dir else None + return OcrReader(backend=backend, cache=cache, dpi=settings.ocr.dpi) + + +def build_checker(settings: Settings, *, enabled: bool = True): + if not enabled: + return NullChecker() + checker = LeanChecker.from_settings(settings.lean) + return checker if checker.available() else NullChecker() + + +def to_row(document: Document, formalization: Formalization) -> ItemRow: + detection = formalization.detection + identifier = item_id( + document.sha256, + detection.start_page, + detection.start_index, + detection.end_page, + detection.end_index, + ) + verify = formalization.verify + return ItemRow( + id=identifier, + source_hash=document.sha256, + kind=detection.kind, + label=detection.label, + summary=detection.summary, + informal=formalization.informal, + start_page=detection.start_page, + start_index=detection.start_index, + end_page=detection.end_page, + end_index=detection.end_index, + page_checksum=document.page(detection.start_page).checksum, + source_text=detection.text, + lean_name=formalization.name, + lean_code=formalization.lean_code, + imports=formalization.imports, + has_proof=formalization.has_proof, + verified=bool(verify and verify.ok), + proved=bool(verify and verify.proved), + verify_mode=verify.mode if verify else verification_mode(detection), + diagnostics=(verify.diagnostics if verify else ""), + confidence=formalization.confidence, + match_method=detection.match_method, + match_score=detection.match_score, + truncated=detection.truncated, + types=formalization.index_types, + ) + + +def store(database: Database, document: Document, + formalizations: Sequence[Formalization]) -> list[str]: + database.upsert_source( + document.sha256, + title=document.title, + path=str(document.path), + pages=len(document), + extractor=document.extractor, + meta=document.meta, + ) + stored: list[str] = [] + for formalization in formalizations: + if formalization.skipped or not formalization.lean_code: + continue + stored.append(database.upsert_item(to_row(document, formalization))) + return stored + + +def ingest( + source: str | Path, + *, + settings: Settings | None = None, + client: Any = None, + database: Database | None = None, + checker: Any = None, + first_page: int | None = None, + last_page: int | None = None, + title: str = "", + replace: bool = True, + store_results: bool = True, + progress: Callable[[str], None] | None = None, +) -> IngestResult: + settings = settings or Settings.from_env() + client = client or build_client(settings.llm) + checker = checker if checker is not None else build_checker(settings) + owns_database = database is None + database = database or Database(settings.database) + + say = progress or (lambda _message: None) + reader = build_ocr(settings, client) if settings.ocr.mode != "text" else None + if reader is not None: + say(f"reading pages with {reader.backend.name} ({settings.ocr.mode} mode)") + document = Document.load( + source, + title=title, + mode=settings.ocr.mode, + ocr=reader, + threshold=settings.ocr.threshold, + first_page=first_page, + last_page=last_page, + progress=progress, + ) + say(f"{document.title}: {len(document)} pages, sha256 {document.sha256[:12]}, " + f"extractor {document.extractor}") + if reader is not None and reader.cache is not None: + say(f" OCR cache: {reader.cache.hits} hit(s), {reader.cache.misses} miss(es)") + + result = IngestResult(document=document, report=DetectionReport()) + try: + if store_results and replace: + if first_page or last_page: + result.removed = database.delete_pages( + document.sha256, first_page or 1, last_page or len(document) + ) + else: + result.removed = database.delete_source(document.sha256) + if result.removed: + say(f"removed {result.removed} item(s) from the previous run") + + say("detecting formalizable passages...") + result.report = detect_document( + client, + document, + first_page=first_page, + last_page=last_page, + max_window_pages=settings.max_window_pages, + progress=progress, + ) + say(result.report.summary()) + + say("formalizing...") + result.formalizations = formalize_all( + client, + result.report.detections, + checker=checker, + document_title=document.title, + repair_attempts=settings.repair_attempts, + progress=progress, + ) + + if store_results: + result.stored = store(database, document, result.formalizations) + say(f"stored {len(result.stored)} item(s) in {database.path}") + finally: + if owns_database: + database.close() + return result 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_ocr.py b/tests/test_ocr.py new file mode 100644 index 0000000..1a0d580 --- /dev/null +++ b/tests/test_ocr.py @@ -0,0 +1,244 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest + +from leanbench.document import Document +from leanbench.ocr import ( + OcrCache, + OcrReader, + VisionLLMBackend, + looks_broken, + text_quality, +) + +# Real extraction of page 90 of the project's test book: the maths font carries a +# custom encoding, so every operator arrives as a stray accent or ligature glyph. +BROKEN = """Proposition 1.11.7 (Chinese remainder theorem). Let a and b be relatively +prime natural numbers. There is an isomorphism of rings +Zab Š Za ˚ Zb; +defined by Œxab 7! .Œxa; Œxb/. +""" + +# What the page actually says, as the vision model reads it. +CLEAN = """Proposition 1.11.7 (Chinese remainder theorem). Let $a$ and $b$ be relatively +prime natural numbers. There is an isomorphism of rings +$$\\mathbb{Z}_{ab} \\cong \\mathbb{Z}_a \\oplus \\mathbb{Z}_b,$$ +defined by $[x]_{ab} \\mapsto ([x]_a, [x]_b)$. +""" + +# The other corruption mode, from page 300 of the same book: the maths is made of +# ordinary ASCII, so a glyph-based check scores it a perfect 1.0. +OPERATORS_AS_LETTERS = """6.2.18. Let I and J be two ideals in a ring R. Show that +I C J D fa C b W a 2 I and b 2 J g +is an ideal in R. +6.2.19. Let I and J be two ideals in a ring R. Show that +IJ D fa1b1 C a2b2 C C asbs W s 1; ai 2 I; bi 2 J g +is an ideal in R, and IJ I \\ J . +""" + +PROSE = """The chapter closes with a discussion of why the theorem matters, and a +remark on its history, with no displayed formulas at all. +""" + + +class FakeBackend: + """Returns a canned transcription and counts how often it was asked.""" + + name = "fake:v1" + + def __init__(self, text: str = CLEAN) -> None: + self.text = text + self.calls = 0 + + def transcribe(self, image: bytes) -> str: + self.calls += 1 + return self.text + + +class ExplodingBackend: + name = "boom:v1" + + def transcribe(self, image: bytes) -> str: + raise RuntimeError("model refused") + + +# -- the quality heuristic ------------------------------------------------- + + +def test_mangled_maths_scores_far_below_clean_text() -> None: + assert text_quality(BROKEN) < 0.75 + assert text_quality(CLEAN) > 0.9 + assert text_quality(PROSE) > 0.9 + + +def test_operators_disguised_as_letters_are_caught() -> None: + """The failure mode that a glyph-only check scores 1.0 and waves through.""" + assert text_quality(OPERATORS_AS_LETTERS) < 0.75 + assert looks_broken(OPERATORS_AS_LETTERS) + + +def test_ordinary_prose_with_numerals_survives() -> None: + assert not looks_broken( + "In chapter 2 we saw that the group has order 2 and index 2 in G.") + + +def test_looks_broken_separates_the_two() -> None: + assert looks_broken(BROKEN) + assert not looks_broken(CLEAN) + assert not looks_broken(PROSE) + + +def test_empty_text_is_the_worst_case() -> None: + assert text_quality("") == 0.0 + assert text_quality(" \n ") == 0.0 + + +def test_a_single_accented_word_is_not_damning() -> None: + # Real prose does contain the odd accent; only a pile of them is evidence. + assert not looks_broken("Consider the Poincare´ conjecture and its history.") + + +# -- the cache ------------------------------------------------------------- + + +def test_cache_round_trips(tmp_path: Path) -> None: + cache = OcrCache(tmp_path) + assert cache.get("hash", 3, 200, "fake:v1") is None + cache.put("hash", 3, 200, "fake:v1", "hello") + assert cache.get("hash", 3, 200, "fake:v1") == "hello" + assert cache.hits == 1 and cache.misses == 1 + + +def test_cache_keys_on_dpi_and_backend(tmp_path: Path) -> None: + cache = OcrCache(tmp_path) + cache.put("hash", 3, 200, "fake:v1", "at 200") + assert cache.get("hash", 3, 400, "fake:v1") is None, "different dpi is a different read" + assert cache.get("hash", 3, 200, "other:v1") is None, "different backend too" + + +def test_cache_can_be_switched_off(tmp_path: Path) -> None: + cache = OcrCache(tmp_path, enabled=False) + cache.put("hash", 1, 200, "fake:v1", "x") + assert cache.get("hash", 1, 200, "fake:v1") is None + + +# -- the reader ------------------------------------------------------------ + + +def test_reader_pays_for_a_page_once(tmp_path: Path, monkeypatch) -> None: + monkeypatch.setattr("leanbench.ocr.render_page", lambda *a, **k: b"png") + backend = FakeBackend() + reader = OcrReader(backend=backend, cache=OcrCache(tmp_path)) + first = reader.page_text("book.pdf", "hash", 7) + second = reader.page_text("book.pdf", "hash", 7) + assert first == second + assert backend.calls == 1, "the second read must come from cache" + + +def test_reader_records_failures_instead_of_raising(tmp_path: Path, monkeypatch) -> None: + monkeypatch.setattr("leanbench.ocr.render_page", lambda *a, **k: b"png") + reader = OcrReader(backend=ExplodingBackend(), cache=OcrCache(tmp_path)) + assert reader.page_text("book.pdf", "hash", 4) is None + assert "model refused" in reader.failures[4] + + +def test_reader_strips_a_markdown_fence(tmp_path: Path, monkeypatch) -> None: + monkeypatch.setattr("leanbench.ocr.render_page", lambda *a, **k: b"png") + reader = OcrReader(backend=FakeBackend("```\nTheorem 1.\n```"), cache=None) + assert reader.page_text("book.pdf", "hash", 1) == "Theorem 1." + + +# -- the vision backend's request shape ------------------------------------ + + +def test_vision_backend_sends_an_image_part() -> None: + captured: dict = {} + + class Client: + def complete(self, messages, **kwargs): + captured["messages"] = messages + captured["kwargs"] = kwargs + return " transcribed " + + backend = VisionLLMBackend(client=Client(), model="qwen-vl-max") + assert backend.transcribe(b"\x89PNG") == "transcribed" + parts = captured["messages"][0]["content"] + assert parts[0]["type"] == "text" + assert parts[1]["type"] == "image_url" + assert parts[1]["image_url"]["url"].startswith("data:image/png;base64,") + assert captured["kwargs"]["model"] == "qwen-vl-max" + assert backend.name == "vision:qwen-vl-max" + + +# -- Document integration -------------------------------------------------- + + +@pytest.fixture() +def mixed_book(tmp_path: Path) -> Path: + """Page 1 is fine, page 2 is font-mangled.""" + path = tmp_path / "mixed.txt" + path.write_text(PROSE + "\f" + BROKEN, encoding="utf-8") + return path + + +def test_auto_mode_only_pays_for_the_broken_page(mixed_book: Path, monkeypatch) -> None: + monkeypatch.setattr("leanbench.ocr.render_page", lambda *a, **k: b"png") + backend = FakeBackend() + reader = OcrReader(backend=backend, cache=None) + document = Document.load(mixed_book, mode="auto", ocr=reader, threshold=0.75) + + assert backend.calls == 1, "the clean page must not be sent to the model" + assert document.meta["ocr"]["pages_replaced"] == [2] + assert document.page(1).text.startswith("The chapter closes") + assert "\\cong" in document.page(2).text + assert "fake:v1" in document.extractor + + +def test_ocr_mode_reads_every_page(mixed_book: Path, monkeypatch) -> None: + monkeypatch.setattr("leanbench.ocr.render_page", lambda *a, **k: b"png") + backend = FakeBackend() + Document.load(mixed_book, mode="ocr", ocr=reader_for(backend)) + assert backend.calls == 2 + + +def reader_for(backend) -> OcrReader: + return OcrReader(backend=backend, cache=None) + + +def test_page_range_bounds_what_gets_read(mixed_book: Path, monkeypatch) -> None: + monkeypatch.setattr("leanbench.ocr.render_page", lambda *a, **k: b"png") + backend = FakeBackend() + document = Document.load(mixed_book, mode="ocr", ocr=reader_for(backend), + first_page=2, last_page=2) + assert backend.calls == 1, "a 25-page trial of a 600-page book costs 25 calls" + assert document.meta["ocr"]["pages_replaced"] == [2] + + +def test_text_mode_never_touches_the_model(mixed_book: Path) -> None: + backend = FakeBackend() + document = Document.load(mixed_book, mode="text", ocr=reader_for(backend)) + assert backend.calls == 0 + assert "Š" in document.page(2).text, "embedded text is left exactly as found" + assert "ocr" not in document.meta + + +def test_a_failed_page_keeps_its_embedded_text(mixed_book: Path, monkeypatch) -> None: + monkeypatch.setattr("leanbench.ocr.render_page", lambda *a, **k: b"png") + reader = OcrReader(backend=ExplodingBackend(), cache=None) + document = Document.load(mixed_book, mode="ocr", ocr=reader) + assert document.meta["ocr"]["pages_failed"] == [1, 2] + assert "Š" in document.page(2).text, "we fall back, we do not lose the page" + + +def test_offsets_stay_usable_after_ocr(mixed_book: Path, monkeypatch) -> None: + """The span contract must hold whichever extractor produced the text.""" + monkeypatch.setattr("leanbench.ocr.render_page", lambda *a, **k: b"png") + document = Document.load(mixed_book, mode="ocr", ocr=reader_for(FakeBackend())) + page = document.page(2) + start = page.text.index("Proposition") + end = page.text.index("rings") + len("rings") + assert document.slice(2, start, 2, end).startswith("Proposition 1.11.7") + # The extractor is part of the offset contract, so it is recorded on the source. + assert "fake:v1" in document.extractor 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") == []