Create parsing/parse_rulebook.py

parsing/parse_rulebook.py

@@ -0,0 +1,46 @@
+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