update test download

Signed-off-by: Michele Dolfi <[email protected]>
DS4SD · Dec 11, 2024 · 9b3ae9b · 9b3ae9b
1 parent 2d0ae7b
commit 9b3ae9b
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/demo/demo_layout_predictor.py b/demo/demo_layout_predictor.py
@@ -96,7 +96,7 @@ def main(args):
     Path(viz_dir).mkdir(parents=True, exist_ok=True)
 
     # Download models from HF
-    download_path = snapshot_download(repo_id="ds4sd/docling-models")
+    download_path = snapshot_download(repo_id="ds4sd/docling-models", revision="v2.0.1")
     artifact_path = os.path.join(download_path, "model_artifacts/layout/beehive_v0.0.5_pt")
 
     # Test the LayoutPredictor

diff --git a/tests/test_layout_predictor.py b/tests/test_layout_predictor.py
@@ -39,7 +39,7 @@ def init() -> dict:
     }
 
     # Download models from HF
-    download_path = snapshot_download(repo_id="ds4sd/docling-models")
+    download_path = snapshot_download(repo_id="ds4sd/docling-models", revision="v2.0.1")
     artifact_path = os.path.join(download_path, "model_artifacts/layout/beehive_v0.0.5_pt")
 
     # Add the missing config keys

diff --git a/tests/test_tf_predictor.py b/tests/test_tf_predictor.py
@@ -468,7 +468,7 @@ def init() -> list[dict]:
     Initialize the testing environment
     """
     # Download models from HF
-    download_path = snapshot_download(repo_id="ds4sd/docling-models")
+    download_path = snapshot_download(repo_id="ds4sd/docling-models", revision="v2.0.1")
     save_dir = os.path.join(download_path, "model_artifacts/tableformer")
 
     # Add the missing config keys