docling 2.69.0__py3-none-any.whl

docling/document_extractor.py

@@ -0,0 +1,327 @@
+import hashlib
+import logging
+import sys
+import threading
+import time
+import warnings
+from collections.abc import Iterable, Iterator
+from concurrent.futures import ThreadPoolExecutor
+from functools import partial
+from pathlib import Path
+from typing import Optional, Type, Union
+
+from pydantic import ConfigDict, model_validator, validate_call
+from typing_extensions import Self
+
+from docling.backend.abstract_backend import AbstractDocumentBackend
+from docling.backend.image_backend import ImageDocumentBackend
+from docling.backend.pypdfium2_backend import PyPdfiumDocumentBackend
+from docling.datamodel.base_models import (
+    BaseFormatOption,
+    ConversionStatus,
+    DoclingComponentType,
+    DocumentStream,
+    ErrorItem,
+    InputFormat,
+)
+from docling.datamodel.document import (
+    InputDocument,
+    _DocumentConversionInput,  # intentionally reused builder
+)
+from docling.datamodel.extraction import ExtractionResult, ExtractionTemplateType
+from docling.datamodel.pipeline_options import PipelineOptions
+from docling.datamodel.settings import (
+    DEFAULT_PAGE_RANGE,
+    DocumentLimits,
+    PageRange,
+    settings,
+)
+from docling.exceptions import ConversionError
+from docling.pipeline.base_extraction_pipeline import BaseExtractionPipeline
+from docling.pipeline.extraction_vlm_pipeline import ExtractionVlmPipeline
+from docling.utils.utils import chunkify
+
+_log = logging.getLogger(__name__)
+_PIPELINE_CACHE_LOCK = threading.Lock()
+
+
+class ExtractionFormatOption(BaseFormatOption):
+    """Per-format configuration for extraction.
+
+    Notes:
+        - `pipeline_cls` must subclass `BaseExtractionPipeline`.
+        - `pipeline_options` is typed as `PipelineOptions` which MUST inherit from
+          `BaseOptions` (as used by `BaseExtractionPipeline`).
+        - `backend` is the document-opening backend used by `_DocumentConversionInput`.
+    """
+
+    pipeline_cls: Type[BaseExtractionPipeline]
+
+    @model_validator(mode="after")
+    def set_optional_field_default(self) -> Self:
+        if self.pipeline_options is None:
+            # `get_default_options` comes from BaseExtractionPipeline
+            self.pipeline_options = self.pipeline_cls.get_default_options()  # type: ignore[assignment]
+        return self
+
+
+def _get_default_extraction_option(fmt: InputFormat) -> ExtractionFormatOption:
+    """Return the default extraction option for a given input format.
+
+    Defaults mirror the converter's *backend* choices, while the pipeline is
+    the VLM extractor. This duplication will be removed when we deduplicate
+    the format registry between convert/extract.
+    """
+    format_to_default_backend: dict[InputFormat, Type[AbstractDocumentBackend]] = {
+        InputFormat.IMAGE: ImageDocumentBackend,
+        InputFormat.PDF: PyPdfiumDocumentBackend,
+    }
+
+    backend = format_to_default_backend.get(fmt)
+    if backend is None:
+        raise RuntimeError(f"No default extraction backend configured for {fmt}")
+
+    return ExtractionFormatOption(
+        pipeline_cls=ExtractionVlmPipeline,
+        backend=backend,
+    )
+
+
+class DocumentExtractor:
+    """Standalone extractor class.
+
+    Public API:
+        - `extract(...) -> ExtractionResult`
+        - `extract_all(...) -> Iterator[ExtractionResult]`
+
+    Implementation intentionally reuses `_DocumentConversionInput` to build
+    `InputDocument` with the correct backend per format.
+    """
+
+    def __init__(
+        self,
+        allowed_formats: Optional[list[InputFormat]] = None,
+        extraction_format_options: Optional[
+            dict[InputFormat, ExtractionFormatOption]
+        ] = None,
+    ) -> None:
+        self.allowed_formats: list[InputFormat] = (
+            allowed_formats if allowed_formats is not None else list(InputFormat)
+        )
+        # Build per-format options with defaults, then apply any user overrides
+        overrides = extraction_format_options or {}
+        self.extraction_format_to_options: dict[InputFormat, ExtractionFormatOption] = {
+            fmt: overrides.get(fmt, _get_default_extraction_option(fmt))
+            for fmt in self.allowed_formats
+        }
+
+        # Cache pipelines by (class, options-hash)
+        self._initialized_pipelines: dict[
+            tuple[Type[BaseExtractionPipeline], str], BaseExtractionPipeline
+        ] = {}
+
+    # ---------------------------- Public API ---------------------------------
+
+    @validate_call(config=ConfigDict(strict=True))
+    def extract(
+        self,
+        source: Union[Path, str, DocumentStream],
+        template: ExtractionTemplateType,
+        headers: Optional[dict[str, str]] = None,
+        raises_on_error: bool = True,
+        max_num_pages: int = sys.maxsize,
+        max_file_size: int = sys.maxsize,
+        page_range: PageRange = DEFAULT_PAGE_RANGE,
+    ) -> ExtractionResult:
+        all_res = self.extract_all(
+            source=[source],
+            headers=headers,
+            raises_on_error=raises_on_error,
+            max_num_pages=max_num_pages,
+            max_file_size=max_file_size,
+            page_range=page_range,
+            template=template,
+        )
+        return next(all_res)
+
+    @validate_call(config=ConfigDict(strict=True))
+    def extract_all(
+        self,
+        source: Iterable[Union[Path, str, DocumentStream]],
+        template: ExtractionTemplateType,
+        headers: Optional[dict[str, str]] = None,
+        raises_on_error: bool = True,
+        max_num_pages: int = sys.maxsize,
+        max_file_size: int = sys.maxsize,
+        page_range: PageRange = DEFAULT_PAGE_RANGE,
+    ) -> Iterator[ExtractionResult]:
+        warnings.warn(
+            "The extract API is currently experimental and may change without prior notice.\n"
+            "Only PDF and image formats are supported.",
+            UserWarning,
+            stacklevel=2,
+        )
+
+        limits = DocumentLimits(
+            max_num_pages=max_num_pages,
+            max_file_size=max_file_size,
+            page_range=page_range,
+        )
+        conv_input = _DocumentConversionInput(
+            path_or_stream_iterator=source, limits=limits, headers=headers
+        )
+
+        ext_res_iter = self._extract(
+            conv_input, raises_on_error=raises_on_error, template=template
+        )
+
+        had_result = False
+        for ext_res in ext_res_iter:
+            had_result = True
+            if raises_on_error and ext_res.status not in {
+                ConversionStatus.SUCCESS,
+                ConversionStatus.PARTIAL_SUCCESS,
+            }:
+                raise ConversionError(
+                    f"Extraction failed for: {ext_res.input.file} with status: {ext_res.status}"
+                )
+            else:
+                yield ext_res
+
+        if not had_result and raises_on_error:
+            raise ConversionError(
+                "Extraction failed because the provided file has no recognizable format or it wasn't in the list of allowed formats."
+            )
+
+    # --------------------------- Internal engine ------------------------------
+
+    def _extract(
+        self,
+        conv_input: _DocumentConversionInput,
+        raises_on_error: bool,
+        template: ExtractionTemplateType,
+    ) -> Iterator[ExtractionResult]:
+        start_time = time.monotonic()
+
+        for input_batch in chunkify(
+            conv_input.docs(self.extraction_format_to_options),
+            settings.perf.doc_batch_size,
+        ):
+            _log.info("Going to extract document batch...")
+            process_func = partial(
+                self._process_document_extraction,
+                raises_on_error=raises_on_error,
+                template=template,
+            )
+
+            if (
+                settings.perf.doc_batch_concurrency > 1
+                and settings.perf.doc_batch_size > 1
+            ):
+                with ThreadPoolExecutor(
+                    max_workers=settings.perf.doc_batch_concurrency
+                ) as pool:
+                    for item in pool.map(
+                        process_func,
+                        input_batch,
+                    ):
+                        yield item
+            else:
+                for item in map(
+                    process_func,
+                    input_batch,
+                ):
+                    elapsed = time.monotonic() - start_time
+                    start_time = time.monotonic()
+                    _log.info(
+                        f"Finished extracting document {item.input.file.name} in {elapsed:.2f} sec."
+                    )
+                    yield item
+
+    def _process_document_extraction(
+        self,
+        in_doc: InputDocument,
+        raises_on_error: bool,
+        template: ExtractionTemplateType,
+    ) -> ExtractionResult:
+        valid = (
+            self.allowed_formats is not None and in_doc.format in self.allowed_formats
+        )
+        if valid:
+            return self._execute_extraction_pipeline(
+                in_doc, raises_on_error=raises_on_error, template=template
+            )
+        else:
+            error_message = f"File format not allowed: {in_doc.file}"
+            if raises_on_error:
+                raise ConversionError(error_message)
+            else:
+                error_item = ErrorItem(
+                    component_type=DoclingComponentType.USER_INPUT,
+                    module_name="",
+                    error_message=error_message,
+                )
+                return ExtractionResult(
+                    input=in_doc, status=ConversionStatus.SKIPPED, errors=[error_item]
+                )
+
+    def _execute_extraction_pipeline(
+        self,
+        in_doc: InputDocument,
+        raises_on_error: bool,
+        template: ExtractionTemplateType,
+    ) -> ExtractionResult:
+        if not in_doc.valid:
+            if raises_on_error:
+                raise ConversionError(f"Input document {in_doc.file} is not valid.")
+            else:
+                return ExtractionResult(input=in_doc, status=ConversionStatus.FAILURE)
+
+        pipeline = self._get_pipeline(in_doc.format)
+        if pipeline is None:
+            if raises_on_error:
+                raise ConversionError(
+                    f"No extraction pipeline could be initialized for {in_doc.file}."
+                )
+            else:
+                return ExtractionResult(input=in_doc, status=ConversionStatus.FAILURE)
+
+        return pipeline.execute(
+            in_doc, raises_on_error=raises_on_error, template=template
+        )
+
+    def _get_pipeline(
+        self, doc_format: InputFormat
+    ) -> Optional[BaseExtractionPipeline]:
+        """Retrieve or initialize a pipeline, reusing instances based on class and options."""
+        fopt = self.extraction_format_to_options.get(doc_format)
+        if fopt is None or fopt.pipeline_options is None:
+            return None
+
+        pipeline_class = fopt.pipeline_cls
+        pipeline_options = fopt.pipeline_options
+        options_hash = self._get_pipeline_options_hash(pipeline_options)
+
+        cache_key = (pipeline_class, options_hash)
+        with _PIPELINE_CACHE_LOCK:
+            if cache_key not in self._initialized_pipelines:
+                _log.info(
+                    f"Initializing extraction pipeline for {pipeline_class.__name__} with options hash {options_hash}"
+                )
+                self._initialized_pipelines[cache_key] = pipeline_class(
+                    pipeline_options=pipeline_options  # type: ignore[arg-type]
+                )
+            else:
+                _log.debug(
+                    f"Reusing cached extraction pipeline for {pipeline_class.__name__} with options hash {options_hash}"
+                )
+
+            return self._initialized_pipelines[cache_key]
+
+    @staticmethod
+    def _get_pipeline_options_hash(pipeline_options: PipelineOptions) -> str:
+        """Generate a stable hash of pipeline options to use as part of the cache key."""
+        options_str = str(pipeline_options.model_dump())
+        return hashlib.md5(
+            options_str.encode("utf-8"), usedforsecurity=False
+        ).hexdigest()

docling/exceptions.py

@@ -0,0 +1,10 @@
+class BaseError(RuntimeError):
+    pass
+
+
+class ConversionError(BaseError):
+    pass
+
+
+class OperationNotAllowed(BaseError):
+    pass

docling/experimental/__init__.py

@@ -0,0 +1,5 @@
+"""Experimental modules for Docling.
+
+This package contains experimental features that are under development
+and may change or be removed in future versions.
+"""

docling/experimental/datamodel/__init__.py

@@ -0,0 +1 @@
+"""Experimental datamodel modules."""

docling/experimental/datamodel/table_crops_layout_options.py

@@ -0,0 +1,13 @@
+"""Internal options for the experimental TableCrops layout model."""
+
+from typing import ClassVar
+
+from docling.datamodel.pipeline_options import BaseLayoutOptions
+
+__all__ = ["TableCropsLayoutOptions"]
+
+
+class TableCropsLayoutOptions(BaseLayoutOptions):
+    """Options for TableCropsLayoutModel (internal-only)."""
+
+    kind: ClassVar[str] = "docling_experimental_table_crops_layout"

docling/experimental/datamodel/threaded_layout_vlm_pipeline_options.py

@@ -0,0 +1,45 @@
+"""Options for the threaded layout+VLM pipeline."""
+
+from typing import Union
+
+from pydantic import model_validator
+
+from docling.datamodel.layout_model_specs import DOCLING_LAYOUT_HERON
+from docling.datamodel.pipeline_options import LayoutOptions, PaginatedPipelineOptions
+from docling.datamodel.pipeline_options_vlm_model import (
+    ApiVlmOptions,
+    InlineVlmOptions,
+    ResponseFormat,
+)
+from docling.datamodel.vlm_model_specs import GRANITEDOCLING_TRANSFORMERS
+
+
+class ThreadedLayoutVlmPipelineOptions(PaginatedPipelineOptions):
+    """Pipeline options for the threaded layout+VLM pipeline."""
+
+    images_scale: float = 2.0
+
+    # VLM configuration (will be enhanced with layout awareness by the pipeline)
+    vlm_options: Union[InlineVlmOptions, ApiVlmOptions] = GRANITEDOCLING_TRANSFORMERS
+
+    # Layout model configuration
+    layout_options: LayoutOptions = LayoutOptions(
+        model_spec=DOCLING_LAYOUT_HERON, skip_cell_assignment=True
+    )
+
+    # Threading and batching controls
+    layout_batch_size: int = 4
+    vlm_batch_size: int = 4
+    batch_timeout_seconds: float = 2.0
+    queue_max_size: int = 50
+
+    @model_validator(mode="after")
+    def validate_response_format(self):
+        """Validate that VLM response format is DOCTAGS (required for this pipeline)."""
+        if self.vlm_options.response_format != ResponseFormat.DOCTAGS:
+            raise ValueError(
+                f"ThreadedLayoutVlmPipeline only supports DOCTAGS response format, "
+                f"but got {self.vlm_options.response_format}. "
+                f"Please set vlm_options.response_format=ResponseFormat.DOCTAGS"
+            )
+        return self

docling/experimental/models/__init__.py

@@ -0,0 +1,3 @@
+"""Experimental models for Docling."""
+
+__all__: list[str] = []

docling/experimental/models/table_crops_layout_model.py

@@ -0,0 +1,114 @@
+"""Internal TableCrops layout model that marks full pages as table clusters."""
+
+from __future__ import annotations
+
+import warnings
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+from docling_core.types.doc import DocItemLabel
+
+from docling.datamodel.accelerator_options import AcceleratorOptions
+from docling.datamodel.base_models import BoundingBox, Cluster, LayoutPrediction, Page
+from docling.datamodel.document import ConversionResult
+from docling.experimental.datamodel.table_crops_layout_options import (
+    TableCropsLayoutOptions,
+)
+from docling.models.base_layout_model import BaseLayoutModel
+
+__all__ = ["TableCropsLayoutModel"]
+
+
+class TableCropsLayoutModel(BaseLayoutModel):
+    """Experimental layout model that treats the full page as a table cluster.
+    This is useful in cases where a Docling pipeline is applied to images of table crops only.
+
+    This model is internal and not part of the stable public interface.
+    """
+
+    def __init__(
+        self,
+        artifacts_path: Optional[Path],
+        accelerator_options: AcceleratorOptions,
+        options: TableCropsLayoutOptions,
+    ):
+        self.options = options
+        self.artifacts_path = artifacts_path
+        self.accelerator_options = accelerator_options
+
+    @classmethod
+    def get_options_type(cls) -> type[TableCropsLayoutOptions]:
+        return TableCropsLayoutOptions
+
+    def predict_layout(
+        self,
+        conv_res: ConversionResult,
+        pages: Sequence[Page],
+    ) -> Sequence[LayoutPrediction]:
+        layout_predictions: list[LayoutPrediction] = []
+
+        for page in pages:
+            if page._backend is None or not page._backend.is_valid():
+                existing_prediction = page.predictions.layout or LayoutPrediction()
+                layout_predictions.append(existing_prediction)
+                continue
+
+            clusters = self._build_page_clusters(page)
+            prediction = LayoutPrediction(clusters=clusters)
+
+            self._update_confidence(conv_res, page, clusters)
+
+            layout_predictions.append(prediction)
+
+        return layout_predictions
+
+    def _build_page_clusters(self, page: Page) -> list[Cluster]:
+        page_size = page.size
+        if page_size is None:
+            return []
+
+        bbox = BoundingBox(
+            l=0.0,
+            t=0.0,
+            r=page_size.width,
+            b=page_size.height,
+        )
+
+        cluster = Cluster(
+            id=0,
+            label=DocItemLabel.TABLE,
+            bbox=bbox,
+            confidence=1.0,
+            cells=[],
+        )
+
+        clusters = [cluster]
+
+        if not self.options.skip_cell_assignment:
+            page_cells = list(page.cells)
+            cluster.cells = page_cells
+
+            if not page_cells and not self.options.keep_empty_clusters:
+                clusters = []
+
+        return clusters
+
+    def _update_confidence(
+        self, conv_res: ConversionResult, page: Page, clusters: list[Cluster]
+    ) -> None:
+        """Populate layout and OCR confidence scores for the page."""
+        with warnings.catch_warnings():
+            warnings.filterwarnings(
+                "ignore",
+                "Mean of empty slice|invalid value encountered in scalar divide",
+                RuntimeWarning,
+                "numpy",
+            )
+
+            conv_res.confidence.pages[page.page_no].layout_score = 1.0
+
+            ocr_cells = [cell for cell in page.cells if cell.from_ocr]
+            ocr_confidence = float(np.mean([cell.confidence for cell in ocr_cells]))
+            conv_res.confidence.pages[page.page_no].ocr_score = ocr_confidence

docling/experimental/pipeline/__init__.py

@@ -0,0 +1 @@
+"""Experimental pipeline modules."""