mirror of
https://github.com/docling-project/docling.git
synced 2025-06-27 05:20:05 +00:00
67 lines
2.1 KiB
Python
67 lines
2.1 KiB
Python
|
from pathlib import Path
|
||
|
|
|
||
|
|
from docling_core.types.doc import CodeItem, TextItem
|
||
|
|
from docling_core.types.doc.labels import CodeLanguageLabel, DocItemLabel
|
||
|
|
|
||
|
|
from docling.backend.docling_parse_backend import DoclingParseDocumentBackend
|
||
|
|
from docling.backend.docling_parse_v2_backend import DoclingParseV2DocumentBackend
|
||
|
|
from docling.datamodel.base_models import InputFormat
|
||
|
|
from docling.datamodel.document import ConversionResult
|
||
|
|
from docling.datamodel.pipeline_options import PdfPipelineOptions
|
||
|
|
from docling.document_converter import DocumentConverter, PdfFormatOption
|
||
|
|
from docling.pipeline.standard_pdf_pipeline import StandardPdfPipeline
|
||
|
|
|
||
|
|
|
||
|
|
def get_converter():
|
||
|
|
|
||
|
|
pipeline_options = PdfPipelineOptions()
|
||
|
|
pipeline_options.generate_page_images = True
|
||
|
|
|
||
|
|
pipeline_options.do_ocr = False
|
||
|
|
pipeline_options.do_table_structure = False
|
||
|
|
pipeline_options.do_code_enrichment = True
|
||
|
|
pipeline_options.do_formula_enrichment = True
|
||
|
|
|
||
|
|
converter = DocumentConverter(
|
||
|
|
format_options={
|
||
|
|
InputFormat.PDF: PdfFormatOption(
|
||
|
|
backend=DoclingParseV2DocumentBackend,
|
||
|
|
pipeline_cls=StandardPdfPipeline,
|
||
|
|
pipeline_options=pipeline_options,
|
||
|
|
)
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
return converter
|
||
|
|
|
||
|
|
|
||
|
|
def test_code_and_formula_conversion():
|
||
|
|
pdf_path = Path("tests/data/code_and_formula.pdf")
|
||
|
|
converter = get_converter()
|
||
|
|
|
||
|
|
print(f"converting {pdf_path}")
|
||
|
|
|
||
|
|
doc_result: ConversionResult = converter.convert(pdf_path)
|
||
|
|
|
||
|
|
results = doc_result.document.texts
|
||
|
|
|
||
|
|
code_blocks = [el for el in results if isinstance(el, CodeItem)]
|
||
|
|
assert len(code_blocks) == 1
|
||
|
|
|
||
|
|
gt = 'public static void print() {\n System.out.println("Java Code");\n}'
|
||
|
|
|
||
|
|
predicted = code_blocks[0].text.strip()
|
||
|
|
assert predicted == gt, f"mismatch in text {predicted=}, {gt=}"
|
||
|
|
assert code_blocks[0].code_language == CodeLanguageLabel.JAVA
|
||
|
|
|
||
|
|
formula_blocks = [
|
||
|
|
el
|
||
|
|
for el in results
|
||
|
|
if isinstance(el, TextItem) and el.label == DocItemLabel.FORMULA
|
||
|
|
]
|
||
|
|
assert len(formula_blocks) == 1
|
||
|
|
|
||
|
|
gt = "a ^ { 2 } + 8 = 1 2"
|
||
|
|
predicted = formula_blocks[0].text
|
||
|
|
assert predicted == gt, f"mismatch in text {predicted=}, {gt=}"
|