lean-bench/leanbench/anchors.py

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