fix(extract-python): ensure pipeline run on the proper device

Author: ClemDoum

extract-core/extract_core/__init__.py

@@ -7,6 +7,7 @@
 from .objects import (
     BaseModel,
     ConversionOutput,
+    Device,
     Error,
     InputDoc,
     MarkdownDoc,
@@ -46,6 +47,7 @@
     "BaseModel",
     "BasePipelineConfig",
     "ConversionOutput",
+    "Device",
     "DoclingPipelineConfig",
     "Error",
     "InputDoc",

extract-core/extract_core/docling_.py

@@ -35,7 +35,7 @@
 from pydantic_core.core_schema import SerializerFunctionWrapHandler
 
 from .configs import BasePipelineConfig, PipelineType
-from .objects import BaseModel, SupportedExt
+from .objects import BaseModel, Device, SupportedExt
 from .utils import all_subclasses
 
 
@@ -161,7 +161,9 @@ def _table_structure_opts_type_adapter() -> TypeAdapter:
 
 
 def _resolve_pipeline_options(
-    pipeline_options: dict[str, Any] | None | PipelineOptions, pipeline_cls: type
+    pipeline_options: dict[str, Any] | None | PipelineOptions,
+    pipeline_cls: type,
+    device: Device,
 ) -> PipelineOptions:
     option_cls = _find_init_arg_type(pipeline_cls, "pipeline_options")
     picture_descr_opts = pipeline_options.get("picture_description_options")
@@ -198,6 +200,8 @@ def _resolve_pipeline_options(
         )
         pipeline_options["table_structure_options"] = table_structure_opts
     pipeline_options = option_cls.model_validate(pipeline_options)
+    if isinstance(pipeline_options, PipelineOptions):
+        pipeline_options.accelerator_options.device = device.to_docling()
     return pipeline_options
 
 
@@ -214,11 +218,13 @@ class DoclingFormatOption(BaseFormatOption):
     pipeline_cls: str
     pipeline_options: dict[str, Any] | None = None
 
-    def to_docling(self) -> BaseFormatOption:  # noqa: ANN201
+    def to_docling(self, device: Device) -> 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)
+        pipeline_opts = _resolve_pipeline_options(
+            self.pipeline_options, pipeline_cls, device
+        )
         pipeline_opts = _validate_pipeline_opts(pipeline_opts)
         return FormatOption(
             pipeline_cls=pipeline_cls,

extract-core/extract_core/objects.py

@@ -9,6 +9,7 @@
 from pathlib import Path
 from typing import Annotated, Any, NoReturn, Self
 
+from docling.datamodel.accelerator_options import AcceleratorDevice
 from icij_common.pydantic_utils import (
     icij_config,
     merge_configs,
@@ -26,6 +27,23 @@ class BaseModel(_BaseModel):
     model_config = base_config
 
 
+class Device(StrEnum):
+    CPU = "cpu"
+    CUDA = "cuda"
+    MPS = "mps"
+
+    def to_docling(self) -> AcceleratorDevice:
+        match self:
+            case Device.CUDA:
+                return AcceleratorDevice.CPU
+            case Device.MPS:
+                return AcceleratorDevice.MPS
+            case Device.CPU:
+                return AcceleratorDevice.CPU
+            case _:
+                raise ValueError(f"unsupported device {self}")
+
+
 class SupportedExt(StrEnum):
     ADOC = ".adoc"
     ASC = ".asc"

extract-core/extract_core/pipeline.py

@@ -4,10 +4,13 @@
 
 from icij_common.registrable import RegistrableFromConfig
 
-from .objects import InputDoc, OutputFormat, Result
+from .objects import Device, InputDoc, OutputFormat, Result
 
 
 class Pipeline(RegistrableFromConfig, ABC):
+    def __init__(self, device: Device = Device.CPU):
+        self._device = device
+
     @abstractmethod
     async def extract_content(
         self, docs: Iterable[InputDoc], output_format: OutputFormat, output_path: Path

extract-core/tests/test_objects.py

@@ -1,10 +1,12 @@
+from docling.datamodel.accelerator_options import AcceleratorDevice, AcceleratorOptions
 from docling.datamodel.base_models import InputFormat
 from docling.datamodel.pipeline_options import (
     TesseractOcrOptions,
     ThreadedPdfPipelineOptions,
 )
 from docling.document_converter import PdfFormatOption
 from extract_core import DoclingPipelineConfig, PipelineConfig
+from extract_core.objects import Device
 from pydantic import TypeAdapter
 
 
@@ -20,6 +22,7 @@ def test_docling_pipeline_config() -> None:
                 "pipeline_options": {
                     "ocr_options": {"kind": "tesserocr", "lang": ["auto"]},
                     "generate_picture_images": True,
+                    "accelerator_options": {"device": "cpu"},
                 },
             }
         },
@@ -30,11 +33,13 @@ 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()
+    pdf_pipeline_options = pdf_options.to_docling(Device.CPU)
+
     expected_options = PdfFormatOption(
         pipeline_options=ThreadedPdfPipelineOptions(
             ocr_options=TesseractOcrOptions(lang=["auto"]),
             generate_picture_images=True,
+            accelerator_options=AcceleratorOptions(device=AcceleratorDevice.CPU),
         )
     )
     assert pdf_pipeline_options.model_dump() == expected_options.model_dump()

extract-python/extract_python/docling_.py

@@ -27,6 +27,7 @@
     Result,
     Status,
 )
+from extract_core.objects import Device
 from icij_common.pydantic_utils import merge_configs
 from icij_common.registrable import FromConfig
 from pydantic import ConfigDict, field_serializer
@@ -41,9 +42,15 @@
 @Pipeline.register(PipelineType.DOCLING)
 class DoclingPipeline(Pipeline):
     def __init__(
-        self, format_options: dict["InputFormat", DoclingFormatOption] | None = None
+        self,
+        format_options: dict["InputFormat", DoclingFormatOption] | None = None,
+        *,
+        device: Device = Device.CPU,
     ):
-        format_options = {k: v.to_docling() for k, v in format_options.items()}
+        super().__init__(device)
+        format_options = {
+            k: v.to_docling(self._device) for k, v in format_options.items()
+        }
         allowed_format = [
             f.to_docling() for f in DoclingPipelineConfig.supported_exts()
         ]
@@ -66,8 +73,13 @@ async def extract_content(
             yield _to_result(res, doc, output_format, output_path=output_path)
 
     @classmethod
-    def _from_config(cls, config: DoclingPipelineConfig) -> FromConfig:
-        return cls(config.format_options)
+    def _from_config(
+        cls,
+        config: DoclingPipelineConfig,
+        *,
+        device: Device = Device.CPU,
+    ) -> FromConfig:
+        return cls(config.format_options, device=device)
 
 
 def _to_docling(docs: Iterable[InputDoc]) -> Iterator["Path | DocumentStream"]:

extract-python/extract_python/marker_.py

@@ -8,6 +8,7 @@
 
 from extract_core import BasePipelineConfig, Pipeline, PipelineType
 from extract_core.objects import (
+    Device,
     InputDoc,
     MarkdownDoc,
     OutputFormat,
@@ -64,7 +65,13 @@ def supported_exts(cls) -> set[SupportedExt]:
 
 @Pipeline.register(PipelineType.MARKER)
 class MarkerPipeline(Pipeline):
-    def __init__(self, marker_config: dict[str, Any] | None = None):
+    def __init__(
+        self,
+        marker_config: dict[str, Any] | None = None,
+        *,
+        device: Device = Device.CPU,
+    ):
+        super().__init__(device)
         if marker_config is None:
             marker_config = dict()
         self._marker_config = marker_config
@@ -82,16 +89,21 @@ async def extract_content(
         renderer = config_parser.get_renderer()
         converter = PdfConverter(
             config=config_parser.generate_config_dict(),
-            artifact_dict=create_model_dict(),
+            artifact_dict=create_model_dict(device=self._device),
             processor_list=config_parser.get_processors(),
             renderer=renderer,
         )
         for doc in docs:
             yield await _process_doc(doc, converter, output_format, output_path)
 
     @classmethod
-    def _from_config(cls, config: MarkerPipelineConfig) -> Self:
-        return cls(config.config)
+    def _from_config(
+        cls,
+        config: MarkerPipelineConfig,
+        *,
+        device: Device = Device.CPU,
+    ) -> Self:
+        return cls(config.config, device=device)
 
 
 @report_recoverable_errors(_MARKER_CONVERSION_ERRORS)

extract-python/extract_python/miner_u.py

@@ -1,4 +1,5 @@
 import json
+import os
 import shutil
 from collections.abc import AsyncGenerator, Callable, Iterable
 from functools import partial
@@ -19,17 +20,21 @@
     Result,
     Status,
 )
+from objects import Device
 
 from .constants import ARTIFACTS, DEFAULT_MD_PAGE_SEP
-from .utils import path_to_artifacts_dirname
+from .utils import path_to_artifacts_dirname, reset_env
 
 _MINER_U_CONVERSION_ERRORS = tuple()
 MDMakeFunction = Callable[[list, str, str], str | None]
 
 
 @Pipeline.register(PipelineType.MINER_U)
 class MinerUPipeline(Pipeline):
-    def __init__(self, config: MinerUConfig, language: str):
+    def __init__(
+        self, config: MinerUConfig, language: str, *, device: Device = Device.CPU
+    ):
+        super().__init__(device)
         self._config = config
         self._language = language
         self._md_make_fn = _parse_md_make_fn(config.backend)
@@ -39,36 +44,43 @@ async def extract_content(
     ) -> AsyncGenerator[Result, None]:
         from mineru.cli.common import aio_do_parse  # noqa: PLC0415
 
-        docs = list(docs)
-        # TODO: exclude files which are not pdf and return an error
-        pdfs_bytes = [d.path.read_bytes() for d in docs]
-        pdfs_names = [d.path.name for d in docs]
-        p_lang_list = [self._language for _ in pdfs_names]
-        # TODO: we should only process valid PDFs
-        with TemporaryDirectory(prefix="mineru-") as workdir:
-            workdir = Path(workdir)  # noqa: PLW2901
-            await aio_do_parse(
-                output_dir=workdir,
-                pdf_file_names=pdfs_names,
-                pdf_bytes_list=pdfs_bytes,
-                p_lang_list=p_lang_list,
-                **self._config.as_parse_kwargs(),
-            )
-            res_paths = [
-                _revert_mineru_output(workdir, pdf_filename=p) for p in pdfs_names
-            ]
-            for doc, res_path in zip(docs, res_paths, strict=True):
-                yield _process_doc(
-                    doc,
-                    md_make_fn=self._md_make_fn,
-                    res_path=res_path,
-                    output_format=output_format,
-                    output_path=output_path,
+        with reset_env():
+            os.environ["MINERU_DEVICE_MODE"] = self._device
+            docs = list(docs)
+            # TODO: exclude files which are not pdf and return an error
+            pdfs_bytes = [d.path.read_bytes() for d in docs]
+            pdfs_names = [d.path.name for d in docs]
+            p_lang_list = [self._language for _ in pdfs_names]
+            # TODO: we should only process valid PDFs
+            with TemporaryDirectory(prefix="mineru-") as workdir:
+                workdir = Path(workdir)  # noqa: PLW2901
+                await aio_do_parse(
+                    output_dir=workdir,
+                    pdf_file_names=pdfs_names,
+                    pdf_bytes_list=pdfs_bytes,
+                    p_lang_list=p_lang_list,
+                    **self._config.as_parse_kwargs(),
                 )
+                res_paths = [
+                    _revert_mineru_output(workdir, pdf_filename=p) for p in pdfs_names
+                ]
+                for doc, res_path in zip(docs, res_paths, strict=True):
+                    yield _process_doc(
+                        doc,
+                        md_make_fn=self._md_make_fn,
+                        res_path=res_path,
+                        output_format=output_format,
+                        output_path=output_path,
+                    )
 
     @classmethod
-    def _from_config(cls, config: MinerUPipelineConfig) -> Self:
-        return cls(config.config, language=config.language)
+    def _from_config(
+        cls,
+        config: MinerUPipelineConfig,
+        *,
+        device: Device = Device.CPU,
+    ) -> Self:
+        return cls(config.config, language=config.language, device=device)
 
 
 def _revert_mineru_output(output_dir: Path, *, pdf_filename: str) -> Path:

extract-python/extract_python/utils.py

@@ -1,6 +1,7 @@
 import os
 from collections.abc import Callable, Generator, Iterable, Iterator
 from contextlib import contextmanager
+from copy import copy
 from functools import wraps
 from itertools import tee
 from pathlib import Path, PurePath
@@ -62,3 +63,13 @@ def chdir(path: Path) -> Generator[None, None, None]:
         yield
     finally:
         os.chdir(cwd)
+
+
+@contextmanager
+def reset_env() -> Generator[None, None, None]:
+    old_env = copy(dict(os.environ))
+    try:
+        yield
+    finally:
+        os.environ.clear()
+        os.environ.update(old_env)

extract-python/tests/conftest.py

@@ -1,13 +1,21 @@
 import os
+import platform
 import shutil
 from pathlib import Path
 
 import pytest
-from extract_core import InputDoc
+from extract_core import Device, InputDoc
 
 from tests import TEST_DATA_DIR
 
 
+@pytest.fixture(scope="session")
+def device() -> Device:
+    if "darwin" in platform.system().lower():
+        return Device.MPS
+    return Device.CPU
+
+
 @pytest.fixture(scope="session")
 def docs() -> list[InputDoc]:
     doc_paths = ("scanned.pdf", "computer_generated.pdf")