docling 2.69.0__py3-none-any.whl

docling/models/vlm_pipeline_models/mlx_model.py

@@ -0,0 +1,325 @@
+import logging
+import sys
+import threading
+import time
+from collections.abc import Iterable
+from pathlib import Path
+from typing import Optional, Union
+
+import numpy as np
+from PIL.Image import Image
+from transformers import StoppingCriteria
+
+from docling.datamodel.accelerator_options import (
+    AcceleratorOptions,
+)
+from docling.datamodel.base_models import (
+    Page,
+    VlmPrediction,
+    VlmPredictionToken,
+    VlmStopReason,
+)
+from docling.datamodel.document import ConversionResult
+from docling.datamodel.pipeline_options_vlm_model import InlineVlmOptions
+from docling.models.base_model import BaseVlmPageModel
+from docling.models.utils.generation_utils import GenerationStopper
+from docling.models.utils.hf_model_download import (
+    HuggingFaceModelDownloadMixin,
+)
+from docling.utils.profiling import TimeRecorder
+
+_log = logging.getLogger(__name__)
+
+# Global lock for MLX model calls - MLX models are not thread-safe
+# All MLX models share this lock to prevent concurrent MLX operations
+_MLX_GLOBAL_LOCK = threading.Lock()
+
+
+class HuggingFaceMlxModel(BaseVlmPageModel, HuggingFaceModelDownloadMixin):
+    def __init__(
+        self,
+        enabled: bool,
+        artifacts_path: Optional[Path],
+        accelerator_options: AcceleratorOptions,
+        vlm_options: InlineVlmOptions,
+    ):
+        self.enabled = enabled
+
+        self.vlm_options = vlm_options
+        self.max_tokens = vlm_options.max_new_tokens
+        self.temperature = vlm_options.temperature
+
+        if self.enabled:
+            try:
+                from mlx_vlm import generate, load, stream_generate  # type: ignore
+                from mlx_vlm.prompt_utils import apply_chat_template  # type: ignore
+                from mlx_vlm.utils import load_config  # type: ignore
+            except ImportError:
+                raise ImportError(
+                    "mlx-vlm is not installed. Please install it via `pip install mlx-vlm` to use MLX VLM models."
+                )
+
+            repo_cache_folder = vlm_options.repo_id.replace("/", "--")
+
+            self.apply_chat_template = apply_chat_template
+            self.stream_generate = stream_generate
+
+            # PARAMETERS:
+            if artifacts_path is None:
+                artifacts_path = self.download_models(
+                    self.vlm_options.repo_id,
+                    revision=self.vlm_options.revision,
+                )
+            elif (artifacts_path / repo_cache_folder).exists():
+                artifacts_path = artifacts_path / repo_cache_folder
+
+            ## Load the model
+            self.vlm_model, self.processor = load(artifacts_path)
+            self.config = load_config(artifacts_path)
+
+            # Validate custom stopping criteria - MLX doesn't support HF StoppingCriteria
+            if self.vlm_options.custom_stopping_criteria:
+                for criteria in self.vlm_options.custom_stopping_criteria:
+                    if isinstance(criteria, StoppingCriteria):
+                        raise ValueError(
+                            f"MLX models do not support HuggingFace StoppingCriteria instances. "
+                            f"Found {type(criteria).__name__}. Use GenerationStopper instead."
+                        )
+                    elif isinstance(criteria, type) and issubclass(
+                        criteria, StoppingCriteria
+                    ):
+                        raise ValueError(
+                            f"MLX models do not support HuggingFace StoppingCriteria classes. "
+                            f"Found {criteria.__name__}. Use GenerationStopper instead."
+                        )
+
+    def __call__(
+        self, conv_res: ConversionResult, page_batch: Iterable[Page]
+    ) -> Iterable[Page]:
+        page_list = list(page_batch)
+        if not page_list:
+            return
+
+        valid_pages = []
+        invalid_pages = []
+
+        for page in page_list:
+            assert page._backend is not None
+            if not page._backend.is_valid():
+                invalid_pages.append(page)
+            else:
+                valid_pages.append(page)
+
+        # Process valid pages in batch
+        if valid_pages:
+            with TimeRecorder(conv_res, f"vlm-mlx-{self.vlm_options.repo_id}"):
+                # Prepare images and prompts for batch processing
+                images = []
+                user_prompts = []
+                pages_with_images = []
+
+                for page in valid_pages:
+                    assert page.size is not None
+                    hi_res_image = page.get_image(
+                        scale=self.vlm_options.scale, max_size=self.vlm_options.max_size
+                    )
+
+                    # Only process pages with valid images
+                    if hi_res_image is not None:
+                        images.append(hi_res_image)
+
+                        # Define prompt structure
+                        user_prompt = self._build_prompt_safe(page)
+
+                        user_prompts.append(user_prompt)
+                        pages_with_images.append(page)
+
+                # Use process_images for the actual inference
+                if images:  # Only if we have valid images
+                    predictions = list(self.process_images(images, user_prompts))
+
+                    # Attach results to pages
+                    for page, prediction in zip(pages_with_images, predictions):
+                        page.predictions.vlm_response = prediction
+
+        # Yield all pages (valid and invalid)
+        for page in invalid_pages:
+            yield page
+        for page in valid_pages:
+            yield page
+
+    def process_images(
+        self,
+        image_batch: Iterable[Union[Image, np.ndarray]],
+        prompt: Union[str, list[str]],
+    ) -> Iterable[VlmPrediction]:
+        """Process raw images without page metadata.
+
+        Args:
+            image_batch: Iterable of PIL Images or numpy arrays
+            prompt: Either:
+                - str: Single prompt used for all images
+                - list[str]: List of prompts (one per image, must match image count)
+
+        Raises:
+            ValueError: If prompt list length doesn't match image count.
+        """
+        # Convert image batch to list for length validation
+        image_list = list(image_batch)
+
+        if len(image_list) == 0:
+            return
+
+        # Handle prompt parameter
+        if isinstance(prompt, str):
+            # Single prompt for all images
+            user_prompts = [prompt] * len(image_list)
+        elif isinstance(prompt, list):
+            # List of prompts (one per image)
+            if len(prompt) != len(image_list):
+                raise ValueError(
+                    f"Number of prompts ({len(prompt)}) must match number of images ({len(image_list)})"
+                )
+            user_prompts = prompt
+        else:
+            raise ValueError(f"prompt must be str or list[str], got {type(prompt)}")
+
+        # MLX models are not thread-safe - use global lock to serialize access
+        with _MLX_GLOBAL_LOCK:
+            _log.debug("MLX model: Acquired global lock for thread safety")
+            for image, user_prompt in zip(image_list, user_prompts):
+                # Convert numpy array to PIL Image if needed
+                if isinstance(image, np.ndarray):
+                    if image.ndim == 3 and image.shape[2] in [3, 4]:
+                        # RGB or RGBA array
+                        from PIL import Image as PILImage
+
+                        image = PILImage.fromarray(image.astype(np.uint8))
+                    elif image.ndim == 2:
+                        # Grayscale array
+                        from PIL import Image as PILImage
+
+                        image = PILImage.fromarray(image.astype(np.uint8), mode="L")
+                    else:
+                        raise ValueError(
+                            f"Unsupported numpy array shape: {image.shape}"
+                        )
+
+                # Ensure image is in RGB mode (handles RGBA, L, etc.)
+                if image.mode != "RGB":
+                    image = image.convert("RGB")
+
+                # Use the MLX chat template approach like in the __call__ method
+                formatted_prompt = self.apply_chat_template(
+                    self.processor, self.config, user_prompt, num_images=1
+                )
+
+                # Stream generate with stop strings and custom stopping criteria support
+                start_time = time.time()
+                _log.debug("start generating ...")
+
+                tokens: list[VlmPredictionToken] = []
+                output = ""
+
+                # Use stream_generate for proper stop string handling
+                for token in self.stream_generate(
+                    self.vlm_model,
+                    self.processor,
+                    formatted_prompt,
+                    [image],  # MLX stream_generate expects list of images
+                    max_tokens=self.max_tokens,
+                    verbose=False,
+                    temp=self.temperature,
+                ):
+                    # Collect token information
+                    if len(token.logprobs.shape) == 1:
+                        tokens.append(
+                            VlmPredictionToken(
+                                text=token.text,
+                                token=token.token,
+                                logprob=token.logprobs[token.token],
+                            )
+                        )
+                    elif (
+                        len(token.logprobs.shape) == 2 and token.logprobs.shape[0] == 1
+                    ):
+                        tokens.append(
+                            VlmPredictionToken(
+                                text=token.text,
+                                token=token.token,
+                                logprob=token.logprobs[0, token.token],
+                            )
+                        )
+                    else:
+                        _log.warning(
+                            f"incompatible shape for logprobs: {token.logprobs.shape}"
+                        )
+
+                    output += token.text
+
+                    # Check for any configured stop strings
+                    if self.vlm_options.stop_strings:
+                        if any(
+                            stop_str in output
+                            for stop_str in self.vlm_options.stop_strings
+                        ):
+                            _log.debug("Stopping generation due to stop string match")
+                            break
+
+                    # Check for custom stopping criteria (GenerationStopper instances)
+                    if self.vlm_options.custom_stopping_criteria:
+                        for criteria in self.vlm_options.custom_stopping_criteria:
+                            # Handle both instances and classes of GenerationStopper
+                            if isinstance(criteria, GenerationStopper):
+                                stopper = criteria
+                            elif isinstance(criteria, type) and issubclass(
+                                criteria, GenerationStopper
+                            ):
+                                stopper = criteria()
+
+                            # Determine the text window to check based on lookback_tokens
+                            lookback_tokens = stopper.lookback_tokens()
+                            # Check only the last N characters worth of text
+                            # This is a simplified approach - in practice, you might want to
+                            # decode the last N tokens from the token list for more accuracy
+                            text_to_check = (
+                                output[-lookback_tokens:]
+                                if len(output) > lookback_tokens
+                                else output
+                            )
+
+                            try:
+                                if stopper.should_stop(text_to_check):
+                                    _log.info(
+                                        f"Stopping generation due to GenerationStopper: {type(stopper).__name__}"
+                                    )
+                                    break
+                            except Exception as e:
+                                _log.warning(
+                                    f"Error in GenerationStopper.should_stop: {e}"
+                                )
+                                continue
+                        else:  # note: for-else idiom
+                            continue  # Only executed if the inner loop didn't break
+                        break  # Break the outer loop if any stopper triggered
+
+                generation_time = time.time() - start_time
+
+                _log.debug(
+                    f"{generation_time:.2f} seconds for {len(tokens)} tokens ({len(tokens) / generation_time:.1f} tokens/sec)."
+                )
+
+                # Apply decode_response to the output before yielding
+                decoded_output = self.vlm_options.decode_response(output)
+                input_prompt = (
+                    formatted_prompt if self.vlm_options.track_input_prompt else None
+                )
+                yield VlmPrediction(
+                    text=decoded_output,
+                    generation_time=generation_time,
+                    generated_tokens=tokens,
+                    num_tokens=len(tokens),
+                    stop_reason=VlmStopReason.UNSPECIFIED,
+                    input_prompt=input_prompt,
+                )
+            _log.debug("MLX model: Released global lock")

docling/models/vlm_pipeline_models/vllm_model.py

@@ -0,0 +1,344 @@
+import logging
+import sys
+import time
+from collections.abc import Iterable
+from pathlib import Path
+from typing import Any, Dict, Optional, Union
+
+import numpy as np
+from PIL.Image import Image
+
+from docling.datamodel.accelerator_options import AcceleratorOptions
+from docling.datamodel.base_models import (
+    Page,
+    VlmPrediction,
+    VlmPredictionToken,
+    VlmStopReason,
+)
+from docling.datamodel.document import ConversionResult
+from docling.datamodel.pipeline_options_vlm_model import (
+    InlineVlmOptions,
+    TransformersPromptStyle,
+)
+from docling.models.base_model import BaseVlmPageModel
+from docling.models.utils.hf_model_download import HuggingFaceModelDownloadMixin
+from docling.utils.accelerator_utils import decide_device
+from docling.utils.profiling import TimeRecorder
+
+_log = logging.getLogger(__name__)
+
+
+class VllmVlmModel(BaseVlmPageModel, HuggingFaceModelDownloadMixin):
+    """
+    vLLM-backed vision-language model that accepts PIL images (or numpy arrays)
+    via vLLM's multi_modal_data, with prompt formatting handled by formulate_prompt().
+    """
+
+    # --------- Allowlist of vLLM args ---------
+    # SamplingParams (runtime generation controls)
+    _VLLM_SAMPLING_KEYS = {
+        # Core
+        "max_tokens",
+        "temperature",
+        "top_p",
+        "top_k",
+        # Penalties
+        "presence_penalty",
+        "frequency_penalty",
+        "repetition_penalty",
+        # Stops / outputs
+        "stop",
+        "stop_token_ids",
+        "skip_special_tokens",
+        "spaces_between_special_tokens",
+        # Search / length
+        "n",
+        "best_of",
+        "length_penalty",
+        "early_stopping",
+        # Misc
+        "logprobs",
+        "prompt_logprobs",
+        "min_p",
+        "seed",
+    }
+
+    # LLM(...) / EngineArgs (engine/load-time controls)
+    _VLLM_ENGINE_KEYS = {
+        # Model/tokenizer/impl
+        "tokenizer",
+        "tokenizer_mode",
+        "download_dir",
+        # Parallelism / memory / lengths
+        "tensor_parallel_size",
+        "pipeline_parallel_size",
+        "gpu_memory_utilization",
+        "max_model_len",
+        "max_num_batched_tokens",
+        "kv_cache_dtype",
+        "dtype",
+        # Quantization (coarse switch)
+        "quantization",
+        # Multimodal limits
+        "limit_mm_per_prompt",
+        # Execution toggles
+        "enforce_eager",
+    }
+
+    def __init__(
+        self,
+        enabled: bool,
+        artifacts_path: Optional[Path],
+        accelerator_options: AcceleratorOptions,
+        vlm_options: InlineVlmOptions,
+    ):
+        self.enabled = enabled
+        self.vlm_options: InlineVlmOptions = vlm_options
+
+        self.llm = None
+        self.sampling_params = None
+        self.processor = None  # used for CHAT templating in formulate_prompt()
+        self.device = "cpu"
+        self.max_new_tokens = vlm_options.max_new_tokens
+        self.temperature = vlm_options.temperature
+
+        if not self.enabled:
+            return
+
+        from transformers import AutoProcessor
+
+        try:
+            from vllm import LLM, SamplingParams
+        except ImportError:
+            if sys.version_info < (3, 14):
+                raise ImportError(
+                    "vllm is not installed. Please install it via `pip install vllm`."
+                )
+            else:
+                raise ImportError(
+                    "vllm is not installed. It is not yet available on Python 3.14."
+                )
+
+        # Device selection
+        self.device = decide_device(
+            accelerator_options.device, supported_devices=vlm_options.supported_devices
+        )
+        _log.debug(f"Available device for VLM: {self.device}")
+
+        # Resolve artifacts path / cache folder
+        repo_cache_folder = vlm_options.repo_id.replace("/", "--")
+        if artifacts_path is None:
+            artifacts_path = self.download_models(
+                self.vlm_options.repo_id, revision=self.vlm_options.revision
+            )
+        elif (artifacts_path / repo_cache_folder).exists():
+            artifacts_path = artifacts_path / repo_cache_folder
+
+        # --------- Strict split & validation of extra_generation_config ---------
+        extra_cfg = self.vlm_options.extra_generation_config
+
+        load_cfg = {k: v for k, v in extra_cfg.items() if k in self._VLLM_ENGINE_KEYS}
+        gen_cfg = {k: v for k, v in extra_cfg.items() if k in self._VLLM_SAMPLING_KEYS}
+
+        unknown = sorted(
+            k
+            for k in extra_cfg.keys()
+            if k not in self._VLLM_ENGINE_KEYS and k not in self._VLLM_SAMPLING_KEYS
+        )
+        if unknown:
+            _log.warning(
+                "Ignoring unknown extra_generation_config keys for vLLM: %s", unknown
+            )
+
+        # --------- Construct LLM kwargs (engine/load-time) ---------
+        llm_kwargs: Dict[str, Any] = {
+            "model": str(artifacts_path),
+            "model_impl": "transformers",
+            "limit_mm_per_prompt": {"image": 1},
+            "revision": self.vlm_options.revision,
+            "trust_remote_code": self.vlm_options.trust_remote_code,
+            **load_cfg,
+        }
+
+        if self.device == "cpu":
+            llm_kwargs.setdefault("enforce_eager", True)
+        else:
+            llm_kwargs.setdefault(
+                "gpu_memory_utilization", 0.3
+            )  # room for other models
+
+        # Quantization (kept as-is; coarse)
+        if self.vlm_options.quantized and self.vlm_options.load_in_8bit:
+            llm_kwargs.setdefault("quantization", "bitsandbytes")
+
+        # Initialize vLLM LLM
+        self.llm = LLM(**llm_kwargs)
+
+        # Initialize processor for prompt templating (needed for CHAT style)
+        self.processor = AutoProcessor.from_pretrained(
+            artifacts_path,
+            trust_remote_code=self.vlm_options.trust_remote_code,
+            revision=self.vlm_options.revision,
+        )
+
+        # --------- SamplingParams (runtime) ---------
+        self.sampling_params = SamplingParams(
+            temperature=self.temperature,
+            max_tokens=self.max_new_tokens,
+            stop=(self.vlm_options.stop_strings or None),
+            **gen_cfg,
+        )
+
+    def __call__(
+        self, conv_res: ConversionResult, page_batch: Iterable[Page]
+    ) -> Iterable[Page]:
+        # If disabled, pass-through
+        if not self.enabled:
+            for page in page_batch:
+                yield page
+            return
+
+        page_list = list(page_batch)
+        if not page_list:
+            return
+
+        # Preserve original order
+        original_order = page_list[:]
+
+        # Separate valid/invalid
+        valid_pages: list[Page] = []
+        invalid_pages: list[Page] = []
+        for page in page_list:
+            assert page._backend is not None
+            if page._backend.is_valid():
+                valid_pages.append(page)
+            else:
+                invalid_pages.append(page)
+
+        if valid_pages:
+            with TimeRecorder(conv_res, "vlm"):
+                images: list[Image] = []
+                user_prompts: list[str] = []
+                pages_with_images: list[Page] = []
+
+                for page in valid_pages:
+                    assert page.size is not None
+                    hi_res_image = page.get_image(
+                        scale=self.vlm_options.scale,
+                        max_size=self.vlm_options.max_size,
+                    )
+                    if hi_res_image is None:
+                        continue
+
+                    images.append(hi_res_image)
+
+                    # Define prompt structure
+                    user_prompt = self._build_prompt_safe(page)
+
+                    user_prompts.append(user_prompt)
+                    pages_with_images.append(page)
+
+                if images:
+                    with TimeRecorder(conv_res, "vlm_inference"):
+                        predictions = list(self.process_images(images, user_prompts))
+                    for page, prediction in zip(pages_with_images, predictions):
+                        page.predictions.vlm_response = prediction
+
+        # Yield in original order
+        for page in original_order:
+            yield page
+
+    def process_images(
+        self,
+        image_batch: Iterable[Union[Image, np.ndarray]],
+        prompt: Union[str, list[str]],
+    ) -> Iterable[VlmPrediction]:
+        """Process images in a single batched vLLM inference call."""
+        import numpy as np
+        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 prompts
+        if isinstance(prompt, str):
+            user_prompts = [prompt] * len(pil_images)
+        elif isinstance(prompt, list):
+            if len(prompt) != len(pil_images):
+                raise ValueError(
+                    f"Number of prompts ({len(prompt)}) must match number of images ({len(pil_images)})"
+                )
+            user_prompts = prompt
+        else:
+            raise ValueError(f"prompt must be str or list[str], got {type(prompt)}")
+
+        # Format prompts
+        prompts: list[str] = [self.formulate_prompt(up) for up in user_prompts]
+
+        # Build vLLM inputs
+        llm_inputs = [
+            {"prompt": p, "multi_modal_data": {"image": im}}
+            for p, im in zip(prompts, pil_images)
+        ]
+
+        # Generate
+        assert self.llm is not None and self.sampling_params is not None
+        start_time = time.time()
+        outputs = self.llm.generate(llm_inputs, sampling_params=self.sampling_params)  # type: ignore
+        generation_time = time.time() - start_time
+
+        # Optional debug
+        if outputs:
+            try:
+                num_tokens_within_batch = len(outputs[0].outputs[0].token_ids)
+                _log.debug(
+                    f"Generated {num_tokens_within_batch} tokens for batch in {generation_time:.2f}s."
+                )
+            except Exception:
+                num_tokens_within_batch = 0
+
+        # Emit predictions
+        for i, output in enumerate(outputs):
+            text = output.outputs[0].text if output.outputs else ""
+            stop_reason = (
+                VlmStopReason.END_OF_SEQUENCE
+                if output.outputs[0].stop_reason
+                else VlmStopReason.LENGTH
+            )
+
+            generated_tokens = [
+                VlmPredictionToken(token=int(t)) for t in output.outputs[0].token_ids
+            ]
+            num_tokens = len(generated_tokens)
+
+            if not self.vlm_options.track_generated_tokens:
+                generated_tokens = []
+
+            input_prompt = prompts[i] if self.vlm_options.track_input_prompt else None
+            _log.debug(f"VLM generated response carries input prompt: {input_prompt}")
+
+            decoded_text = self.vlm_options.decode_response(text)
+            yield VlmPrediction(
+                text=decoded_text,
+                generation_time=generation_time,
+                num_tokens=num_tokens,
+                stop_reason=stop_reason,
+                generated_tokens=generated_tokens,
+                input_prompt=input_prompt,
+            )