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

extract_python/docling_.py

@@ -1,16 +1,20 @@
+import asyncio
 import shutil
 import tempfile
 from collections.abc import AsyncGenerator, Iterable, Iterator
 from pathlib import Path
+from typing import Any, Self
 
 from docling.datamodel.base_models import InputFormat
 from docling.datamodel.document import ConversionResult
-from docling.document_converter import DocumentConverter
+from docling.datamodel.pipeline_options import PipelineOptions
+from docling.document_converter import DocumentConverter, FormatOption
 
 # TODO: this is long to load improve it
 from docling_core.types.doc import ImageRefMode
 from docling_core.types.io import DocumentStream
 from extract_core import (
+    BaseModel,
     DoclingFormatOption,
     DoclingPipelineConfig,
     Error,
@@ -23,7 +27,11 @@ 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
+from pydantic_core.core_schema import SerializerFunctionWrapHandler
 
 from .constants import ARTIFACTS, DEFAULT_MD_PAGE_SEP
 from .utils import chdir, map_and_preserve, path_to_artifacts_dirname
@@ -34,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()
         ]
@@ -49,12 +63,23 @@ class DoclingPipeline(Pipeline):
     ) -> AsyncGenerator[Result, None]:
         docs, path_or_streams = map_and_preserve(_to_docling, docs)
         outputs = self._converter.convert_all(path_or_streams, raises_on_error=False)
-        for doc, res in zip(docs, outputs, strict=True):
+
+        sentinel = object()
+        while True:
+            res = await asyncio.to_thread(next, outputs, sentinel)
+            if res is sentinel:
+                return
+            doc = next(docs)
             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"]:
@@ -128,3 +153,52 @@ def _to_markdown_doc(
         shutil.move(tmp_dir, md_dir)
     pages = PageIndexes.from_page_end_indices(end_indices)
     return MarkdownDoc(path=Path(md_dir_name), pages=pages)
+
+
+class SerializableFormatOptions(DoclingFormatOption):
+    # Utility class to serialize Python format options into a JSON which can be
+    # correctly deserialized into a docling FormatOption
+    # via DoclingFormatOption.to_docling
+    model_config = merge_configs(
+        BaseModel.model_config, ConfigDict(polymorphic_serialization=True)
+    )
+
+    pipeline_options: PipelineOptions | None = None
+
+    @classmethod
+    def from_docling(cls, format_opts: FormatOption) -> Self:
+        return cls(
+            pipeline_cls=format_opts.pipeline_cls.__name__,
+            pipeline_options=format_opts.pipeline_options,
+            backend=format_opts.backend.__name__,
+            backend_options=format_opts.backend_options,
+        )
+
+    @field_serializer("pipeline_options", mode="wrap")
+    def _serialize_pipeline_opts(
+        self, v: PipelineOptions | None, handler: SerializerFunctionWrapHandler
+    ) -> Any:
+        if v is None:
+            return handler(v)
+        serialized = handler(v)
+        picture_desc_opts = getattr(v, "picture_description_options", None)
+        if picture_desc_opts is not None:
+            if "picture_description_options" not in serialized:
+                serialized["picture_description_options"] = dict()
+            serialized["picture_description_options"]["kind"] = picture_desc_opts.kind
+        ocr_opts = getattr(v, "ocr_options", None)
+        if ocr_opts is not None:
+            if "ocr_options" not in serialized:
+                serialized["ocr_options"] = dict()
+            serialized["ocr_options"]["kind"] = ocr_opts.kind
+        layout_opts = getattr(v, "layout_options", None)
+        if layout_opts is not None:
+            if "layout_options" not in serialized:
+                serialized["layout_options"] = dict()
+            serialized["layout_opts"]["kind"] = layout_opts.kind
+        table_structure_opts = getattr(v, "table_structure_options", None)
+        if table_structure_opts is not None:
+            if "table_structure_options" not in serialized:
+                serialized["table_structure_options"] = dict()
+            serialized["table_structure_options"]["kind"] = table_structure_opts.kind
+        return serialized

extract_python/marker_.py

@@ -1,3 +1,4 @@
+import asyncio
 import gc
 from collections.abc import AsyncGenerator, Iterable
 from copy import deepcopy
@@ -7,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,
@@ -63,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
@@ -81,20 +89,25 @@ 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,
         )
         for doc in docs:
-            yield _process_doc(doc, converter, output_format, output_path)
+            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)
-def _process_doc(
+async def _process_doc(
     doc: InputDoc,
     converter: "PdfConverter",
     output_format: OutputFormat,
@@ -102,7 +115,7 @@ def _process_doc(
 ) -> Result:
     from marker.output import text_from_rendered  # noqa: PLC0415
 
-    rendered = converter(str(doc.path))
+    rendered = await asyncio.to_thread(converter, str(doc.path))
     content, _, images = text_from_rendered(rendered)
     match output_format:
         case OutputFormat.MARKDOWN:

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.5.dist-info → extract_python-0.5.8.dist-info}/METADATA

@@ -1,13 +1,13 @@
 Metadata-Version: 2.4
 Name: extract-python
-Version: 0.5.5
+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
 Project-URL: Issues, https://github.com/ICIJ/extract-python/issues
 Author-email: Clément Doumouro <cdoumouro@icij.org>
 Requires-Python: <3.14,>=3.11
-Requires-Dist: extract-core~=0.1
+Requires-Dist: extract-core~=0.5.5
 Requires-Dist: icij-common~=0.8.2
 Provides-Extra: benches
 Requires-Dist: html2image~=2.0.7; extra == 'benches'

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.5.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=C4WP1AJrvS2n-KytlGc_1CShjdTGM077I6b9tvw4NhY,4727
-extract_python/marker_.py,sha256=mLJA1m9G4JQtBs1wz8rmshdbaH81DhIwkRzDKZPJH8A,5058
-extract_python/miner_u.py,sha256=jjHqHx7-2w0LSxYNcjvgWoLDTXsv_y1eeyteSfXqjk4,5771
-extract_python/utils.py,sha256=NiYf65iCF7QO4loh7u4t38Ww3eVJUdBpWStL4eX_DqE,1781
-extract_python-0.5.5.dist-info/METADATA,sha256=iENRXysGcLOtZ3mJNPpGwHiixhkYJYaYQ-sj8j26q2o,1216
-extract_python-0.5.5.dist-info/WHEEL,sha256=mffPy8wBnZQn2VnJUU5jE99KsxaSfiyMHV9Yt0aLVxs,87
-extract_python-0.5.5.dist-info/RECORD,,