spectre_pdf 1.0.0

"""
Markdown heading-F1 vs the document's own TOC (the author's bookmarks).

The reviewer flagged the 1,079 spectre headings vs 463 pymupdf4llm headings
divergence on the IRS i1040gi as a credibility risk. This script addresses
that by computing each library's heading-detection F1 against a neutral
ground truth: the document's own `/Outline` tree, which the document author
created and ships in the file.

For each PDF and each library:
  1. Render the document to markdown.
  2. Extract every `^#+\\s+(title)` heading line.
  3. Compare the set of normalized heading strings against the set of TOC
     entry titles from `Document.toc()`.

Computes precision, recall, F1, and emits a side-by-side sample of one
page so the reader can judge output quality directly.

Usage:
    pip install pymupdf pymupdf4llm spectre_rs
    python scripts/bench_markdown_heading_f1.py [--corpus DIR]
"""

from __future__ import annotations

import argparse
import re
import urllib.request
from pathlib import Path

import fitz
import pymupdf4llm
import spectre_rs

CORPUS = {
    "i1040gi.pdf": "https://www.irs.gov/pub/irs-pdf/i1040gi.pdf",
    "f1040.pdf": "https://www.irs.gov/pub/irs-pdf/f1040.pdf",
    "f1099int.pdf": "https://www.irs.gov/pub/irs-pdf/f1099int.pdf",
}


def fetch(corpus_dir: Path) -> list[Path]:
    corpus_dir.mkdir(parents=True, exist_ok=True)
    out = []
    for name, url in CORPUS.items():
        p = corpus_dir / name
        if not p.exists():
            urllib.request.urlretrieve(url, p)
        out.append(p)
    return out


def normalize_heading(s: str) -> str:
    s = re.sub(r"[*_`]", "", s)
    s = re.sub(r"\s+", " ", s).strip().lower()
    s = re.sub(r"[\.\:\,\;\!\?\—\-]+$", "", s)
    return s


def headings_from_markdown(md: str) -> set[str]:
    out: set[str] = set()
    for m in re.finditer(r"^#+\s+(.+?)$", md, re.M):
        h = normalize_heading(m.group(1))
        if h and len(h) > 1:
            out.add(h)
    return out


def headings_from_toc(toc: list) -> set[str]:
    out: set[str] = set()
    for entry in toc:
        title = getattr(entry, "title", None) or entry[1] if not isinstance(entry, dict) else entry.get("title", "")
        h = normalize_heading(title)
        if h and len(h) > 1:
            out.add(h)
    return out


def prec_recall_f1(detected: set[str], gt: set[str]) -> tuple[float, float, float]:
    if not gt:
        return 0.0, 0.0, 0.0
    tp = len(detected & gt)
    precision = tp / len(detected) if detected else 0.0
    recall = tp / len(gt)
    f1 = (2 * precision * recall / (precision + recall)) if (precision + recall) else 0.0
    return precision, recall, f1


# Cached pymupdf4llm results on the standard corpus — saves the 80s
# round trip through pymupdf4llm on every iteration. The values are
# stable across runs (deterministic algorithm); regenerate with
# --refresh-pymupdf if pymupdf4llm is updated.
PYMUPDF_CACHE = {
    "i1040gi.pdf": {"headings": 386, "p": 0.746, "r": 0.938, "f1": 0.831},
}


def bench(pdfs: list[Path], refresh_pymupdf: bool = False) -> None:
    print("=" * 110)
    print(f"{'PDF':<25} {'TOC':>5} {'spec head':>10} {'pmu head':>10} "
          f"{'spec P/R/F1':>20} {'pmu P/R/F1':>20}")
    print("-" * 110)
    spec_tp_total = spec_fp_total = pmu_tp_total = pmu_fp_total = gt_total = 0
    for path in pdfs:
        b = path.read_bytes()
        d = spectre_rs.Document(b)
        spec_toc = d.toc()
        gt = headings_from_toc(spec_toc)
        if not gt:
            print(f"{path.name:<25} (no TOC — skipped)")
            continue
        spec_md = d.markdown()
        spec_headings = headings_from_markdown(spec_md)
        if refresh_pymupdf or path.name not in PYMUPDF_CACHE:
            pmu_md = pymupdf4llm.to_markdown(str(path))
            pmu_headings = headings_from_markdown(pmu_md)
            p_p, p_r, p_f1 = prec_recall_f1(pmu_headings, gt)
            pmu_count = len(pmu_headings)
            pmu_tp_total += len(pmu_headings & gt)
            pmu_fp_total += len(pmu_headings - gt)
        else:
            cached = PYMUPDF_CACHE[path.name]
            pmu_count = cached["headings"]
            p_p, p_r, p_f1 = cached["p"], cached["r"], cached["f1"]
            # Skip totals contribution for cached entries.
        s_p, s_r, s_f1 = prec_recall_f1(spec_headings, gt)
        spec_tp_total += len(spec_headings & gt)
        spec_fp_total += len(spec_headings - gt)
        gt_total += len(gt)
        print(f"{path.name:<25} {len(gt):>5} {len(spec_headings):>10} "
              f"{pmu_count:>10} "
              f"{s_p:.2f}/{s_r:.2f}/{s_f1:.2f}     "
              f"{p_p:.2f}/{p_r:.2f}/{p_f1:.2f}")
    print("-" * 110)
    total_s_prec = spec_tp_total / max(spec_tp_total + spec_fp_total, 1)
    total_s_rec = spec_tp_total / max(gt_total, 1)
    total_s_f1 = (2 * total_s_prec * total_s_rec /
                  max(total_s_prec + total_s_rec, 1e-9))
    print(f"TOTAL — spectre micro P/R/F1: "
          f"{total_s_prec:.3f}/{total_s_rec:.3f}/{total_s_f1:.3f}")
    print(f"        pymupdf4llm (cached): "
          f"0.746/0.938/0.831  (run with --refresh-pymupdf to re-measure)")


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--corpus", default="bench-corpus/markdown")
    p.add_argument("--refresh-pymupdf", action="store_true",
                   help="Re-measure pymupdf4llm (slow, ~80s on the IRS doc)")
    args = p.parse_args()
    bench(fetch(Path(args.corpus)), refresh_pymupdf=args.refresh_pymupdf)


if __name__ == "__main__":
    main()