spectre_pdf 1.0.0

"""Settle the blocks-accuracy question.

The original parity harness scored spectre's blocks at ~0.43 using
pymupdf's `page.get_text("blocks")` as the reference. That number is
unmeasurable as quality unless we also know what pymupdf itself
scores against an independent reference on the same metric.

This script does two things:

  1. Treats pymupdf's `page.get_text("dict")` block partition as the
     ground truth. The "dict" output gives a hierarchical
     blocks → lines → spans tree with explicit block ids per span, so
     it's a clean partition of the document's visible characters.

  2. Computes a **pair-counting F1** for each candidate against that
     reference:

       - For every pair of tokens (a, b) emitted by the candidate, the
         candidate either places them in the same block or different
         blocks. Same for the reference.
       - TP = pair same in both. FP = same in candidate, different in
         ref. FN = different in candidate, same in ref. (TN ignored.)
       - F1 = 2·TP / (2·TP + FP + FN). Adjusted Rand-style; standard
         metric for partition comparison.

  3. Reports **two** F1s per PDF: pymupdf-blocks vs pymupdf-dict
     (internal baseline) and spectre vs pymupdf-dict (the contested
     number). Side-by-side average plus per-PDF deltas.

Usage:

    python scripts/blocks_baseline.py <corpus-dir>
"""
from __future__ import annotations

import statistics
import sys
from collections import defaultdict
from pathlib import Path

import fitz  # pymupdf
import spectre_rs


def tokenize_block_text(text: str) -> list[str]:
    """Lowercased whitespace-split tokens. Same normalization both sides."""
    return [t for t in text.lower().split() if t]


def pymupdf_dict_partition(pdf_bytes: bytes) -> tuple[list[str], list[int]]:
    """Use pymupdf's hierarchical dict output as the ground-truth
    partition. Returns (tokens, block_id_per_token)."""
    tokens: list[str] = []
    block_ids: list[int] = []
    bid = 0
    with fitz.open(stream=pdf_bytes, filetype="pdf") as doc:
        for page in doc:
            data = page.get_text("dict")
            for block in data.get("blocks", []):
                if block.get("type") != 0:  # type 1 = image; skip
                    continue
                block_text_parts: list[str] = []
                for line in block.get("lines", []):
                    for span in line.get("spans", []):
                        block_text_parts.append(span.get("text", ""))
                joined = " ".join(block_text_parts)
                toks = tokenize_block_text(joined)
                if not toks:
                    continue
                tokens.extend(toks)
                block_ids.extend([bid] * len(toks))
                bid += 1
    return tokens, block_ids


def pymupdf_blocks_partition(pdf_bytes: bytes) -> tuple[list[str], list[int]]:
    tokens: list[str] = []
    block_ids: list[int] = []
    bid = 0
    with fitz.open(stream=pdf_bytes, filetype="pdf") as doc:
        for page in doc:
            for block in page.get_text("blocks"):
                # block[4] is the text; block[6] is block_no in pymupdf.
                if not isinstance(block[4], str):
                    continue
                toks = tokenize_block_text(block[4])
                if not toks:
                    continue
                tokens.extend(toks)
                block_ids.extend([bid] * len(toks))
                bid += 1
    return tokens, block_ids


def spectre_blocks_partition(pdf_bytes: bytes) -> tuple[list[str], list[int]]:
    tokens: list[str] = []
    block_ids: list[int] = []
    bid = 0
    doc = spectre_rs.Document(pdf_bytes)
    for block in doc.blocks():
        toks = tokenize_block_text(block.text)
        if not toks:
            continue
        tokens.extend(toks)
        block_ids.extend([bid] * len(toks))
        bid += 1
    return tokens, block_ids


def pair_f1(
    cand_tokens: list[str],
    cand_block_ids: list[int],
    ref_tokens: list[str],
    ref_block_ids: list[int],
) -> tuple[float, int, int, int]:
    """Pair-counting F1 between the candidate and reference partitions.

    Tokens are aligned by walking both sides and matching the first
    occurrence of each token in the reference. Tokens that don't appear
    in the reference are excluded from the comparison (you can't fairly
    judge boundaries for tokens the reference doesn't carry).

    Returns (f1, tp, fp, fn).
    """
    # Reference: which block does each occurrence of a token belong to?
    # Build a queue of (block_id, position_in_ref) per token.
    ref_index: dict[str, list[int]] = defaultdict(list)
    for pos, tok in enumerate(ref_tokens):
        ref_index[tok].append(pos)

    # Walk the candidate; for each token, consume the next ref occurrence.
    consumed: dict[str, int] = defaultdict(int)
    aligned: list[tuple[int, int]] = []  # (cand_block, ref_block)
    for tok, cb in zip(cand_tokens, cand_block_ids):
        positions = ref_index.get(tok)
        if not positions:
            continue
        idx = consumed[tok]
        if idx >= len(positions):
            continue  # candidate has more of this token than reference
        ref_pos = positions[idx]
        consumed[tok] += 1
        aligned.append((cb, ref_block_ids[ref_pos]))

    n = len(aligned)
    if n < 2:
        return (1.0 if n == 0 else 0.0, 0, 0, 0)

    # Pair counts. To keep cost manageable on large docs, group by
    # (cand_block, ref_block) and use closed-form pair sums.
    cand_groups: dict[int, list[int]] = defaultdict(list)
    ref_groups: dict[int, list[int]] = defaultdict(list)
    pair_counts: dict[tuple[int, int], int] = defaultdict(int)
    for cb, rb in aligned:
        cand_groups[cb].append(rb)
        ref_groups[rb].append(cb)
        pair_counts[(cb, rb)] += 1

    def npairs(k: int) -> int:
        return k * (k - 1) // 2

    same_cand_pairs = sum(npairs(len(v)) for v in cand_groups.values())
    same_ref_pairs = sum(npairs(len(v)) for v in ref_groups.values())
    same_both = sum(npairs(c) for c in pair_counts.values())

    tp = same_both
    fp = same_cand_pairs - same_both
    fn = same_ref_pairs - same_both
    if tp == 0 and fp == 0 and fn == 0:
        return (1.0, tp, fp, fn)
    prec = tp / (tp + fp) if (tp + fp) else 0.0
    rec = tp / (tp + fn) if (tp + fn) else 0.0
    f1 = (2 * prec * rec / (prec + rec)) if (prec + rec) else 0.0
    return (f1, tp, fp, fn)


def main() -> int:
    if len(sys.argv) < 2:
        print(__doc__, file=sys.stderr)
        return 2
    p = Path(sys.argv[1])
    pdfs = sorted(p.glob("*.pdf")) if p.is_dir() else [p]
    if not pdfs:
        print("no PDFs", file=sys.stderr)
        return 2

    print(
        f"{'pdf':<22}{'pymupdf F1':>14}{'spectre F1':>14}{'Δ':>8}"
    )
    print("-" * 60)
    pmu_scores: list[float] = []
    spectre_scores: list[float] = []
    for path in pdfs:
        bytes_ = path.read_bytes()
        try:
            ref_t, ref_b = pymupdf_dict_partition(bytes_)
        except Exception as e:  # noqa: BLE001
            print(f"{path.name:<22}ref err: {e}")
            continue
        try:
            cand_t, cand_b = pymupdf_blocks_partition(bytes_)
            pmu_f1, *_ = pair_f1(cand_t, cand_b, ref_t, ref_b)
        except Exception as e:  # noqa: BLE001
            pmu_f1 = float("nan")
            print(f"{path.name:<22}pymupdf err: {e}")
        try:
            sp_t, sp_b = spectre_blocks_partition(bytes_)
            sp_f1, *_ = pair_f1(sp_t, sp_b, ref_t, ref_b)
        except Exception as e:  # noqa: BLE001
            sp_f1 = float("nan")
            print(f"{path.name:<22}spectre err: {e}")
        delta = sp_f1 - pmu_f1
        print(
            f"{path.name:<22}{pmu_f1:>14.3f}{sp_f1:>14.3f}{delta:>+8.3f}"
        )
        if pmu_f1 == pmu_f1:  # not NaN
            pmu_scores.append(pmu_f1)
        if sp_f1 == sp_f1:
            spectre_scores.append(sp_f1)

    print("-" * 60)
    if pmu_scores and spectre_scores:
        print(
            f"{'MEAN ({} PDFs)'.format(len(spectre_scores)):<22}"
            f"{statistics.mean(pmu_scores):>14.3f}"
            f"{statistics.mean(spectre_scores):>14.3f}"
            f"{statistics.mean(spectre_scores) - statistics.mean(pmu_scores):>+8.3f}"
        )
        print(
            f"{'MEDIAN':<22}"
            f"{statistics.median(pmu_scores):>14.3f}"
            f"{statistics.median(spectre_scores):>14.3f}"
            f"{statistics.median(spectre_scores) - statistics.median(pmu_scores):>+8.3f}"
        )
        print()
        print(
            f"reference = pymupdf get_text('dict')'s block partition. "
            f"pymupdf-blocks-vs-dict is the internal-baseline (how much pymupdf "
            f"agrees with itself across two of its own output modes)."
        )
    return 0


if __name__ == "__main__":
    sys.exit(main())