extract-python 0.1.0__py3-none-any.whl

extract_python/__init__.py

@@ -0,0 +1,41 @@
+from .objects import InputDoc, OutputFormat, Status
+from .pipeline import Pipeline, PipelineConfig, PipelineType
+
+try:
+    from .docling_ import (
+        DOCLING_DEFAULT_ARTIFACTS_PATH,
+        DoclingPipeline,
+        DoclingPipelineConfig,
+    )
+except ImportError:
+    DOCKING_DEFAULT_ARTIFACTS_PATH, DoclingPipeline, DoclingPipelineConfig = (
+        None,
+        None,
+        None,
+    )
+
+try:
+    from .marker_ import MarkerPipeline, MarkerPipelineConfig
+except ImportError:
+    MarkerPipeline, MarkerPipelineConfig = None, None
+
+
+try:
+    from .miner_u import MinerUPipeline, MinerUPipelineConfig
+except ImportError:
+    MinerUPipeline, MinerUPipelineConfig = None, None
+
+
+__all__ = [
+    "DoclingPipeline",
+    "DoclingPipelineConfig",
+    "InputDoc",
+    "DOCLING_DEFAULT_ARTIFACTS_PATH",
+    "MarkerPipeline",
+    "MarkerPipelineConfig",
+    "OutputFormat",
+    "Pipeline",
+    "PipelineType",
+    "PipelineConfig",
+    "Status",
+]

extract_python/constants.py

@@ -0,0 +1,6 @@
+ARTIFACTS = "artifacts"
+CPU_GROUP = "cpu"
+MINER_U_GROUP = "miner-u"
+EXTRACT_CONTENT_TASK = "extract-content"
+EXTRACT_CONTENT_MINER_U_TASK = "extract-content-miner-u"
+DEFAULT_MD_PAGE_SEP = '<div style="page-break-after: always;"></div>'

extract_python/docling_.py

@@ -0,0 +1,233 @@
+import shutil
+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 docling.backend.abstract_backend import AbstractDocumentBackend
+from docling.datamodel.base_models import 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_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 .constants import ARTIFACTS, CPU_GROUP, DEFAULT_MD_PAGE_SEP
+from .objects import (
+    BaseModel,
+    Error,
+    InputDoc,
+    MarkdownDoc,
+    OutputFormat,
+    PageIndexes,
+    Result,
+    Status,
+)
+from .pipeline import Pipeline, PipelineConfig, PipelineType
+from .utils import all_subclasses, chdir, map_and_preserve, path_to_artifacts_dirname
+
+DOCLING_DEFAULT_ARTIFACTS_PATH = Path.home().joinpath(".cache", "docling", "models")
+
+
+class _PdfPipelineOptions(PdfPipelineOptions):
+    generate_picture_images: bool = Field(default=True, frozen=True)
+
+    @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,
+        )
+
+
+@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
+    }
+
+
+T = TypeVar("T")
+
+
+def _find_subcls(cls: type[T], name: str) -> type[T]:
+    for c in all_subclasses(cls):
+        if c.__name__ == name:
+            return c
+    raise ValueError(f"unknown {cls.__name__} subclass {name}")
+
+
+@PipelineConfig.register()
+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()
+        }
+
+
+DEFAULT_FORMAT_OPTIONS = DoclingPipelineConfig().to_format_options()
+
+
+@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)
+
+    async def extract_content(
+        self, docs: Iterable[InputDoc], output_format: OutputFormat, output_path: Path
+    ) -> 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):
+            yield _to_result(res, doc, output_format, output_path=output_path)
+
+    @classmethod
+    def _from_config(cls, config: DoclingPipelineConfig) -> FromConfig:
+        return cls(config.to_format_options())
+
+
+def _to_docling(docs: Iterable[InputDoc]) -> Iterator[Path | DocumentStream]:
+    for d in docs:
+        yield d.to_docling()
+
+
+def _to_result(
+    res: ConversionResult,
+    input_document: InputDoc,
+    output_format: OutputFormat,
+    output_path: Path,
+    **kwargs,
+) -> Result:
+    output_path.mkdir(parents=True, exist_ok=True)
+    output = None
+    status = Status.from_docling(res.status)
+    if status.allows_conversion:
+        match output_format:
+            case OutputFormat.MARKDOWN:
+                output = _to_markdown_doc(res, output_path, **kwargs)
+            case _:
+                raise NotImplementedError(f"unsupported output format {output_format}")
+    errors = [Error.from_docling(e) for e in res.errors]
+    input_doc = input_document.without_content()
+    return Result(input=input_doc, status=status, errors=errors, output=output)
+
+
+def _to_markdown_doc(
+    res: ConversionResult,
+    output_path: Path,
+    page_sep: str = DEFAULT_MD_PAGE_SEP,
+    **kwargs,
+) -> MarkdownDoc:
+    # TODO: Should we add a hash to avoid collision between files with same names
+    #  nested in the tree structured
+    md_dir_name = path_to_artifacts_dirname(res.input.file)
+    md_dir = output_path / md_dir_name
+    if md_dir.exists():
+        raise FileExistsError(f"directory {md_dir} already exists")
+    # Let's avoid issue of duplicated input file names flattened top level
+    md_filename = md_dir_name + OutputFormat.MARKDOWN
+    total_length = 0
+    n_pages = len(res.pages)
+
+    with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as td:
+        tmp_dir = Path(td)
+        page_path = Path("page.md")
+        # We do a chdir to bypass a Docling bug which only allows to maintain relative
+        # image ref when saving the markdown to a relative path
+        with (tmp_dir / md_filename).open("w") as f, chdir(tmp_dir):
+            end_indices = []
+            for page_i in range(n_pages):
+                res.document.save_as_markdown(
+                    page_path,
+                    page_no=page_i + 1,
+                    image_mode=ImageRefMode.REFERENCED,
+                    artifacts_dir=Path(ARTIFACTS),
+                    **kwargs,
+                )
+                content = page_path.read_text()
+                if page_i > 0:
+                    content += "\n"
+                if page_i < n_pages - 1:
+                    content += page_sep
+                total_length += len(content)
+                end_indices.append(total_length)
+                f.write(content)
+                f.flush()
+                page_path.unlink()
+        shutil.move(tmp_dir, md_dir)
+    pages = PageIndexes.from_page_end_indices(end_indices)
+    return MarkdownDoc(path=Path(md_dir_name), pages=pages)

extract_python/marker_.py

@@ -0,0 +1,118 @@
+import gc
+from collections.abc import AsyncGenerator, Iterable
+from copy import deepcopy
+from pathlib import Path
+from typing import Any, ClassVar, Self
+
+from marker.config.parser import ConfigParser
+from marker.converters.pdf import PdfConverter
+from marker.models import create_model_dict
+from marker.output import text_from_rendered
+from marker.renderers.markdown import MarkdownRenderer
+from PIL.Image import Image
+from pydantic import Field
+
+from .constants import ARTIFACTS, CPU_GROUP
+from .objects import (
+    InputDoc,
+    MarkdownDoc,
+    OutputFormat,
+    PageIndexes,
+    Result,
+    Status,
+)
+from .pipeline import Pipeline, PipelineConfig, PipelineType
+from .utils import path_to_artifacts_dirname, report_recoverable_errors
+
+
+@PipelineConfig.register()
+class MarkerPipelineConfig(PipelineConfig):
+    pipeline: PipelineType = Field(frozen=True, default=PipelineType.MARKER)
+    task_group: ClassVar[str] = Field(frozen=True, default=CPU_GROUP)
+
+    config: dict[str, Any] = dict()
+
+
+_MARKER_CONVERSION_ERRORS = tuple()
+
+
+@Pipeline.register(PipelineType.MARKER)
+class MarkerPipeline(Pipeline):
+    def __init__(self, marker_config: dict[str, Any] | None = None):
+        if marker_config is None:
+            marker_config = dict()
+        self._marker_config = marker_config
+
+    async def extract_content(
+        self, docs: Iterable[InputDoc], output_format: OutputFormat, output_path: Path
+    ) -> AsyncGenerator[Result, None]:
+        config = deepcopy(self._marker_config)
+        config["output_format"] = output_format.to_marker()
+        config_parser = ConfigParser(config)
+        renderer = config_parser.get_renderer()
+        converter = PdfConverter(
+            config=config_parser.generate_config_dict(),
+            artifact_dict=create_model_dict(),
+            processor_list=config_parser.get_processors(),
+            renderer=renderer,
+        )
+        for doc in docs:
+            yield _process_doc(doc, converter, output_format, output_path)
+
+    @classmethod
+    def _from_config(cls, config: MarkerPipelineConfig) -> Self:
+        return cls(config.config)
+
+
+@report_recoverable_errors(_MARKER_CONVERSION_ERRORS)
+def _process_doc(
+    doc: InputDoc,
+    converter: PdfConverter,
+    output_format: OutputFormat,
+    output_path: Path,
+) -> Result:
+    rendered = converter(str(doc.path))
+    content, _, images = text_from_rendered(rendered)
+    match output_format:
+        case OutputFormat.MARKDOWN:
+            output = _to_markdown_doc(doc, content, images, output_path)
+        case _:
+            raise NotImplementedError(f"unsupported output format {output_format}")
+    input_doc = doc.without_content()
+    return Result(input=input_doc, status=Status.SUCCESS, output=output)
+
+
+def _to_markdown_doc(
+    input_doc: InputDoc, content: str, images: dict[str, Image], output_path: Path
+) -> MarkdownDoc:
+    # TODO: Should we add a hash to avoid collision between files with same names
+    #  nested in the tree structured
+    md_dir_name = path_to_artifacts_dirname(input_doc.path)
+    md_dir = output_path / md_dir_name
+    artifacts_dir = md_dir / ARTIFACTS
+    artifacts_dir.mkdir(parents=True)
+    for im_name, im in images.items():
+        im.save(artifacts_dir / im_name)
+    del images
+    gc.collect()
+    page_sep = MarkdownRenderer.page_separator
+    content = content.split(page_sep)
+    n_pages = len(content)
+    md_path = (output_path / md_dir_name / md_dir_name).with_suffix(
+        OutputFormat.MARKDOWN.value
+    )
+    total_length = 0
+    end_indices = []
+    with md_path.open("w", encoding="utf-8") as f:
+        for page_i, page_content in enumerate(content):
+            content = page_content
+            if page_i > 0:
+                content += "\n"
+            if page_i < n_pages - 1:
+                content += page_sep
+            total_length += len(content)
+            end_indices.append(total_length)
+            f.write(content)
+            f.flush()
+    pages = PageIndexes.from_page_end_indices(end_indices)
+    return MarkdownDoc(path=Path(md_dir_name), pages=pages)

extract_python/miner_u.py

@@ -0,0 +1,212 @@
+import json
+import shutil
+from collections.abc import AsyncGenerator, Callable, Iterable
+from copy import copy
+from enum import StrEnum
+from functools import partial
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import Any, ClassVar, Self
+
+from mineru.backend.pipeline.pipeline_middle_json_mkcontent import (
+    union_make as pipeline_union_make,
+)
+from mineru.backend.vlm.vlm_middle_json_mkcontent import union_make as vlm_union_make
+from mineru.cli.common import aio_do_parse
+from mineru.utils.enum_class import MakeMode
+from pydantic import Field
+from pydantic_extra_types.language_code import LanguageAlpha2
+
+from .constants import ARTIFACTS, DEFAULT_MD_PAGE_SEP, MINER_U_GROUP
+from .objects import (
+    BaseModel,
+    ConversionOutput,
+    InputDoc,
+    OutputFormat,
+    PageIndexes,
+    Result,
+    Status,
+)
+from .pipeline import Pipeline, PipelineConfig, PipelineType
+from .utils import path_to_artifacts_dirname
+
+_MINER_U_CONVERSION_ERRORS = tuple()
+MDMakeFunction = Callable[[list, str, str], str | None]
+
+
+class MinerUBackend(StrEnum):
+    PIPELINE = "pipeline"
+    VLM = "vlm"
+
+
+class MinerUConfig(BaseModel):
+    backend: MinerUBackend = MinerUBackend.PIPELINE
+    enable_formula_extraction: bool = True
+    enable_table_extraction: bool = True
+    # TODO: use enum or literal here
+    parse_method: str = "auto"
+
+    default_kwargs: ClassVar[dict] = {
+        "server_url": None,
+        # We don't dump md directly we process, we dump the middle json in order to be
+        # able to get page indexes
+        "parse_method": "auto",
+        "dump_md": False,
+        "dump_middle_json": True,
+        "f_draw_layout_bbox": False,
+        "f_draw_span_bbox": False,
+        "f_dump_model_output": False,  # might be useful for debug though
+        "f_dump_orig_pdf": False,
+        "f_dump_content_list": False,  # might be useful for debug though
+        "start_page_id": 0,
+        "f_make_md_mode": MakeMode.MM_MD,
+        "image_analysis": True,
+        "end_page_id": None,
+        "client_side_output_generation": False,
+    }
+
+    def as_parse_kwargs(self) -> dict[str, Any]:
+        kwargs = copy(self.default_kwargs)
+        kwargs["backend"] = self.backend
+        kwargs["parse_method"] = self.parse_method
+        kwargs["formula_enable"] = self.enable_formula_extraction
+        kwargs["table_enable"] = self.enable_table_extraction
+        return kwargs
+
+
+@PipelineConfig.register()  # noqa: F821
+class MinerUPipelineConfig(PipelineConfig):  # noqa: F821
+    pipeline: PipelineType = Field(frozen=True, default=PipelineType.MINER_U)
+    task_group: ClassVar[str] = Field(frozen=True, default=MINER_U_GROUP)
+
+    config: MinerUConfig = Field(frozen=True, default=MinerUConfig())
+    language: LanguageAlpha2 = Field(frozen=True, default="en")
+
+
+@Pipeline.register(PipelineType.MINER_U)
+class MinerUPipeline(Pipeline):
+    def __init__(self, config: MinerUConfig, language: str):
+        self._config = config
+        self._language = language
+        self._md_make_fn = _parse_md_make_fn(config.backend)
+
+    async def extract_content(
+        self, docs: Iterable[InputDoc], output_format: OutputFormat, output_path: Path
+    ) -> AsyncGenerator[Result, None]:
+        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 _revert_mineru_output(output_dir: Path, *, pdf_filename: str) -> Path:
+    output_path = output_dir / pdf_filename
+    if not output_path.exists():
+        msg = f"couldn't find result for {pdf_filename}"
+        raise FileNotFoundError(msg)
+    dirs = [p for p in output_path.iterdir() if p.is_dir()]
+    if len(dirs) != 1:
+        msg = f"expected exactly one result directory, found: {dirs}"
+        raise ValueError(msg)
+    return output_dir / dirs[0]
+
+
+def _parse_md_make_fn(backend: MinerUBackend) -> MDMakeFunction:
+    match backend:
+        case MinerUBackend.PIPELINE:
+            return pipeline_union_make
+        case MinerUBackend.VLM:
+            return vlm_union_make
+        case _:
+            raise ValueError(f"Unsupported backend: {backend}")
+
+
+def _process_doc(
+    doc: InputDoc,
+    *,
+    md_make_fn: MDMakeFunction,
+    res_path: Path,
+    output_format: OutputFormat,
+    output_path: Path,
+) -> Result:
+    md_dir_name = path_to_artifacts_dirname(doc.path)
+    md_dir = Path(output_path) / md_dir_name
+    md_dir.mkdir(parents=True, exist_ok=False)
+    artifacts_dir = md_dir / ARTIFACTS
+    md_path = (md_dir / md_dir_name).with_suffix(OutputFormat.MARKDOWN.value)
+    # Fail early
+    match output_format:
+        case OutputFormat.MARKDOWN:
+            im_rel_dir = artifacts_dir.relative_to(md_dir)
+            dump_content_fn = partial(
+                _dump_md_content,
+                md_make_fn=md_make_fn,
+                output_path=output_path,
+                md_path=md_path,
+                im_dir=im_rel_dir,
+            )
+        case _:
+            raise NotImplementedError(f"unsupported output format {output_format}")
+    middle_json_path = res_path / f"{doc.path.name}_middle.json"
+    middle_json = json.loads(middle_json_path.read_text())
+    pdf_info = middle_json["pdf_info"]
+    shutil.move(res_path / "images", artifacts_dir)
+    output = dump_content_fn(pdf_info)
+    input_doc = doc.without_content()
+    return Result(input=input_doc, status=Status.SUCCESS, output=output)
+
+
+def _dump_md_content(
+    pdf_info: list[dict],
+    *,
+    md_make_fn: MDMakeFunction,
+    page_sep: str = DEFAULT_MD_PAGE_SEP,
+    output_path: Path,
+    md_path: Path,
+    im_dir: Path,
+    md_make_mode: str = MakeMode.MM_MD,
+) -> ConversionOutput:
+    total_length = 0
+    end_indices = []
+    with md_path.open("w") as f:
+        n_pages = len(pdf_info)
+        for page_i, page in enumerate(pdf_info):
+            content = md_make_fn([page], md_make_mode, str(im_dir))
+            if page_i > 0:
+                content += "\n"
+            if page_i < n_pages - 1:
+                content += page_sep
+            total_length += len(content)
+            end_indices.append(total_length)
+            f.write(content)
+            f.flush()
+    end_indices = PageIndexes.from_page_end_indices(end_indices)
+    output_path = md_path.parent.relative_to(output_path)
+    output = ConversionOutput(path=output_path, pages=end_indices)
+    return output

extract_python/objects.py

@@ -0,0 +1,254 @@
+from __future__ import annotations
+
+import logging
+import os
+import traceback
+import uuid
+from abc import ABC
+from enum import StrEnum
+from functools import cache
+from io import BytesIO
+from pathlib import Path
+from typing import Annotated, Any, NoReturn, Self
+
+from icij_common.pydantic_utils import (
+    icij_config,
+    merge_configs,
+    no_enum_values_config,
+    safe_copy,
+)
+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.document import InputDocument
+    from docling_core.types.io import DocumentStream
+except ImportError:
+    ConversionStatus, ErrorItem, InputFormat = None, None, None
+    InputDocument = None
+    DocumentStream = None
+
+logger = logging.getLogger(__name__)
+base_config = merge_configs(icij_config(), no_enum_values_config())
+
+
+class BaseModel(_BaseModel):
+    model_config = base_config
+
+
+class SupportedExt(StrEnum):
+    PDF = ".pdf"
+
+    def to_docling(self) -> InputFormat:
+        return InputFormat(self.value[1:])
+
+
+class OutputFormat(StrEnum):
+    MARKDOWN = ".md"
+
+    @property
+    def suffix(self) -> str:
+        return self.value[1:]
+
+    def to_marker(self) -> str:
+        match self:
+            case OutputFormat.MARKDOWN:
+                return "markdown"
+            case _:
+                raise ValueError(f"{self} is unsupported by marker")
+
+
+class Status(StrEnum):
+    FAILURE = "failure"
+    SUCCESS = "success"
+    PARTIAL_SUCCESS = "partial_success"
+
+    @classmethod
+    def from_docling(cls, v: Any) -> Self:
+        from docling.datamodel.base_models import ConversionStatus  # noqa: PLC0415
+
+        if v is ConversionStatus.SUCCESS:
+            return cls.SUCCESS
+        if v is ConversionStatus.PARTIAL_SUCCESS:
+            return cls.PARTIAL_SUCCESS
+        if isinstance(v, ConversionStatus):
+            return cls.FAILURE
+        raise TypeError(f"can't convert {v!r} to {cls.__name__!r}")
+
+    @property
+    def allows_conversion(self) -> bool:
+        return self is Status.SUCCESS or self is Status.PARTIAL_SUCCESS
+
+
+class Error(BaseModel):
+    id: str
+    title: str
+    detail: str
+
+    @classmethod
+    def from_exception(cls, exception: BaseException) -> Self:
+        title = exception.__class__.__name__
+        trace_lines = traceback.format_exception(
+            None, value=exception, tb=exception.__traceback__
+        )
+        detail = f"{exception}\n{''.join(trace_lines)}"
+        error_id = f"{_id_title(title)}-{uuid.uuid4().hex}"
+        error = cls(id=error_id, title=title, detail=detail)
+        return error
+
+    @classmethod
+    def from_docling(cls, docling_error: ErrorItem) -> Self:
+        title = "DoclingConversionError"
+        error_id = f"{_id_title(title)}-{uuid.uuid4().hex}"
+        detail = (
+            f"error in module {docling_error.module_name} of"
+            f" {docling_error.component_type}:\n{docling_error.error_message}"
+        )
+        return cls(id=error_id, title=title, detail=detail)
+
+
+def _id_title(title: str) -> str:
+    id_title = []
+    for i, letter in enumerate(title):
+        if i and letter.isupper():
+            id_title.append("-")
+        id_title.append(letter.lower())
+    return "".join(id_title)
+
+
+class InputDoc(BaseModel):
+    ext: SupportedExt
+    path: Path
+    content: bytes | None = None
+
+    @classmethod
+    def from_path(cls, path: str | Path) -> Self:
+        if isinstance(path, str):
+            path = Path(path)
+        ext = SupportedExt(path.suffix)
+        return cls(path=path, ext=ext)
+
+    def to_docling(self) -> Path | DocumentStream:
+        if self.content is not None:
+            return DocumentStream(name=str(self.path), stream=BytesIO(self.content))
+        if not self.path.suffix:
+            return DocumentStream(
+                name=str(self.path), stream=BytesIO(self.path.read_bytes())
+            )
+        return self.path
+
+    def without_content(self) -> Self:
+        return safe_copy(self, update={"content": None})
+
+
+class PageIndexes(RootModel[list[tuple[int, int]]]):
+    # Stores page end index
+    @classmethod
+    def from_page_end_indices(cls, lengths: list[int]) -> Self:
+        return [
+            ((lengths[p - 1] if p > 0 else 0), lengths[p]) for p in range(len(lengths))
+        ]
+
+
+class ConversionOutput(BaseModel):
+    path: Path
+    pages: PageIndexes = []
+
+
+class MarkdownDoc(ConversionOutput):
+    @classmethod
+    @property
+    @cache
+    def _valid_conversion_statuses(cls) -> set[ConversionStatus]:
+        from docling.datamodel.base_models import ConversionStatus  # noqa: PLC0415
+
+        return {ConversionStatus.SUCCESS, ConversionStatus.PARTIAL_SUCCESS}
+
+
+def _input_should_not_have_content(value: InputDoc) -> InputDoc:
+    if value.content is not None:
+        raise ValueError(f"response input can't have content, but got {value}")
+    return value
+
+
+class _BaseResult(BaseModel, ABC):
+    input: InputDoc
+    status: Status
+    errors: list[Error] = []
+
+
+class Result(_BaseResult):
+    # TODO: we could also use generics here when we add more output formats
+    output: ConversionOutput | None
+
+    def to_response(self) -> ResponseResult:
+        return ResponseResult(
+            input=self.input.without_content(),
+            status=self.status,
+            errors=self.errors,
+            output_path=self.output.path,
+        )
+
+
+class ResponseResult(_BaseResult):
+    input: Annotated[InputDoc, AfterValidator(func=_input_should_not_have_content)]
+    output_path: Path
+
+
+class ExtractionResponse(BaseModel):
+    results: list[ResponseResult]
+
+
+_INPUT_DOCS_ADAPTER = TypeAdapter(list[InputDoc | Path])
+
+
+def parse_extraction_request(
+    docs: str | list[dict | str], *, data_dir: Path
+) -> list[InputDoc]:
+    if isinstance(docs, str):
+        logger.debug("exploring files in %s", data_dir.absolute())
+        docs_dir = Path(data_dir) / docs
+        docs = _as_input_docs(docs_dir)
+        msg = "found %s"
+        if len(docs) > 10:
+            msg = msg + ", and more..."
+        logger.debug("found %s", docs[:10])
+        return docs
+    docs = _INPUT_DOCS_ADAPTER.validate_python(docs)
+    if not docs:
+        return []
+    if isinstance(docs[0], Path):
+        doc_meta = []
+        unknown_exts = []
+        for doc in docs:
+            _, ext = os.path.splitext(str(doc))
+            if not ext:
+                unknown_exts.append(doc)
+            else:
+                doc_meta.append(InputDoc.from_path(path=doc.relative_to(data_dir)))
+        if unknown_exts:
+            raise ValueError(f"found files with unknown extensions {unknown_exts}")
+        return doc_meta
+    return docs
+
+
+def _raise(err: OSError) -> NoReturn:
+    raise err
+
+
+def _as_input_docs(
+    docs_dir: Path, *, supported_ext: set[str] | None = None
+) -> list[InputDoc]:
+    if supported_ext is None:
+        supported_ext = {v.value for v in SupportedExt}
+    docs = []
+    for root, _, files in os.walk(docs_dir, onerror=_raise):
+        root = Path(root)  # noqa: PLW2901
+        for f in files:
+            ext = Path(f).suffix
+            if not ext or ext not in supported_ext:
+                continue
+            docs.append(InputDoc.from_path(path=root / f))
+    docs = sorted(docs, key=lambda x: x.path)
+    return docs

extract_python/pipeline.py

@@ -0,0 +1,36 @@
+from abc import ABC, abstractmethod
+from collections.abc import AsyncGenerator, Iterable
+from enum import StrEnum
+from pathlib import Path
+from typing import ClassVar
+
+from icij_common.pydantic_utils import icij_config, merge_configs, no_enum_values_config
+from icij_common.registrable import RegistrableConfig, RegistrableFromConfig
+from pydantic import Field
+
+from .objects import InputDoc, OutputFormat, Result
+
+StructuredContent = str
+
+
+class PipelineType(StrEnum):
+    DOCLING = "docling"
+    MARKER = "marker"
+    MINER_U = "miner_u"
+
+
+class PipelineConfig(RegistrableConfig, ABC):
+    # TODO: move this icij_config() to RegistrableConfig
+    model_config = merge_configs(icij_config(), no_enum_values_config())
+
+    registry_key: ClassVar[str] = Field(frozen=True, default="pipeline")
+    pipeline: PipelineType
+
+    task_group: ClassVar[str] = Field(frozen=True)
+
+
+class Pipeline(RegistrableFromConfig, ABC):
+    @abstractmethod
+    async def extract_content(
+        self, docs: Iterable[InputDoc], output_format: OutputFormat, output_path: Path
+    ) -> AsyncGenerator[Result, None]: ...

extract_python/utils.py

@@ -0,0 +1,70 @@
+import os
+from collections.abc import Callable, Generator, Iterable, Iterator
+from contextlib import contextmanager
+from functools import wraps
+from itertools import tee
+from pathlib import Path, PurePath
+from typing import Protocol, TypeVar
+
+from .objects import Error, InputDoc, Result, Status
+
+R = TypeVar("R")
+T = TypeVar("T")
+
+
+def map_and_preserve(
+    fn: Callable[[Iterable[T]], Iterator[R]], inputs: Iterable[T]
+) -> tuple[Iterable[T], Iterator[R]]:
+    save_inputs, function_inputs = tee(inputs)
+    outputs = iter(fn(function_inputs))
+    return save_inputs, outputs
+
+
+def all_subclasses(cls: type[T]) -> set[type[T]]:
+    return set(cls.__subclasses__()).union(
+        [s for c in cls.__subclasses__() for s in all_subclasses(c)]
+    )
+
+
+def path_to_artifacts_dirname(path: PurePath, sep: str = "_") -> str:
+    dirname = f"{path.name[: -len(path.suffix)]}"
+    ext = path.suffix
+    if ext:
+        dirname += sep + ext[1:]
+    return dirname
+
+
+class DocProcessingFn(Protocol):
+    def __call__(self, doc: InputDoc, *arg, **kwargs) -> Result: ...
+
+
+def report_recoverable_errors(
+    recoverable_errors: tuple[type[Exception], ...] = tuple(),
+) -> Callable[[DocProcessingFn], DocProcessingFn]:
+    def make_decorator(f: DocProcessingFn) -> DocProcessingFn:
+        @wraps(f)
+        def wrapped(doc: InputDoc, *args, **kwargs) -> Result:
+            try:
+                return f(doc, *args, **kwargs)
+            except recoverable_errors as e:
+                error = Error.from_exception(e)
+                return Result(
+                    input=doc.without_content(),
+                    status=Status.FAILURE,
+                    errors=[error],
+                    output=None,
+                )
+
+        return wrapped
+
+    return make_decorator
+
+
+@contextmanager
+def chdir(path: Path) -> Generator[None, None, None]:
+    cwd = Path.cwd()
+    try:
+        os.chdir(path)
+        yield
+    finally:
+        os.chdir(cwd)

extract_python-0.1.0.dist-info/METADATA

@@ -0,0 +1,24 @@
+Metadata-Version: 2.4
+Name: extract-python
+Version: 0.1.0
+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: icij-common~=0.8.2
+Provides-Extra: benches
+Requires-Dist: html2image~=2.0.7; extra == 'benches'
+Requires-Dist: markdown2>=2.5.4; extra == 'benches'
+Requires-Dist: notebook>=7.4.5; extra == 'benches'
+Requires-Dist: pypdfium2>=4.30.0; extra == 'benches'
+Provides-Extra: docling
+Requires-Dist: docling-slim[feat-ocr-easyocr,feat-ocr-mac,feat-ocr-tesserocr,standard]~=2.96; extra == 'docling'
+Provides-Extra: marker
+Requires-Dist: marker-pdf~=1.10; extra == 'marker'
+Provides-Extra: mineru
+Requires-Dist: mineru[mlx]~=3.2; (sys_platform == 'darwin') and extra == 'mineru'
+Requires-Dist: mineru[pipeline,vlm]~=3.2; extra == 'mineru'
+Requires-Dist: pydantic-extra-types[pycountry]~=2.11; extra == 'mineru'
+Requires-Dist: six~=1.17; extra == 'mineru'

extract_python-0.1.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=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,,

extract_python-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.30.1
+Root-Is-Purelib: true
+Tag: py3-none-any