docling 2.69.0__py3-none-any.whl

docling/models/stages/table_structure/table_structure_model.py

@@ -0,0 +1,305 @@
+import copy
+import warnings
+from collections.abc import Iterable, Sequence
+from pathlib import Path
+from typing import Optional
+
+import numpy
+from docling_core.types.doc import BoundingBox, DocItemLabel, TableCell
+from docling_core.types.doc.page import (
+    BoundingRectangle,
+    TextCellUnit,
+)
+from PIL import ImageDraw
+
+from docling.datamodel.accelerator_options import AcceleratorDevice, AcceleratorOptions
+from docling.datamodel.base_models import Page, Table, TableStructurePrediction
+from docling.datamodel.document import ConversionResult
+from docling.datamodel.pipeline_options import (
+    TableFormerMode,
+    TableStructureOptions,
+)
+from docling.datamodel.settings import settings
+from docling.models.base_table_model import BaseTableStructureModel
+from docling.models.utils.hf_model_download import download_hf_model
+from docling.utils.accelerator_utils import decide_device
+from docling.utils.profiling import TimeRecorder
+
+
+class TableStructureModel(BaseTableStructureModel):
+    _model_repo_folder = "docling-project--docling-models"
+    _model_path = "model_artifacts/tableformer"
+
+    def __init__(
+        self,
+        enabled: bool,
+        artifacts_path: Optional[Path],
+        options: TableStructureOptions,
+        accelerator_options: AcceleratorOptions,
+    ):
+        self.options = options
+        self.do_cell_matching = self.options.do_cell_matching
+        self.mode = self.options.mode
+
+        self.enabled = enabled
+        if self.enabled:
+            if artifacts_path is None:
+                artifacts_path = self.download_models() / self._model_path
+            else:
+                # will become the default in the future
+                if (artifacts_path / self._model_repo_folder).exists():
+                    artifacts_path = (
+                        artifacts_path / self._model_repo_folder / self._model_path
+                    )
+                elif (artifacts_path / self._model_path).exists():
+                    warnings.warn(
+                        "The usage of artifacts_path containing directly "
+                        f"{self._model_path} is deprecated. Please point "
+                        "the artifacts_path to the parent containing "
+                        f"the {self._model_repo_folder} folder.",
+                        DeprecationWarning,
+                        stacklevel=3,
+                    )
+                    artifacts_path = artifacts_path / self._model_path
+
+            if self.mode == TableFormerMode.ACCURATE:
+                artifacts_path = artifacts_path / "accurate"
+            else:
+                artifacts_path = artifacts_path / "fast"
+
+            # Third Party
+            import docling_ibm_models.tableformer.common as c
+            from docling_ibm_models.tableformer.data_management.tf_predictor import (
+                TFPredictor,
+            )
+
+            device = decide_device(accelerator_options.device)
+
+            # Disable MPS here, until we know why it makes things slower.
+            if device == AcceleratorDevice.MPS.value:
+                device = AcceleratorDevice.CPU.value
+
+            self.tm_config = c.read_config(f"{artifacts_path}/tm_config.json")
+            self.tm_config["model"]["save_dir"] = artifacts_path
+            self.tm_model_type = self.tm_config["model"]["type"]
+
+            self.tf_predictor = TFPredictor(
+                self.tm_config, device, accelerator_options.num_threads
+            )
+            self.scale = 2.0  # Scale up table input images to 144 dpi
+
+    @classmethod
+    def get_options_type(cls) -> type[TableStructureOptions]:
+        return TableStructureOptions
+
+    @staticmethod
+    def download_models(
+        local_dir: Optional[Path] = None, force: bool = False, progress: bool = False
+    ) -> Path:
+        return download_hf_model(
+            repo_id="docling-project/docling-models",
+            revision="v2.3.0",
+            local_dir=local_dir,
+            force=force,
+            progress=progress,
+        )
+
+    def draw_table_and_cells(
+        self,
+        conv_res: ConversionResult,
+        page: Page,
+        tbl_list: Iterable[Table],
+        show: bool = False,
+    ):
+        assert page._backend is not None
+        assert page.size is not None
+
+        image = (
+            page._backend.get_page_image()
+        )  # make new image to avoid drawing on the saved ones
+
+        scale_x = image.width / page.size.width
+        scale_y = image.height / page.size.height
+
+        draw = ImageDraw.Draw(image)
+
+        for table_element in tbl_list:
+            x0, y0, x1, y1 = table_element.cluster.bbox.as_tuple()
+            y0 *= scale_y
+            y1 *= scale_y
+            x0 *= scale_x
+            x1 *= scale_x
+
+            draw.rectangle([(x0, y0), (x1, y1)], outline="red")
+
+            for cell in table_element.cluster.cells:
+                x0, y0, x1, y1 = cell.rect.to_bounding_box().as_tuple()
+                x0 *= scale_x
+                x1 *= scale_x
+                y0 *= scale_y
+                y1 *= scale_y
+
+                draw.rectangle([(x0, y0), (x1, y1)], outline="green")
+
+            for tc in table_element.table_cells:
+                if tc.bbox is not None:
+                    x0, y0, x1, y1 = tc.bbox.as_tuple()
+                    x0 *= scale_x
+                    x1 *= scale_x
+                    y0 *= scale_y
+                    y1 *= scale_y
+
+                    if tc.column_header:
+                        width = 3
+                    else:
+                        width = 1
+                    draw.rectangle([(x0, y0), (x1, y1)], outline="blue", width=width)
+                    draw.text(
+                        (x0 + 3, y0 + 3),
+                        text=f"{tc.start_row_offset_idx}, {tc.start_col_offset_idx}",
+                        fill="black",
+                    )
+        if show:
+            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"table_struct_page_{page.page_no:05}.png"
+            image.save(str(out_file), format="png")
+
+    def predict_tables(
+        self,
+        conv_res: ConversionResult,
+        pages: Sequence[Page],
+    ) -> Sequence[TableStructurePrediction]:
+        pages = list(pages)
+        predictions: list[TableStructurePrediction] = []
+
+        for page in pages:
+            assert page._backend is not None
+            if not page._backend.is_valid():
+                existing_prediction = (
+                    page.predictions.tablestructure or TableStructurePrediction()
+                )
+                page.predictions.tablestructure = existing_prediction
+                predictions.append(existing_prediction)
+                continue
+
+            with TimeRecorder(conv_res, "table_structure"):
+                assert page.predictions.layout is not None
+                assert page.size is not None
+
+                table_prediction = TableStructurePrediction()
+                page.predictions.tablestructure = table_prediction
+
+                in_tables = [
+                    (
+                        cluster,
+                        [
+                            round(cluster.bbox.l) * self.scale,
+                            round(cluster.bbox.t) * self.scale,
+                            round(cluster.bbox.r) * self.scale,
+                            round(cluster.bbox.b) * self.scale,
+                        ],
+                    )
+                    for cluster in page.predictions.layout.clusters
+                    if cluster.label
+                    in [DocItemLabel.TABLE, DocItemLabel.DOCUMENT_INDEX]
+                ]
+                if not in_tables:
+                    predictions.append(table_prediction)
+                    continue
+
+                page_input = {
+                    "width": page.size.width * self.scale,
+                    "height": page.size.height * self.scale,
+                    "image": numpy.asarray(page.get_image(scale=self.scale)),
+                }
+
+                for table_cluster, tbl_box in in_tables:
+                    # Check if word-level cells are available from backend:
+                    sp = page._backend.get_segmented_page()
+                    if sp is not None:
+                        tcells = sp.get_cells_in_bbox(
+                            cell_unit=TextCellUnit.WORD,
+                            bbox=table_cluster.bbox,
+                        )
+                        if len(tcells) == 0:
+                            # In case word-level cells yield empty
+                            tcells = table_cluster.cells
+                    else:
+                        # Otherwise - we use normal (line/phrase) cells
+                        tcells = table_cluster.cells
+                    tokens = []
+                    for c in tcells:
+                        # Only allow non empty strings (spaces) into the cells of a table
+                        if len(c.text.strip()) > 0:
+                            new_cell = copy.deepcopy(c)
+                            new_cell.rect = BoundingRectangle.from_bounding_box(
+                                new_cell.rect.to_bounding_box().scaled(scale=self.scale)
+                            )
+                            tokens.append(
+                                {
+                                    "id": new_cell.index,
+                                    "text": new_cell.text,
+                                    "bbox": new_cell.rect.to_bounding_box().model_dump(),
+                                }
+                            )
+                    page_input["tokens"] = tokens
+
+                    tf_output = self.tf_predictor.multi_table_predict(
+                        page_input, [tbl_box], do_matching=self.do_cell_matching
+                    )
+                    table_out = tf_output[0]
+                    table_cells = []
+                    for element in table_out["tf_responses"]:
+                        if not self.do_cell_matching:
+                            the_bbox = BoundingBox.model_validate(
+                                element["bbox"]
+                            ).scaled(1 / self.scale)
+                            text_piece = page._backend.get_text_in_rect(the_bbox)
+                            element["bbox"]["token"] = text_piece
+
+                        tc = TableCell.model_validate(element)
+                        if tc.bbox is not None:
+                            tc.bbox = tc.bbox.scaled(1 / self.scale)
+                        table_cells.append(tc)
+
+                    assert "predict_details" in table_out
+
+                    # Retrieving cols/rows, after post processing:
+                    num_rows = table_out["predict_details"].get("num_rows", 0)
+                    num_cols = table_out["predict_details"].get("num_cols", 0)
+                    otsl_seq = (
+                        table_out["predict_details"]
+                        .get("prediction", {})
+                        .get("rs_seq", [])
+                    )
+
+                    tbl = Table(
+                        otsl_seq=otsl_seq,
+                        table_cells=table_cells,
+                        num_rows=num_rows,
+                        num_cols=num_cols,
+                        id=table_cluster.id,
+                        page_no=page.page_no,
+                        cluster=table_cluster,
+                        label=table_cluster.label,
+                    )
+
+                    table_prediction.table_map[table_cluster.id] = tbl
+
+                if settings.debug.visualize_tables:
+                    self.draw_table_and_cells(
+                        conv_res,
+                        page,
+                        page.predictions.tablestructure.table_map.values(),
+                    )
+
+                predictions.append(table_prediction)
+
+        return predictions

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

@@ -0,0 +1,45 @@
+import logging
+from pathlib import Path
+from typing import Optional
+
+_log = logging.getLogger(__name__)
+
+
+def download_hf_model(
+    repo_id: str,
+    local_dir: Optional[Path] = None,
+    force: bool = False,
+    progress: bool = False,
+    revision: Optional[str] = None,
+) -> Path:
+    from huggingface_hub import snapshot_download
+    from huggingface_hub.utils import disable_progress_bars
+
+    if not progress:
+        disable_progress_bars()
+    download_path = snapshot_download(
+        repo_id=repo_id,
+        force_download=force,
+        local_dir=local_dir,
+        revision=revision,
+    )
+
+    return Path(download_path)
+
+
+class HuggingFaceModelDownloadMixin:
+    @staticmethod
+    def download_models(
+        repo_id: str,
+        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,
+            revision=revision,
+        )

docling/models/vlm_pipeline_models/__init__.py

@@ -0,0 +1 @@
+

docling/models/vlm_pipeline_models/api_vlm_model.py

@@ -0,0 +1,180 @@
+from collections.abc import Iterable
+from concurrent.futures import ThreadPoolExecutor
+from typing import Union
+
+import numpy as np
+from PIL.Image import Image
+
+from docling.datamodel.base_models import Page, VlmPrediction, VlmStopReason
+from docling.datamodel.document import ConversionResult
+from docling.datamodel.pipeline_options_vlm_model import ApiVlmOptions
+from docling.exceptions import OperationNotAllowed
+from docling.models.base_model import BaseVlmPageModel
+from docling.models.utils.generation_utils import GenerationStopper
+from docling.utils.api_image_request import (
+    api_image_request,
+    api_image_request_streaming,
+)
+from docling.utils.profiling import TimeRecorder
+
+
+class ApiVlmModel(BaseVlmPageModel):
+    # Override the vlm_options type annotation from BaseVlmPageModel
+    vlm_options: ApiVlmOptions  # type: ignore[assignment]
+
+    def __init__(
+        self,
+        enabled: bool,
+        enable_remote_services: bool,
+        vlm_options: ApiVlmOptions,
+    ):
+        self.enabled = enabled
+        self.vlm_options = vlm_options
+        if self.enabled:
+            if not enable_remote_services:
+                raise OperationNotAllowed(
+                    "Connections to remote services is only allowed when set explicitly. "
+                    "pipeline_options.enable_remote_services=True, or using the CLI "
+                    "--enable-remote-services."
+                )
+
+            self.timeout = self.vlm_options.timeout
+            self.concurrency = self.vlm_options.concurrency
+            self.params = {
+                **self.vlm_options.params,
+                "temperature": self.vlm_options.temperature,
+            }
+
+    def __call__(
+        self, conv_res: ConversionResult, page_batch: Iterable[Page]
+    ) -> Iterable[Page]:
+        page_list = list(page_batch)
+        if not page_list:
+            return
+
+        original_order = page_list[:]
+        valid_pages = []
+
+        for page in page_list:
+            assert page._backend is not None
+            if page._backend.is_valid():
+                valid_pages.append(page)
+
+        # Process valid pages in batch
+        if valid_pages:
+            with TimeRecorder(conv_res, "vlm"):
+                # Prepare images and prompts for batch processing
+                images = []
+                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)
+                        prompt = self._build_prompt_safe(page)
+                        prompts.append(prompt)
+                        pages_with_images.append(page)
+
+                # Use process_images for the actual inference
+                if images:  # Only if we have valid images
+                    with TimeRecorder(conv_res, "vlm_inference"):
+                        predictions = list(self.process_images(images, prompts))
+
+                    # Attach results to pages
+                    for page, prediction in zip(pages_with_images, predictions):
+                        page.predictions.vlm_response = prediction
+
+        # Yield pages preserving 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 raw images without page metadata."""
+        images = list(image_batch)
+
+        # Handle prompt parameter
+        if isinstance(prompt, str):
+            prompts = [prompt] * len(images)
+        elif isinstance(prompt, list):
+            if len(prompt) != len(images):
+                raise ValueError(
+                    f"Prompt list length ({len(prompt)}) must match image count ({len(images)})"
+                )
+            prompts = prompt
+
+        def _process_single_image(image_prompt_pair):
+            image, prompt_text = image_prompt_pair
+
+            # Convert numpy array to PIL Image if needed
+            if isinstance(image, np.ndarray):
+                if image.ndim == 3 and image.shape[2] in [3, 4]:
+                    from PIL import Image as PILImage
+
+                    image = PILImage.fromarray(image.astype(np.uint8))
+                elif image.ndim == 2:
+                    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
+            if image.mode != "RGB":
+                image = image.convert("RGB")
+
+            stop_reason = VlmStopReason.UNSPECIFIED
+
+            if self.vlm_options.custom_stopping_criteria:
+                # Instantiate any GenerationStopper classes before passing to streaming
+                instantiated_stoppers = []
+                for criteria in self.vlm_options.custom_stopping_criteria:
+                    if isinstance(criteria, GenerationStopper):
+                        instantiated_stoppers.append(criteria)
+                    elif isinstance(criteria, type) and issubclass(
+                        criteria, GenerationStopper
+                    ):
+                        instantiated_stoppers.append(criteria())
+                    # Skip non-GenerationStopper criteria (should have been caught in validation)
+
+                # Streaming path with early abort support
+                page_tags, num_tokens = api_image_request_streaming(
+                    image=image,
+                    prompt=prompt_text,
+                    url=self.vlm_options.url,
+                    timeout=self.timeout,
+                    headers=self.vlm_options.headers,
+                    generation_stoppers=instantiated_stoppers,
+                    **self.params,
+                )
+            else:
+                # Non-streaming fallback (existing behavior)
+                page_tags, num_tokens, stop_reason = api_image_request(
+                    image=image,
+                    prompt=prompt_text,
+                    url=self.vlm_options.url,
+                    timeout=self.timeout,
+                    headers=self.vlm_options.headers,
+                    **self.params,
+                )
+
+            page_tags = self.vlm_options.decode_response(page_tags)
+            input_prompt = prompt_text if self.vlm_options.track_input_prompt else None
+            return VlmPrediction(
+                text=page_tags,
+                num_tokens=num_tokens,
+                stop_reason=stop_reason,
+                input_prompt=input_prompt,
+            )
+
+        with ThreadPoolExecutor(max_workers=self.concurrency) as executor:
+            yield from executor.map(_process_single_image, zip(images, prompts))