fix: Use device_map for transformer models (#124)

Signed-off-by: Christoph Auer <[email protected]>

Author: cau-git

docling_ibm_models/code_formula_model/code_formula_predictor.py

@@ -89,8 +89,8 @@ def __init__(
             self._tokenizer = AutoTokenizer.from_pretrained(
                 artifacts_path, use_fast=True, padding_side="left"
             )
-            self._model = SamOPTForCausalLM.from_pretrained(artifacts_path).to(
-                self._device
+            self._model = SamOPTForCausalLM.from_pretrained(
+                artifacts_path, device_map=self._device
             )
             self._model.eval()
 

docling_ibm_models/document_figure_classifier_model/document_figure_classifier_predictor.py

@@ -90,8 +90,9 @@ def __init__(
             torch.set_num_threads(self._num_threads)
 
         with _model_init_lock:
-            model = AutoModelForImageClassification.from_pretrained(artifacts_path)
-            self._model = model.to(device)
+            self._model = AutoModelForImageClassification.from_pretrained(
+                artifacts_path, device_map=device
+            )
             self._model.eval()
 
             self._image_processor = transforms.Compose(

docling_ibm_models/layoutmodel/layout_predictor.py

@@ -84,8 +84,8 @@ def __init__(
         # Use lock to prevent threading issues during model initialization
         with _model_init_lock:
             self._model = AutoModelForObjectDetection.from_pretrained(
-                artifact_path, config=self._model_config
-            ).to(self._device)
+                artifact_path, config=self._model_config, device_map=self._device
+            )
             self._model.eval()
 
         # Set classes map

pyproject.toml

@@ -43,6 +43,7 @@ dependencies = [
   'transformers (>=4.42.0,<5.0.0)',
   'numpy (>=1.24.4,<3.0.0)',
   "rtree>=1.0.0",
+  'accelerate (>=1.2.1,<2.0.0)',
 ]
 
 [project.urls]