social-data-attribution-w2/scripts/docs/extract_pdf_text.py

#!/usr/bin/env python3
"""Text-only PDF extraction using PyMuPDF.

- isolates PyMuPDF-powered parsing so downstream tooling can reuse it without notebook glue
- strips image blocks (opinionated: figure bodies are omitted intentionally)
- emits both plain text and LaTeX-friendly streams so users can choose whichever feeds their LMs best
"""

import argparse
from pathlib import Path
from typing import Iterable, Sequence

import fitz  # type: ignore


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Extract PDF body text without figures"
    )
    parser.add_argument("--pdf_path", type=Path, required=True)
    parser.add_argument("--output_dir", type=Path, default=Path("papers"))
    parser.add_argument(
        "--text_name", type=str, default=None, help="Override plain-text filename"
    )
    parser.add_argument(
        "--latex_name", type=str, default=None, help="Override LaTeX filename"
    )
    parser.add_argument(
        "--min_chars",
        type=int,
        default=12,
        help="Skip text blocks shorter than this many characters to avoid noise",
    )
    return parser.parse_args()


LATEX_REPLACEMENTS: tuple[tuple[str, str], ...] = (
    ("\\", r"\\textbackslash{}"),
    ("{", r"\\{"),
    ("}", r"\\}"),
    ("%", r"\\%"),
    ("$", r"\\$"),
    ("#", r"\\#"),
    ("_", r"\\_"),
    ("&", r"\\&"),
    ("~", r"\\textasciitilde{}"),
    ("^", r"\\textasciicircum{}"),
)


def latex_escape(text: str) -> str:
    escaped = text
    for old, new in LATEX_REPLACEMENTS:
        escaped = escaped.replace(old, new)
    return escaped


def _sorted_text_blocks(blocks: Sequence[Sequence]) -> Iterable[str]:
    filtered = []
    for block in blocks:
        if len(block) < 7:
            continue
        block_type = block[6]
        text = (block[4] or "").strip()
        if block_type != 0 or len(text) == 0:
            continue
        filtered.append((round(block[1], 2), round(block[0], 2), text))
    filtered.sort(key=lambda item: (item[0], item[1]))
    for _, _, text in filtered:
        yield text


def collect_pages(doc: fitz.Document, min_chars: int) -> list[list[str]]:
    pages: list[list[str]] = []
    for page in doc:
        page_blocks = [
            chunk
            for chunk in _sorted_text_blocks(page.get_text("blocks"))
            if len(chunk) >= min_chars
        ]
        if not page_blocks:
            continue
        pages.append(page_blocks)
    return pages


def render_plain_text(pages: list[list[str]]) -> str:
    return "\n\n\n".join("\n\n".join(blocks) for blocks in pages)


def render_latex(pages: list[list[str]]) -> str:
    sections = []
    for idx, blocks in enumerate(pages, start=1):
        content = "\n\n".join(latex_escape(block) for block in blocks)
        sections.append(f"\\section*{{Page {idx}}}\n{content}")
    body = "\n\n".join(sections)
    return (
        "\\documentclass{article}\n\\usepackage[utf8]{inputenc}\n\\begin{document}\n"
        + body
        + "\n\\end{document}\n"
    )


def main() -> None:
    args = parse_args()
    pdf_path = args.pdf_path.expanduser().resolve()
    if not pdf_path.exists():
        raise FileNotFoundError(f"PDF not found: {pdf_path}")

    target_dir = args.output_dir.expanduser().resolve()
    target_dir.mkdir(parents=True, exist_ok=True)

    text_name = args.text_name or f"{pdf_path.stem}_body.txt"
    latex_name = args.latex_name or f"{pdf_path.stem}_body.tex"

    with fitz.open(pdf_path) as document:
        pages = collect_pages(document, args.min_chars)

    if not pages:
        raise RuntimeError(f"No text blocks found in {pdf_path}")

    plain_text = render_plain_text(pages)
    latex_text = render_latex(pages)

    text_path = target_dir / text_name
    latex_path = target_dir / latex_name

    text_path.write_text(plain_text, encoding="utf-8")
    latex_path.write_text(latex_text, encoding="utf-8")

    print(f"Wrote plain text to {text_path}")
    print(f"Wrote LaTeX text to {latex_path}")


if __name__ == "__main__":
    main()