docling 2.69.0__py3-none-any.whl

docling/models/extraction/nuextract_transformers_model.py

@@ -0,0 +1,305 @@
+import logging
+import sys
+import time
+from collections.abc import Iterable
+from pathlib import Path
+from typing import Any, Optional, Union
+
+import numpy as np
+from PIL.Image import Image
+from transformers import AutoModelForImageTextToText, AutoProcessor, GenerationConfig
+
+from docling.datamodel.accelerator_options import (
+    AcceleratorOptions,
+)
+from docling.datamodel.base_models import VlmPrediction, VlmStopReason
+from docling.datamodel.pipeline_options_vlm_model import InlineVlmOptions
+from docling.models.base_model import BaseVlmModel
+from docling.models.utils.hf_model_download import (
+    HuggingFaceModelDownloadMixin,
+)
+from docling.utils.accelerator_utils import decide_device
+
+_log = logging.getLogger(__name__)
+
+
+# Source code from https://huggingface.co/numind/NuExtract-2.0-8B
+def process_all_vision_info(messages, examples=None):
+    """
+    Process vision information from both messages and in-context examples, supporting batch processing.
+
+    Args:
+        messages: List of message dictionaries (single input) OR list of message lists (batch input)
+        examples: Optional list of example dictionaries (single input) OR list of example lists (batch)
+
+    Returns:
+        A flat list of all images in the correct order:
+        - For single input: example images followed by message images
+        - For batch input: interleaved as (item1 examples, item1 input, item2 examples, item2 input, etc.)
+        - Returns None if no images were found
+    """
+    try:
+        from qwen_vl_utils import fetch_image, process_vision_info
+    except ImportError:
+        raise ImportError(
+            "qwen-vl-utils is required for NuExtractTransformersModel. "
+            "Please install it with: pip install qwen-vl-utils"
+        )
+
+    from qwen_vl_utils import fetch_image, process_vision_info
+
+    # Helper function to extract images from examples
+    def extract_example_images(example_item):
+        if not example_item:
+            return []
+
+        # Handle both list of examples and single example
+        examples_to_process = (
+            example_item if isinstance(example_item, list) else [example_item]
+        )
+        images = []
+
+        for example in examples_to_process:
+            if (
+                isinstance(example.get("input"), dict)
+                and example["input"].get("type") == "image"
+            ):
+                images.append(fetch_image(example["input"]))
+
+        return images
+
+    # Normalize inputs to always be batched format
+    is_batch = messages and isinstance(messages[0], list)
+    messages_batch = messages if is_batch else [messages]
+    is_batch_examples = (
+        examples
+        and isinstance(examples, list)
+        and (isinstance(examples[0], list) or examples[0] is None)
+    )
+    examples_batch = (
+        examples
+        if is_batch_examples
+        else ([examples] if examples is not None else None)
+    )
+
+    # Ensure examples batch matches messages batch if provided
+    if examples and len(examples_batch) != len(messages_batch):
+        if not is_batch and len(examples_batch) == 1:
+            # Single example set for a single input is fine
+            pass
+        else:
+            raise ValueError("Examples batch length must match messages batch length")
+
+    # Process all inputs, maintaining correct order
+    all_images = []
+    for i, message_group in enumerate(messages_batch):
+        # Get example images for this input
+        if examples and i < len(examples_batch):
+            input_example_images = extract_example_images(examples_batch[i])
+            all_images.extend(input_example_images)
+
+        # Get message images for this input
+        input_message_images = process_vision_info(message_group)[0] or []
+        all_images.extend(input_message_images)
+
+    return all_images if all_images else None
+
+
+class NuExtractTransformersModel(BaseVlmModel, HuggingFaceModelDownloadMixin):
+    def __init__(
+        self,
+        enabled: bool,
+        artifacts_path: Optional[Path],
+        accelerator_options: AcceleratorOptions,
+        vlm_options: InlineVlmOptions,
+    ):
+        self.enabled = enabled
+        self.vlm_options = vlm_options
+
+        if self.enabled:
+            import torch
+
+            self.device = decide_device(
+                accelerator_options.device,
+                supported_devices=vlm_options.supported_devices,
+            )
+            _log.debug(f"Available device for NuExtract VLM: {self.device}")
+
+            self.max_new_tokens = vlm_options.max_new_tokens
+            self.temperature = vlm_options.temperature
+
+            repo_cache_folder = vlm_options.repo_id.replace("/", "--")
+
+            if artifacts_path is None:
+                artifacts_path = self.download_models(
+                    repo_id=self.vlm_options.repo_id,
+                    revision=self.vlm_options.revision,
+                )
+            elif (artifacts_path / repo_cache_folder).exists():
+                artifacts_path = artifacts_path / repo_cache_folder
+
+            self.processor = AutoProcessor.from_pretrained(
+                artifacts_path,
+                trust_remote_code=vlm_options.trust_remote_code,
+                use_fast=True,
+            )
+            self.processor.tokenizer.padding_side = "left"
+
+            self.vlm_model = AutoModelForImageTextToText.from_pretrained(
+                artifacts_path,
+                device_map=self.device,
+                dtype=self.vlm_options.torch_dtype,
+                _attn_implementation=(
+                    "flash_attention_2"
+                    if self.device.startswith("cuda")
+                    and accelerator_options.cuda_use_flash_attention2
+                    else "sdpa"
+                ),
+                trust_remote_code=vlm_options.trust_remote_code,
+            )
+            if sys.version_info < (3, 14):
+                self.vlm_model = torch.compile(self.vlm_model)  # type: ignore
+            else:
+                self.vlm_model.eval()
+
+            # Load generation config
+            self.generation_config = GenerationConfig.from_pretrained(artifacts_path)
+
+    def process_images(
+        self,
+        image_batch: Iterable[Union[Image, np.ndarray]],
+        prompt: Union[str, list[str]],
+    ) -> Iterable[VlmPrediction]:
+        """
+        Batched inference for NuExtract VLM using the specialized input format.
+
+        Args:
+            image_batch: Iterable of PIL Images or numpy arrays
+            prompt: Either:
+                - str: Single template used for all images
+                - list[str]: List of templates (one per image, must match image count)
+        """
+        import torch
+        from PIL import Image as PILImage
+
+        # Normalize images to RGB PIL
+        pil_images: list[Image] = []
+        for img in image_batch:
+            if isinstance(img, np.ndarray):
+                if img.ndim == 3 and img.shape[2] in (3, 4):
+                    pil_img = PILImage.fromarray(img.astype(np.uint8))
+                elif img.ndim == 2:
+                    pil_img = PILImage.fromarray(img.astype(np.uint8), mode="L")
+                else:
+                    raise ValueError(f"Unsupported numpy array shape: {img.shape}")
+            else:
+                pil_img = img
+            if pil_img.mode != "RGB":
+                pil_img = pil_img.convert("RGB")
+            pil_images.append(pil_img)
+
+        if not pil_images:
+            return
+
+        # Normalize templates (1 per image)
+        if isinstance(prompt, str):
+            templates = [prompt] * len(pil_images)
+        else:
+            if len(prompt) != len(pil_images):
+                raise ValueError(
+                    f"Number of templates ({len(prompt)}) must match number of images ({len(pil_images)})"
+                )
+            templates = prompt
+
+        # Construct NuExtract input format
+        inputs = []
+        for pil_img, template in zip(pil_images, templates):
+            input_item = {
+                "document": {"type": "image", "image": pil_img},
+                "template": template,
+            }
+            inputs.append(input_item)
+
+        # Create messages structure for batch processing
+        messages = [
+            [
+                {
+                    "role": "user",
+                    "content": [x["document"]],
+                }
+            ]
+            for x in inputs
+        ]
+
+        # Apply chat template to each example individually
+        texts = [
+            self.processor.tokenizer.apply_chat_template(
+                messages[i],
+                template=x["template"],
+                tokenize=False,
+                add_generation_prompt=True,
+            )
+            for i, x in enumerate(inputs)
+        ]
+
+        # Process vision inputs using qwen-vl-utils
+        image_inputs = process_all_vision_info(messages)
+
+        # Process with the processor
+        processor_inputs = self.processor(
+            text=texts,
+            images=image_inputs,
+            padding=True,
+            return_tensors="pt",
+            **self.vlm_options.extra_processor_kwargs,
+        )
+        processor_inputs = {k: v.to(self.device) for k, v in processor_inputs.items()}
+
+        # Generate
+        gen_kwargs = {
+            **processor_inputs,
+            "max_new_tokens": self.max_new_tokens,
+            "generation_config": self.generation_config,
+            **self.vlm_options.extra_generation_config,
+        }
+        if self.temperature > 0:
+            gen_kwargs["do_sample"] = True
+            gen_kwargs["temperature"] = self.temperature
+        else:
+            gen_kwargs["do_sample"] = False
+
+        start_time = time.time()
+        with torch.inference_mode():
+            generated_ids = self.vlm_model.generate(**gen_kwargs)
+        generation_time = time.time() - start_time
+
+        # Trim generated sequences
+        input_len = processor_inputs["input_ids"].shape[1]
+        trimmed_sequences = generated_ids[:, input_len:]
+
+        # Decode with the processor/tokenizer
+        decoded_texts: list[str] = self.processor.batch_decode(
+            trimmed_sequences,
+            skip_special_tokens=True,
+            clean_up_tokenization_spaces=False,
+        )
+
+        # Optional logging
+        num_tokens = None
+        if generated_ids.shape[0] > 0:  # type: ignore
+            # Todo: confirm num tokens is actually from first item, code was already like this
+            num_tokens = int(generated_ids[0].shape[0])
+            _log.debug(
+                f"Generated {num_tokens} tokens in {generation_time:.2f}s "
+                f"for batch size {generated_ids.shape[0]}."  # type: ignore
+            )
+
+        for text in decoded_texts:
+            # Apply decode_response to the output text
+            decoded_text = self.vlm_options.decode_response(text)
+            yield VlmPrediction(
+                text=decoded_text,
+                generation_time=generation_time,
+                num_tokens=num_tokens,
+                stop_reason=VlmStopReason.UNSPECIFIED,
+            )

docling/models/factories/__init__.py

@@ -0,0 +1,47 @@
+import logging
+from functools import lru_cache
+
+from docling.models.factories.layout_factory import LayoutFactory
+from docling.models.factories.ocr_factory import OcrFactory
+from docling.models.factories.picture_description_factory import (
+    PictureDescriptionFactory,
+)
+from docling.models.factories.table_factory import TableStructureFactory
+
+logger = logging.getLogger(__name__)
+
+
+@lru_cache
+def get_ocr_factory(allow_external_plugins: bool = False) -> OcrFactory:
+    factory = OcrFactory()
+    factory.load_from_plugins(allow_external_plugins=allow_external_plugins)
+    logger.info("Registered ocr engines: %r", factory.registered_kind)
+    return factory
+
+
+@lru_cache
+def get_picture_description_factory(
+    allow_external_plugins: bool = False,
+) -> PictureDescriptionFactory:
+    factory = PictureDescriptionFactory()
+    factory.load_from_plugins(allow_external_plugins=allow_external_plugins)
+    logger.info("Registered picture descriptions: %r", factory.registered_kind)
+    return factory
+
+
+@lru_cache
+def get_layout_factory(allow_external_plugins: bool = False) -> LayoutFactory:
+    factory = LayoutFactory()
+    factory.load_from_plugins(allow_external_plugins=allow_external_plugins)
+    logger.info("Registered layout engines: %r", factory.registered_kind)
+    return factory
+
+
+@lru_cache
+def get_table_structure_factory(
+    allow_external_plugins: bool = False,
+) -> TableStructureFactory:
+    factory = TableStructureFactory()
+    factory.load_from_plugins(allow_external_plugins=allow_external_plugins)
+    logger.info("Registered table structure engines: %r", factory.registered_kind)
+    return factory

docling/models/factories/base_factory.py

@@ -0,0 +1,122 @@
+import enum
+import logging
+from abc import ABCMeta
+from typing import Generic, Optional, Type, TypeVar
+
+from pluggy import PluginManager
+from pydantic import BaseModel
+
+from docling.datamodel.pipeline_options import BaseOptions
+from docling.models.base_model import BaseModelWithOptions
+
+A = TypeVar("A", bound=BaseModelWithOptions)
+
+
+logger = logging.getLogger(__name__)
+
+
+class FactoryMeta(BaseModel):
+    kind: str
+    plugin_name: str
+    module: str
+
+
+class BaseFactory(Generic[A], metaclass=ABCMeta):
+    default_plugin_name = "docling"
+
+    def __init__(self, plugin_attr_name: str, plugin_name=default_plugin_name):
+        self.plugin_name = plugin_name
+        self.plugin_attr_name = plugin_attr_name
+
+        self._classes: dict[Type[BaseOptions], Type[A]] = {}
+        self._meta: dict[Type[BaseOptions], FactoryMeta] = {}
+
+    @property
+    def registered_kind(self) -> list[str]:
+        return [opt.kind for opt in self._classes.keys()]
+
+    def get_enum(self) -> enum.Enum:
+        return enum.Enum(
+            self.plugin_attr_name + "_enum",
+            names={kind: kind for kind in self.registered_kind},
+            type=str,
+            module=__name__,
+        )
+
+    @property
+    def classes(self):
+        return self._classes
+
+    @property
+    def registered_meta(self):
+        return self._meta
+
+    def create_instance(self, options: BaseOptions, **kwargs) -> A:
+        try:
+            _cls = self._classes[type(options)]
+            return _cls(options=options, **kwargs)
+        except KeyError:
+            raise RuntimeError(self._err_msg_on_class_not_found(options.kind))
+
+    def create_options(self, kind: str, *args, **kwargs) -> BaseOptions:
+        for opt_cls, _ in self._classes.items():
+            if opt_cls.kind == kind:
+                return opt_cls(*args, **kwargs)
+        raise RuntimeError(self._err_msg_on_class_not_found(kind))
+
+    def _err_msg_on_class_not_found(self, kind: str):
+        msg = []
+
+        for opt, cls in self._classes.items():
+            msg.append(f"\t{opt.kind!r} => {cls!r}")
+
+        msg_str = "\n".join(msg)
+
+        return f"No class found with the name {kind!r}, known classes are:\n{msg_str}"
+
+    def register(self, cls: Type[A], plugin_name: str, plugin_module_name: str):
+        opt_type = cls.get_options_type()
+
+        if opt_type in self._classes:
+            raise ValueError(
+                f"{opt_type.kind!r} already registered to class {self._classes[opt_type]!r}"
+            )
+
+        self._classes[opt_type] = cls
+        self._meta[opt_type] = FactoryMeta(
+            kind=opt_type.kind, plugin_name=plugin_name, module=plugin_module_name
+        )
+
+    def load_from_plugins(
+        self, plugin_name: Optional[str] = None, allow_external_plugins: bool = False
+    ):
+        plugin_name = plugin_name or self.plugin_name
+
+        plugin_manager = PluginManager(plugin_name)
+        plugin_manager.load_setuptools_entrypoints(plugin_name)
+
+        for plugin_name, plugin_module in plugin_manager.list_name_plugin():
+            plugin_module_name = str(plugin_module.__name__)  # type: ignore
+
+            if not allow_external_plugins and not plugin_module_name.startswith(
+                "docling."
+            ):
+                logger.warning(
+                    f"The plugin {plugin_name} will not be loaded because Docling is being executed with allow_external_plugins=false."
+                )
+                continue
+
+            attr = getattr(plugin_module, self.plugin_attr_name, None)
+
+            if callable(attr):
+                logger.info("Loading plugin %r", plugin_name)
+
+                config = attr()
+                self.process_plugin(config, plugin_name, plugin_module_name)
+
+    def process_plugin(self, config, plugin_name: str, plugin_module_name: str):
+        for item in config[self.plugin_attr_name]:
+            try:
+                self.register(item, plugin_name, plugin_module_name)
+            except ValueError:
+                logger.warning("%r already registered", item)

docling/models/factories/layout_factory.py

@@ -0,0 +1,7 @@
+from docling.models.base_layout_model import BaseLayoutModel
+from docling.models.factories.base_factory import BaseFactory
+
+
+class LayoutFactory(BaseFactory[BaseLayoutModel]):
+    def __init__(self, *args, **kwargs):
+        super().__init__("layout_engines", *args, **kwargs)

docling/models/factories/ocr_factory.py

@@ -0,0 +1,11 @@
+import logging
+
+from docling.models.base_ocr_model import BaseOcrModel
+from docling.models.factories.base_factory import BaseFactory
+
+logger = logging.getLogger(__name__)
+
+
+class OcrFactory(BaseFactory[BaseOcrModel]):
+    def __init__(self, *args, **kwargs):
+        super().__init__("ocr_engines", *args, **kwargs)

docling/models/factories/picture_description_factory.py

@@ -0,0 +1,11 @@
+import logging
+
+from docling.models.factories.base_factory import BaseFactory
+from docling.models.picture_description_base_model import PictureDescriptionBaseModel
+
+logger = logging.getLogger(__name__)
+
+
+class PictureDescriptionFactory(BaseFactory[PictureDescriptionBaseModel]):
+    def __init__(self, *args, **kwargs):
+        super().__init__("picture_description", *args, **kwargs)

docling/models/factories/table_factory.py

@@ -0,0 +1,7 @@
+from docling.models.base_table_model import BaseTableStructureModel
+from docling.models.factories.base_factory import BaseFactory
+
+
+class TableStructureFactory(BaseFactory[BaseTableStructureModel]):
+    def __init__(self, *args, **kwargs):
+        super().__init__("table_structure_engines", *args, **kwargs)

docling/models/picture_description_base_model.py

@@ -0,0 +1,149 @@
+from abc import abstractmethod
+from collections.abc import Iterable
+from pathlib import Path
+from typing import List, Optional, Type, Union
+
+from docling_core.types.doc import (
+    DescriptionMetaField,
+    DoclingDocument,
+    NodeItem,
+    PictureClassificationLabel,
+    PictureItem,
+    PictureMeta,
+)
+from docling_core.types.doc.document import PictureDescriptionData
+from PIL import Image
+
+from docling.datamodel.accelerator_options import AcceleratorOptions
+from docling.datamodel.pipeline_options import (
+    PictureDescriptionBaseOptions,
+)
+from docling.models.base_model import (
+    BaseItemAndImageEnrichmentModel,
+    BaseModelWithOptions,
+    ItemAndImageEnrichmentElement,
+)
+
+
+class PictureDescriptionBaseModel(
+    BaseItemAndImageEnrichmentModel, BaseModelWithOptions
+):
+    images_scale: float = 2.0
+
+    def __init__(
+        self,
+        *,
+        enabled: bool,
+        enable_remote_services: bool,
+        artifacts_path: Optional[Union[Path, str]],
+        options: PictureDescriptionBaseOptions,
+        accelerator_options: AcceleratorOptions,
+    ):
+        self.enabled = enabled
+        self.options = options
+        self.provenance = "not-implemented"
+
+    def is_processable(self, doc: DoclingDocument, element: NodeItem) -> bool:
+        return self.enabled and isinstance(element, PictureItem)
+
+    def _annotate_images(self, images: Iterable[Image.Image]) -> Iterable[str]:
+        raise NotImplementedError
+
+    def __call__(
+        self,
+        doc: DoclingDocument,
+        element_batch: Iterable[ItemAndImageEnrichmentElement],
+    ) -> Iterable[NodeItem]:
+        if not self.enabled:
+            for element in element_batch:
+                yield element.item
+            return
+
+        images: List[Image.Image] = []
+        elements: List[PictureItem] = []
+        for el in element_batch:
+            assert isinstance(el.item, PictureItem)
+            describe_image = True
+            # Don't describe the image if it's smaller than the threshold
+            if len(el.item.prov) > 0:
+                prov = el.item.prov[0]  # PictureItems have at most a single provenance
+                page = doc.pages.get(prov.page_no)
+                if page is not None:
+                    page_area = page.size.width * page.size.height
+                    if page_area > 0:
+                        area_fraction = prov.bbox.area() / page_area
+                        if area_fraction < self.options.picture_area_threshold:
+                            describe_image = False
+            if describe_image and not _passes_classification(
+                el.item.meta,
+                self.options.classification_allow,
+                self.options.classification_deny,
+                self.options.classification_min_confidence,
+            ):
+                describe_image = False
+            if describe_image:
+                elements.append(el.item)
+                images.append(el.image)
+
+        outputs = self._annotate_images(images)
+
+        for item, output in zip(elements, outputs):
+            # FIXME: annotations is deprecated, remove once all consumers use meta.classification
+            item.annotations.append(
+                PictureDescriptionData(text=output, provenance=self.provenance)
+            )
+
+            # Store classification in the new meta field
+            if item.meta is None:
+                item.meta = PictureMeta()
+            item.meta.description = DescriptionMetaField(
+                text=output,
+                created_by=self.provenance,
+            )
+
+            yield item
+
+    @classmethod
+    @abstractmethod
+    def get_options_type(cls) -> Type[PictureDescriptionBaseOptions]:
+        pass
+
+
+def _passes_classification(
+    meta: Optional[PictureMeta],
+    allow: Optional[List[PictureClassificationLabel]],
+    deny: Optional[List[PictureClassificationLabel]],
+    min_confidence: float,
+) -> bool:
+    if not allow and not deny:
+        return True
+    predicted = None
+    if meta and meta.classification:
+        predicted = meta.classification.predictions
+    if not predicted:
+        return allow is None
+    if deny:
+        deny_set = {_label_value(label) for label in deny}
+        for entry in predicted:
+            if _meets_confidence(entry.confidence, min_confidence) and (
+                entry.class_name in deny_set
+            ):
+                return False
+    if allow:
+        allow_set = {_label_value(label) for label in allow}
+        return any(
+            _meets_confidence(entry.confidence, min_confidence)
+            and entry.class_name in allow_set
+            for entry in predicted
+        )
+    return True
+
+
+def _label_value(label: Union[PictureClassificationLabel, str]) -> str:
+    return label.value if isinstance(label, PictureClassificationLabel) else str(label)
+
+
+def _meets_confidence(confidence: Optional[float], min_confidence: float) -> bool:
+    return min_confidence <= 0 or (
+        confidence is not None and confidence >= min_confidence
+    )