File size: 1,376 Bytes
f656538 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 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 |