From faf66aff573c6b3ee18a3e94cfc9fbd513cdba41 Mon Sep 17 00:00:00 2001 From: Michele Dolfi Date: Thu, 12 Sep 2024 14:47:06 +0200 Subject: [PATCH] use typer for the docling CLI Signed-off-by: Michele Dolfi --- docling/cli/__init__.py | 0 examples/convert.py => docling/cli/main.py | 134 +++++++++++++-------- poetry.lock | 35 +++++- pyproject.toml | 5 + 4 files changed, 124 insertions(+), 50 deletions(-) create mode 100644 docling/cli/__init__.py rename examples/convert.py => docling/cli/main.py (60%) diff --git a/docling/cli/__init__.py b/docling/cli/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/examples/convert.py b/docling/cli/main.py similarity index 60% rename from examples/convert.py rename to docling/cli/main.py index 7f986551..acb21715 100644 --- a/examples/convert.py +++ b/docling/cli/main.py @@ -1,9 +1,14 @@ -import argparse +import importlib import json import logging import time +import warnings +from enum import Enum from pathlib import Path -from typing import Iterable +from typing import Annotated, Iterable, List, Optional + +import typer +from pydantic import AnyUrl from docling.backend.docling_parse_backend import DoclingParseDocumentBackend from docling.backend.pypdfium2_backend import PyPdfiumDocumentBackend @@ -11,14 +16,39 @@ from docling.datamodel.base_models import ConversionStatus, PipelineOptions from docling.datamodel.document import ConversionResult, DocumentConversionInput from docling.document_converter import DocumentConverter -_log = logging.getLogger(__name__) +warnings.filterwarnings(action="ignore", category=UserWarning, module="pydantic|torch") +warnings.filterwarnings(action="ignore", category=FutureWarning, module="easyocr") -from enum import Enum +_log = logging.getLogger(__name__) +from rich.console import Console + +err_console = Console(stderr=True) + + +app = typer.Typer( + name="Docling", + no_args_is_help=True, + add_completion=False, + pretty_exceptions_enable=False, +) + + +def version_callback(value: bool): + if value: + docling_version = importlib.metadata.version("docling") + docling_core_version = importlib.metadata.version("docling-core") + docling_ibm_models_version = importlib.metadata.version("docling-ibm-models") + docling_parse_version = importlib.metadata.version("docling-parse") + print(f"Docling version: {docling_version}") + print(f"Docling Core version: {docling_core_version}") + print(f"Docling IBM Models version: {docling_ibm_models_version}") + print(f"Docling Parse version: {docling_parse_version}") + raise typer.Exit() # Define an enum for the backend options -class Backend(Enum): - PDFIUM = "pdfium" +class Backend(str, Enum): + PYPDFIUM2 = "pypdfium2" DOCLING = "docling" @@ -26,7 +56,6 @@ def export_documents( conv_results: Iterable[ConversionResult], output_dir: Path, ): - output_dir.mkdir(parents=True, exist_ok=True) success_count = 0 failure_count = 0 @@ -62,13 +91,52 @@ def export_documents( f"Processed {success_count + failure_count} docs, of which {failure_count} failed" ) - return success_count, failure_count - -def main(pdf, ocr, backend): +@app.command(no_args_is_help=True) +def convert( + input_files: Annotated[ + List[Path], + typer.Argument( + ..., + metavar="file", + help="PDF files to convert. Directories are also accepted.", + ), + ], + 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, + output: Annotated[ + Path, typer.Option(..., help="Output directory where results are saved.") + ] = Path("."), + version: Annotated[ + Optional[bool], + typer.Option( + "--version", + callback=version_callback, + is_eager=True, + help="Show version information.", + ), + ] = None, +): logging.basicConfig(level=logging.INFO) - input_doc_paths = [Path(pdf)] + input_doc_paths: List[Path] = [] + for source in input_files: + if not source.exists(): + err_console.print( + f"[red]Error: The input file {source} does not exist.[/red]" + ) + raise typer.Abort() + elif source.is_dir(): + input_doc_paths.extend(list(source.glob("**/*.pdf", case_sensitive=False))) + else: + input_doc_paths.append(source) ########################################################################### @@ -77,7 +145,7 @@ def main(pdf, ocr, backend): # Uncomment one section at the time to see the differences in the output. doc_converter = None - if backend == Backend.PDFIUM.value and not ocr: # PyPdfium without OCR + if backend == Backend.PYPDFIUM2 and not ocr: # PyPdfium without OCR pipeline_options = PipelineOptions() pipeline_options.do_ocr = False pipeline_options.do_table_structure = True @@ -88,7 +156,7 @@ def main(pdf, ocr, backend): pdf_backend=PyPdfiumDocumentBackend, ) - elif backend == Backend.PDFIUM.value and ocr: # PyPdfium with OCR + elif backend == Backend.PYPDFIUM2.value and ocr: # PyPdfium with OCR pipeline_options = PipelineOptions() pipeline_options.do_ocr = False pipeline_options.do_table_structure = True @@ -121,8 +189,6 @@ def main(pdf, ocr, backend): pdf_backend=DoclingParseDocumentBackend, ) - else: - return ########################################################################### # Define input files @@ -131,44 +197,14 @@ def main(pdf, ocr, backend): start_time = time.time() conv_results = doc_converter.convert(input) - success_count, failure_count = export_documents( - conv_results, output_dir=Path("./scratch") - ) + + output.mkdir(parents=True, exist_ok=True) + export_documents(conv_results, output_dir=output) end_time = time.time() - start_time _log.info(f"All documents were converted in {end_time:.2f} seconds.") - if failure_count > 0: - raise RuntimeError( - f"The example failed converting {failure_count} on {len(input_doc_paths)}." - ) - if __name__ == "__main__": - - # Create an argument parser - parser = argparse.ArgumentParser(description="Process PDF files with optional OCR.") - - # Add arguments - parser.add_argument( - "--pdf", - type=str, - default="./tests/data/2206.01062.pdf", - help="Path to the PDF file.", - ) - parser.add_argument( - "--ocr", type=bool, default=False, help="Enable OCR (True or False)." - ) - parser.add_argument( - "--backend", - type=lambda b: Backend[b.upper()], - choices=list(Backend), - default=Backend.DOCLING, - help="Select backend (pdfium or docling). Default is docling.", - ) - - # Parse the arguments - args = parser.parse_args() - - main(args.pdf, args.ocr, args.backend.value) + app() diff --git a/poetry.lock b/poetry.lock index e711babc..3f30f438 100644 --- a/poetry.lock +++ b/poetry.lock @@ -5684,6 +5684,17 @@ numpy = ">=1.14,<3" docs = ["matplotlib", "numpydoc (==1.1.*)", "sphinx", "sphinx-book-theme", "sphinx-remove-toctrees"] test = ["pytest", "pytest-cov"] +[[package]] +name = "shellingham" +version = "1.5.4" +description = "Tool to Detect Surrounding Shell" +optional = false +python-versions = ">=3.7" +files = [ + {file = "shellingham-1.5.4-py2.py3-none-any.whl", hash = "sha256:7ecfff8f2fd72616f7481040475a65b2bf8af90a56c89140852d1120324e8686"}, + {file = "shellingham-1.5.4.tar.gz", hash = "sha256:8dbca0739d487e5bd35ab3ca4b36e11c4078f3a234bfce294b0a0291363404de"}, +] + [[package]] name = "simplejson" version = "3.19.3" @@ -6584,6 +6595,11 @@ files = [ {file = "triton-3.0.0-1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:34e509deb77f1c067d8640725ef00c5cbfcb2052a1a3cb6a6d343841f92624eb"}, {file = "triton-3.0.0-1-cp38-cp38-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:bcbf3b1c48af6a28011a5c40a5b3b9b5330530c3827716b5fbf6d7adcc1e53e9"}, {file = "triton-3.0.0-1-cp39-cp39-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6e5727202f7078c56f91ff13ad0c1abab14a0e7f2c87e91b12b6f64f3e8ae609"}, + {file = "triton-3.0.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:39b052da883351fdf6be3d93cedae6db3b8e3988d3b09ed221bccecfa9612230"}, + {file = "triton-3.0.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:cd34f19a8582af96e6291d4afce25dac08cb2a5d218c599163761e8e0827208e"}, + {file = "triton-3.0.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0d5e10de8c011adeb7c878c6ce0dd6073b14367749e34467f1cff2bde1b78253"}, + {file = "triton-3.0.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e8903767951bf86ec960b4fe4e21bc970055afc65e9d57e916d79ae3c93665e3"}, + {file = "triton-3.0.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:41004fb1ae9a53fcb3e970745feb87f0e3c94c6ce1ba86e95fa3b8537894bef7"}, ] [package.dependencies] @@ -6617,6 +6633,23 @@ rfc3986 = ">=1.4.0" tqdm = ">=4.14" urllib3 = ">=1.26.0" +[[package]] +name = "typer" +version = "0.12.5" +description = "Typer, build great CLIs. Easy to code. Based on Python type hints." +optional = false +python-versions = ">=3.7" +files = [ + {file = "typer-0.12.5-py3-none-any.whl", hash = "sha256:62fe4e471711b147e3365034133904df3e235698399bc4de2b36c8579298d52b"}, + {file = "typer-0.12.5.tar.gz", hash = "sha256:f592f089bedcc8ec1b974125d64851029c3b1af145f04aca64d69410f0c9b722"}, +] + +[package.dependencies] +click = ">=8.0.0" +rich = ">=10.11.0" +shellingham = ">=1.3.0" +typing-extensions = ">=3.7.4.3" + [[package]] name = "types-requests" version = "2.32.0.20240907" @@ -7169,4 +7202,4 @@ examples = ["langchain-huggingface", "langchain-milvus", "langchain-text-splitte [metadata] lock-version = "2.0" python-versions = "^3.10" -content-hash = "5ce8fc1e245442e355b967430e211b1378fed2e9fd20d2ddbea47f0e9f1dfcd5" +content-hash = "e9eaa2a2de2ef321d274bbe245a290826c8604a730e9e65149421a6d2cfe2202" diff --git a/pyproject.toml b/pyproject.toml index 0542e679..274dae3d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -60,6 +60,7 @@ torchvision = [ {version = "^0", optional = true, markers = "sys_platform != 'darwin' or platform_machine != 'x86_64'"}, {version = "~0.17.2", optional = true, markers = "sys_platform == 'darwin' and platform_machine == 'x86_64'"} ] +typer = "^0.12.5" [tool.poetry.group.dev.dependencies] black = {extras = ["jupyter"], version = "^24.4.2"} @@ -94,6 +95,10 @@ examples = [ "langchain-text-splitters", ] + +[tool.poetry.scripts] +docling = "docling.cli.main:app" + [build-system] requires = ["poetry-core"] build-backend = "poetry.core.masonry.api"