lean-bench/leanbench/document.py

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