diff --git a/docling/cli/__init__.py b/docling/cli/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/docling/cli/main.py b/docling/cli/main.py new file mode 100644 index 0000000..95551b1 --- /dev/null +++ b/docling/cli/main.py @@ -0,0 +1,257 @@ +import importlib +import json +import logging +import time +import warnings +from enum import Enum +from pathlib import Path +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 +from docling.datamodel.base_models import ConversionStatus, PipelineOptions +from docling.datamodel.document import ConversionResult, DocumentConversionInput +from docling.document_converter import DocumentConverter + +warnings.filterwarnings(action="ignore", category=UserWarning, module="pydantic|torch") +warnings.filterwarnings(action="ignore", category=FutureWarning, module="easyocr") + +_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(str, Enum): + PYPDFIUM2 = "pypdfium2" + DOCLING = "docling" + + +def export_documents( + conv_results: Iterable[ConversionResult], + output_dir: Path, + export_json: bool, + export_md: bool, + export_txt: bool, + export_doctags: bool, +): + + success_count = 0 + failure_count = 0 + + for conv_res in conv_results: + if conv_res.status == ConversionStatus.SUCCESS: + success_count += 1 + doc_filename = conv_res.input.file.stem + + # Export Deep Search document JSON format: + if export_json: + 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())) + + # 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()) + + # 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()) + + # 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()) + + else: + _log.warning(f"Document {conv_res.input.file} failed to convert.") + failure_count += 1 + + _log.info( + f"Processed {success_count + failure_count} docs, of which {failure_count} failed" + ) + + +@app.command(no_args_is_help=True) +def convert( + input_sources: Annotated[ + List[Path], + typer.Argument( + ..., + metavar="source", + help="PDF files to convert. Directories are also accepted.", + ), + ], + 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, + 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: List[Path] = [] + for source in input_sources: + 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) + + ########################################################################### + + # The following sections contain a combination of PipelineOptions + # and PDF Backends for various configurations. + # Uncomment one section at the time to see the differences in the output. + + doc_converter = None + if backend == Backend.PYPDFIUM2 and not ocr: # PyPdfium without OCR + pipeline_options = PipelineOptions() + pipeline_options.do_ocr = False + pipeline_options.do_table_structure = True + pipeline_options.table_structure_options.do_cell_matching = False + + doc_converter = DocumentConverter( + pipeline_options=pipeline_options, + pdf_backend=PyPdfiumDocumentBackend, + ) + + elif backend == Backend.PYPDFIUM2.value and ocr: # PyPdfium with OCR + pipeline_options = PipelineOptions() + pipeline_options.do_ocr = False + pipeline_options.do_table_structure = True + pipeline_options.table_structure_options.do_cell_matching = True + + doc_converter = DocumentConverter( + pipeline_options=pipeline_options, + pdf_backend=PyPdfiumDocumentBackend, + ) + + elif backend == Backend.DOCLING.value and not ocr: # Docling Parse without OCR + pipeline_options = PipelineOptions() + pipeline_options.do_ocr = False + pipeline_options.do_table_structure = True + pipeline_options.table_structure_options.do_cell_matching = True + + doc_converter = DocumentConverter( + pipeline_options=pipeline_options, + pdf_backend=DoclingParseDocumentBackend, + ) + + elif backend == Backend.DOCLING.value and ocr: # Docling Parse with OCR + pipeline_options = PipelineOptions() + pipeline_options.do_ocr = True + pipeline_options.do_table_structure = True + pipeline_options.table_structure_options.do_cell_matching = True + + doc_converter = DocumentConverter( + pipeline_options=pipeline_options, + pdf_backend=DoclingParseDocumentBackend, + ) + + ########################################################################### + + # Define input files + input = DocumentConversionInput.from_paths(input_doc_paths) + + start_time = time.time() + + conv_results = doc_converter.convert(input) + + output.mkdir(parents=True, exist_ok=True) + export_documents( + conv_results, + output_dir=output, + export_json=export_json, + export_md=export_md, + export_txt=export_txt, + export_doctags=export_doctags, + ) + + end_time = time.time() - start_time + + _log.info(f"All documents were converted in {end_time:.2f} seconds.") + + +if __name__ == "__main__": + app() diff --git a/poetry.lock b/poetry.lock index 9492e56..b8459ad 100644 --- a/poetry.lock +++ b/poetry.lock @@ -5699,6 +5699,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" @@ -6637,6 +6648,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-pytz" version = "2024.1.0.20240417" @@ -7200,4 +7228,4 @@ examples = ["langchain-huggingface", "langchain-milvus", "langchain-text-splitte [metadata] lock-version = "2.0" python-versions = "^3.10" -content-hash = "b881ea7a3504555707e0778c7c25631cbb353b78da04bd724852c7d34f39d46d" +content-hash = "f50b5f6158b688cc25f80253e3cec8e60d852d66a90fe8eb96798ea3c2372019" diff --git a/pyproject.toml b/pyproject.toml index 9842831..9bb0f94 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"} @@ -95,6 +96,10 @@ examples = [ "langchain-text-splitters", ] + +[tool.poetry.scripts] +docling = "docling.cli.main:app" + [build-system] requires = ["poetry-core"] build-backend = "poetry.core.masonry.api" diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000..48f863a --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,20 @@ +from typer.testing import CliRunner + +from docling.cli.main import app + +runner = CliRunner() + + +def test_cli_help(): + result = runner.invoke(app, ["--help"]) + assert result.exit_code == 0 + + +def test_cli_version(): + result = runner.invoke(app, ["--version"]) + assert result.exit_code == 0 + + +def test_cli_convert(): + result = runner.invoke(app, ["./tests/data/2305.03393v1-pg9.pdf"]) + assert result.exit_code == 0