Models refactoring, add stella, BGE, Linq

mteb/abstasks/AbsTaskBitextMining.py

@@ -105,6 +105,7 @@ def _evaluate_subset(
         evaluator = BitextMiningEvaluator(
             data_split,
             task_name=self.metadata.name,
+            task_type=self.metadata.type,
             pair_columns=pairs,  # type: ignore
             **kwargs,
         )

mteb/abstasks/AbsTaskClassification.py

@@ -149,6 +149,7 @@ def _evaluate_subset(
                     eval_split["text"],  # type: ignore
                     eval_split["label"],  # type: ignore
                     task_name=self.metadata.name,
+                    task_type=self.metadata.type,
                     encode_kwargs=encode_kwargs,
                     **params,
                 )
@@ -159,6 +160,7 @@ def _evaluate_subset(
                     eval_split["text"],  # type: ignore
                     eval_split["label"],  # type: ignore
                     task_name=self.metadata.name,
+                    task_type=self.metadata.type,
                     encode_kwargs=encode_kwargs,
                     **params,
                 )
@@ -169,6 +171,7 @@ def _evaluate_subset(
                     eval_split["text"],  # type: ignore
                     eval_split["label"],  # type: ignore
                     task_name=self.metadata.name,
+                    task_type=self.metadata.type,
                     encode_kwargs=encode_kwargs,
                     **params,
                 )

mteb/abstasks/AbsTaskClustering.py

@@ -64,6 +64,7 @@ def _evaluate_subset(
                 cluster_set["sentences"],  # type: ignore
                 cluster_set["labels"],  # type: ignore
                 task_name=self.metadata.name,
+                task_type=self.metadata.type,
                 **kwargs,
             )
             metrics = evaluator(model, encode_kwargs=encode_kwargs)

mteb/abstasks/AbsTaskClusteringFast.py

@@ -178,6 +178,7 @@ def _evaluate_subset(
             downsampled_dataset["sentences"],  # type: ignore
             model=model,
             prompt_name=self.metadata.name,
+            task_type=self.metadata.type,
             **encode_kwargs,
         )
 

mteb/abstasks/AbsTaskInstructionRetrieval.py

@@ -463,6 +463,7 @@ def evaluate(
         retriever = InstructionRetrievalEvaluator(
             retriever=model,
             task_name=self.metadata.name,
+            task_type=self.metadata.type,
             encode_kwargs=encode_kwargs,
             **kwargs,
         )

mteb/abstasks/AbsTaskMultilabelClassification.py

@@ -166,6 +166,7 @@ def _evaluate_subset(
             unique_train_sentences,
             model=model,
             prompt_name=self.metadata.name,
+            task_type=self.metadata.type,
             **encode_kwargs,
         )
         unique_train_embeddings = dict(
@@ -184,7 +185,11 @@ def _evaluate_subset(
             logger.warning("Couldn't subsample, continuing with the entire test set.")
 
         X_test = model_encode(
-            test_text, model=model, prompt_name=self.metadata.name, **encode_kwargs
+            test_text,
+            model=model,
+            prompt_name=self.metadata.name,
+            task_type=self.metadata.type,
+            **encode_kwargs,
         )
         for i_experiment, sample_indices in enumerate(train_samples):
             logger.info(

mteb/abstasks/AbsTaskPairClassification.py

@@ -65,6 +65,7 @@ def _evaluate_subset(
             data_split["sentence2"],
             data_split["labels"],
             task_name=self.metadata.name,
+            task_type=self.metadata.type,
             **kwargs,
         )
         scores = evaluator.compute_metrics(model, encode_kwargs=encode_kwargs)

mteb/abstasks/AbsTaskReranking.py

@@ -54,6 +54,7 @@ def _evaluate_subset(
         evaluator = RerankingEvaluator(
             data_split,
             task_name=self.metadata.name,
+            task_type=self.metadata.type,
             encode_kwargs=encode_kwargs,
             **kwargs,
         )

mteb/abstasks/AbsTaskRetrieval.py

@@ -274,6 +274,7 @@ def evaluate(
         retriever = RetrievalEvaluator(
             retriever=model,
             task_name=self.metadata.name,
+            task_type=self.metadata.type,
             encode_kwargs=encode_kwargs,
             **kwargs,
         )

mteb/abstasks/AbsTaskSTS.py

@@ -58,6 +58,7 @@ def normalize(x):
             data_split["sentence2"],
             normalized_scores,
             task_name=self.metadata.name,
+            task_type=self.metadata.type,
             **kwargs,
         )
         scores = evaluator(model, encode_kwargs=encode_kwargs)

mteb/abstasks/AbsTaskSummarization.py

@@ -68,6 +68,7 @@ def _evaluate_subset(
             texts=data_split["text"],
             gold_scores=normalized_scores,
             task_name=self.metadata.name,
+            task_type=self.metadata.type,
             **kwargs,
         )
         scores = evaluator(model, encode_kwargs=encode_kwargs)

mteb/evaluation/evaluators/BitextMiningEvaluator.py

@@ -25,6 +25,7 @@ def __init__(
         self,
         sentences: Dataset,
         task_name: str | None = None,
+        task_type: str | None = None,
         pair_columns: list[tuple[str, str]] = DEFAULT_PAIR,
         **kwargs,
     ):
@@ -38,6 +39,7 @@ def __init__(
             else sentences["gold"]
         )
         self.task_name = task_name
+        self.task_type = task_type
 
     def __call__(self, model: Encoder, *, encode_kwargs: dict[str, Any] = {}):
         scores = self.compute_metrics(model, encode_kwargs=encode_kwargs)
@@ -56,6 +58,7 @@ def compute_metrics(self, model: Encoder, encode_kwargs: dict[str, Any] = {}):
                 self.sentences[sub],
                 model=model,
                 prompt_name=self.task_name,
+                task_type=self.task_type,
                 **encode_kwargs,
             )
 

mteb/evaluation/evaluators/ClassificationEvaluator.py

@@ -6,11 +6,7 @@
 import numpy as np
 import torch
 from sklearn.linear_model import LogisticRegression
-from sklearn.metrics import (
-    accuracy_score,
-    average_precision_score,
-    f1_score,
-)
+from sklearn.metrics import accuracy_score, average_precision_score, f1_score
 from sklearn.neighbors import KNeighborsClassifier
 from torch import Tensor
 
@@ -34,6 +30,7 @@ def __init__(
         sentences_test,
         y_test,
         task_name: str | None = None,
+        task_type: str | None = None,
         k: int = 1,
         encode_kwargs: dict[str, Any] = {},
         limit: int | None = None,
@@ -51,6 +48,7 @@ def __init__(
         self.y_test = y_test
 
         self.task_name = task_name
+        self.task_type = task_type
         self.encode_kwargs = encode_kwargs
 
         if "batch_size" not in self.encode_kwargs:
@@ -67,13 +65,15 @@ def __call__(self, model, test_cache=None):
             self.sentences_train,
             model=model,
             prompt_name=self.task_name,
+            task_type=self.task_type,
             **self.encode_kwargs,
         )
         if test_cache is None:
             X_test = model_encode(
                 self.sentences_test,
                 model=model,
                 prompt_name=self.task_name,
+                task_type=self.task_type,
                 **self.encode_kwargs,
             )
             test_cache = X_test
@@ -109,6 +109,7 @@ def __init__(
         sentences_test,
         y_test,
         task_name: str,
+        task_type: str,
         k: int = 1,
         encode_kwargs: dict[str, Any] = {},
         limit: int | None = None,
@@ -127,6 +128,7 @@ def __init__(
         self.y_test = y_test
 
         self.task_name = task_name
+        self.task_type = task_type
         self.encode_kwargs = encode_kwargs
 
         if "batch_size" not in self.encode_kwargs:
@@ -143,6 +145,7 @@ def __call__(self, model: Encoder, test_cache=None):
             self.sentences_train,
             model=model,
             prompt_name=self.task_name,
+            task_type=self.task_type,
             **self.encode_kwargs,
         )
 
@@ -151,6 +154,7 @@ def __call__(self, model: Encoder, test_cache=None):
                 self.sentences_test,
                 model=model,
                 prompt_name=self.task_name,
+                task_type=self.task_type,
                 **self.encode_kwargs,
             )
             test_cache = X_test
@@ -261,6 +265,7 @@ def __init__(
         sentences_test,
         y_test,
         task_name: str,
+        task_type: str,
         max_iter: int = 100,
         encode_kwargs: dict[str, Any] = {},
         limit: int | None = None,
@@ -284,6 +289,7 @@ def __init__(
 
         self.max_iter = max_iter
         self.task_name = task_name
+        self.task_type = task_type
 
     def __call__(self, model, test_cache=None):
         scores = {}
@@ -297,13 +303,15 @@ def __call__(self, model, test_cache=None):
             self.sentences_train,
             model=model,
             prompt_name=self.task_name,
+            task_type=self.task_type,
             **self.encode_kwargs,
         )
         if test_cache is None:
             X_test = model_encode(
                 self.sentences_test,
                 model=model,
                 prompt_name=self.task_name,
+                task_type=self.task_type,
                 **self.encode_kwargs,
             )
             test_cache = X_test

mteb/evaluation/evaluators/ClusteringEvaluator.py

@@ -21,6 +21,7 @@ def __init__(
         sentences,
         labels,
         task_name: str | None = None,
+        task_type: str | None = None,
         clustering_batch_size: int = 500,
         limit: int | None = None,
         **kwargs,
@@ -33,6 +34,7 @@ def __init__(
         self.labels = labels
         self.clustering_batch_size = clustering_batch_size
         self.task_name = task_name
+        self.task_type = task_type
 
     def __call__(self, model: Encoder, *, encode_kwargs: dict[str, Any] = {}):
         if "batch_size" not in encode_kwargs:
@@ -42,6 +44,7 @@ def __call__(self, model: Encoder, *, encode_kwargs: dict[str, Any] = {}):
             self.sentences,
             model=model,
             prompt_name=self.task_name,
+            task_type=self.task_type,
             **encode_kwargs,
         )
 

mteb/evaluation/evaluators/InstructionRetrievalEvaluator.py

@@ -2,9 +2,7 @@
 
 import logging
 
-from .RetrievalEvaluator import (
-    RetrievalEvaluator,
-)
+from .RetrievalEvaluator import RetrievalEvaluator
 
 logger = logging.getLogger(__name__)
 
@@ -35,5 +33,6 @@ def __call__(
                 instructions=instructions,
                 request_qid=qid,
                 prompt_name=self.task_name,
+                task_type=self.task_type,
                 **kwargs,
             )

mteb/evaluation/evaluators/PairClassificationEvaluator.py

@@ -43,6 +43,7 @@ def __init__(
         sentences2,
         labels,
         task_name: str | None = None,
+        task_type: str | None = None,
         limit: int | None = None,
         **kwargs,
     ):
@@ -55,6 +56,7 @@ def __init__(
         self.sentences2 = sentences2
         self.labels = labels
         self.task_name = task_name
+        self.task_type = task_type
 
         assert len(self.sentences1) == len(self.sentences2)
         assert len(self.sentences1) == len(self.labels)
@@ -94,6 +96,7 @@ def compute_metrics(
             sentences,
             model=model,
             prompt_name=self.task_name,
+            task_type=self.task_type,
             **encode_kwargs,
         )
         emb_dict = dict(zip(sentences, embeddings))

mteb/evaluation/evaluators/RerankingEvaluator.py

@@ -34,6 +34,7 @@ def __init__(
         self,
         samples,
         task_name: str | None = None,
+        task_type: str | None = None,
         mrr_at_k: int = 10,
         name: str = "",
         similarity_fct=cos_sim,
@@ -53,6 +54,7 @@ def __init__(
         self.similarity_fct = similarity_fct
         self.use_batched_encoding = use_batched_encoding
         self.task_name = task_name
+        self.task_type = task_type
         self.k_values = k_values
         self.evaluator_type = evaluator_type
         self.encode_kwargs = encode_kwargs
@@ -104,6 +106,7 @@ def compute_metrics_batched(self, model: Encoder | EncoderWithQueryCorpusEncode)
                 encode_queries_func(
                     [sample["query"] for sample in self.samples],
                     prompt_name=self.task_name,
+                    task_type=self.task_type,
                     **self.encode_kwargs,
                 )
             )
@@ -116,6 +119,7 @@ def compute_metrics_batched(self, model: Encoder | EncoderWithQueryCorpusEncode)
                 all_query_flattened,
                 encode_queries_func,
                 prompt_name=self.task_name,
+                task_type=self.task_type,
                 **self.encode_kwargs,
             )
         else:
@@ -210,6 +214,7 @@ def _encode_candidates_batched(
             all_docs,
             encode_corpus_func,
             prompt_name=self.task_name,
+            task_type=self.task_type,
             **self.encode_kwargs,
         )
 
@@ -307,7 +312,10 @@ def _encode_candidates_miracl_batched(self, all_query_embs, encode_corpus_func):
 
         all_docs_embs = np.asarray(
             encode_corpus_func(
-                all_docs, prompt_name=self.task_name, **self.encode_kwargs
+                all_docs,
+                prompt_name=self.task_name,
+                task_type=self.task_type,
+                **self.encode_kwargs,
             )
         )
 
@@ -422,6 +430,7 @@ def _encode_unique_texts(
         all_texts: list[str],
         encode_fn: Callable,
         prompt_name: str | None,
+        task_type: str | None,
         **encode_kwargs: Any,
     ):
         index_map, all_unique_texts, all_texts_indexes = {}, [], []
@@ -435,7 +444,12 @@ def _encode_unique_texts(
             f"A total on {len(all_texts) - len(all_unique_texts)}/{len(all_texts)} duplicate texts were found during encoding. Only encoding unique text and duplicating embeddings across."
         )
         all_unique_texts_embs = np.asarray(
-            encode_fn(all_unique_texts, prompt_name=prompt_name, **encode_kwargs)
+            encode_fn(
+                all_unique_texts,
+                prompt_name=prompt_name,
+                task_type=task_type,
+                **encode_kwargs,
+            )
         )
         return all_unique_texts_embs[all_texts_indexes]