From 9b3ae9ba395900645a03523a15a1ff1d7cecb501 Mon Sep 17 00:00:00 2001 From: Michele Dolfi Date: Wed, 11 Dec 2024 09:46:14 +0100 Subject: [PATCH] update test download Signed-off-by: Michele Dolfi --- demo/demo_layout_predictor.py | 2 +- tests/test_layout_predictor.py | 2 +- tests/test_tf_predictor.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/demo/demo_layout_predictor.py b/demo/demo_layout_predictor.py index 506125b..e8a567c 100644 --- 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 index d8a35a1..5045ae6 100644 --- 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 index 8a0e1c2..daee028 100644 --- 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