extract-python 0.1.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

extract_python/docling_.py

@@ -3,28 +3,23 @@ import tempfile
 from collections.abc import AsyncGenerator, Iterable, Iterator
 from functools import cache
 from pathlib import Path
-from typing import Any, ClassVar, Literal, TypeVar
+from typing import Annotated, ClassVar, TypeVar
 
-from docling.backend.abstract_backend import AbstractDocumentBackend
-from docling.datamodel.base_models import InputFormat
+from docling.datamodel.base_models import FormatToExtensions, InputFormat
 from docling.datamodel.document import ConversionResult
 from docling.datamodel.pipeline_options import (
     EasyOcrOptions,
     PdfPipelineOptions,
     PipelineOptions,
-    VlmPipelineOptions,
 )
-from docling.document_converter import DocumentConverter, FormatOption
-from docling.models.factories import get_ocr_factory
-from docling.pipeline.base_pipeline import BasePipeline
+from docling.document_converter import DocumentConverter, FormatOption, PdfFormatOption
 from docling_core.types.doc import ImageRefMode
 from docling_core.types.io import DocumentStream
 from icij_common.registrable import FromConfig
-from pydantic import Field, model_validator
+from pydantic import AfterValidator, Field
 
 from .constants import ARTIFACTS, CPU_GROUP, DEFAULT_MD_PAGE_SEP
 from .objects import (
-    BaseModel,
     Error,
     InputDoc,
     MarkdownDoc,
@@ -32,6 +27,7 @@ from .objects import (
     PageIndexes,
     Result,
     Status,
+    SupportedExt,
 )
 from .pipeline import Pipeline, PipelineConfig, PipelineType
 from .utils import all_subclasses, chdir, map_and_preserve, path_to_artifacts_dirname
@@ -39,73 +35,27 @@ from .utils import all_subclasses, chdir, map_and_preserve, path_to_artifacts_di
 DOCLING_DEFAULT_ARTIFACTS_PATH = Path.home().joinpath(".cache", "docling", "models")
 
 
-class _PdfPipelineOptions(PdfPipelineOptions):
-    generate_picture_images: bool = Field(default=True, frozen=True)
+def _validate_pipeline_opts(opts: PipelineOptions) -> None:
+    if isinstance(opts, PdfPipelineOptions) and not opts.generate_picture_images:
+        msg = "generate_picture_images should be set to true"
+        raise ValueError(msg)
 
-    @model_validator(mode="before")
-    @classmethod
-    def validate_ocr_options(cls, data: Any) -> Any:
-        if isinstance(data, dict):
-            ocr_options = data.get("ocr_options")
-            if not isinstance(ocr_options, dict):
-                return data
-            allow_external_plugins = ocr_options.get("allow_external_plugins", False)
-            ocr_factory = get_ocr_factory(allow_external_plugins=allow_external_plugins)
-            kind = ocr_options.pop("kind")
-            data["ocr_options"] = ocr_factory.create_options(kind=kind, **ocr_options)
-        return data
-
-
-OptionsByPipeline = list[
-    tuple[Literal["pdf"], _PdfPipelineOptions]
-    | tuple[Literal["vlm"], VlmPipelineOptions]
-]
-
-
-def _default_pipeline_options() -> OptionsByPipeline:
-    pipeline_options = _PdfPipelineOptions(ocr_options=EasyOcrOptions())
-    return [("pdf", pipeline_options), ("vlm", VlmPipelineOptions())]
-
-
-class DoclingFormatOption(BaseModel):
-    pipeline_cls: str
-    backend_cls: str
-
-    def to_docling(
-        self, pipeline_options: dict[Literal["pdf", "vlm"], PipelineOptions]
-    ) -> FormatOption:
-        pipeline_cls = _find_subcls(BasePipeline, self.pipeline_cls)
-        backend_cls = _find_subcls(AbstractDocumentBackend, self.backend_cls)
-        if "vlm" in self.pipeline_cls.lower():
-            pipeline_options = pipeline_options.get("vlm")
-            if pipeline_options is not None:
-                pipeline_options = VlmPipelineOptions.model_validate(pipeline_options)
-        elif "pdf" in self.pipeline_cls.lower():
-            pipeline_options = pipeline_options.get("pdf")
-            if pipeline_options is not None:
-                pipeline_options = _PdfPipelineOptions.model_validate(pipeline_options)
-        else:
-            raise ValueError(
-                f"invalid pipeline_cls: {pipeline_cls}, expected a VLM or PDF pipeline"
-            )
-        return FormatOption(
-            pipeline_cls=pipeline_cls,
-            pipeline_options=pipeline_options,
-            backend=backend_cls,
-        )
 
+def _validate_options(
+    data: dict[InputFormat, FormatOption],
+) -> dict[InputFormat, FormatOption]:
+    for opts in data.values():
+        _validate_pipeline_opts(opts.pipeline_options)
+    return data
 
-@cache
-def _default_format_options() -> dict[InputFormat, DoclingFormatOption]:
-    supported_fmt = {InputFormat.PDF}
-    return {
-        fmt: DoclingFormatOption(
-            pipeline_cls=opt.pipeline_cls.__name__, backend_cls=opt.backend.__name__
-        )
-        for fmt, opt in DocumentConverter().format_to_options.items()
-        if fmt in supported_fmt
-    }
 
+_DEFAULT_FORMAT_OPTS = {
+    InputFormat.PDF: PdfFormatOption(
+        pipeline_options=PdfPipelineOptions(
+            ocr_options=EasyOcrOptions(), generate_picture_images=True
+        )
+    ),
+}
 
 T = TypeVar("T")
 
@@ -122,30 +72,37 @@ class DoclingPipelineConfig(PipelineConfig):
     pipeline: PipelineType = Field(frozen=True, default=PipelineType.DOCLING)
     task_group: ClassVar[str] = Field(frozen=True, default=CPU_GROUP)
 
-    pipeline_options: OptionsByPipeline = Field(
-        default_factory=_default_pipeline_options
-    )
-    format_options: dict[InputFormat, DoclingFormatOption] = Field(
-        default_factory=_default_format_options
-    )
-
-    def to_format_options(self) -> dict[InputFormat, FormatOption]:
-        pipeline_options = dict(self.pipeline_options)
-        return {
-            InputFormat(f): opt.to_docling(pipeline_options)
-            for f, opt in self.format_options.items()
-        }
+    format_options: Annotated[
+        dict[InputFormat, FormatOption] | None, AfterValidator(_validate_options)
+    ] = _DEFAULT_FORMAT_OPTS
 
+    _unsupported_input_formats: ClassVar[set[InputFormat]] = {
+        InputFormat.AUDIO,
+        InputFormat.METS_GBS,
+        InputFormat.VTT,
+    }
 
-DEFAULT_FORMAT_OPTIONS = DoclingPipelineConfig().to_format_options()
+    @classmethod
+    @cache
+    def supported_exts(cls) -> set[SupportedExt]:
+        supported = set()
+        for f in InputFormat:
+            if f in cls._unsupported_input_formats:
+                continue
+            for ext in FormatToExtensions[f]:
+                supported.add(SupportedExt(f".{ext.lower()}"))
+        return supported
 
 
 @Pipeline.register(PipelineType.DOCLING)
 class DoclingPipeline(Pipeline):
     def __init__(self, format_options: dict[InputFormat, FormatOption] | None = None):
-        if format_options is None:
-            format_options = DEFAULT_FORMAT_OPTIONS
-        self._converter = DocumentConverter(format_options=format_options)
+        allowed_format = [
+            f.to_docling() for f in DoclingPipelineConfig.supported_exts()
+        ]
+        self._converter = DocumentConverter(
+            allowed_formats=allowed_format, format_options=format_options
+        )
 
     async def extract_content(
         self, docs: Iterable[InputDoc], output_format: OutputFormat, output_path: Path
@@ -157,7 +114,7 @@ class DoclingPipeline(Pipeline):
 
     @classmethod
     def _from_config(cls, config: DoclingPipelineConfig) -> FromConfig:
-        return cls(config.to_format_options())
+        return cls(config.format_options)
 
 
 def _to_docling(docs: Iterable[InputDoc]) -> Iterator[Path | DocumentStream]:

extract_python/marker_.py

@@ -1,6 +1,7 @@
 import gc
 from collections.abc import AsyncGenerator, Iterable
 from copy import deepcopy
+from functools import cache
 from pathlib import Path
 from typing import Any, ClassVar, Self
 
@@ -20,6 +21,7 @@ from .objects import (
     PageIndexes,
     Result,
     Status,
+    SupportedExt,
 )
 from .pipeline import Pipeline, PipelineConfig, PipelineType
 from .utils import path_to_artifacts_dirname, report_recoverable_errors
@@ -32,6 +34,33 @@ class MarkerPipelineConfig(PipelineConfig):
 
     config: dict[str, Any] = dict()
 
+    @classmethod
+    @cache
+    def supported_exts(cls) -> set[SupportedExt]:
+        # Subset of https://documentation.datalab.to/docs/common/supportedfiletypes
+        return {
+            SupportedExt.PDF,
+            SupportedExt.XLS,
+            SupportedExt.XLSX,
+            SupportedExt.XLSM,
+            SupportedExt.CSV,
+            SupportedExt.ODS,
+            SupportedExt.DOC,
+            SupportedExt.DOCX,
+            SupportedExt.ODT,
+            SupportedExt.PPT,
+            SupportedExt.PPTX,
+            SupportedExt.ODP,
+            SupportedExt.HTLM,
+            SupportedExt.EPUB,
+            SupportedExt.PNG,
+            SupportedExt.JPG,
+            SupportedExt.JPEG,
+            SupportedExt.WEBP,
+            SupportedExt.GIF,
+            SupportedExt.TIFF,
+        }
+
 
 _MARKER_CONVERSION_ERRORS = tuple()
 

extract_python/miner_u.py

@@ -3,7 +3,7 @@ import shutil
 from collections.abc import AsyncGenerator, Callable, Iterable
 from copy import copy
 from enum import StrEnum
-from functools import partial
+from functools import cache, partial
 from pathlib import Path
 from tempfile import TemporaryDirectory
 from typing import Any, ClassVar, Self
@@ -26,6 +26,7 @@ from .objects import (
     PageIndexes,
     Result,
     Status,
+    SupportedExt,
 )
 from .pipeline import Pipeline, PipelineConfig, PipelineType
 from .utils import path_to_artifacts_dirname
@@ -82,6 +83,16 @@ class MinerUPipelineConfig(PipelineConfig):  # noqa: F821
     config: MinerUConfig = Field(frozen=True, default=MinerUConfig())
     language: LanguageAlpha2 = Field(frozen=True, default="en")
 
+    @classmethod
+    @cache
+    def supported_exts(cls) -> set[SupportedExt]:
+        return {
+            SupportedExt.PDF,
+            SupportedExt.DOCX,
+            SupportedExt.PPTX,
+            SupportedExt.XLSX,
+        }
+
 
 @Pipeline.register(PipelineType.MINER_U)
 class MinerUPipeline(Pipeline):

extract_python/objects.py

@@ -21,7 +21,12 @@ from pydantic import AfterValidator, RootModel, TypeAdapter
 from pydantic import BaseModel as _BaseModel
 
 try:
-    from docling.datamodel.base_models import ConversionStatus, ErrorItem, InputFormat
+    from docling.datamodel.base_models import (
+        ConversionStatus,
+        ErrorItem,
+        FormatToExtensions,
+        InputFormat,
+    )
     from docling.datamodel.document import InputDocument
     from docling_core.types.io import DocumentStream
 except ImportError:
@@ -33,15 +38,78 @@ logger = logging.getLogger(__name__)
 base_config = merge_configs(icij_config(), no_enum_values_config())
 
 
+@cache
+def _ext_to_docling_input_format() -> dict:
+    from .docling_ import DoclingPipelineConfig  # noqa: PLC0415
+
+    mapping = dict()
+    supported = DoclingPipelineConfig.supported_exts()
+    for input_f, exts in FormatToExtensions.items():
+        for ext in exts:
+            try:
+                ext = SupportedExt(f".{ext.lower()}")  # noqa: PLW2901
+            except ValueError:
+                continue
+            if ext in supported:
+                mapping[ext] = input_f
+    return mapping
+
+
 class BaseModel(_BaseModel):
     model_config = base_config
 
 
 class SupportedExt(StrEnum):
+    ADOC = ".adoc"
+    ASC = ".asc"
+    ASCIIDOC = ".asciidoc"
+    BMP = ".bmp"
+    CSV = ".csv"
+    DOC = ".doc"
+    DOCX = ".docx"
+    DOTX = ".dotx"
+    DOTM = ".dotm"
+    DOCM = ".docm"
+    EPUB = ".epub"
+    GIF = ".gif"
+    HTLM = ".html"
+    HTM = ".htm"
+    JPEG = ".jpeg"
+    JPG = ".jpg"
+    JSON = ".json"
+    LATEX = ".latex"
+    MD = ".md"
+    NXML = ".nxml"
+    ODP = ".odp"
+    ODS = ".ods"
+    ODT = ".odt"
     PDF = ".pdf"
+    PNG = ".png"
+    PPSX = ".ppsx"
+    PPT = ".ppt"
+    PPTM = ".pptm"
+    PPSM = ".ppsm"
+    POTX = ".potx"
+    POTM = ".potm"
+    PPTX = ".pptx"
+    QMD = ".qmd"
+    RMD = ".rmd"
+    TEX = ".tex"
+    TIF = ".tif"
+    TIFF = ".tiff"
+    TXT = ".txt"
+    TEXT = ".text"
+    WEBP = ".webp"
+    XBRL = ".xbrl"
+    XHTML = ".xhtml"
+    XLS = ".xls"
+    XLSM = ".xlsm"
+    XLSX = ".xlsx"
+    XLTX = ".xltx"
+    XML = ".xml"
 
     def to_docling(self) -> InputFormat:
-        return InputFormat(self.value[1:])
+        return _ext_to_docling_input_format()[self]
 
 
 class OutputFormat(StrEnum):

extract_python/pipeline.py

@@ -8,7 +8,7 @@ from icij_common.pydantic_utils import icij_config, merge_configs, no_enum_value
 from icij_common.registrable import RegistrableConfig, RegistrableFromConfig
 from pydantic import Field
 
-from .objects import InputDoc, OutputFormat, Result
+from .objects import InputDoc, OutputFormat, Result, SupportedExt
 
 StructuredContent = str
 
@@ -28,6 +28,10 @@ class PipelineConfig(RegistrableConfig, ABC):
 
     task_group: ClassVar[str] = Field(frozen=True)
 
+    @classmethod
+    @abstractmethod
+    def supported_exts(cls) -> set[SupportedExt]: ...
+
 
 class Pipeline(RegistrableFromConfig, ABC):
     @abstractmethod

{extract_python-0.1.0.dist-info → extract_python-0.3.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: extract-python
-Version: 0.1.0
+Version: 0.3.0
 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.3.0.dist-info/RECORD

@@ -0,0 +1,11 @@
+extract_python/__init__.py,sha256=Y-lcFbJd5sX3wt1UnyYAbm7H81Sgg7p2fG07CuuJUco,945
+extract_python/constants.py,sha256=JmAkjXyQMYwpTod9DCLc11zOc3caK2j_2ji_r3hGZws,236
+extract_python/docling_.py,sha256=00g7RV33iftjvMLoKaEs2lUZP3LslSCKMpGeSys1Suc,6616
+extract_python/marker_.py,sha256=3Q8H-TeM2_GenB6OOqIuytqgI1VE93Ek99_kW0cJHEw,4905
+extract_python/miner_u.py,sha256=WdaftyINZdnALqSuu1qKaZJKKyIHVRn-wBke-Na78O0,7747
+extract_python/objects.py,sha256=kCxg6m7j01Z1sNAGi8vniokMnw6Ry0gU2lXXH2uan8A,8744
+extract_python/pipeline.py,sha256=VhDvfCxMEKvhFbMA-yxWO7FEeErDoLQCHiTRNnrbI8Y,1204
+extract_python/utils.py,sha256=kXe0CyT6zGYDTOvxs3BJYxq2-cgJZ0d_IT3hdMbuXa8,1943
+extract_python-0.3.0.dist-info/METADATA,sha256=MwFN5PsmkUEv8sbhuS6joh7r17W6rrikpJb4Yr-rdKk,1132
+extract_python-0.3.0.dist-info/WHEEL,sha256=mffPy8wBnZQn2VnJUU5jE99KsxaSfiyMHV9Yt0aLVxs,87
+extract_python-0.3.0.dist-info/RECORD,,

extract_python-0.1.0.dist-info/RECORD

@@ -1,11 +0,0 @@
-extract_python/__init__.py,sha256=Y-lcFbJd5sX3wt1UnyYAbm7H81Sgg7p2fG07CuuJUco,945
-extract_python/constants.py,sha256=JmAkjXyQMYwpTod9DCLc11zOc3caK2j_2ji_r3hGZws,236
-extract_python/docling_.py,sha256=FMDsPVz05sGMPvIOX90lOLygWp6nC5DEjRfgx_ESPJ4,8530
-extract_python/marker_.py,sha256=z3PkUUStC-E78HhqByMwJ7re6-I7YUQzSxWToegHrUQ,4060
-extract_python/miner_u.py,sha256=f5pvLvay1ThBXNOI1R276aWSWsk5mhIPzWVjCy2u_lw,7493
-extract_python/objects.py,sha256=gTyGA5gaMAmW5P_PbAO2LNMqtP69CxlknebBFTojiwQ,7322
-extract_python/pipeline.py,sha256=qUgGar1rlYQgNz78BcUT1nQRsG3hy5UwpCl0e-0V77I,1098
-extract_python/utils.py,sha256=kXe0CyT6zGYDTOvxs3BJYxq2-cgJZ0d_IT3hdMbuXa8,1943
-extract_python-0.1.0.dist-info/METADATA,sha256=wyYMrleKk9yUU1UaTYT0EsGpw_e3qbE8LOBanyLv0Qg,1132
-extract_python-0.1.0.dist-info/WHEEL,sha256=mffPy8wBnZQn2VnJUU5jE99KsxaSfiyMHV9Yt0aLVxs,87
-extract_python-0.1.0.dist-info/RECORD,,