docling/tests/test_e2e_conversion.py
Christoph Auer 15282547cb update test cases
Signed-off-by: Christoph Auer <cau@zurich.ibm.com>
2025-03-12 11:04:48 +01:00

61 lines
1.7 KiB
Python

from pathlib import Path
from docling.backend.docling_parse_v3_backend import DoclingParseV3DocumentBackend
from docling.datamodel.base_models import InputFormat
from docling.datamodel.document import ConversionResult
from docling.datamodel.pipeline_options import AcceleratorDevice, PdfPipelineOptions
from docling.document_converter import DocumentConverter, PdfFormatOption
from .verify_utils import verify_conversion_result_v1, verify_conversion_result_v2
GENERATE_V1 = True
GENERATE_V2 = True
def get_pdf_paths():
# Define the directory you want to search
directory = Path("./tests/data/pdf/")
# List all PDF files in the directory and its subdirectories
pdf_files = sorted(directory.rglob("*.pdf"))
return pdf_files
def get_converter():
pipeline_options = PdfPipelineOptions()
pipeline_options.do_ocr = False
pipeline_options.do_table_structure = True
pipeline_options.table_structure_options.do_cell_matching = True
pipeline_options.accelerator_options.device = AcceleratorDevice.CPU
converter = DocumentConverter(
format_options={
InputFormat.PDF: PdfFormatOption(
pipeline_options=pipeline_options, backend=DoclingParseV3DocumentBackend
)
}
)
return converter
def test_e2e_pdfs_conversions():
pdf_paths = get_pdf_paths()
converter = get_converter()
for pdf_path in pdf_paths:
print(f"converting {pdf_path}")
doc_result: ConversionResult = converter.convert(pdf_path)
verify_conversion_result_v1(
input_path=pdf_path, doc_result=doc_result, generate=GENERATE_V1
)
verify_conversion_result_v2(
input_path=pdf_path, doc_result=doc_result, generate=GENERATE_V2
)