add back large

embeddings-benchmark · Jun 7, 2024 · c64c144 · c64c144
1 parent 3e94983
commit c64c144
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 14 deletions.
diff --git a/scripts/mteb_english_cluster_spearman.py b/scripts/mteb_english_cluster_spearman.py
@@ -8,7 +8,7 @@
 from scipy import stats
 
 from mteb import get_model_meta
-from mteb.models.e5_models import e5_mult_base, e5_mult_small
+from mteb.models.e5_models import e5_mult_base, e5_mult_small, e5_mult_large
 from mteb.MTEBResults import MTEBResults
 
 logging.basicConfig(level=logging.INFO)
@@ -31,7 +31,7 @@
 
 TASK_LIST = [x + ".v2" for x in TASK_LIST_CLUSTERING] + TASK_LIST_CLUSTERING
 
-MODELS = [e5_mult_small, e5_mult_base, e5_mult_small]
+MODELS = [e5_mult_small, e5_mult_base, e5_mult_large]
 
 for model in MODELS:
     model_name = model.name

diff --git a/scripts/run_mteb_english_cluster.py b/scripts/run_mteb_english_cluster.py
@@ -5,6 +5,7 @@
 import logging
 
 from mteb import MTEB, get_model
+from mteb.models.e5_models import e5_mult_base, e5_mult_small, e5_mult_large
 
 logging.basicConfig(level=logging.INFO)
 
@@ -26,15 +27,22 @@
 
 TASK_LIST = [x + ".v2" for x in TASK_LIST_CLUSTERING] + TASK_LIST_CLUSTERING
 
-model = get_model(
-    model_name="intfloat/multilingual-e5-base",
-    revision="d13f1b27baf31030b7fd040960d60d909913633f",
-)
-
-eval_splits = ["test"]
-evaluation = MTEB(
-    tasks=TASK_LIST_CLUSTERING, task_langs=["en"]
-)  # Remove "en" for running all languages
-evaluation.run(
-    model, output_folder="results", eval_splits=eval_splits, overwrite_results=True
-)
+MODELS = [
+    e5_mult_small, 
+    # e5_mult_base, 
+    e5_mult_large,
+]
+
+for model in MODELS:
+    model_name = model.name
+    revision = model.revision
+
+    model = get_model(model_name=model_name, revision=revision)
+
+    eval_splits = ["test"]
+    evaluation = MTEB(
+        tasks=TASK_LIST_CLUSTERING, task_langs=["en"]
+    )  # Remove "en" for running all languages
+    evaluation.run(
+        model, output_folder="results", eval_splits=eval_splits, overwrite_results=True
+    )