feat!: Docling v2 (#117)
--------- Signed-off-by: Christoph Auer <cau@zurich.ibm.com> Signed-off-by: Maxim Lysak <mly@zurich.ibm.com> Signed-off-by: Michele Dolfi <dol@zurich.ibm.com> Signed-off-by: Panos Vagenas <35837085+vagenas@users.noreply.github.com> Co-authored-by: Maxim Lysak <mly@zurich.ibm.com> Co-authored-by: Michele Dolfi <dol@zurich.ibm.com> Co-authored-by: Panos Vagenas <35837085+vagenas@users.noreply.github.com>
This commit is contained in:
Christoph Auer
GitHub
@@ -5,22 +5,27 @@ import time
|
||||
import warnings
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from typing import Annotated, Iterable, List, Optional
|
||||
from typing import Annotated, Dict, Iterable, List, Optional
|
||||
|
||||
import typer
|
||||
from docling_core.utils.file import resolve_file_source
|
||||
|
||||
from docling.backend.docling_parse_backend import DoclingParseDocumentBackend
|
||||
from docling.backend.pypdfium2_backend import PyPdfiumDocumentBackend
|
||||
from docling.datamodel.base_models import ConversionStatus
|
||||
from docling.datamodel.document import ConversionResult, DocumentConversionInput
|
||||
from docling.datamodel.base_models import (
|
||||
ConversionStatus,
|
||||
FormatToExtensions,
|
||||
InputFormat,
|
||||
OutputFormat,
|
||||
)
|
||||
from docling.datamodel.document import ConversionResult
|
||||
from docling.datamodel.pipeline_options import (
|
||||
EasyOcrOptions,
|
||||
PipelineOptions,
|
||||
OcrOptions,
|
||||
PdfPipelineOptions,
|
||||
TesseractCliOcrOptions,
|
||||
TesseractOcrOptions,
|
||||
)
|
||||
from docling.document_converter import DocumentConverter
|
||||
from docling.document_converter import DocumentConverter, FormatOption, PdfFormatOption
|
||||
|
||||
warnings.filterwarnings(action="ignore", category=UserWarning, module="pydantic|torch")
|
||||
warnings.filterwarnings(action="ignore", category=FutureWarning, module="easyocr")
|
||||
@@ -87,28 +92,28 @@ def export_documents(
|
||||
fname = output_dir / f"{doc_filename}.json"
|
||||
with fname.open("w") as fp:
|
||||
_log.info(f"writing JSON output to {fname}")
|
||||
fp.write(json.dumps(conv_res.render_as_dict()))
|
||||
fp.write(json.dumps(conv_res.document.export_to_dict()))
|
||||
|
||||
# Export Text format:
|
||||
if export_txt:
|
||||
fname = output_dir / f"{doc_filename}.txt"
|
||||
with fname.open("w") as fp:
|
||||
_log.info(f"writing Text output to {fname}")
|
||||
fp.write(conv_res.render_as_text())
|
||||
fp.write(conv_res.document.export_to_markdown(strict_text=True))
|
||||
|
||||
# Export Markdown format:
|
||||
if export_md:
|
||||
fname = output_dir / f"{doc_filename}.md"
|
||||
with fname.open("w") as fp:
|
||||
_log.info(f"writing Markdown output to {fname}")
|
||||
fp.write(conv_res.render_as_markdown())
|
||||
fp.write(conv_res.document.export_to_markdown())
|
||||
|
||||
# Export Document Tags format:
|
||||
if export_doctags:
|
||||
fname = output_dir / f"{doc_filename}.doctags"
|
||||
with fname.open("w") as fp:
|
||||
_log.info(f"writing Doc Tags output to {fname}")
|
||||
fp.write(conv_res.render_as_doctags())
|
||||
fp.write(conv_res.document.export_to_document_tokens())
|
||||
|
||||
else:
|
||||
_log.warning(f"Document {conv_res.input.file} failed to convert.")
|
||||
@@ -129,44 +134,31 @@ def convert(
|
||||
help="PDF files to convert. Can be local file / directory paths or URL.",
|
||||
),
|
||||
],
|
||||
export_json: Annotated[
|
||||
bool,
|
||||
typer.Option(
|
||||
..., "--json/--no-json", help="If enabled the document is exported as JSON."
|
||||
),
|
||||
] = False,
|
||||
export_md: Annotated[
|
||||
bool,
|
||||
typer.Option(
|
||||
..., "--md/--no-md", help="If enabled the document is exported as Markdown."
|
||||
),
|
||||
] = True,
|
||||
export_txt: Annotated[
|
||||
bool,
|
||||
typer.Option(
|
||||
..., "--txt/--no-txt", help="If enabled the document is exported as Text."
|
||||
),
|
||||
] = False,
|
||||
export_doctags: Annotated[
|
||||
bool,
|
||||
typer.Option(
|
||||
...,
|
||||
"--doctags/--no-doctags",
|
||||
help="If enabled the document is exported as Doc Tags.",
|
||||
),
|
||||
] = False,
|
||||
from_formats: List[InputFormat] = typer.Option(
|
||||
None,
|
||||
"--from",
|
||||
help="Specify input formats to convert from. Defaults to all formats.",
|
||||
),
|
||||
to_formats: List[OutputFormat] = typer.Option(
|
||||
None, "--to", help="Specify output formats. Defaults to Markdown."
|
||||
),
|
||||
ocr: Annotated[
|
||||
bool,
|
||||
typer.Option(
|
||||
..., help="If enabled, the bitmap content will be processed using OCR."
|
||||
),
|
||||
] = True,
|
||||
backend: Annotated[
|
||||
Backend, typer.Option(..., help="The PDF backend to use.")
|
||||
] = Backend.DOCLING,
|
||||
ocr_engine: Annotated[
|
||||
OcrEngine, typer.Option(..., help="The OCR engine to use.")
|
||||
] = OcrEngine.EASYOCR,
|
||||
abort_on_error: Annotated[
|
||||
bool,
|
||||
typer.Option(
|
||||
...,
|
||||
"--abort-on-error/--no-abort-on-error",
|
||||
help="If enabled, the bitmap content will be processed using OCR.",
|
||||
),
|
||||
] = False,
|
||||
output: Annotated[
|
||||
Path, typer.Option(..., help="Output directory where results are saved.")
|
||||
] = Path("."),
|
||||
@@ -182,6 +174,9 @@ def convert(
|
||||
):
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
|
||||
if from_formats is None:
|
||||
from_formats = [e for e in InputFormat]
|
||||
|
||||
input_doc_paths: List[Path] = []
|
||||
for src in input_sources:
|
||||
source = resolve_file_source(source=src)
|
||||
@@ -191,48 +186,54 @@ def convert(
|
||||
)
|
||||
raise typer.Abort()
|
||||
elif source.is_dir():
|
||||
input_doc_paths.extend(list(source.glob("**/*.pdf")))
|
||||
input_doc_paths.extend(list(source.glob("**/*.PDF")))
|
||||
for fmt in from_formats:
|
||||
for ext in FormatToExtensions[fmt]:
|
||||
input_doc_paths.extend(list(source.glob(f"**/*.{ext}")))
|
||||
input_doc_paths.extend(list(source.glob(f"**/*.{ext.upper()}")))
|
||||
else:
|
||||
input_doc_paths.append(source)
|
||||
|
||||
match backend:
|
||||
case Backend.PYPDFIUM2:
|
||||
do_cell_matching = ocr # only do cell matching when OCR enabled
|
||||
pdf_backend = PyPdfiumDocumentBackend
|
||||
case Backend.DOCLING:
|
||||
do_cell_matching = True
|
||||
pdf_backend = DoclingParseDocumentBackend
|
||||
case _:
|
||||
raise RuntimeError(f"Unexpected backend type {backend}")
|
||||
if to_formats is None:
|
||||
to_formats = [OutputFormat.MARKDOWN]
|
||||
|
||||
export_json = OutputFormat.JSON in to_formats
|
||||
export_md = OutputFormat.MARKDOWN in to_formats
|
||||
export_txt = OutputFormat.TEXT in to_formats
|
||||
export_doctags = OutputFormat.DOCTAGS in to_formats
|
||||
|
||||
match ocr_engine:
|
||||
case OcrEngine.EASYOCR:
|
||||
ocr_options = EasyOcrOptions()
|
||||
ocr_options: OcrOptions = EasyOcrOptions()
|
||||
case OcrEngine.TESSERACT_CLI:
|
||||
ocr_options = TesseractCliOcrOptions()
|
||||
case OcrEngine.TESSERACT:
|
||||
ocr_options = TesseractOcrOptions()
|
||||
case _:
|
||||
raise RuntimeError(f"Unexpected backend type {backend}")
|
||||
raise RuntimeError(f"Unexpected OCR engine type {ocr_engine}")
|
||||
|
||||
pipeline_options = PipelineOptions(
|
||||
pipeline_options = PdfPipelineOptions(
|
||||
do_ocr=ocr,
|
||||
ocr_options=ocr_options,
|
||||
do_table_structure=True,
|
||||
)
|
||||
pipeline_options.table_structure_options.do_cell_matching = do_cell_matching
|
||||
doc_converter = DocumentConverter(
|
||||
pipeline_options=pipeline_options,
|
||||
pdf_backend=pdf_backend,
|
||||
)
|
||||
pipeline_options.table_structure_options.do_cell_matching = True # do_cell_matching
|
||||
|
||||
# Define input files
|
||||
input = DocumentConversionInput.from_paths(input_doc_paths)
|
||||
format_options: Dict[InputFormat, FormatOption] = {
|
||||
InputFormat.PDF: PdfFormatOption(
|
||||
pipeline_options=pipeline_options,
|
||||
backend=DoclingParseDocumentBackend, # pdf_backend
|
||||
)
|
||||
}
|
||||
doc_converter = DocumentConverter(
|
||||
allowed_formats=from_formats,
|
||||
format_options=format_options,
|
||||
)
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
conv_results = doc_converter.convert(input)
|
||||
conv_results = doc_converter.convert_all(
|
||||
input_doc_paths, raises_on_error=abort_on_error
|
||||
)
|
||||
|
||||
output.mkdir(parents=True, exist_ok=True)
|
||||
export_documents(
|
||||
|
||||
Reference in New Issue
Block a user