[v2] add similarity_fn in ModelMeta

mteb/evaluation/evaluators/model_classes.py

@@ -321,8 +321,11 @@ def _full_corpus_search(
             query_embeddings = torch.as_tensor(query_embeddings).to(device)
             sub_corpus_embeddings = torch.as_tensor(sub_corpus_embeddings).to(device)
 
+            # TODO: If the model is not DRESModel
             score_function = (
-                self.model.similarity if hasattr(self.model, "similarity") else cos_sim
+                self.model.similarity
+                if hasattr(self.model, "similarity")
+                else self.model.model.mteb_model_meta.get_evaluation_function()
             )
 
             with torch.inference_mode():

mteb/model_meta.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+from enum import Enum
 from functools import partial
 from typing import TYPE_CHECKING, Any, Callable, Literal
 
@@ -9,6 +10,7 @@
 from mteb.abstasks.AbsTask import AbsTask
 from mteb.abstasks.TaskMetadata import STR_DATE, STR_URL
 from mteb.encoder_interface import Encoder
+from mteb.evaluation.evaluators.utils import cos_sim, dot_score
 
 from .languages import ISO_LANGUAGE_SCRIPT
 
@@ -51,6 +53,11 @@ def get_loader_name(
     return loader.__name__
 
 
+class EvaluationFunction(str, Enum):
+    DOT_PRODUCT: str = "dot_score"
+    COSINE: str = "cos_sim"
+
+
 class ModelMeta(BaseModel):
     """The model metadata object.
 
@@ -101,13 +108,27 @@ class ModelMeta(BaseModel):
     public_training_code: bool | None = None
     framework: list[FRAMEWORKS] = []
     reference: STR_URL | None = None
-    similarity_fn_name: DISTANCE_METRICS | None = None
+    similarity_fn_name: DISTANCE_METRICS | EvaluationFunction | None = None
     use_instructions: bool | None = None
     training_datasets: dict[str, list[str]] | None = None
     adapted_from: str | None = None
     superseded_by: str | None = None
     citation: str | None = None
 
+    def get_evaluation_function(self) -> callable:
+        if (
+            self.similarity_fn_name == "cosine"
+            or self.similarity_fn_name == EvaluationFunction.COSINE
+        ):
+            return cos_sim
+        elif (
+            self.similarity_fn_name == "dot"
+            or self.similarity_fn_name == EvaluationFunction.DOT_PRODUCT
+        ):
+            return dot_score
+        else:
+            raise ValueError(f"Unknown similarity function {self.similarity_fn_name}")
+
     def to_dict(self):
         dict_repr = self.model_dump()
         loader = dict_repr.pop("loader", None)

mteb/models/e5_models.py

@@ -3,7 +3,7 @@
 from functools import partial
 
 from mteb.encoder_interface import PromptType
-from mteb.model_meta import ModelMeta, sentence_transformers_loader
+from mteb.model_meta import EvaluationFunction, ModelMeta, sentence_transformers_loader
 
 E5_PAPER_RELEASE_DATE = "2024-02-08"
 XLMR_LANGUAGES = [
@@ -149,7 +149,7 @@
     license="mit",
     max_tokens=512,
     reference="https://huggingface.co/intfloat/multilingual-e5-small",
-    similarity_fn_name="cosine",
+    similarity_fn_name=EvaluationFunction.COSINE,
     framework=["Sentence Transformers", "PyTorch"],
     use_instructions=True,
     citation=MULTILINGUAL_E5_CITATION,

mteb/models/sentence_transformer_wrapper.py

@@ -10,6 +10,7 @@
 
 from mteb.encoder_interface import PromptType
 
+from ..evaluation import dot_distance
 from .wrapper import Wrapper
 
 logger = logging.getLogger(__name__)
@@ -59,7 +60,7 @@ def __init__(
         if isinstance(self.model, CrossEncoder):
             self.predict = self._predict
 
-        if hasattr(self.model, "similarity"):
+        if hasattr(self.model, "similarity") and not hasattr(self, "similarity"):
             self.similarity = self.model.similarity
 
     def encode(
@@ -125,3 +126,8 @@ def _predict(
             convert_to_numpy=True,
             **kwargs,
         )
+
+
+class SentenceTransformerWrapperDotSimilarity(SentenceTransformerWrapper):
+    def similarity(self, embedding1: np.ndarray, embedding2: np.ndarray) -> float:
+        return dot_distance(embedding1, embedding2)

mteb/models/sentence_transformers_models.py

@@ -2,7 +2,12 @@
 
 from __future__ import annotations
 
+from functools import partial
+
 from mteb.model_meta import ModelMeta
+from mteb.models.sentence_transformer_wrapper import (
+    SentenceTransformerWrapperDotSimilarity,
+)
 
 paraphrase_langs = [
     "ara_Arab",
@@ -375,6 +380,11 @@
 )
 
 contriever = ModelMeta(
+    loader=partial(
+        SentenceTransformerWrapperDotSimilarity,
+        model="facebook/contriever-msmarco",
+        revision="abe8c1493371369031bcb1e02acb754cf4e162fa",
+    ),
     name="facebook/contriever-msmarco",
     languages=["eng-Latn"],
     open_weights=True,