docling 2.69.0__py3-none-any.whl

docling/pipeline/base_pipeline.py

@@ -0,0 +1,326 @@
+import functools
+import logging
+import time
+import traceback
+from abc import ABC, abstractmethod
+from collections.abc import Iterable
+from pathlib import Path
+from typing import Any, Callable, List, Optional
+
+from docling_core.types.doc import NodeItem
+
+from docling.backend.abstract_backend import (
+    AbstractDocumentBackend,
+    PaginatedDocumentBackend,
+)
+from docling.backend.pdf_backend import PdfDocumentBackend
+from docling.datamodel.base_models import (
+    ConversionStatus,
+    DoclingComponentType,
+    ErrorItem,
+    Page,
+)
+from docling.datamodel.document import ConversionResult, InputDocument
+from docling.datamodel.pipeline_options import (
+    ConvertPipelineOptions,
+    PdfPipelineOptions,
+    PipelineOptions,
+)
+from docling.datamodel.settings import settings
+from docling.models.base_model import GenericEnrichmentModel
+from docling.models.factories import get_picture_description_factory
+from docling.models.picture_description_base_model import PictureDescriptionBaseModel
+from docling.models.stages.picture_classifier.document_picture_classifier import (
+    DocumentPictureClassifier,
+    DocumentPictureClassifierOptions,
+)
+from docling.utils.profiling import ProfilingScope, TimeRecorder
+from docling.utils.utils import chunkify
+
+_log = logging.getLogger(__name__)
+
+
+class BasePipeline(ABC):
+    def __init__(self, pipeline_options: PipelineOptions):
+        self.pipeline_options = pipeline_options
+        self.keep_images = False
+        self.build_pipe: List[Callable] = []
+        self.enrichment_pipe: List[GenericEnrichmentModel[Any]] = []
+
+        self.artifacts_path: Optional[Path] = None
+        if pipeline_options.artifacts_path is not None:
+            self.artifacts_path = Path(pipeline_options.artifacts_path).expanduser()
+        elif settings.artifacts_path is not None:
+            self.artifacts_path = Path(settings.artifacts_path).expanduser()
+
+        if self.artifacts_path is not None and not self.artifacts_path.is_dir():
+            raise RuntimeError(
+                f"The value of {self.artifacts_path=} is not valid. "
+                "When defined, it must point to a folder containing all models required by the pipeline."
+            )
+
+    def execute(self, in_doc: InputDocument, raises_on_error: bool) -> ConversionResult:
+        conv_res = ConversionResult(input=in_doc)
+
+        _log.info(f"Processing document {in_doc.file.name}")
+        try:
+            with TimeRecorder(
+                conv_res, "pipeline_total", scope=ProfilingScope.DOCUMENT
+            ):
+                # These steps are building and assembling the structure of the
+                # output DoclingDocument.
+                conv_res = self._build_document(conv_res)
+                conv_res = self._assemble_document(conv_res)
+                # From this stage, all operations should rely only on conv_res.output
+                conv_res = self._enrich_document(conv_res)
+                conv_res.status = self._determine_status(conv_res)
+        except Exception as e:
+            conv_res.status = ConversionStatus.FAILURE
+            if not raises_on_error:
+                error_item = ErrorItem(
+                    component_type=DoclingComponentType.PIPELINE,
+                    module_name=self.__class__.__name__,
+                    error_message=str(e),
+                )
+                conv_res.errors.append(error_item)
+            else:
+                raise RuntimeError(f"Pipeline {self.__class__.__name__} failed") from e
+        finally:
+            self._unload(conv_res)
+
+        return conv_res
+
+    @abstractmethod
+    def _build_document(self, conv_res: ConversionResult) -> ConversionResult:
+        pass
+
+    def _assemble_document(self, conv_res: ConversionResult) -> ConversionResult:
+        return conv_res
+
+    def _enrich_document(self, conv_res: ConversionResult) -> ConversionResult:
+        def _prepare_elements(
+            conv_res: ConversionResult, model: GenericEnrichmentModel[Any]
+        ) -> Iterable[NodeItem]:
+            for doc_element, _level in conv_res.document.iterate_items():
+                prepared_element = model.prepare_element(
+                    conv_res=conv_res, element=doc_element
+                )
+                if prepared_element is not None:
+                    yield prepared_element
+
+        with TimeRecorder(conv_res, "doc_enrich", scope=ProfilingScope.DOCUMENT):
+            for model in self.enrichment_pipe:
+                for element_batch in chunkify(
+                    _prepare_elements(conv_res, model),
+                    model.elements_batch_size,
+                ):
+                    for element in model(
+                        doc=conv_res.document, element_batch=element_batch
+                    ):  # Must exhaust!
+                        pass
+
+        return conv_res
+
+    @abstractmethod
+    def _determine_status(self, conv_res: ConversionResult) -> ConversionStatus:
+        pass
+
+    def _unload(self, conv_res: ConversionResult):
+        pass
+
+    @classmethod
+    @abstractmethod
+    def get_default_options(cls) -> PipelineOptions:
+        pass
+
+    @classmethod
+    @abstractmethod
+    def is_backend_supported(cls, backend: AbstractDocumentBackend):
+        pass
+
+
+class ConvertPipeline(BasePipeline):
+    def __init__(self, pipeline_options: ConvertPipelineOptions):
+        super().__init__(pipeline_options)
+        self.pipeline_options: ConvertPipelineOptions
+
+        # ------ Common enrichment models working on all backends
+
+        # Picture description model
+        if (
+            picture_description_model := self._get_picture_description_model(
+                artifacts_path=self.artifacts_path
+            )
+        ) is None:
+            raise RuntimeError(
+                f"The specified picture description kind is not supported: {pipeline_options.picture_description_options.kind}."
+            )
+
+        self.enrichment_pipe = [
+            # Document Picture Classifier
+            DocumentPictureClassifier(
+                enabled=pipeline_options.do_picture_classification,
+                artifacts_path=self.artifacts_path,
+                options=DocumentPictureClassifierOptions(),
+                accelerator_options=pipeline_options.accelerator_options,
+            ),
+            # Document Picture description
+            picture_description_model,
+        ]
+
+    def _get_picture_description_model(
+        self, artifacts_path: Optional[Path] = None
+    ) -> Optional[PictureDescriptionBaseModel]:
+        factory = get_picture_description_factory(
+            allow_external_plugins=self.pipeline_options.allow_external_plugins
+        )
+        return factory.create_instance(
+            options=self.pipeline_options.picture_description_options,
+            enabled=self.pipeline_options.do_picture_description,
+            enable_remote_services=self.pipeline_options.enable_remote_services,
+            artifacts_path=artifacts_path,
+            accelerator_options=self.pipeline_options.accelerator_options,
+        )
+
+    @classmethod
+    @abstractmethod
+    def get_default_options(cls) -> ConvertPipelineOptions:
+        pass
+
+
+class PaginatedPipeline(ConvertPipeline):  # TODO this is a bad name.
+    def __init__(self, pipeline_options: ConvertPipelineOptions):
+        super().__init__(pipeline_options)
+        self.keep_backend = False
+
+    def _apply_on_pages(
+        self, conv_res: ConversionResult, page_batch: Iterable[Page]
+    ) -> Iterable[Page]:
+        for model in self.build_pipe:
+            page_batch = model(conv_res, page_batch)
+
+        yield from page_batch
+
+    def _build_document(self, conv_res: ConversionResult) -> ConversionResult:
+        if not isinstance(conv_res.input._backend, PaginatedDocumentBackend):
+            raise RuntimeError(
+                f"The selected backend {type(conv_res.input._backend).__name__} for {conv_res.input.file} is not a paginated backend. "
+                f"Can not convert this with a paginated PDF pipeline. "
+                f"Please check your format configuration on DocumentConverter."
+            )
+            # conv_res.status = ConversionStatus.FAILURE
+            # return conv_res
+
+        total_elapsed_time = 0.0
+        with TimeRecorder(conv_res, "doc_build", scope=ProfilingScope.DOCUMENT):
+            for i in range(conv_res.input.page_count):
+                start_page, end_page = conv_res.input.limits.page_range
+                if (start_page - 1) <= i <= (end_page - 1):
+                    conv_res.pages.append(Page(page_no=i + 1))
+
+            try:
+                total_pages_processed = 0
+                # Iterate batches of pages (page_batch_size) in the doc
+                for page_batch in chunkify(
+                    conv_res.pages, settings.perf.page_batch_size
+                ):
+                    start_batch_time = time.monotonic()
+
+                    # 1. Initialise the page resources
+                    init_pages = map(
+                        functools.partial(self.initialize_page, conv_res), page_batch
+                    )
+
+                    # 2. Run pipeline stages
+                    pipeline_pages = self._apply_on_pages(conv_res, init_pages)
+
+                    for p in pipeline_pages:  # Must exhaust!
+                        # Cleanup cached images
+                        if not self.keep_images:
+                            p._image_cache = {}
+
+                        # Cleanup page backends
+                        if not self.keep_backend and p._backend is not None:
+                            p._backend.unload()
+                        if (
+                            isinstance(self.pipeline_options, PdfPipelineOptions)
+                            and not self.pipeline_options.generate_parsed_pages
+                        ):
+                            del p.parsed_page
+                            p.parsed_page = None
+
+                    end_batch_time = time.monotonic()
+                    total_elapsed_time += end_batch_time - start_batch_time
+                    if (
+                        self.pipeline_options.document_timeout is not None
+                        and total_elapsed_time > self.pipeline_options.document_timeout
+                    ):
+                        _log.warning(
+                            f"Document processing time ({total_elapsed_time:.3f} seconds) exceeded the specified timeout of {self.pipeline_options.document_timeout:.3f} seconds"
+                        )
+                        conv_res.status = ConversionStatus.PARTIAL_SUCCESS
+                        break
+                    total_pages_processed += len(page_batch)
+                    _log.debug(
+                        f"Finished converting pages {total_pages_processed}/{len(conv_res.pages)} time={end_batch_time:.3f}"
+                    )
+
+            except Exception as e:
+                conv_res.status = ConversionStatus.FAILURE
+                trace = "\n".join(
+                    traceback.format_exception(type(e), e, e.__traceback__)
+                )
+                _log.warning(
+                    f"Encountered an error during conversion of document {conv_res.input.document_hash}:\n"
+                    f"{trace}"
+                )
+                raise e
+
+            # Filter out uninitialized pages (those with size=None) that may remain
+            # after timeout or processing failures to prevent assertion errors downstream
+            initial_page_count = len(conv_res.pages)
+            conv_res.pages = [page for page in conv_res.pages if page.size is not None]
+
+            if len(conv_res.pages) < initial_page_count:
+                _log.info(
+                    f"Filtered out {initial_page_count - len(conv_res.pages)} uninitialized pages "
+                    f"due to timeout or processing failures"
+                )
+
+        return conv_res
+
+    def _unload(self, conv_res: ConversionResult) -> ConversionResult:
+        for page in conv_res.pages:
+            if page._backend is not None:
+                page._backend.unload()
+
+        if conv_res.input._backend:
+            conv_res.input._backend.unload()
+
+        return conv_res
+
+    def _determine_status(self, conv_res: ConversionResult) -> ConversionStatus:
+        status = conv_res.status
+        if status in [
+            ConversionStatus.PENDING,
+            ConversionStatus.STARTED,
+        ]:  # preserves ConversionStatus.PARTIAL_SUCCESS
+            status = ConversionStatus.SUCCESS
+
+        for page in conv_res.pages:
+            if page._backend is None or not page._backend.is_valid():
+                conv_res.errors.append(
+                    ErrorItem(
+                        component_type=DoclingComponentType.DOCUMENT_BACKEND,
+                        module_name=type(page._backend).__name__,
+                        error_message=f"Page {page.page_no} failed to parse.",
+                    )
+                )
+                status = ConversionStatus.PARTIAL_SUCCESS
+
+        return status
+
+    # Initialise and load resources for a page
+    @abstractmethod
+    def initialize_page(self, conv_res: ConversionResult, page: Page) -> Page:
+        pass

docling/pipeline/extraction_vlm_pipeline.py

@@ -0,0 +1,207 @@
+import inspect
+import json
+import logging
+from typing import Optional
+
+from PIL.Image import Image
+from pydantic import BaseModel
+
+from docling.backend.abstract_backend import PaginatedDocumentBackend
+from docling.backend.pdf_backend import PdfDocumentBackend
+from docling.datamodel.base_models import ConversionStatus, ErrorItem, VlmStopReason
+from docling.datamodel.document import InputDocument
+from docling.datamodel.extraction import (
+    ExtractedPageData,
+    ExtractionResult,
+    ExtractionTemplateType,
+)
+from docling.datamodel.pipeline_options import (
+    PipelineOptions,
+    VlmExtractionPipelineOptions,
+)
+from docling.datamodel.settings import settings
+from docling.models.extraction.nuextract_transformers_model import (
+    NuExtractTransformersModel,
+)
+from docling.pipeline.base_extraction_pipeline import BaseExtractionPipeline
+from docling.utils.accelerator_utils import decide_device
+
+_log = logging.getLogger(__name__)
+
+
+class ExtractionVlmPipeline(BaseExtractionPipeline):
+    def __init__(self, pipeline_options: VlmExtractionPipelineOptions):
+        super().__init__(pipeline_options)
+
+        # Initialize VLM model with default options
+        self.accelerator_options = pipeline_options.accelerator_options
+        self.pipeline_options: VlmExtractionPipelineOptions
+
+        # Create VLM model instance
+        self.vlm_model = NuExtractTransformersModel(
+            enabled=True,
+            artifacts_path=self.artifacts_path,  # Will download automatically
+            accelerator_options=self.accelerator_options,
+            vlm_options=pipeline_options.vlm_options,
+        )
+
+    def _extract_data(
+        self,
+        ext_res: ExtractionResult,
+        template: Optional[ExtractionTemplateType] = None,
+    ) -> ExtractionResult:
+        """Extract data using the VLM model."""
+        try:
+            # Get images from input document using the backend
+            images = self._get_images_from_input(ext_res.input)
+            if not images:
+                ext_res.status = ConversionStatus.FAILURE
+                ext_res.errors.append(
+                    ErrorItem(
+                        component_type="extraction_pipeline",
+                        module_name=self.__class__.__name__,
+                        error_message="No images found in document",
+                    )
+                )
+                return ext_res
+
+            # Use provided template or default prompt
+            if template is not None:
+                prompt = self._serialize_template(template)
+            else:
+                prompt = "Extract all text and structured information from this document. Return as JSON."
+
+            # Process all images with VLM model
+            start_page, end_page = ext_res.input.limits.page_range
+            for i, image in enumerate(images):
+                # Calculate the actual page number based on the filtered range
+                page_number = start_page + i
+                try:
+                    predictions = list(self.vlm_model.process_images([image], prompt))
+
+                    if predictions:
+                        # Parse the extracted text as JSON if possible, otherwise use as-is
+                        extracted_text = predictions[0].text
+                        extracted_data = None
+                        vlm_stop_reason: VlmStopReason = predictions[0].stop_reason
+                        if (
+                            vlm_stop_reason == VlmStopReason.LENGTH
+                            or vlm_stop_reason == VlmStopReason.STOP_SEQUENCE
+                        ):
+                            ext_res.status = ConversionStatus.PARTIAL_SUCCESS
+
+                        try:
+                            extracted_data = json.loads(extracted_text)
+                        except (json.JSONDecodeError, ValueError):
+                            # If not valid JSON, keep extracted_data as None
+                            pass
+
+                        # Create page data with proper structure
+                        page_data = ExtractedPageData(
+                            page_no=page_number,
+                            extracted_data=extracted_data,
+                            raw_text=extracted_text,  # Always populate raw_text
+                        )
+                        ext_res.pages.append(page_data)
+                    else:
+                        # Add error page data
+                        page_data = ExtractedPageData(
+                            page_no=page_number,
+                            extracted_data=None,
+                            errors=["No extraction result from VLM model"],
+                        )
+                        ext_res.pages.append(page_data)
+
+                except Exception as e:
+                    _log.error(f"Error processing page {page_number}: {e}")
+                    page_data = ExtractedPageData(
+                        page_no=page_number, extracted_data=None, errors=[str(e)]
+                    )
+                    ext_res.pages.append(page_data)
+
+        except Exception as e:
+            _log.error(f"Error during extraction: {e}")
+            ext_res.errors.append(
+                ErrorItem(
+                    component_type="extraction_pipeline",
+                    module_name=self.__class__.__name__,
+                    error_message=str(e),
+                )
+            )
+
+        return ext_res
+
+    def _determine_status(self, ext_res: ExtractionResult) -> ConversionStatus:
+        """Determine the status based on extraction results."""
+        if ext_res.pages and not any(page.errors for page in ext_res.pages):
+            return (
+                ConversionStatus.PARTIAL_SUCCESS
+                if ext_res.status == ConversionStatus.PARTIAL_SUCCESS
+                else ConversionStatus.SUCCESS
+            )
+        else:
+            return ConversionStatus.FAILURE
+
+    def _get_images_from_input(self, input_doc: InputDocument) -> list[Image]:
+        """Extract images from input document using the backend."""
+        images = []
+
+        try:
+            backend = input_doc._backend
+
+            assert isinstance(backend, PdfDocumentBackend)
+            # Use the backend's pagination interface
+            page_count = backend.page_count()
+
+            # Respect page range limits, following the same pattern as PaginatedPipeline
+            start_page, end_page = input_doc.limits.page_range
+            _log.info(
+                f"Processing pages {start_page}-{end_page} of {page_count} total pages for extraction"
+            )
+
+            for page_num in range(page_count):
+                # Only process pages within the specified range (0-based indexing)
+                if start_page - 1 <= page_num <= end_page - 1:
+                    try:
+                        page_backend = backend.load_page(page_num)
+                        if page_backend.is_valid():
+                            # Get page image at a reasonable scale
+                            page_image = page_backend.get_page_image(
+                                scale=self.pipeline_options.vlm_options.scale
+                            )
+                            images.append(page_image)
+                        else:
+                            _log.warning(f"Page {page_num + 1} backend is not valid")
+                    except Exception as e:
+                        _log.error(f"Error loading page {page_num + 1}: {e}")
+
+        except Exception as e:
+            _log.error(f"Error getting images from input document: {e}")
+
+        return images
+
+    def _serialize_template(self, template: ExtractionTemplateType) -> str:
+        """Serialize template to string based on its type."""
+        if isinstance(template, str):
+            return template
+        elif isinstance(template, dict):
+            return json.dumps(template, indent=2)
+        elif isinstance(template, BaseModel):
+            return template.model_dump_json(indent=2)
+        elif inspect.isclass(template) and issubclass(template, BaseModel):
+            from polyfactory.factories.pydantic_factory import ModelFactory
+
+            class ExtractionTemplateFactory(ModelFactory[template]):  # type: ignore
+                __use_examples__ = True  # prefer Field(examples=...) when present
+                __use_defaults__ = True  # use field defaults instead of random values
+                __check_model__ = (
+                    True  # setting the value to avoid deprecation warnings
+                )
+
+            return ExtractionTemplateFactory.build().model_dump_json(indent=2)  # type: ignore
+        else:
+            raise ValueError(f"Unsupported template type: {type(template)}")
+
+    @classmethod
+    def get_default_options(cls) -> PipelineOptions:
+        return VlmExtractionPipelineOptions()