docling 2.53.0__py3-none-any.whl → 2.55.0__py3-none-any.whl

docling/models/utils/generation_utils.py

@@ -0,0 +1,157 @@
+import logging
+import re
+import sys
+from abc import abstractmethod
+from typing import List
+
+from transformers import StoppingCriteria
+
+_log = logging.getLogger(__name__)
+
+
+class GenerationStopper:
+    """
+    Base interface for stopping logic.
+    - should_stop(s): True to stop given the current decoded text window.
+    - lookback_tokens(): how many tokens should be considered (default: sys.maxsize).
+    """
+
+    @abstractmethod
+    def should_stop(self, s: str) -> bool:
+        pass
+
+    def lookback_tokens(self) -> int:
+        return sys.maxsize
+
+
+class DocTagsRepetitionStopper(GenerationStopper):
+    """
+    Detects repetitive <tag>...<loc_x><loc_y><loc_w><loc_h>text</tag> blocks,
+    but only when repeats are **consecutive** and both tag & inner text are identical.
+
+    Performance:
+    - Heavy check runs every N calls (default 32).
+    - Only decodes the last LOOKBACK_TOKENS tokens per sequence (default 200).
+    """
+
+    def __init__(self, *, N: int = 32, lookback_tokens: int = 200):
+        self.N = max(1, int(N))
+        self._lookback_tokens = max(1, int(lookback_tokens))
+        self._call_count = 0
+
+        # <tag> ... <loc_x><loc_y><loc_w><loc_h> text ... </tag>
+        self._PATTERN = re.compile(
+            r"""
+            <(?P<tag>[a-zA-Z0-9_]+)>\s*
+            (?P<prefix>.*?)?
+            <loc_(?P<x>\d+)><loc_(?P<y>\d+)><loc_(?P<w>\d+)><loc_(?P<h>\d+)>
+            (?P<text>.*?)
+            </(?P=tag)>
+            """,
+            re.DOTALL | re.VERBOSE,
+        )
+
+    # --- small helper ---
+    def _regular(self, vals: List[int]) -> bool:
+        """3+ strictly increasing values with ~regular spacing (±20%)."""
+        if len(vals) < 3:
+            return False
+        diffs = [b - a for a, b in zip(vals, vals[1:])]
+        if any(d <= 0 for d in diffs):
+            return False
+        mean = sum(diffs) / len(diffs)
+        tol = 0.2 * mean
+        return all(abs(d - mean) <= tol for d in diffs)
+
+    def should_stop(self, s: str) -> bool:
+        """
+        Trip only on **consecutive** runs (no other matched blocks between) of ≥3 items
+        with the same <tag> and identical inner text, where within that run we see:
+          - any exact duplicate (x,y,w,h), or
+          - stable X/W with regular Y progression, or
+          - stable Y/H with regular X progression.
+        """
+        # Stream matches and evaluate runs on-the-fly to stay compact and fast.
+        prev_tag = prev_text = None
+        run = []  # list of (x,y,w,h)
+
+        def run_repetitive(boxes: List[tuple]) -> bool:
+            if len(boxes) < 3:
+                return False
+            # duplicates?
+            if len(set(boxes)) < len(boxes):
+                return True
+            xs, ys, ws, hs = zip(*boxes)
+            x_stable = all(x == xs[0] for x in xs)
+            y_stable = all(y == ys[0] for y in ys)
+            w_stable = all(w == ws[0] for w in ws)
+            h_stable = all(h == hs[0] for h in hs)
+            # horizontal (down the page): X/W stable, Y regular
+            if (x_stable or w_stable) and self._regular(list(ys)):
+                return True
+            # vertical (across): Y/H stable, X regular
+            if (y_stable or h_stable) and self._regular(list(xs)):
+                return True
+            return False
+
+        for m in self._PATTERN.finditer(s):
+            tag, text = m.group("tag"), m.group("text")
+            box = (
+                int(m.group("x")),
+                int(m.group("y")),
+                int(m.group("w")),
+                int(m.group("h")),
+            )
+
+            if prev_tag == tag and prev_text == text:
+                run.append(box)  # consecutive same-tag+text
+            else:
+                # evaluate previous run before starting a new one
+                if run_repetitive(run):
+                    return True
+                prev_tag, prev_text = tag, text
+                run = [box]
+
+        # check the last run
+        return run_repetitive(run)
+
+
+class HFStoppingCriteriaWrapper(StoppingCriteria):
+    """
+    Adapts any GenerationStopper to HuggingFace Transformers.
+    Decodes exactly min(seq_len, stopper.lookback_tokens()) tokens from the end.
+    """
+
+    def __init__(
+        self,
+        tokenizer,
+        stopper: GenerationStopper,
+        *,
+        skip_special_tokens: bool = False,
+    ):
+        self.tokenizer = tokenizer
+        self.stopper = stopper
+        self.skip_special_tokens = skip_special_tokens
+
+    def __call__(self, input_ids, scores, **kwargs) -> bool:
+        lb = max(1, int(self.stopper.lookback_tokens()))
+        for seq in input_ids:  # (batch, seq_len)
+            window = seq[-lb:]  # slicing handles lb > len(seq)
+            try:
+                text = self.tokenizer.decode(
+                    window, skip_special_tokens=self.skip_special_tokens
+                )
+            except Exception as e:
+                _log.info(f"Decoding failed for stopping check: {e}")
+                continue
+
+            try:
+                if self.stopper.should_stop(text):
+                    _log.info(
+                        "HF wrapper: stopping due to TextStopper.should_stop==True"
+                    )
+                    return True
+            except Exception as e:
+                _log.info(f"Error in TextStopper.should_stop: {e}")
+                continue
+        return False

docling/models/utils/hf_model_download.py

@@ -34,7 +34,12 @@ class HuggingFaceModelDownloadMixin:
         local_dir: Optional[Path] = None,
         force: bool = False,
         progress: bool = False,
+        revision: Optional[str] = None,
     ) -> Path:
         return download_hf_model(
-            repo_id=repo_id, local_dir=local_dir, force=force, progress=progress
+            repo_id=repo_id,
+            local_dir=local_dir,
+            force=force,
+            progress=progress,
+            revision=revision,
         )

docling/models/vlm_models_inline/hf_transformers_model.py

@@ -7,7 +7,7 @@ from typing import Any, Optional, Union
 
 import numpy as np
 from PIL.Image import Image
-from transformers import StoppingCriteriaList, StopStringCriteria
+from transformers import StoppingCriteria, StoppingCriteriaList, StopStringCriteria
 
 from docling.datamodel.accelerator_options import (
     AcceleratorOptions,
@@ -20,6 +20,10 @@ from docling.datamodel.pipeline_options_vlm_model import (
     TransformersPromptStyle,
 )
 from docling.models.base_model import BaseVlmPageModel
+from docling.models.utils.generation_utils import (
+    GenerationStopper,
+    HFStoppingCriteriaWrapper,
+)
 from docling.models.utils.hf_model_download import (
     HuggingFaceModelDownloadMixin,
 )
@@ -75,7 +79,9 @@ class HuggingFaceTransformersVlmModel(BaseVlmPageModel, HuggingFaceModelDownload
             repo_cache_folder = vlm_options.repo_id.replace("/", "--")
 
             if artifacts_path is None:
-                artifacts_path = self.download_models(self.vlm_options.repo_id)
+                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
 
@@ -106,6 +112,7 @@ class HuggingFaceTransformersVlmModel(BaseVlmPageModel, HuggingFaceModelDownload
             self.processor = AutoProcessor.from_pretrained(
                 artifacts_path,
                 trust_remote_code=vlm_options.trust_remote_code,
+                revision=vlm_options.revision,
             )
             self.processor.tokenizer.padding_side = "left"
 
@@ -120,11 +127,14 @@ class HuggingFaceTransformersVlmModel(BaseVlmPageModel, HuggingFaceModelDownload
                     else "sdpa"
                 ),
                 trust_remote_code=vlm_options.trust_remote_code,
+                revision=vlm_options.revision,
             )
             self.vlm_model = torch.compile(self.vlm_model)  # type: ignore
 
             # Load generation config
-            self.generation_config = GenerationConfig.from_pretrained(artifacts_path)
+            self.generation_config = GenerationConfig.from_pretrained(
+                artifacts_path, revision=vlm_options.revision
+            )
 
     def __call__(
         self, conv_res: ConversionResult, page_batch: Iterable[Page]
@@ -196,7 +206,7 @@ class HuggingFaceTransformersVlmModel(BaseVlmPageModel, HuggingFaceModelDownload
         import torch
         from PIL import Image as PILImage
 
-        # -- Normalize images to RGB PIL (SmolDocling & friends accept PIL/np via processor)
+        # -- Normalize images to RGB PIL
         pil_images: list[Image] = []
         for img in image_batch:
             if isinstance(img, np.ndarray):
@@ -247,24 +257,74 @@ class HuggingFaceTransformersVlmModel(BaseVlmPageModel, HuggingFaceModelDownload
         inputs = {k: v.to(self.device) for k, v in inputs.items()}
 
         # -- Optional stopping criteria
-        stopping_criteria = None
+        stopping_criteria_list: StoppingCriteriaList = StoppingCriteriaList()
+
+        # Add string-based stopping criteria
         if self.vlm_options.stop_strings:
-            stopping_criteria = StoppingCriteriaList(
-                [
-                    StopStringCriteria(
-                        stop_strings=self.vlm_options.stop_strings,
-                        tokenizer=self.processor.tokenizer,
-                    )
-                ]
+            stopping_criteria_list.append(
+                StopStringCriteria(
+                    stop_strings=self.vlm_options.stop_strings,
+                    tokenizer=self.processor.tokenizer,
+                )
             )
 
+        # Add custom stopping criteria
+        if self.vlm_options.custom_stopping_criteria:
+            for criteria in self.vlm_options.custom_stopping_criteria:
+                # If it's a class (not an instance), determine the type and handle accordingly
+                if isinstance(criteria, type):
+                    # Check if it's a GenerationStopper class
+                    if issubclass(criteria, GenerationStopper):
+                        # Instantiate GenerationStopper and wrap it
+                        stopper_instance = criteria()
+                        wrapped_criteria = HFStoppingCriteriaWrapper(
+                            self.processor.tokenizer, stopper_instance
+                        )
+                        stopping_criteria_list.append(wrapped_criteria)
+                    elif issubclass(criteria, StoppingCriteria):
+                        # It's a StoppingCriteria class, instantiate with tokenizer
+                        criteria_instance = criteria(self.processor.tokenizer)
+                        stopping_criteria_list.append(criteria_instance)
+                elif isinstance(criteria, GenerationStopper):
+                    # Wrap GenerationStopper instances in HFStoppingCriteriaWrapper
+                    wrapped_criteria = HFStoppingCriteriaWrapper(
+                        self.processor.tokenizer, criteria
+                    )
+                    stopping_criteria_list.append(wrapped_criteria)
+                else:
+                    # If it's already an instance of StoppingCriteria, use it directly
+                    stopping_criteria_list.append(criteria)
+
+        stopping_criteria = (
+            StoppingCriteriaList(stopping_criteria_list)
+            if stopping_criteria_list
+            else None
+        )
+
+        # -- Filter out decoder-specific keys from extra_generation_config
+        decoder_keys = {
+            "skip_special_tokens",
+            "clean_up_tokenization_spaces",
+            "spaces_between_special_tokens",
+        }
+        generation_config = {
+            k: v
+            for k, v in self.vlm_options.extra_generation_config.items()
+            if k not in decoder_keys
+        }
+        decoder_config = {
+            k: v
+            for k, v in self.vlm_options.extra_generation_config.items()
+            if k in decoder_keys
+        }
+
         # -- Generate (Image-Text-to-Text class expects these inputs from processor)
         gen_kwargs = {
             **inputs,
             "max_new_tokens": self.max_new_tokens,
             "use_cache": self.use_cache,
             "generation_config": self.generation_config,
-            **self.vlm_options.extra_generation_config,
+            **generation_config,
         }
         if self.temperature > 0:
             gen_kwargs["do_sample"] = True
@@ -293,7 +353,8 @@ class HuggingFaceTransformersVlmModel(BaseVlmPageModel, HuggingFaceModelDownload
             )
 
         decoded_texts: list[str] = decode_fn(
-            trimmed_sequences, skip_special_tokens=False
+            trimmed_sequences,
+            **decoder_config,
         )
 
         # -- Clip off pad tokens from decoded texts

docling/models/vlm_models_inline/mlx_model.py

@@ -1,4 +1,5 @@
 import logging
+import sys
 import threading
 import time
 from collections.abc import Iterable
@@ -7,6 +8,7 @@ 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,
@@ -15,6 +17,7 @@ from docling.datamodel.base_models import Page, VlmPrediction, VlmPredictionToke
 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,
 )
@@ -60,6 +63,7 @@ class HuggingFaceMlxModel(BaseVlmPageModel, HuggingFaceModelDownloadMixin):
             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
@@ -68,6 +72,22 @@ class HuggingFaceMlxModel(BaseVlmPageModel, HuggingFaceModelDownloadMixin):
             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]:
@@ -192,7 +212,7 @@ class HuggingFaceMlxModel(BaseVlmPageModel, HuggingFaceModelDownloadMixin):
                     self.processor, self.config, user_prompt, num_images=1
                 )
 
-                # Stream generate with stop strings support
+                # Stream generate with stop strings and custom stopping criteria support
                 start_time = time.time()
                 _log.debug("start generating ...")
 
@@ -244,6 +264,43 @@ class HuggingFaceMlxModel(BaseVlmPageModel, HuggingFaceModelDownloadMixin):
                             _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(