docling 1.19.0__py3-none-any.whl → 2.1.0__py3-none-any.whl

docling/pipeline/base_pipeline.py

@@ -0,0 +1,190 @@
+import functools
+import logging
+import time
+import traceback
+from abc import ABC, abstractmethod
+from typing import Callable, Iterable, List
+
+from docling_core.types.doc import DoclingDocument, NodeItem
+
+from docling.backend.abstract_backend import AbstractDocumentBackend
+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 PipelineOptions
+from docling.datamodel.settings import settings
+from docling.models.base_model import BaseEnrichmentModel
+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.build_pipe: List[Callable] = []
+        self.enrichment_pipe: List[BaseEnrichmentModel] = []
+
+    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:
+            # These steps are building and assembling the structure of the
+            # output DoclingDocument
+            conv_res = self._build_document(in_doc, conv_res)
+            conv_res = self._assemble_document(in_doc, conv_res)
+            # From this stage, all operations should rely only on conv_res.output
+            conv_res = self._enrich_document(in_doc, conv_res)
+            conv_res.status = self._determine_status(in_doc, conv_res)
+        except Exception as e:
+            conv_res.status = ConversionStatus.FAILURE
+            if raises_on_error:
+                raise e
+
+        return conv_res
+
+    @abstractmethod
+    def _build_document(
+        self, in_doc: InputDocument, conv_res: ConversionResult
+    ) -> ConversionResult:
+        pass
+
+    def _assemble_document(
+        self, in_doc: InputDocument, conv_res: ConversionResult
+    ) -> ConversionResult:
+        return conv_res
+
+    def _enrich_document(
+        self, in_doc: InputDocument, conv_res: ConversionResult
+    ) -> ConversionResult:
+
+        def _filter_elements(
+            doc: DoclingDocument, model: BaseEnrichmentModel
+        ) -> Iterable[NodeItem]:
+            for element, _level in doc.iterate_items():
+                if model.is_processable(doc=doc, element=element):
+                    yield element
+
+        for model in self.enrichment_pipe:
+            for element_batch in chunkify(
+                _filter_elements(conv_res.document, model),
+                settings.perf.elements_batch_size,
+            ):
+                # TODO: currently we assume the element itself is modified, because
+                # we don't have an interface to save the element back to the document
+                for element in model(
+                    doc=conv_res.document, element_batch=element_batch
+                ):  # Must exhaust!
+                    pass
+
+        return conv_res
+
+    @abstractmethod
+    def _determine_status(
+        self, in_doc: InputDocument, conv_res: ConversionResult
+    ) -> ConversionStatus:
+        pass
+
+    @classmethod
+    @abstractmethod
+    def get_default_options(cls) -> PipelineOptions:
+        pass
+
+    @classmethod
+    @abstractmethod
+    def is_backend_supported(cls, backend: AbstractDocumentBackend):
+        pass
+
+    # def _apply_on_elements(self, element_batch: Iterable[NodeItem]) -> Iterable[Any]:
+    #    for model in self.build_pipe:
+    #        element_batch = model(element_batch)
+    #
+    #    yield from element_batch
+
+
+class PaginatedPipeline(BasePipeline):  # TODO this is a bad name.
+
+    def _apply_on_pages(self, page_batch: Iterable[Page]) -> Iterable[Page]:
+        for model in self.build_pipe:
+            page_batch = model(page_batch)
+
+        yield from page_batch
+
+    def _build_document(
+        self, in_doc: InputDocument, conv_res: ConversionResult
+    ) -> ConversionResult:
+
+        if not isinstance(in_doc._backend, PdfDocumentBackend):
+            raise RuntimeError(
+                f"The selected backend {type(in_doc._backend).__name__} for {in_doc.file} is not a PDF backend. "
+                f"Can not convert this with a PDF pipeline. "
+                f"Please check your format configuration on DocumentConverter."
+            )
+            # conv_res.status = ConversionStatus.FAILURE
+            # return conv_res
+
+        for i in range(0, in_doc.page_count):
+            conv_res.pages.append(Page(page_no=i))
+
+        try:
+            # Iterate batches of pages (page_batch_size) in the doc
+            for page_batch in chunkify(conv_res.pages, settings.perf.page_batch_size):
+                start_pb_time = time.time()
+
+                # 1. Initialise the page resources
+                init_pages = map(
+                    functools.partial(self.initialize_page, in_doc), page_batch
+                )
+
+                # 2. Run pipeline stages
+                pipeline_pages = self._apply_on_pages(init_pages)
+
+                for p in pipeline_pages:  # Must exhaust!
+                    pass
+
+                end_pb_time = time.time() - start_pb_time
+                _log.debug(f"Finished converting page batch time={end_pb_time:.3f}")
+
+        except Exception as e:
+            conv_res.status = ConversionStatus.FAILURE
+            trace = "\n".join(traceback.format_exception(e))
+            _log.warning(
+                f"Encountered an error during conversion of document {in_doc.document_hash}:\n"
+                f"{trace}"
+            )
+            raise e
+
+        finally:
+            # Always unload the PDF backend, even in case of failure
+            if in_doc._backend:
+                in_doc._backend.unload()
+
+        return conv_res
+
+    def _determine_status(
+        self, in_doc: InputDocument, conv_res: ConversionResult
+    ) -> ConversionStatus:
+        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, doc: InputDocument, page: Page) -> Page:
+        pass

docling/pipeline/simple_pipeline.py

@@ -0,0 +1,59 @@
+import logging
+
+from docling.backend.abstract_backend import (
+    AbstractDocumentBackend,
+    DeclarativeDocumentBackend,
+)
+from docling.datamodel.base_models import ConversionStatus
+from docling.datamodel.document import ConversionResult, InputDocument
+from docling.datamodel.pipeline_options import PipelineOptions
+from docling.pipeline.base_pipeline import BasePipeline
+
+_log = logging.getLogger(__name__)
+
+
+class SimplePipeline(BasePipeline):
+    """SimpleModelPipeline.
+
+    This class is used at the moment for formats / backends
+    which produce straight DoclingDocument output.
+    """
+
+    def __init__(self, pipeline_options: PipelineOptions):
+        super().__init__(pipeline_options)
+
+    def _build_document(
+        self, in_doc: InputDocument, conv_res: ConversionResult
+    ) -> ConversionResult:
+
+        if not isinstance(in_doc._backend, DeclarativeDocumentBackend):
+            raise RuntimeError(
+                f"The selected backend {type(in_doc._backend).__name__} for {in_doc.file} is not a declarative backend. "
+                f"Can not convert this with simple pipeline. "
+                f"Please check your format configuration on DocumentConverter."
+            )
+            # conv_res.status = ConversionStatus.FAILURE
+            # return conv_res
+
+        # Instead of running a page-level pipeline to build up the document structure,
+        # the backend is expected to be of type DeclarativeDocumentBackend, which can output
+        # a DoclingDocument straight.
+
+        conv_res.document = in_doc._backend.convert()
+        return conv_res
+
+    def _determine_status(
+        self, in_doc: InputDocument, conv_res: ConversionResult
+    ) -> ConversionStatus:
+        # This is called only if the previous steps didn't raise.
+        # Since we don't have anything else to evaluate, we can
+        # safely return SUCCESS.
+        return ConversionStatus.SUCCESS
+
+    @classmethod
+    def get_default_options(cls) -> PipelineOptions:
+        return PipelineOptions()
+
+    @classmethod
+    def is_backend_supported(cls, backend: AbstractDocumentBackend):
+        return isinstance(backend, DeclarativeDocumentBackend)

docling/pipeline/standard_pdf_pipeline.py

@@ -0,0 +1,198 @@
+import logging
+from pathlib import Path
+from typing import Optional
+
+from docling_core.types.doc import DocItem, ImageRef, PictureItem, TableItem
+
+from docling.backend.abstract_backend import AbstractDocumentBackend
+from docling.backend.pdf_backend import PdfDocumentBackend
+from docling.datamodel.base_models import AssembledUnit, Page
+from docling.datamodel.document import ConversionResult, InputDocument
+from docling.datamodel.pipeline_options import (
+    EasyOcrOptions,
+    PdfPipelineOptions,
+    TesseractCliOcrOptions,
+    TesseractOcrOptions,
+)
+from docling.models.base_ocr_model import BaseOcrModel
+from docling.models.ds_glm_model import GlmModel, GlmOptions
+from docling.models.easyocr_model import EasyOcrModel
+from docling.models.layout_model import LayoutModel
+from docling.models.page_assemble_model import PageAssembleModel, PageAssembleOptions
+from docling.models.page_preprocessing_model import (
+    PagePreprocessingModel,
+    PagePreprocessingOptions,
+)
+from docling.models.table_structure_model import TableStructureModel
+from docling.models.tesseract_ocr_cli_model import TesseractOcrCliModel
+from docling.models.tesseract_ocr_model import TesseractOcrModel
+from docling.pipeline.base_pipeline import PaginatedPipeline
+
+_log = logging.getLogger(__name__)
+
+
+class StandardPdfPipeline(PaginatedPipeline):
+    _layout_model_path = "model_artifacts/layout/beehive_v0.0.5_pt"
+    _table_model_path = "model_artifacts/tableformer"
+
+    def __init__(self, pipeline_options: PdfPipelineOptions):
+        super().__init__(pipeline_options)
+        self.pipeline_options: PdfPipelineOptions
+
+        if pipeline_options.artifacts_path is None:
+            self.artifacts_path = self.download_models_hf()
+        else:
+            self.artifacts_path = Path(pipeline_options.artifacts_path)
+
+        keep_images = (
+            self.pipeline_options.generate_page_images
+            or self.pipeline_options.generate_picture_images
+            or self.pipeline_options.generate_table_images
+        )
+
+        self.glm_model = GlmModel(options=GlmOptions())
+
+        if (ocr_model := self.get_ocr_model()) is None:
+            raise RuntimeError(
+                f"The specified OCR kind is not supported: {pipeline_options.ocr_options.kind}."
+            )
+
+        self.build_pipe = [
+            # Pre-processing
+            PagePreprocessingModel(
+                options=PagePreprocessingOptions(
+                    images_scale=pipeline_options.images_scale
+                )
+            ),
+            # OCR
+            ocr_model,
+            # Layout model
+            LayoutModel(
+                artifacts_path=self.artifacts_path
+                / StandardPdfPipeline._layout_model_path
+            ),
+            # Table structure model
+            TableStructureModel(
+                enabled=pipeline_options.do_table_structure,
+                artifacts_path=self.artifacts_path
+                / StandardPdfPipeline._table_model_path,
+                options=pipeline_options.table_structure_options,
+            ),
+            # Page assemble
+            PageAssembleModel(options=PageAssembleOptions(keep_images=keep_images)),
+        ]
+
+        self.enrichment_pipe = [
+            # Other models working on `NodeItem` elements in the DoclingDocument
+        ]
+
+    @staticmethod
+    def download_models_hf(
+        local_dir: Optional[Path] = None, force: bool = False
+    ) -> Path:
+        from huggingface_hub import snapshot_download
+
+        download_path = snapshot_download(
+            repo_id="ds4sd/docling-models",
+            force_download=force,
+            local_dir=local_dir,
+            revision="v2.0.1",
+        )
+
+        return Path(download_path)
+
+    def get_ocr_model(self) -> Optional[BaseOcrModel]:
+        if isinstance(self.pipeline_options.ocr_options, EasyOcrOptions):
+            return EasyOcrModel(
+                enabled=self.pipeline_options.do_ocr,
+                options=self.pipeline_options.ocr_options,
+            )
+        elif isinstance(self.pipeline_options.ocr_options, TesseractCliOcrOptions):
+            return TesseractOcrCliModel(
+                enabled=self.pipeline_options.do_ocr,
+                options=self.pipeline_options.ocr_options,
+            )
+        elif isinstance(self.pipeline_options.ocr_options, TesseractOcrOptions):
+            return TesseractOcrModel(
+                enabled=self.pipeline_options.do_ocr,
+                options=self.pipeline_options.ocr_options,
+            )
+        return None
+
+    def initialize_page(self, doc: InputDocument, page: Page) -> Page:
+        page._backend = doc._backend.load_page(page.page_no)  # type: ignore
+        if page._backend is not None and page._backend.is_valid():
+            page.size = page._backend.get_size()
+
+        return page
+
+    def _assemble_document(
+        self, in_doc: InputDocument, conv_res: ConversionResult
+    ) -> ConversionResult:
+        all_elements = []
+        all_headers = []
+        all_body = []
+
+        for p in conv_res.pages:
+            if p.assembled is not None:
+                for el in p.assembled.body:
+                    all_body.append(el)
+                for el in p.assembled.headers:
+                    all_headers.append(el)
+                for el in p.assembled.elements:
+                    all_elements.append(el)
+
+        conv_res.assembled = AssembledUnit(
+            elements=all_elements, headers=all_headers, body=all_body
+        )
+
+        conv_res.document = self.glm_model(conv_res)
+
+        # Generate page images in the output
+        if self.pipeline_options.generate_page_images:
+            for page in conv_res.pages:
+                assert page.image is not None
+                page_no = page.page_no + 1
+                conv_res.document.pages[page_no].image = ImageRef.from_pil(
+                    page.image, dpi=int(72 * self.pipeline_options.images_scale)
+                )
+
+        # Generate images of the requested element types
+        if (
+            self.pipeline_options.generate_picture_images
+            or self.pipeline_options.generate_table_images
+        ):
+            scale = self.pipeline_options.images_scale
+            for element, _level in conv_res.document.iterate_items():
+                if not isinstance(element, DocItem) or len(element.prov) == 0:
+                    continue
+                if (
+                    isinstance(element, PictureItem)
+                    and self.pipeline_options.generate_picture_images
+                ) or (
+                    isinstance(element, TableItem)
+                    and self.pipeline_options.generate_table_images
+                ):
+                    page_ix = element.prov[0].page_no - 1
+                    page = conv_res.pages[page_ix]
+                    assert page.size is not None
+                    assert page.image is not None
+
+                    crop_bbox = (
+                        element.prov[0]
+                        .bbox.scaled(scale=scale)
+                        .to_top_left_origin(page_height=page.size.height * scale)
+                    )
+
+                    cropped_im = page.image.crop(crop_bbox.as_tuple())
+                    element.image = ImageRef.from_pil(cropped_im, dpi=int(72 * scale))
+
+        return conv_res
+
+    @classmethod
+    def get_default_options(cls) -> PdfPipelineOptions:
+        return PdfPipelineOptions()
+
+    @classmethod
+    def is_backend_supported(cls, backend: AbstractDocumentBackend):
+        return isinstance(backend, PdfDocumentBackend)

docling/utils/export.py

@@ -1,9 +1,10 @@
 import logging
 from typing import Any, Dict, Iterable, List, Tuple, Union
 
-from docling_core.types.doc.base import BaseCell, BaseText, Ref, Table, TableCell
+from docling_core.types.doc import BoundingBox, CoordOrigin
+from docling_core.types.legacy_doc.base import BaseCell, BaseText, Ref, Table
 
-from docling.datamodel.base_models import BoundingBox, CoordOrigin, OcrCell
+from docling.datamodel.base_models import OcrCell
 from docling.datamodel.document import ConversionResult, Page
 
 _log = logging.getLogger(__name__)
@@ -40,7 +41,7 @@ def generate_multimodal_pages(
     end_ix = 0
     doc_items: List[Tuple[int, Union[BaseCell, BaseText]]] = []
 
-    doc = doc_result.output
+    doc = doc_result.legacy_document
 
     def _process_page_segments(doc_items: list[Tuple[int, BaseCell]], page: Page):
         segments = []

docling/utils/layout_utils.py

@@ -2,6 +2,7 @@ import copy
 import logging
 
 import networkx as nx
+from docling_core.types.doc import DocItemLabel
 
 logger = logging.getLogger("layout_utils")
 
@@ -370,7 +371,7 @@ def adapt_bboxes(raw_cells, clusters, orphan_cell_indices):
             "Treating cluster " + str(ix) + ", type " + str(new_cluster["type"])
         )
         logger.debug("  with cells: " + str(new_cluster["cell_ids"]))
-        if len(cluster["cell_ids"]) == 0 and cluster["type"] != "Picture":
+        if len(cluster["cell_ids"]) == 0 and cluster["type"] != DocItemLabel.PICTURE:
             logger.debug("  Empty non-picture, removed")
             continue  ## Skip this former cluster, now without cells.
         new_bbox = adapt_bbox(raw_cells, new_cluster, orphan_cell_indices)
@@ -380,14 +381,14 @@ def adapt_bboxes(raw_cells, clusters, orphan_cell_indices):
 
 
 def adapt_bbox(raw_cells, cluster, orphan_cell_indices):
-    if not (cluster["type"] in ["Table", "Picture"]):
+    if not (cluster["type"] in [DocItemLabel.TABLE, DocItemLabel.PICTURE]):
         ## A text-like cluster. The bbox only needs to be around the text cells:
         logger.debug("    Initial bbox: " + str(cluster["bbox"]))
         new_bbox = surrounding_list(
             [raw_cells[cid]["bbox"] for cid in cluster["cell_ids"]]
         )
         logger.debug("  New bounding box:" + str(new_bbox))
-    if cluster["type"] == "Picture":
+    if cluster["type"] == DocItemLabel.PICTURE:
         ## We only make the bbox completely comprise included text cells:
         logger.debug("  Picture")
         if len(cluster["cell_ids"]) != 0:
@@ -587,7 +588,7 @@ def set_orphan_as_text(
     max_id = -1
     figures = []
     for cluster in cluster_predictions:
-        if cluster["type"] == "Picture":
+        if cluster["type"] == DocItemLabel.PICTURE:
             figures.append(cluster)
 
         if cluster["id"] > max_id:
@@ -638,13 +639,13 @@ def set_orphan_as_text(
             # if fig_flag == False and raw_cells[orph_id]["text"] not in line_orphans:
             if fig_flag == False and lines_detector == False:
                 # get class from low confidence detections if not set as text:
-                class_type = "Text"
+                class_type = DocItemLabel.TEXT
 
                 for cluster in cluster_predictions_low:
                     intersection = compute_intersection(
                         orph_cell["bbox"], cluster["bbox"]
                     )
-                    class_type = "Text"
+                    class_type = DocItemLabel.TEXT
                     if (
                         cluster["confidence"] > 0.1
                         and bb_iou(cluster["bbox"], orph_cell["bbox"]) > 0.4
@@ -718,7 +719,9 @@ def merge_cells(cluster_predictions):
                     if cluster["id"] == node:
                         lines.append(cluster)
                         cluster_predictions.remove(cluster)
-            new_merged_cluster = build_cluster_from_lines(lines, "Text", max_id)
+            new_merged_cluster = build_cluster_from_lines(
+                lines, DocItemLabel.TEXT, max_id
+            )
             cluster_predictions.append(new_merged_cluster)
     return cluster_predictions
 
@@ -753,9 +756,9 @@ def clean_up_clusters(
                 # remove clusters that might appear inside tables, or images (such as pdf cells in graphs)
                 elif img_table == True:
                     if (
-                        cluster_1["type"] == "Text"
-                        and cluster_2["type"] == "Picture"
-                        or cluster_2["type"] == "Table"
+                        cluster_1["type"] == DocItemLabel.TEXT
+                        and cluster_2["type"] == DocItemLabel.PICTURE
+                        or cluster_2["type"] == DocItemLabel.TABLE
                     ):
                         if bb_iou(cluster_1["bbox"], cluster_2["bbox"]) > 0.5:
                             DuplicateDeletedClusterIDs.append(cluster_1["id"])
@@ -771,7 +774,10 @@ def clean_up_clusters(
                             DuplicateDeletedClusterIDs.append(cluster_1["id"])
             # remove tables that have one pdf cell
             if one_cell_table == True:
-                if cluster_1["type"] == "Table" and len(cluster_1["cell_ids"]) < 2:
+                if (
+                    cluster_1["type"] == DocItemLabel.TABLE
+                    and len(cluster_1["cell_ids"]) < 2
+                ):
                     DuplicateDeletedClusterIDs.append(cluster_1["id"])
 
     DuplicateDeletedClusterIDs = list(set(DuplicateDeletedClusterIDs))