[v2] fix contriever (add similarity_fn_name to ST wrapper)

mteb/models/sentence_transformer_wrapper.py

@@ -21,6 +21,7 @@ def __init__(
         model: str | SentenceTransformer | CrossEncoder,
         revision: str | None = None,
         model_prompts: dict[str, str] | None = None,
+        similarity_fn_name: str | None = None,
         **kwargs,
     ) -> None:
         """Wrapper for SentenceTransformer models.
@@ -32,6 +33,7 @@ def __init__(
                 First priority is given to the composed prompt of task name + prompt type (query or passage), then to the specific task prompt,
                 then to the composed prompt of task type + prompt type, then to the specific task type prompt,
                 and finally to the specific prompt type.
+            similarity_fn_name: A similarity function to use.
             **kwargs: Additional arguments to pass to the SentenceTransformer model.
         """
         if isinstance(model, str):
@@ -59,7 +61,9 @@ def __init__(
         if isinstance(self.model, CrossEncoder):
             self.predict = self._predict
 
-        if hasattr(self.model, "similarity"):
+        if similarity_fn_name:
+            self.similarity = self.get_similarity_function(similarity_fn_name)
+        elif hasattr(self.model, "similarity"):
             self.similarity = self.model.similarity
 
     def encode(

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 (
+    SentenceTransformerWrapper,
+)
 
 paraphrase_langs = [
     "ara_Arab",
@@ -375,6 +380,12 @@
 )
 
 contriever = ModelMeta(
+    loader=partial(
+        SentenceTransformerWrapper,
+        model="facebook/contriever-msmarco",
+        revision="abe8c1493371369031bcb1e02acb754cf4e162fa",
+        similarity_fn_name="dot",
+    ),
     name="facebook/contriever-msmarco",
     languages=["eng-Latn"],
     open_weights=True,

mteb/models/wrapper.py

@@ -3,9 +3,12 @@
 import logging
 from typing import Callable, get_args
 
+import numpy as np
+
 import mteb
 from mteb.abstasks.TaskMetadata import TASK_TYPE
 from mteb.encoder_interface import PromptType
+from mteb.evaluation.evaluators.utils import cos_sim, dot_score
 
 logger = logging.getLogger(__name__)
 
@@ -64,6 +67,18 @@ def get_prompt_name(
         )
         return None
 
+    @staticmethod
+    def get_similarity_function(
+        similarity_fn_name: str,
+    ) -> Callable[[np.ndarray, np.ndarray], np.ndarray]:
+        if similarity_fn_name == "cosine":
+            return cos_sim
+        if similarity_fn_name == "dot":
+            return dot_score
+        raise ValueError(
+            "Invalid similarity function. Should be one of ['cosine', 'dot']"
+        )
+
     @staticmethod
     def validate_task_to_prompt_name(
         task_to_prompt_name: dict[str, str] | None,