scripts/extract_pdf_text.py

from __future__ import annotations

"""
Extract text from every page of a fragrance literature PDF using PyMuPDF.

Outputs a JSONL file where each line is:
    {"source": "filename", "page": 1, "text": "..."}
"""

import argparse
import json
import logging
from pathlib import Path

import fitz

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger("extract_pdf_text")


def extract_text_from_pdf(pdf_path: Path) -> list[dict[str, str]]:
    doc = fitz.open(str(pdf_path))
    records = []
    for i in range(len(doc)):
        page = doc.load_page(i)
        text = page.get_text()
        records.append({
            "source": pdf_path.stem,
            "page": i + 1,
            "text": text,
        })
    doc.close()
    return records


def main() -> int:
    parser = argparse.ArgumentParser(description="Extract text from PDF pages")
    parser.add_argument("--input", required=True, help="PDF file or directory of PDFs")
    parser.add_argument("--output", required=True, help="Output JSONL path")
    args = parser.parse_args()

    input_path = Path(args.input)
    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    if input_path.is_dir():
        pdf_files = sorted(input_path.glob("*.pdf"))
    else:
        pdf_files = [input_path]

    with output_path.open("w", encoding="utf-8") as out:
        for pdf_file in pdf_files:
            logger.info("Extracting %s", pdf_file.name)
            records = extract_text_from_pdf(pdf_file)
            for record in records:
                out.write(json.dumps(record, ensure_ascii=False) + "\n")

    logger.info("Done: %s", output_path)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())