From 4eb10b5adf7cb0b75bf9f852d894e37924b8af40 Mon Sep 17 00:00:00 2001
From: Akshay Ballal <arballal95@gmail.com>
Date: Sat, 9 Nov 2024 14:58:56 +0100
Subject: [PATCH] fix tests

---
 Cargo.lock                     | 8 ++++----
 Cargo.toml                     | 2 +-
 tests/model_tests/conftest.py  | 4 +++-
 tests/model_tests/test_bert.py | 3 ++-
 4 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index e585397..dde5e63 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2689,9 +2689,9 @@ checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d"
 
 [[package]]
 name = "ort"
-version = "2.0.0-rc.6"
+version = "2.0.0-rc.8"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a5f95fe501e1cb81dec2f66ee3129025759b602817aa2c77ff421390c418cc34"
+checksum = "11826e6118cc42fea0cb2b102f7d006c1bb339cb167f8badb5fb568616438234"
 dependencies = [
  "half",
  "libloading",
@@ -2702,9 +2702,9 @@ dependencies = [
 
 [[package]]
 name = "ort-sys"
-version = "2.0.0-rc.6"
+version = "2.0.0-rc.8"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b4174960a7b93a17564a05b26e05889f0dea9ee70e68db5841f27b40c0c9804e"
+checksum = "c4780a8b8681e653b2bed85c7f0e2c6e8547224c3e983e5ad27bf0457e012407"
 dependencies = [
  "flate2",
  "pkg-config",
diff --git a/Cargo.toml b/Cargo.toml
index 8f96b8d..2c195f2 100644
--- 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
index 588af64..16392de 100644
--- 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
index 6e808d4..7626be5 100644
--- 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