docling 2.69.0__py3-none-any.whl

docling/models/plugins/defaults.py

@@ -0,0 +1,60 @@
+def ocr_engines():
+    from docling.models.stages.ocr.auto_ocr_model import OcrAutoModel
+    from docling.models.stages.ocr.easyocr_model import EasyOcrModel
+    from docling.models.stages.ocr.ocr_mac_model import OcrMacModel
+    from docling.models.stages.ocr.rapid_ocr_model import RapidOcrModel
+    from docling.models.stages.ocr.tesseract_ocr_cli_model import TesseractOcrCliModel
+    from docling.models.stages.ocr.tesseract_ocr_model import TesseractOcrModel
+
+    return {
+        "ocr_engines": [
+            OcrAutoModel,
+            EasyOcrModel,
+            OcrMacModel,
+            RapidOcrModel,
+            TesseractOcrModel,
+            TesseractOcrCliModel,
+        ]
+    }
+
+
+def picture_description():
+    from docling.models.stages.picture_description.picture_description_api_model import (
+        PictureDescriptionApiModel,
+    )
+    from docling.models.stages.picture_description.picture_description_vlm_model import (
+        PictureDescriptionVlmModel,
+    )
+
+    return {
+        "picture_description": [
+            PictureDescriptionVlmModel,
+            PictureDescriptionApiModel,
+        ]
+    }
+
+
+def layout_engines():
+    from docling.experimental.models.table_crops_layout_model import (
+        TableCropsLayoutModel,
+    )
+    from docling.models.stages.layout.layout_model import LayoutModel
+
+    return {
+        "layout_engines": [
+            LayoutModel,
+            TableCropsLayoutModel,
+        ]
+    }
+
+
+def table_structure_engines():
+    from docling.models.stages.table_structure.table_structure_model import (
+        TableStructureModel,
+    )
+
+    return {
+        "table_structure_engines": [
+            TableStructureModel,
+        ]
+    }

docling/models/stages/code_formula/code_formula_model.py

@@ -0,0 +1,342 @@
+import re
+from collections.abc import Iterable
+from pathlib import Path
+from typing import List, Literal, Optional, Tuple, Union
+
+import numpy as np
+from docling_core.types.doc import (
+    CodeItem,
+    DocItemLabel,
+    DoclingDocument,
+    NodeItem,
+    TextItem,
+)
+from docling_core.types.doc.labels import CodeLanguageLabel
+from PIL import Image
+from pydantic import BaseModel
+
+from docling.datamodel.accelerator_options import AcceleratorDevice, AcceleratorOptions
+from docling.datamodel.base_models import ItemAndImageEnrichmentElement
+from docling.models.base_model import BaseItemAndImageEnrichmentModel
+from docling.models.utils.hf_model_download import download_hf_model
+from docling.utils.accelerator_utils import decide_device
+
+
+class CodeFormulaModelOptions(BaseModel):
+    """
+    Configuration options for the CodeFormulaModel.
+
+    Attributes
+    ----------
+    kind : str
+        Type of the model. Fixed value "code_formula".
+    do_code_enrichment : bool
+        True if code enrichment is enabled, False otherwise.
+    do_formula_enrichment : bool
+        True if formula enrichment is enabled, False otherwise.
+    """
+
+    kind: Literal["code_formula"] = "code_formula"
+    do_code_enrichment: bool = True
+    do_formula_enrichment: bool = True
+
+
+class CodeFormulaModel(BaseItemAndImageEnrichmentModel):
+    """
+    Model for processing and enriching documents with code and formula predictions.
+
+    Attributes
+    ----------
+    enabled : bool
+        True if the model is enabled, False otherwise.
+    options : CodeFormulaModelOptions
+        Configuration options for the CodeFormulaModel.
+    code_formula_model : CodeFormulaPredictor
+        The predictor model for code and formula processing.
+
+    Methods
+    -------
+    __init__(self, enabled, artifacts_path, accelerator_options, code_formula_options)
+        Initializes the CodeFormulaModel with the given configuration options.
+    is_processable(self, doc, element)
+        Determines if a given element in a document can be processed by the model.
+    __call__(self, doc, element_batch)
+        Processes the given batch of elements and enriches them with predictions.
+    """
+
+    _model_repo_folder = "docling-project--CodeFormulaV2"
+    elements_batch_size = 5
+    images_scale = 1.67  # = 120 dpi, aligned with training data resolution
+    expansion_factor = 0.18
+
+    def __init__(
+        self,
+        enabled: bool,
+        artifacts_path: Optional[Path],
+        options: CodeFormulaModelOptions,
+        accelerator_options: AcceleratorOptions,
+    ):
+        """
+        Initializes the CodeFormulaModel with the given configuration.
+
+        Parameters
+        ----------
+        enabled : bool
+            True if the model is enabled, False otherwise.
+        artifacts_path : Path
+            Path to the directory containing the model artifacts.
+        options : CodeFormulaModelOptions
+            Configuration options for the model.
+        accelerator_options : AcceleratorOptions
+            Options specifying the device and number of threads for acceleration.
+        """
+        self.enabled = enabled
+        self.options = options
+
+        if self.enabled:
+            self.device = decide_device(
+                accelerator_options.device,
+                supported_devices=[
+                    AcceleratorDevice.CPU,
+                    AcceleratorDevice.CUDA,
+                    AcceleratorDevice.XPU,
+                ],
+            )
+
+            if artifacts_path is None:
+                artifacts_path = self.download_models()
+            else:
+                artifacts_path = artifacts_path / self._model_repo_folder
+
+            from transformers import AutoModelForImageTextToText, AutoProcessor
+
+            self._processor = AutoProcessor.from_pretrained(
+                artifacts_path,
+            )
+            self._model_max_length = self._processor.tokenizer.model_max_length
+            self._model = AutoModelForImageTextToText.from_pretrained(
+                artifacts_path, device_map=self.device
+            )
+            self._model.eval()
+
+    @staticmethod
+    def download_models(
+        local_dir: Optional[Path] = None,
+        force: bool = False,
+        progress: bool = False,
+    ) -> Path:
+        return download_hf_model(
+            repo_id="docling-project/CodeFormulaV2",
+            revision="main",
+            local_dir=local_dir,
+            force=force,
+            progress=progress,
+        )
+
+    def is_processable(self, doc: DoclingDocument, element: NodeItem) -> bool:
+        """
+        Determines if a given element in a document can be processed by the model.
+
+        Parameters
+        ----------
+        doc : DoclingDocument
+            The document being processed.
+        element : NodeItem
+            The element within the document to check.
+
+        Returns
+        -------
+        bool
+            True if the element can be processed, False otherwise.
+        """
+        return self.enabled and (
+            (isinstance(element, CodeItem) and self.options.do_code_enrichment)
+            or (
+                isinstance(element, TextItem)
+                and element.label == DocItemLabel.FORMULA
+                and self.options.do_formula_enrichment
+            )
+        )
+
+    def _extract_code_language(self, input_string: str) -> Tuple[str, Optional[str]]:
+        """Extracts a programming language from the beginning of a string.
+
+        This function checks if the input string starts with a pattern of the form
+        ``<_some_language_>``. If it does, it extracts the language string and returns
+        a tuple of (remainder, language). Otherwise, it returns the original string
+        and `None`.
+
+        Args:
+            input_string (str): The input string, which may start with ``<_language_>``.
+
+        Returns:
+            Tuple[str, Optional[str]]:
+                A tuple where:
+                - The first element is either:
+                    - The remainder of the string (everything after ``<_language_>``),
+                    if a match is found; or
+                    - The original string, if no match is found.
+                - The second element is the extracted language if a match is found;
+                otherwise, `None`.
+        """
+        pattern = r"^<_([^_>]+)_>\s*(.*)"
+        match = re.match(pattern, input_string, flags=re.DOTALL)
+        if match:
+            language = str(match.group(1))  # the captured programming language
+            remainder = str(match.group(2))  # everything after the <_language_>
+            return remainder, language
+        else:
+            return input_string, None
+
+    def _get_code_language_enum(self, value: Optional[str]) -> CodeLanguageLabel:
+        """
+        Converts a string to a corresponding `CodeLanguageLabel` enum member.
+
+        If the provided string does not match any value in `CodeLanguageLabel`,
+        it defaults to `CodeLanguageLabel.UNKNOWN`.
+
+        Args:
+            value (Optional[str]): The string representation of the code language or None.
+
+        Returns:
+            CodeLanguageLabel: The corresponding enum member if the value is valid,
+            otherwise `CodeLanguageLabel.UNKNOWN`.
+        """
+        if not isinstance(value, str):
+            return CodeLanguageLabel.UNKNOWN
+
+        try:
+            return CodeLanguageLabel(value)
+        except ValueError:
+            return CodeLanguageLabel.UNKNOWN
+
+    def _get_prompt(self, label: str) -> str:
+        """
+        Constructs the prompt for the model based on the input label.
+
+        Parameters
+        ----------
+        label : str
+            The type of input, either 'code' or 'formula'.
+
+        Returns
+        -------
+        str
+            The constructed prompt including necessary tokens and query.
+
+        Raises
+        ------
+        NotImplementedError
+            If the label is not 'code' or 'formula'.
+        """
+        if label == "code":
+            query = "<code>"
+        elif label == "formula":
+            query = "<formula>"
+        else:
+            raise NotImplementedError("Label must be either code or formula")
+
+        messages = [
+            {
+                "role": "user",
+                "content": [{"type": "image"}, {"type": "text", "text": query}],
+            },
+        ]
+
+        prompt = self._processor.apply_chat_template(
+            messages, add_generation_prompt=True
+        )
+
+        return prompt
+
+    def _post_process(self, texts: list[str]) -> list[str]:
+        """
+        Processes a list of text strings by truncating at '<end_of_utterance>' and
+        removing a predefined set of unwanted substrings.
+
+        Parameters
+        ----------
+        texts : list[str]
+            A list of strings to be post-processed.
+
+        Returns
+        -------
+        list[str]
+            A list of cleaned strings with specified substrings removed and truncated at
+                '<end_of_utterance>' if present.
+        """
+        to_remove = ["</code>", "</formula>", "<loc_0><loc_0><loc_500><loc_500>"]
+
+        def clean_text(text: str) -> str:
+            idx = text.find("<end_of_utterance>")
+            if idx != -1:
+                text = text[:idx]
+
+            for token in to_remove:
+                if token in text:
+                    text = text.replace(token, "")
+            return text.lstrip()
+
+        return [clean_text(t) for t in texts]
+
+    def __call__(
+        self,
+        doc: DoclingDocument,
+        element_batch: Iterable[ItemAndImageEnrichmentElement],
+    ) -> Iterable[NodeItem]:
+        """
+        Processes the given batch of elements and enriches them with predictions.
+
+        Parameters
+        ----------
+        doc : DoclingDocument
+            The document being processed.
+        element_batch : Iterable[ItemAndImageEnrichmentElement]
+            A batch of elements to be processed.
+
+        Returns
+        -------
+        Iterable[Any]
+            An iterable of enriched elements.
+        """
+        if not self.enabled:
+            for element in element_batch:
+                yield element.item
+            return
+
+        labels: List[str] = []
+        images: List[Union[Image.Image, np.ndarray]] = []
+        elements: List[TextItem] = []
+        for el in element_batch:
+            elements.append(el.item)  # type: ignore[arg-type]
+            labels.append(el.item.label)  # type: ignore[attr-defined]
+            images.append(el.image)
+
+        prompts = [self._get_prompt(label) for label in labels]
+        inputs = self._processor(
+            text=prompts,
+            images=images,
+            return_tensors="pt",
+        )
+        inputs = inputs.to(self.device)
+
+        gen_kwargs = dict(
+            max_new_tokens=self._model_max_length - inputs.input_ids.shape[1],
+            use_cache=True,
+            do_sample=False,
+        )
+
+        generated_ids = self._model.generate(**inputs, **gen_kwargs)
+
+        outputs = self._processor.batch_decode(
+            generated_ids[:, inputs.input_ids.shape[1] :], skip_special_tokens=False
+        )
+        outputs = self._post_process(outputs)
+
+        for item, output in zip(elements, outputs):
+            if isinstance(item, CodeItem):
+                output, code_language = self._extract_code_language(output)
+                item.code_language = self._get_code_language_enum(code_language)
+            item.text = output
+
+            yield item

docling/models/stages/layout/layout_model.py

@@ -0,0 +1,249 @@
+import copy
+import logging
+import warnings
+from collections.abc import Sequence
+from pathlib import Path
+from typing import List, Optional, Union
+
+import numpy as np
+from docling_core.types.doc import DocItemLabel
+from PIL import Image
+
+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.datamodel.layout_model_specs import DOCLING_LAYOUT_V2, LayoutModelConfig
+from docling.datamodel.pipeline_options import LayoutOptions
+from docling.datamodel.settings import settings
+from docling.models.base_layout_model import BaseLayoutModel
+from docling.models.utils.hf_model_download import download_hf_model
+from docling.utils.accelerator_utils import decide_device
+from docling.utils.layout_postprocessor import LayoutPostprocessor
+from docling.utils.profiling import TimeRecorder
+from docling.utils.visualization import draw_clusters
+
+_log = logging.getLogger(__name__)
+
+
+class LayoutModel(BaseLayoutModel):
+    TEXT_ELEM_LABELS = [
+        DocItemLabel.TEXT,
+        DocItemLabel.FOOTNOTE,
+        DocItemLabel.CAPTION,
+        DocItemLabel.CHECKBOX_UNSELECTED,
+        DocItemLabel.CHECKBOX_SELECTED,
+        DocItemLabel.SECTION_HEADER,
+        DocItemLabel.PAGE_HEADER,
+        DocItemLabel.PAGE_FOOTER,
+        DocItemLabel.CODE,
+        DocItemLabel.LIST_ITEM,
+        DocItemLabel.FORMULA,
+    ]
+    PAGE_HEADER_LABELS = [DocItemLabel.PAGE_HEADER, DocItemLabel.PAGE_FOOTER]
+
+    TABLE_LABELS = [DocItemLabel.TABLE, DocItemLabel.DOCUMENT_INDEX]
+    FIGURE_LABEL = DocItemLabel.PICTURE
+    FORMULA_LABEL = DocItemLabel.FORMULA
+    CONTAINER_LABELS = [DocItemLabel.FORM, DocItemLabel.KEY_VALUE_REGION]
+
+    def __init__(
+        self,
+        artifacts_path: Optional[Path],
+        accelerator_options: AcceleratorOptions,
+        options: LayoutOptions,
+    ):
+        from docling_ibm_models.layoutmodel.layout_predictor import LayoutPredictor
+
+        self.options = options
+
+        device = decide_device(accelerator_options.device)
+        layout_model_config = options.model_spec
+        model_repo_folder = layout_model_config.model_repo_folder
+        model_path = layout_model_config.model_path
+
+        if artifacts_path is None:
+            artifacts_path = (
+                self.download_models(layout_model_config=layout_model_config)
+                / model_path
+            )
+        else:
+            if (artifacts_path / model_repo_folder).exists():
+                artifacts_path = artifacts_path / model_repo_folder / model_path
+            elif (artifacts_path / model_path).exists():
+                warnings.warn(
+                    "The usage of artifacts_path containing directly "
+                    f"{model_path} is deprecated. Please point "
+                    "the artifacts_path to the parent containing "
+                    f"the {model_repo_folder} folder.",
+                    DeprecationWarning,
+                    stacklevel=3,
+                )
+                artifacts_path = artifacts_path / model_path
+
+        self.layout_predictor = LayoutPredictor(
+            artifact_path=str(artifacts_path),
+            device=device,
+            num_threads=accelerator_options.num_threads,
+        )
+
+    @classmethod
+    def get_options_type(cls) -> type[LayoutOptions]:
+        return LayoutOptions
+
+    @staticmethod
+    def download_models(
+        local_dir: Optional[Path] = None,
+        force: bool = False,
+        progress: bool = False,
+        layout_model_config: LayoutModelConfig = LayoutOptions().model_spec,  # use default
+    ) -> Path:
+        return download_hf_model(
+            repo_id=layout_model_config.repo_id,
+            revision=layout_model_config.revision,
+            local_dir=local_dir,
+            force=force,
+            progress=progress,
+        )
+
+    def draw_clusters_and_cells_side_by_side(
+        self, conv_res, page, clusters, mode_prefix: str, show: bool = False
+    ):
+        """
+        Draws a page image side by side with clusters filtered into two categories:
+        - Left: Clusters excluding FORM, KEY_VALUE_REGION, and PICTURE.
+        - Right: Clusters including FORM, KEY_VALUE_REGION, and PICTURE.
+        Includes label names and confidence scores for each cluster.
+        """
+        scale_x = page.image.width / page.size.width
+        scale_y = page.image.height / page.size.height
+
+        # Filter clusters for left and right images
+        exclude_labels = {
+            DocItemLabel.FORM,
+            DocItemLabel.KEY_VALUE_REGION,
+            DocItemLabel.PICTURE,
+        }
+        left_clusters = [c for c in clusters if c.label not in exclude_labels]
+        right_clusters = [c for c in clusters if c.label in exclude_labels]
+        # Create a deep copy of the original image for both sides
+        left_image = page.image.copy()
+        right_image = page.image.copy()
+
+        # Draw clusters on both images
+        draw_clusters(left_image, left_clusters, scale_x, scale_y)
+        draw_clusters(right_image, right_clusters, scale_x, scale_y)
+        # Combine the images side by side
+        combined_width = left_image.width * 2
+        combined_height = left_image.height
+        combined_image = Image.new("RGB", (combined_width, combined_height))
+        combined_image.paste(left_image, (0, 0))
+        combined_image.paste(right_image, (left_image.width, 0))
+        if show:
+            combined_image.show()
+        else:
+            out_path: Path = (
+                Path(settings.debug.debug_output_path)
+                / f"debug_{conv_res.input.file.stem}"
+            )
+            out_path.mkdir(parents=True, exist_ok=True)
+            out_file = out_path / f"{mode_prefix}_layout_page_{page.page_no:05}.png"
+            combined_image.save(str(out_file), format="png")
+
+    def predict_layout(
+        self,
+        conv_res: ConversionResult,
+        pages: Sequence[Page],
+    ) -> Sequence[LayoutPrediction]:
+        # Convert to list to ensure predictable iteration
+        pages = list(pages)
+
+        # Separate valid and invalid pages
+        valid_pages = []
+        valid_page_images: List[Union[Image.Image, np.ndarray]] = []
+
+        for page in pages:
+            assert page._backend is not None
+            if not page._backend.is_valid():
+                continue
+
+            assert page.size is not None
+            page_image = page.get_image(scale=1.0)
+            assert page_image is not None
+
+            valid_pages.append(page)
+            valid_page_images.append(page_image)
+
+        # Process all valid pages with batch prediction
+        batch_predictions = []
+        if valid_page_images:
+            with TimeRecorder(conv_res, "layout"):
+                batch_predictions = self.layout_predictor.predict_batch(  # type: ignore[attr-defined]
+                    valid_page_images
+                )
+
+        # Process each page with its predictions
+        layout_predictions: list[LayoutPrediction] = []
+        valid_page_idx = 0
+        for page in pages:
+            assert page._backend is not None
+            if not page._backend.is_valid():
+                existing_prediction = page.predictions.layout or LayoutPrediction()
+                page.predictions.layout = existing_prediction
+                layout_predictions.append(existing_prediction)
+                continue
+
+            page_predictions = batch_predictions[valid_page_idx]
+            valid_page_idx += 1
+
+            clusters = []
+            for ix, pred_item in enumerate(page_predictions):
+                label = DocItemLabel(
+                    pred_item["label"].lower().replace(" ", "_").replace("-", "_")
+                )  # Temporary, until docling-ibm-model uses docling-core types
+                cluster = Cluster(
+                    id=ix,
+                    label=label,
+                    confidence=pred_item["confidence"],
+                    bbox=BoundingBox.model_validate(pred_item),
+                    cells=[],
+                )
+                clusters.append(cluster)
+
+            if settings.debug.visualize_raw_layout:
+                self.draw_clusters_and_cells_side_by_side(
+                    conv_res, page, clusters, mode_prefix="raw"
+                )
+
+            # Apply postprocessing
+            processed_clusters, processed_cells = LayoutPostprocessor(
+                page, clusters, self.options
+            ).postprocess()
+            # Note: LayoutPostprocessor updates page.cells and page.parsed_page internally
+
+            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 = float(
+                    np.mean([c.confidence for c in processed_clusters])
+                )
+
+                conv_res.confidence.pages[page.page_no].ocr_score = float(
+                    np.mean([c.confidence for c in processed_cells if c.from_ocr])
+                )
+
+            prediction = LayoutPrediction(clusters=processed_clusters)
+            page.predictions.layout = prediction
+
+            if settings.debug.visualize_layout:
+                self.draw_clusters_and_cells_side_by_side(
+                    conv_res, page, processed_clusters, mode_prefix="postprocessed"
+                )
+
+            layout_predictions.append(prediction)
+
+        return layout_predictions