lean-bench/leanbench/ocr.py

"""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",
]