fix: upgrade to code formula model v1.0.1 (#75)

Signed-off-by: Matteo Omenetti <[email protected]>
DS4SD · Feb 5, 2025 · d6a3549 · d6a3549
1 parent a47673b
commit d6a3549
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 2 deletions.
diff --git a/docling_ibm_models/code_formula_model/code_formula_predictor.py b/docling_ibm_models/code_formula_model/code_formula_predictor.py
@@ -207,6 +207,7 @@ def predict(
                 temperature=temperature,
                 max_new_tokens=4096 - prompt_ids.shape[1],
                 use_cache=True,
+                no_repeat_ngram_size=300,
             )
         else:
             with torch.autocast(device_type=self._device, dtype=torch.bfloat16):
@@ -217,6 +218,7 @@ def predict(
                     temperature=temperature,
                     max_new_tokens=4096 - prompt_ids.shape[1],
                     use_cache=True,
+                    no_repeat_ngram_size=300,
                 )
 
         outputs = self._tokenizer.batch_decode(

diff --git a/docling_ibm_models/code_formula_model/models/sam_opt_image_processor.py b/docling_ibm_models/code_formula_model/models/sam_opt_image_processor.py
@@ -28,4 +28,7 @@ def __call__(self, image):
         return image
 
 
-AutoImageProcessor.register(SamOptImageProcessor, SamOptImageProcessor)
+AutoImageProcessor.register(
+    config_class="SamOptImageProcessor",
+    slow_image_processor_class=SamOptImageProcessor,
+)
diff --git a/tests/test_code_formula_predictor.py b/tests/test_code_formula_predictor.py
@@ -37,7 +37,7 @@ def init() -> dict:
     }
 
     # Download models from HF
-    artifact_path = snapshot_download(repo_id="ds4sd/CodeFormula", revision="v1.0.0")
+    artifact_path = snapshot_download(repo_id="ds4sd/CodeFormula", revision="v1.0.1")
 
     init["artifact_path"] = artifact_path