Update using last version of mteb benchmark.

bukosabino · Jun 12, 2024 · 7ffd34d · 7ffd34d
1 parent 50f8b48
commit 7ffd34d
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 6 deletions.
diff --git a/evaluation/embeddings_model/mteb_benchmark.py b/evaluation/embeddings_model/mteb_benchmark.py
@@ -1,4 +1,4 @@
-from mteb import MTEB
+import mteb
 from sentence_transformers import SentenceTransformer
 
 
@@ -16,10 +16,17 @@
 # model_name = "sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
 # model_name = "intfloat/multilingual-e5-small"
 # model_name = "intfloat/multilingual-e5-base"
-# model_name = "intfloat/multilingual-e5-large"
-model_name = "intfloat/multilingual-e5-large-instruct"
+model_name = "intfloat/multilingual-e5-large"
+# model_name = "intfloat/multilingual-e5-large-instruct"
 
+try:
+    model = SentenceTransformer(model_name, device='cuda')
+    print("Loaded model embedding using GPU")
+except:
+    model = SentenceTransformer(model_name, device='cpu')
+    print("Loaded model embedding using CPU")
 
-model = SentenceTransformer(model_name)
-evaluation = MTEB(task_langs=["es"])
+tasks = mteb.get_tasks(languages=["spa"])  # Spanish
+print(tasks)
+evaluation = mteb.MTEB(tasks=tasks)
 results = evaluation.run(model, output_folder=f"results/{model_name}")
diff --git a/requirements.txt b/requirements.txt
@@ -30,5 +30,5 @@ black==23.9.1
 isort==5.12.0
 
 # Evaluation
-mteb==1.1.1
+mteb==1.12.25
 # ragas==0.1.0rc1