feature(extract-python): add a debug cli to extract-python

Author: ClemDoum

extract-core/extract_core/configs.py

@@ -1,7 +1,7 @@
 from abc import ABC, abstractmethod
 from enum import StrEnum
 from pathlib import Path
-from typing import ClassVar
+from typing import ClassVar, Self
 
 from icij_common.pydantic_utils import icij_config, merge_configs, no_enum_values_config
 from icij_common.registrable import RegistrableConfig
@@ -29,6 +29,9 @@ class BasePipelineConfig(RegistrableConfig, ABC):
     @abstractmethod
     def supported_exts(cls) -> set[SupportedExt]: ...
 
+    @abstractmethod
+    def to_device(self) -> Self: ...
+
 
 class ResultBufferConfig(BaseModel):
     max_size: ByteSize = "500MiB"

extract-core/extract_core/docling_.py

@@ -3,7 +3,7 @@
 from collections import defaultdict
 from copy import deepcopy
 from functools import cache
-from typing import Annotated, Any, ClassVar, get_type_hints
+from typing import Annotated, Any, ClassVar, Self, get_type_hints
 
 from docling.datamodel.backend_options import BackendOptions, BaseBackendOptions
 from docling.datamodel.base_models import (
@@ -159,7 +159,6 @@ def _table_structure_opts_type_adapter() -> TypeAdapter:
 def _resolve_pipeline_options(
     pipeline_options: dict[str, Any] | None,
     pipeline_cls: type,
-    device: Device,
 ) -> PipelineOptions:
     option_cls = _find_init_arg_type(pipeline_cls, "pipeline_options")
     if pipeline_options is None:
@@ -198,10 +197,6 @@ def _resolve_pipeline_options(
         )
         pipeline_options["table_structure_options"] = table_structure_opts
     pipeline_options = option_cls.model_validate(pipeline_options)
-    accelerator_opts = getattr(pipeline_options, "accelerator_options", None)
-    if accelerator_opts is not None:
-        accelerator_opts.device = device.to_docling()
-        pipeline_options.accelerator_options = accelerator_opts
     return pipeline_options
 
 
@@ -218,13 +213,11 @@ class DoclingFormatOption(BaseFormatOption):
     pipeline_cls: str
     pipeline_options: dict[str, Any] | None = None
 
-    def to_docling(self, device: Device) -> BaseFormatOption:  # noqa: ANN201
+    def to_docling(self) -> BaseFormatOption:  # noqa: ANN201
         from docling.document_converter import FormatOption  # noqa: PLC0415
 
         pipeline_cls = _resolve_pipeline_cls(self.pipeline_cls)
-        pipeline_opts = _resolve_pipeline_options(
-            self.pipeline_options, pipeline_cls, device
-        )
+        pipeline_opts = _resolve_pipeline_options(self.pipeline_options, pipeline_cls)
         pipeline_opts = _validate_pipeline_opts(pipeline_opts)
         return FormatOption(
             pipeline_cls=pipeline_cls,
@@ -448,3 +441,23 @@ def supported_exts(cls) -> set[SupportedExt]:
     def with_setting(self, settings: DoclingSettings) -> "DoclingPipelineConfig":
         update = {"settings": settings}
         return safe_copy(self, update=update)
+
+    def to_device(self) -> Self:
+        device = self.device.to_docling()
+        new_format_opts = dict()
+        for fmt, fmt_opts in self.format_options.items():
+            pipeline_opts = fmt_opts.pipeline_options
+            if pipeline_opts is None:
+                pipeline_opts = {"generate_picture_images": True}
+            accelerator_opts = pipeline_opts.get("accelerator_options", dict())
+            accelerator_opts["device"] = device
+            pipeline_opts["accelerator_options"] = accelerator_opts
+            update = {"pipeline_options": pipeline_opts}
+            new_format_opts[fmt] = safe_copy(fmt_opts, update=update)
+
+        return DoclingPipelineConfig(
+            device=self.device,
+            format_options=new_format_opts,
+            settings=self.settings,
+            result_buffer=self.result_buffer,
+        )

extract-core/extract_core/marker_.py

@@ -1,5 +1,6 @@
+from copy import deepcopy
 from functools import cache
-from typing import Any, ClassVar
+from typing import Any, ClassVar, Self
 
 from pydantic import Field
 
@@ -38,3 +39,6 @@ def supported_exts(cls) -> set[SupportedExt]:
             SupportedExt.GIF,
             SupportedExt.TIFF,
         }
+
+    def to_device(self) -> Self:
+        return deepcopy(self)

extract-core/extract_core/miner_u.py

@@ -1,7 +1,8 @@
 from collections.abc import Callable
+from copy import deepcopy
 from enum import StrEnum
 from functools import cache
-from typing import ClassVar, Literal
+from typing import ClassVar, Literal, Self
 
 from mineru.config import VlmConfig
 from mineru.types import Tier
@@ -44,3 +45,6 @@ def supported_exts(cls) -> set[SupportedExt]:
             SupportedExt.PPTX,
             SupportedExt.XLSX,
         }
+
+    def to_device(self) -> Self:
+        return deepcopy(self)

extract-core/extract_core/pipeline.py

@@ -12,7 +12,7 @@
 
 class Pipeline[C: BasePipelineConfig](RegistrableFromConfig, ABC):
     def __init__(self, config: C):
-        self._config = config
+        self._config = config.to_device()
         self._device = self._config.device
 
     @abstractmethod

extract-core/tests/test_docling.py

@@ -47,14 +47,18 @@ def test_format_option_ser() -> None:
     assert deserialized == config
 
 
-def test_docling_format_options_should_resolve_gpu_accelerator() -> None:
+def test_docling_pipeline_should_resolve_cuda_accelerator() -> None:
     # Given
-    device = Device.CUDA
     format_opts = DoclingFormatOption(
-        backend="PdfDocumentBackend", pipeline_cls="VlmPipeline"
+        backend="PdfDocumentBackend",
+        pipeline_cls="VlmPipeline",
+    )
+    config = DoclingPipelineConfig(
+        device=Device.CUDA, format_options={InputFormat.PDF: format_opts}
     )
     # When
-    as_docling = format_opts.to_docling(device)
+    config = config.to_device()
     # Then
-    accelerator_opts = as_docling.pipeline_options.accelerator_options
-    assert accelerator_opts.device is AcceleratorDevice.CUDA
+    pdf_fmt_opts = config.format_options[InputFormat.PDF]
+    device = pdf_fmt_opts.pipeline_options["accelerator_options"]["device"]
+    assert device is AcceleratorDevice.CUDA

extract-core/tests/test_objects.py

@@ -7,7 +7,7 @@
 )
 from docling.document_converter import PdfFormatOption
 from extract_core import DoclingPipelineConfig, PipelineConfig
-from extract_core.objects import Device, Status
+from extract_core.objects import Status
 from pydantic import TypeAdapter
 
 
@@ -24,6 +24,7 @@ def test_docling_pipeline_config() -> None:
     ta = TypeAdapter(PipelineConfig)
     config = {
         "pipeline": "docling",
+        "device": "cpu",
         "format_options": {
             "pdf": {
                 "pipeline_cls": "StandardPdfPipeline",
@@ -42,7 +43,7 @@ def test_docling_pipeline_config() -> None:
     assert isinstance(pipeline_config, DoclingPipelineConfig)
     format_options = pipeline_config.format_options
     pdf_options = format_options[InputFormat.PDF]
-    pdf_pipeline_options = pdf_options.to_docling(Device.CPU)
+    pdf_pipeline_options = pdf_options.to_docling()
 
     expected_options = PdfFormatOption(
         pipeline_options=ThreadedPdfPipelineOptions(

extract-python/extract_python/__init__.py

@@ -5,20 +5,20 @@
 try:
     from .docling_ import DOCLING_DEFAULT_ARTIFACTS_PATH, DoclingPipeline
 except ModuleNotFoundError:
-    logger.exception("docling is not available, might be optional...")
+    logger.warning("docling is not available, might be optional...")
     DOCKING_DEFAULT_ARTIFACTS_PATH, DoclingPipeline = None, None
 
 try:
     from .marker_ import MarkerPipeline
 except ModuleNotFoundError:
-    logger.exception("marker is not available, might be optional...")
+    logger.warning("marker is not available, might be optional...")
     MarkerPipeline = None
 
 
 try:
     from .miner_u import MinerUPipeline
 except ModuleNotFoundError:
-    logger.exception("mineru is not available, might be optional...")
+    logger.warning("mineru is not available, might be optional...")
     MinerUPipeline = None
 
 

extract-python/extract_python/__main__.py

@@ -0,0 +1,9 @@
+from .cli import cli_app
+
+
+def main() -> None:
+    cli_app()
+
+
+if __name__ == "__main__":
+    main()

extract-python/extract_python/cli/__init__.py

@@ -0,0 +1,50 @@
+import importlib
+import logging
+import os
+from typing import Annotated
+
+import typer
+from icij_common.logging_utils import setup_loggers
+
+import extract_python
+
+from .doc import configs_app
+from .utils import AsyncTyper
+
+cli_app = AsyncTyper(
+    context_settings={"help_option_names": ["-h", "--help"]},
+    pretty_exceptions_enable=False,
+)
+cli_app.add_typer(configs_app)
+
+
+def version_callback(value: bool) -> None:  # noqa: FBT001
+    if value:
+        package_version = importlib.metadata.version(extract_python.__name__)
+        print(package_version)
+        raise typer.Exit()
+
+
+def pretty_exc_callback(value: bool) -> None:  # noqa: FBT001
+    if not value:
+        os.environ["TYPER_STANDARD_TRACEBACK"] = "1"
+
+
+@cli_app.callback()
+def main(
+    version: Annotated[  # noqa: ARG001
+        bool | None,
+        typer.Option("--version", callback=version_callback, is_eager=True),
+    ] = None,
+    *,
+    pretty_exceptions: Annotated[  # noqa: ARG001
+        bool,
+        typer.Option(
+            "--pretty-exceptions", callback=pretty_exc_callback, is_eager=True
+        ),
+    ] = False,
+) -> None:
+    setup_loggers(
+        ["__main__", extract_python.__name__, "docling", "mineru", "marker"],
+        level=logging.DEBUG,
+    )