parsing/parse_rulebook.py

from pathlib import Path

from docling.document_converter import DocumentConverter, PdfFormatOption
from docling.datamodel.base_models import InputFormat
from docling.datamodel.pipeline_options import PdfPipelineOptions, RapidOcrOptions

# Assumes this file lives at <repo_root>/parsing/parse_rulebook.py
PROJECT_ROOT = Path(__file__).resolve().parent.parent
RULEBOOK_DIR = PROJECT_ROOT / "rulebooks"

OUTPUT_DIR = PROJECT_ROOT / "parsed"
OUTPUT_DIR.mkdir(exist_ok=True)

def parse(file_name: str, out_name: str, ocr = False):

    pdf_path =  RULEBOOK_DIR / file_name
    output_file = OUTPUT_DIR / f"{out_name}.md"

    if not pdf_path.exists():
        raise FileNotFoundError(f"Could not find {pdf_path}")

    pipeline_options = PdfPipelineOptions()
    pipeline_options.do_ocr = ocr

    if ocr:
        pipeline_options.ocr_options = RapidOcrOptions(
            lang=["english"],
            force_full_page_ocr=True,
            backend="onnxruntime",
        )
    converter = DocumentConverter(
    format_options={
        InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options)
        }
    )

    doc = converter.convert(str(pdf_path)).document

    markdown = doc.export_to_markdown()

    with open(output_file, "w", encoding = "utf-8" ) as f:
        f.write(markdown)

    print(f"Saved parsed rulebook to {output_file}")

    return output_file