extract-python 0.5.7__py3-none-any.whl → 0.5.8__py3-none-any.whl

extract_python/docling_.py

@@ -27,6 +27,7 @@ from extract_core import (
     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 @@ DOCLING_DEFAULT_ARTIFACTS_PATH = Path.home().joinpath(".cache", "docling", "mode
 @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 @@ class DoclingPipeline(Pipeline):
             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/marker_.py

@@ -8,6 +8,7 @@ from typing import TYPE_CHECKING, Any, ClassVar, Self
 
 from extract_core import BasePipelineConfig, Pipeline, PipelineType
 from extract_core.objects import (
+    Device,
     InputDoc,
     MarkdownDoc,
     OutputFormat,
@@ -64,7 +65,13 @@ _MARKER_CONVERSION_ERRORS = tuple()
 
 @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,7 +89,7 @@ class MarkerPipeline(Pipeline):
         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,
         )
@@ -90,8 +97,13 @@ class MarkerPipeline(Pipeline):
             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/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,9 +20,10 @@ from extract_core import (
     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]
@@ -29,7 +31,10 @@ 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 @@ class MinerUPipeline(Pipeline):
     ) -> 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/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-0.5.7.dist-info → extract_python-0.5.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: extract-python
-Version: 0.5.7
+Version: 0.5.8
 Summary: Structured content extraction
 Project-URL: Homepage, https://github.com/ICIJ/extract-python
 Project-URL: Repository, https://github.com/ICIJ/extract-python

extract_python-0.5.8.dist-info/RECORD

@@ -0,0 +1,9 @@
+extract_python/__init__.py,sha256=CrqmcyLwD2JgtQNuGRIQ8wr1cWdlKkgMlCz_2reaPJo,470
+extract_python/constants.py,sha256=Hxl2Bc-GJX71gjFgj3U7tRRA_nHhneoT6YPOx9CjYsc,94
+extract_python/docling_.py,sha256=5q3JGJa8Jg8T8WL8MvufFkyILbD3G5mqtkV7HQed9zI,7632
+extract_python/marker_.py,sha256=CIqRvHbdm6Itb7RiOVZWBDx5IvXzBGHjmHsWQFXIMyw,5331
+extract_python/miner_u.py,sha256=5_KhSq5weRy4AmFrt0OZEDhVsnxzHN9TtM3L68mk-8I,6174
+extract_python/utils.py,sha256=HL-84NkjfJEiWp8GPRaJIiBL2Cywp4ABN41EkxYYnPI,2004
+extract_python-0.5.8.dist-info/METADATA,sha256=GbJOVU_ytNS46Z0M49YR_d7OVxQW6tP1SYSwl_LUr9E,1218
+extract_python-0.5.8.dist-info/WHEEL,sha256=mffPy8wBnZQn2VnJUU5jE99KsxaSfiyMHV9Yt0aLVxs,87
+extract_python-0.5.8.dist-info/RECORD,,

extract_python-0.5.7.dist-info/RECORD

@@ -1,9 +0,0 @@
-extract_python/__init__.py,sha256=CrqmcyLwD2JgtQNuGRIQ8wr1cWdlKkgMlCz_2reaPJo,470
-extract_python/constants.py,sha256=Hxl2Bc-GJX71gjFgj3U7tRRA_nHhneoT6YPOx9CjYsc,94
-extract_python/docling_.py,sha256=TXXSayirZtA82VM9-2WrGqxTKodX6KyHUDZHURuvd7k,7382
-extract_python/marker_.py,sha256=i_zWNISTI--E7Y-71DYQ5TlsTN7I9T5vT-3szPqMkWE,5110
-extract_python/miner_u.py,sha256=jjHqHx7-2w0LSxYNcjvgWoLDTXsv_y1eeyteSfXqjk4,5771
-extract_python/utils.py,sha256=NiYf65iCF7QO4loh7u4t38Ww3eVJUdBpWStL4eX_DqE,1781
-extract_python-0.5.7.dist-info/METADATA,sha256=4jp__6WVmbXwrGGpf7ExGdJWO1-bqcaN5r411xq--OY,1218
-extract_python-0.5.7.dist-info/WHEEL,sha256=mffPy8wBnZQn2VnJUU5jE99KsxaSfiyMHV9Yt0aLVxs,87
-extract_python-0.5.7.dist-info/RECORD,,