remove hf_vlm_model and add extra_generation_args

Signed-off-by: Michele Dolfi <dol@zurich.ibm.com>

docling/datamodel/pipeline_options_vlm_model.py

@@ -46,7 +46,8 @@ class InlineVlmOptions(BaseVlmOptions):
     scale: float = 2.0
 
     temperature: float = 0.0
-    stop_strings: list[str] = []
+    stop_strings: List[str] = []
+    extra_generation_config: Dict[str, Any] = {}
 
     use_kv_cache: bool = True
     max_new_tokens: int = 4096

docling/datamodel/vlm_model_specs.py

@@ -97,6 +97,7 @@ PHI4_TRANSFORMERS = InlineVlmOptions(
     supported_devices=[AcceleratorDevice.CPU, AcceleratorDevice.CUDA],
     scale=2.0,
     temperature=0.0,
+    extra_generation_config=dict(num_logits_to_keep=0),
 )
 
 # Qwen

docling/models/code_formula_model.py

@@ -19,6 +19,7 @@ from pydantic import BaseModel
 from docling.datamodel.accelerator_options import AcceleratorOptions
 from docling.datamodel.base_models import ItemAndImageEnrichmentElement
 from docling.models.base_model import BaseItemAndImageEnrichmentModel
+from docling.models.utils.hf_model_download import download_hf_model
 from docling.utils.accelerator_utils import decide_device
 
 
@@ -117,20 +118,14 @@ class CodeFormulaModel(BaseItemAndImageEnrichmentModel):
         force: bool = False,
         progress: bool = False,
     ) -> Path:
-        from huggingface_hub import snapshot_download
+        return download_hf_model(
-        from huggingface_hub.utils import disable_progress_bars
-
-        if not progress:
-            disable_progress_bars()
-        download_path = snapshot_download(
             repo_id="ds4sd/CodeFormula",
-            force_download=force,
-            local_dir=local_dir,
             revision="v1.0.2",
+            local_dir=local_dir,
+            force=force,
+            progress=progress,
         )
 
-        return Path(download_path)
-
     def is_processable(self, doc: DoclingDocument, element: NodeItem) -> bool:
         """
         Determines if a given element in a document can be processed by the model.

docling/models/document_picture_classifier.py

@@ -15,6 +15,7 @@ from pydantic import BaseModel
 
 from docling.datamodel.accelerator_options import AcceleratorOptions
 from docling.models.base_model import BaseEnrichmentModel
+from docling.models.utils.hf_model_download import download_hf_model
 from docling.utils.accelerator_utils import decide_device
 
 
@@ -105,20 +106,14 @@ class DocumentPictureClassifier(BaseEnrichmentModel):
     def download_models(
         local_dir: Optional[Path] = None, force: bool = False, progress: bool = False
     ) -> Path:
-        from huggingface_hub import snapshot_download
+        return download_hf_model(
-        from huggingface_hub.utils import disable_progress_bars
-
-        if not progress:
-            disable_progress_bars()
-        download_path = snapshot_download(
             repo_id="ds4sd/DocumentFigureClassifier",
-            force_download=force,
-            local_dir=local_dir,
             revision="v1.0.1",
+            local_dir=local_dir,
+            force=force,
+            progress=progress,
         )
 
-        return Path(download_path)
-
     def is_processable(self, doc: DoclingDocument, element: NodeItem) -> bool:
         """
         Determines if the given element can be processed by the classifier.

docling/models/hf_vlm_model.py

@@ -1,28 +0,0 @@
-import logging
-from pathlib import Path
-from typing import Optional
-
-_log = logging.getLogger(__name__)
-
-
-class HuggingFaceVlmModel:
-    @staticmethod
-    def download_models(
-        repo_id: str,
-        local_dir: Optional[Path] = None,
-        force: bool = False,
-        progress: bool = False,
-    ) -> 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="v0.0.1",
-        )
-
-        return Path(download_path)

docling/models/layout_model.py

@@ -15,6 +15,7 @@ from docling.datamodel.base_models import BoundingBox, Cluster, LayoutPrediction
 from docling.datamodel.document import ConversionResult
 from docling.datamodel.settings import settings
 from docling.models.base_model import BasePageModel
+from docling.models.utils.hf_model_download import download_hf_model
 from docling.utils.accelerator_utils import decide_device
 from docling.utils.layout_postprocessor import LayoutPostprocessor
 from docling.utils.profiling import TimeRecorder
@@ -83,20 +84,14 @@ class LayoutModel(BasePageModel):
         force: bool = False,
         progress: bool = False,
     ) -> Path:
-        from huggingface_hub import snapshot_download
+        return download_hf_model(
-        from huggingface_hub.utils import disable_progress_bars
-
-        if not progress:
-            disable_progress_bars()
-        download_path = snapshot_download(
             repo_id="ds4sd/docling-models",
-            force_download=force,
+            revision="v2.2.0",
             local_dir=local_dir,
-            revision="v2.1.0",
+            force=force,
+            progress=progress,
         )
 
-        return Path(download_path)
-
     def draw_clusters_and_cells_side_by_side(
         self, conv_res, page, clusters, mode_prefix: str, show: bool = False
     ):

docling/models/picture_description_vlm_model.py

@@ -10,10 +10,15 @@ from docling.datamodel.pipeline_options import (
     PictureDescriptionVlmOptions,
 )
 from docling.models.picture_description_base_model import PictureDescriptionBaseModel
+from docling.models.utils.hf_model_download import (
+    HuggingFaceModelDownloadMixin,
+)
 from docling.utils.accelerator_utils import decide_device
 
 
-class PictureDescriptionVlmModel(PictureDescriptionBaseModel):
+class PictureDescriptionVlmModel(
+    PictureDescriptionBaseModel, HuggingFaceModelDownloadMixin
+):
     @classmethod
     def get_options_type(cls) -> Type[PictureDescriptionBaseOptions]:
         return PictureDescriptionVlmOptions
@@ -66,26 +71,6 @@ class PictureDescriptionVlmModel(PictureDescriptionBaseModel):
 
             self.provenance = f"{self.options.repo_id}"
 
-    @staticmethod
-    def download_models(
-        repo_id: str,
-        local_dir: Optional[Path] = None,
-        force: bool = False,
-        progress: bool = False,
-    ) -> 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,
-        )
-
-        return Path(download_path)
-
     def _annotate_images(self, images: Iterable[Image.Image]) -> Iterable[str]:
         from transformers import GenerationConfig
 

docling/models/table_structure_model.py

@@ -22,6 +22,7 @@ from docling.datamodel.pipeline_options import (
 )
 from docling.datamodel.settings import settings
 from docling.models.base_model import BasePageModel
+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
 
@@ -89,20 +90,14 @@ class TableStructureModel(BasePageModel):
     def download_models(
         local_dir: Optional[Path] = None, force: bool = False, progress: bool = False
     ) -> Path:
-        from huggingface_hub import snapshot_download
+        return download_hf_model(
-        from huggingface_hub.utils import disable_progress_bars
-
-        if not progress:
-            disable_progress_bars()
-        download_path = snapshot_download(
             repo_id="ds4sd/docling-models",
-            force_download=force,
-            local_dir=local_dir,
             revision="v2.2.0",
+            local_dir=local_dir,
+            force=force,
+            progress=progress,
         )
 
-        return Path(download_path)
-
     def draw_table_and_cells(
         self,
         conv_res: ConversionResult,

docling/models/utils/hf_model_download.py

@@ -0,0 +1,40 @@
+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,
+    ) -> Path:
+        return download_hf_model(
+            repo_id=repo_id, local_dir=local_dir, force=force, progress=progress
+        )

docling/models/vlm_models_inline/hf_transformers_causallm_model.py

@@ -12,14 +12,18 @@ from docling.datamodel.base_models import Page, VlmPrediction
 from docling.datamodel.document import ConversionResult
 from docling.datamodel.pipeline_options_vlm_model import InlineVlmOptions
 from docling.models.base_model import BasePageModel
-from docling.models.hf_vlm_model import HuggingFaceVlmModel
+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 HuggingFaceVlmModel_AutoModelForCausalLM(BasePageModel):
+class HuggingFaceVlmModel_AutoModelForCausalLM(
+    BasePageModel, HuggingFaceModelDownloadMixin
+):
     def __init__(
         self,
         enabled: bool,
@@ -62,9 +66,7 @@ class HuggingFaceVlmModel_AutoModelForCausalLM(BasePageModel):
             repo_cache_folder = vlm_options.repo_id.replace("/", "--")
 
             if artifacts_path is None:
-                artifacts_path = HuggingFaceVlmModel.download_models(
+                artifacts_path = self.download_models(self.vlm_options.repo_id)
-                    self.vlm_options.repo_id
-                )
             elif (artifacts_path / repo_cache_folder).exists():
                 artifacts_path = artifacts_path / repo_cache_folder
 
@@ -128,7 +130,7 @@ class HuggingFaceVlmModel_AutoModelForCausalLM(BasePageModel):
                         use_cache=self.use_cache,  # Enables KV caching which can improve performance
                         temperature=self.temperature,
                         generation_config=self.generation_config,
-                        num_logits_to_keep=1,
+                        **self.vlm_options.extra_generation_config,
                     )
                     generate_ids = generate_ids[:, inputs["input_ids"].shape[1] :]
 

docling/models/vlm_models_inline/hf_transformers_vision2seq_model.py

@@ -11,14 +11,18 @@ from docling.datamodel.base_models import Page, VlmPrediction
 from docling.datamodel.document import ConversionResult
 from docling.datamodel.pipeline_options_vlm_model import InlineVlmOptions
 from docling.models.base_model import BasePageModel
-from docling.models.hf_vlm_model import HuggingFaceVlmModel
+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 HuggingFaceVlmModel_AutoModelForVision2Seq(BasePageModel):
+class HuggingFaceVlmModel_AutoModelForVision2Seq(
+    BasePageModel, HuggingFaceModelDownloadMixin
+):
     def __init__(
         self,
         enabled: bool,
@@ -52,10 +56,7 @@ class HuggingFaceVlmModel_AutoModelForVision2Seq(BasePageModel):
 
             # PARAMETERS:
             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)
-                artifacts_path = HuggingFaceVlmModel.download_models(
-                    self.vlm_options.repo_id
-                )
             elif (artifacts_path / repo_cache_folder).exists():
                 artifacts_path = artifacts_path / repo_cache_folder
 

docling/models/vlm_models_inline/mlx_model.py

@@ -11,13 +11,15 @@ 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 BasePageModel
-from docling.models.hf_vlm_model import HuggingFaceVlmModel
+from docling.models.utils.hf_model_download import (
+    HuggingFaceModelDownloadMixin,
+)
 from docling.utils.profiling import TimeRecorder
 
 _log = logging.getLogger(__name__)
 
 
-class HuggingFaceMlxModel(BasePageModel):
+class HuggingFaceMlxModel(BasePageModel, HuggingFaceModelDownloadMixin):
     def __init__(
         self,
         enabled: bool,
@@ -48,12 +50,8 @@ class HuggingFaceMlxModel(BasePageModel):
 
             # PARAMETERS:
             if artifacts_path is None:
-                _log.debug(
+                artifacts_path = self.download_models(
-                    f"before HuggingFaceVlmModel.download_models: {self.vlm_options.repo_id}"
-                )
-                artifacts_path = HuggingFaceVlmModel.download_models(
                     self.vlm_options.repo_id,
-                    progress=True,
                 )
             elif (artifacts_path / repo_cache_folder).exists():
                 artifacts_path = artifacts_path / repo_cache_folder

docling/utils/model_downloader.py

@@ -14,10 +14,10 @@ from docling.datamodel.vlm_model_specs import (
 from docling.models.code_formula_model import CodeFormulaModel
 from docling.models.document_picture_classifier import DocumentPictureClassifier
 from docling.models.easyocr_model import EasyOcrModel
-from docling.models.hf_vlm_model import HuggingFaceVlmModel
 from docling.models.layout_model import LayoutModel
 from docling.models.picture_description_vlm_model import PictureDescriptionVlmModel
 from docling.models.table_structure_model import TableStructureModel
+from docling.models.utils.hf_model_download import download_hf_model
 
 _log = logging.getLogger(__name__)
 
@@ -77,7 +77,7 @@ def download_models(
 
     if with_smolvlm:
         _log.info("Downloading SmolVlm model...")
-        PictureDescriptionVlmModel.download_models(
+        download_hf_model(
             repo_id=smolvlm_picture_description.repo_id,
             local_dir=output_dir / smolvlm_picture_description.repo_cache_folder,
             force=force,
@@ -86,7 +86,7 @@ def download_models(
 
     if with_smoldocling:
         _log.info("Downloading SmolDocling model...")
-        HuggingFaceVlmModel.download_models(
+        download_hf_model(
             repo_id=SMOLDOCLING_TRANSFORMERS.repo_id,
             local_dir=output_dir / SMOLDOCLING_TRANSFORMERS.repo_cache_folder,
             force=force,
@@ -95,7 +95,7 @@ def download_models(
 
     if with_smoldocling_mlx:
         _log.info("Downloading SmolDocling MLX model...")
-        HuggingFaceVlmModel.download_models(
+        download_hf_model(
             repo_id=SMOLDOCLING_MLX.repo_id,
             local_dir=output_dir / SMOLDOCLING_MLX.repo_cache_folder,
             force=force,
@@ -104,7 +104,7 @@ def download_models(
 
     if with_granite_vision:
         _log.info("Downloading Granite Vision model...")
-        PictureDescriptionVlmModel.download_models(
+        download_hf_model(
             repo_id=granite_picture_description.repo_id,
             local_dir=output_dir / granite_picture_description.repo_cache_folder,
             force=force,