fix tests

StarlightSearch · Nov 9, 2024 · 4eb10b5 · 4eb10b5
1 parent cd84765
commit 4eb10b5
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 7 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -20,7 +20,7 @@ pdf-extract = "0.7.7"
 candle-nn = { git = "https://github.com/huggingface/candle.git", version = "0.7.2" }
 candle-transformers = { git = "https://github.com/huggingface/candle.git", version = "0.7.2" }
 candle-core = { git = "https://github.com/huggingface/candle.git", version = "0.7.2" }
-ort = {version = "=2.0.0-rc.6", features = ["cuda"]}
+ort = {version = "=2.0.0-rc.8", features = ["cuda"]}
 strum = "0.26.1"
 strum_macros = "0.26"
 

diff --git a/tests/model_tests/conftest.py b/tests/model_tests/conftest.py
@@ -9,6 +9,8 @@
     ColpaliModel
 )
 
+from embed_anything import ONNXModel
+
 
 @pytest.fixture
 def clip_model() -> EmbeddingModel:
@@ -88,7 +90,7 @@ def openai_model() -> EmbeddingModel:
 
 @pytest.fixture
 def onnx_model() -> EmbeddingModel: 
-    model = EmbeddingModel.from_pretrained_onnx(WhichModel.Bert, "BGESmallENV15Q")
+    model = EmbeddingModel.from_pretrained_onnx(WhichModel.Bert, ONNXModel.AllMiniLML6V2Q)
     return model
 
 @pytest.fixture

diff --git a/tests/model_tests/test_bert.py b/tests/model_tests/test_bert.py
@@ -5,6 +5,7 @@
     embed_query,
     embed_file,
     embed_directory,
+    ONNXModel
 )
 
 import os
@@ -45,7 +46,7 @@ def test_bert_model_creation():
     assert model is not None
 
 def test_onnx_model_creation():
-    model = EmbeddingModel.from_pretrained_onnx(WhichModel.Bert, "BGESmallENV15Q")
+    model = EmbeddingModel.from_pretrained_onnx(WhichModel.Bert, ONNXModel.AllMiniLML6V2Q)
     assert model is not None
 
 @model_fixture_parametrize