mteb 2.1.4__py3-none-any.whl → 2.7.2__py3-none-any.whl

mteb/models/search_wrappers.py

@@ -14,6 +14,7 @@ from mteb.types import (
     Array,
     BatchedInput,
     CorpusDatasetType,
+    EncodeKwargs,
     PromptType,
     QueryDatasetType,
     RetrievalOutputType,
@@ -21,6 +22,7 @@ from mteb.types import (
 )
 
 from .models_protocols import CrossEncoderProtocol, EncoderProtocol
+from .search_encoder_index.search_backend_protocol import IndexEncoderSearchProtocol
 
 logger = logging.getLogger(__name__)
 
@@ -28,13 +30,19 @@ logger = logging.getLogger(__name__)
 class SearchEncoderWrapper:
     """Wrapper for Encoder models to be used in search tasks."""
 
-    corpus_chunk_size = 50_000
     task_corpus: CorpusDatasetType | None
 
-    def __init__(self, model: EncoderProtocol):
+    def __init__(
+        self,
+        model: EncoderProtocol,
+        corpus_chunk_size: int = 50_000,
+        index_backend: IndexEncoderSearchProtocol | None = None,
+    ) -> None:
         self.model = model
         self.task_corpus = None
         self.mteb_model_meta = model.mteb_model_meta
+        self.corpus_chunk_size = corpus_chunk_size
+        self.index_backend = index_backend
 
     def index(
         self,
@@ -43,7 +51,7 @@ class SearchEncoderWrapper:
         task_metadata: TaskMetadata,
         hf_split: str,
         hf_subset: str,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
     ) -> None:
         """Index the corpus for retrieval.
 
@@ -56,6 +64,22 @@ class SearchEncoderWrapper:
         """
         # Always retain corpus for potential reranking or fallback flows
         self.task_corpus = corpus
+        if self.index_backend is not None:
+            all_doc_embeddings = self.model.encode(
+                create_dataloader(
+                    corpus,
+                    task_metadata,
+                    prompt_type=PromptType.document,
+                    **encode_kwargs,
+                ),
+                task_metadata=task_metadata,
+                hf_split=hf_split,
+                hf_subset=hf_subset,
+                prompt_type=PromptType.document,
+                **encode_kwargs,
+            )
+
+            self.index_backend.add_documents(all_doc_embeddings, corpus["id"])
 
     def search(
         self,
@@ -65,7 +89,7 @@ class SearchEncoderWrapper:
         hf_split: str,
         hf_subset: str,
         top_k: int,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         top_ranked: TopRankedDocumentsType | None = None,
     ) -> RetrievalOutputType:
         """Search the corpus for the given queries.
@@ -90,7 +114,7 @@ class SearchEncoderWrapper:
             queries,
             task_metadata,
             prompt_type=PromptType.query,
-            batch_size=encode_kwargs.get("batch_size", 32),
+            **encode_kwargs,
         )
 
         query_embeddings = self.model.encode(
@@ -105,32 +129,79 @@ class SearchEncoderWrapper:
 
         if top_ranked is not None:
             logger.info("Reranking pre-ranked documents...")
-            result_heaps = self._rerank_documents(
-                query_idx_to_id=query_idx_to_id,
-                query_embeddings=query_embeddings,
-                top_ranked=top_ranked,
-                top_k=top_k,
-                task_metadata=task_metadata,
-                hf_subset=hf_subset,
-                hf_split=hf_split,
-                encode_kwargs=encode_kwargs,
-            )
+            if self.index_backend is None:
+                result_heaps = self._rerank_documents(
+                    query_idx_to_id=query_idx_to_id,
+                    query_embeddings=query_embeddings,
+                    top_ranked=top_ranked,
+                    top_k=top_k,
+                    task_metadata=task_metadata,
+                    hf_subset=hf_subset,
+                    hf_split=hf_split,
+                    encode_kwargs=encode_kwargs,
+                )
+            else:
+                cos_scores_top_k_values, cos_scores_top_k_idx = (
+                    self.index_backend.search(
+                        query_embeddings,
+                        top_k,
+                        similarity_fn=self.model.similarity,
+                        top_ranked=top_ranked,
+                        query_idx_to_id=query_idx_to_id,
+                    )
+                )
+                result_heaps = {qid: [] for qid in query_idx_to_id.values()}
+                for query_itr in range(len(query_embeddings)):
+                    result_heaps = self._rerank_sort_results(
+                        result_heaps=result_heaps,
+                        query_id=query_idx_to_id[query_itr],
+                        ranked_ids=top_ranked[query_idx_to_id[query_itr]],
+                        scores_top_k_idx=torch.tensor(
+                            [cos_scores_top_k_idx[query_itr]]
+                        ),
+                        scores_top_k_values=torch.tensor(
+                            [cos_scores_top_k_values[query_itr]]
+                        ),
+                    )
+                self.index_backend.clear()
         else:
             logger.info("Performing full corpus search...")
-            result_heaps = self._full_corpus_search(
-                query_idx_to_id=query_idx_to_id,
-                query_embeddings=query_embeddings,
-                task_metadata=task_metadata,
-                hf_subset=hf_subset,
-                hf_split=hf_split,
-                top_k=top_k,
-                encode_kwargs=encode_kwargs,
-            )
+            if self.index_backend is None:
+                result_heaps = self._full_corpus_search(
+                    query_idx_to_id=query_idx_to_id,
+                    query_embeddings=query_embeddings,
+                    task_metadata=task_metadata,
+                    hf_subset=hf_subset,
+                    hf_split=hf_split,
+                    top_k=top_k,
+                    encode_kwargs=encode_kwargs,
+                )
+            else:
+                cos_scores_top_k_values, cos_scores_top_k_idx = (
+                    self.index_backend.search(
+                        query_embeddings,
+                        top_k,
+                        similarity_fn=self.model.similarity,
+                        top_ranked=None,
+                        query_idx_to_id=None,
+                    )
+                )
+                result_heaps = {qid: [] for qid in query_idx_to_id.values()}
+                result_heaps = self._sort_full_corpus_results(
+                    result_heaps=result_heaps,
+                    query_idx_to_id=query_idx_to_id,
+                    query_embeddings=query_embeddings,
+                    cos_scores_top_k_idx=cos_scores_top_k_idx,
+                    cos_scores_top_k_values=cos_scores_top_k_values,
+                    sub_corpus_ids=self.task_corpus["id"],
+                    top_k=top_k,
+                )
+                self.index_backend.clear()
 
         # Reset the task corpus dataloader to None to free up memory
         self.task_corpus = None
 
-        results = {qid: {} for qid in query_idx_to_id.values()}
+        results: RetrievalOutputType = {qid: {} for qid in query_idx_to_id.values()}
         for qid in result_heaps:
             for score, corpus_id in result_heaps[qid]:
                 results[qid][corpus_id] = score
@@ -145,16 +216,22 @@ class SearchEncoderWrapper:
         hf_subset: str,
         hf_split: str,
         top_k: int,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
     ) -> dict[str, list[tuple[float, str]]]:
-        logger.info("Encoding Corpus in batches... Warning: This might take a while!")
+        logger.info("Encoding Corpus in batches (this might take a while)...")
+        if self.task_corpus is None:
+            raise ValueError("Corpus must be indexed before searching.")
+
         itr = range(0, len(self.task_corpus), self.corpus_chunk_size)
 
-        result_heaps = {qid: [] for qid in query_idx_to_id.values()}
+        result_heaps: dict[str, list[tuple[float, str]]] = {
+            qid: [] for qid in query_idx_to_id.values()
+        }
         for batch_num, corpus_start_idx in enumerate(itr):
             logger.info(f"Encoding Batch {batch_num + 1}/{len(itr)}...")
             corpus_end_idx = min(
-                corpus_start_idx + self.corpus_chunk_size, len(self.task_corpus)
+                corpus_start_idx + self.corpus_chunk_size,
+                len(self.task_corpus),
             )
             sub_corpus = self.task_corpus.select(
                 range(corpus_start_idx, corpus_end_idx)
@@ -165,7 +242,7 @@ class SearchEncoderWrapper:
                     sub_corpus,
                     task_metadata,
                     prompt_type=PromptType.document,
-                    batch_size=encode_kwargs.get("batch_size", 32),
+                    **encode_kwargs,
                 ),
                 task_metadata=task_metadata,
                 hf_split=hf_split,
@@ -179,8 +256,8 @@ class SearchEncoderWrapper:
             scores = self.model.similarity(query_embeddings, sub_corpus_embeddings)
 
             # get top-k values
-            cos_scores_top_k_values, cos_scores_top_k_idx = torch.topk(
-                torch.tensor(scores),
+            cos_scores_top_k_values_tensor, cos_scores_top_k_idx_tensor = torch.topk(
+                torch.as_tensor(scores),
                 min(
                     top_k + 1,
                     len(scores[1]) if len(scores) > 1 else len(scores[-1]),
@@ -188,22 +265,49 @@ class SearchEncoderWrapper:
                 dim=1,
                 largest=True,
             )
-            cos_scores_top_k_idx = cos_scores_top_k_idx.cpu().tolist()
-            cos_scores_top_k_values = cos_scores_top_k_values.cpu().tolist()
-
-            for query_itr in range(len(query_embeddings)):
-                query_id = query_idx_to_id[query_itr]
-                for sub_corpus_id, score in zip(
-                    cos_scores_top_k_idx[query_itr],
-                    cos_scores_top_k_values[query_itr],
-                ):
-                    corpus_id = sub_corpus_ids[sub_corpus_id]
-                    if len(result_heaps[query_id]) < top_k:
-                        # push item on the heap
-                        heapq.heappush(result_heaps[query_id], (score, corpus_id))
-                    else:
-                        # If item is larger than the smallest in the heap, push it on the heap then pop the smallest element
-                        heapq.heappushpop(result_heaps[query_id], (score, corpus_id))
+            cos_scores_top_k_idx = cos_scores_top_k_idx_tensor.cpu().tolist()
+            cos_scores_top_k_values = cos_scores_top_k_values_tensor.cpu().tolist()
+
+            sub_corpus_ids = list(sub_corpus_ids)
+            result_heaps = self._sort_full_corpus_results(
+                result_heaps=result_heaps,
+                query_idx_to_id=query_idx_to_id,
+                query_embeddings=query_embeddings,
+                cos_scores_top_k_idx=cos_scores_top_k_idx,
+                cos_scores_top_k_values=cos_scores_top_k_values,
+                sub_corpus_ids=sub_corpus_ids,
+                top_k=top_k,
+            )
+        return result_heaps
+
+    def _sort_full_corpus_results(
+        self,
+        result_heaps: dict[str, list[tuple[float, str]]],
+        query_idx_to_id: dict[int, str],
+        query_embeddings: Array,
+        cos_scores_top_k_idx: list[list[int]],
+        cos_scores_top_k_values: list[list[float]],
+        sub_corpus_ids: list[str],
+        top_k: int,
+    ) -> dict[str, list[tuple[float, str]]]:
+        """Sort the heaps into descending order lists.
+
+        Returns:
+            A dictionary mapping query IDs to a sorted list of tuples, each containing a relevance score and a document ID.
+        """
+        for query_itr in range(len(query_embeddings)):
+            query_id = query_idx_to_id[query_itr]
+            for sub_corpus_id, score in zip(
+                cos_scores_top_k_idx[query_itr],
+                cos_scores_top_k_values[query_itr],
+            ):
+                corpus_id = sub_corpus_ids[sub_corpus_id]
+                if len(result_heaps[query_id]) < top_k:
+                    # push item on the heap
+                    heapq.heappush(result_heaps[query_id], (score, corpus_id))
+                else:
+                    # If item is larger than the smallest in the heap, push it on the heap then pop the smallest element
+                    heapq.heappushpop(result_heaps[query_id], (score, corpus_id))
         return result_heaps
 
     def _rerank_documents(
@@ -215,14 +319,18 @@ class SearchEncoderWrapper:
         task_metadata: TaskMetadata,
         hf_subset: str,
         hf_split: str,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
     ) -> dict[str, list[tuple[float, str]]]:
         """Rerank documents based on pre-ranked documents.
 
         Returns:
             A dictionary mapping query IDs to a list of tuples, each containing a relevance score and a document ID.
         """
-        result_heaps = {qid: [] for qid in query_idx_to_id.values()}
+        if self.task_corpus is None:
+            raise ValueError("Corpus must be indexed before searching.")
+        result_heaps: dict[str, list[tuple[float, str]]] = {
+            qid: [] for qid in query_idx_to_id.values()
+        }
         doc_id_to_idx = {doc["id"]: idx for idx, doc in enumerate(self.task_corpus)}
 
         all_doc_embeddings = self.model.encode(
@@ -230,7 +338,7 @@ class SearchEncoderWrapper:
                 self.task_corpus,
                 task_metadata,
                 prompt_type=PromptType.document,
-                batch_size=encode_kwargs.get("batch_size", 32),
+                **encode_kwargs,
             ),
             task_metadata=task_metadata,
             hf_split=hf_split,
@@ -243,7 +351,8 @@ class SearchEncoderWrapper:
         for query_idx, query_embedding in enumerate(query_embeddings):
             query_id = query_idx_to_id[query_idx]
             if query_id not in top_ranked:
-                logger.warning(f"No pre-ranked documents found for query {query_id}")
+                msg = f"No pre-ranked documents found for query {query_id}"
+                logger.warning(msg)
                 continue
 
             ranked_ids = top_ranked[query_id]
@@ -278,14 +387,34 @@ class SearchEncoderWrapper:
             scores_top_k_values = scores_top_k_values.cpu()
             scores_top_k_idx = scores_top_k_idx.cpu()
 
-            # Build result heap
-            for doc_idx, score in zip(
-                scores_top_k_idx[0].tolist(),
-                scores_top_k_values[0].tolist(),
-            ):
-                corpus_id = ranked_ids[doc_idx]
-                heapq.heappush(result_heaps[query_id], (score, corpus_id))
+            result_heaps = self._rerank_sort_results(
+                result_heaps=result_heaps,
+                query_id=query_id,
+                ranked_ids=ranked_ids,
+                scores_top_k_idx=scores_top_k_idx,
+                scores_top_k_values=scores_top_k_values,
+            )
+        return result_heaps
+
+    def _rerank_sort_results(
+        self,
+        result_heaps: dict[str, list[tuple[float, str]]],
+        query_id: str,
+        ranked_ids: list[str],
+        scores_top_k_idx: torch.Tensor,
+        scores_top_k_values: torch.Tensor,
+    ) -> dict[str, list[tuple[float, str]]]:
+        """Sort the heap into descending order list.
 
+        Returns:
+            A sorted list of tuples, each containing a relevance score and a document ID.
+        """
+        for doc_idx, score in zip(
+            scores_top_k_idx[0].tolist(),
+            scores_top_k_values[0].tolist(),
+        ):
+            corpus_id = ranked_ids[doc_idx]
+            heapq.heappush(result_heaps[query_id], (score, corpus_id))
         return result_heaps
 
     def encode(
@@ -342,7 +471,7 @@ class SearchCrossEncoderWrapper:
         task_metadata: TaskMetadata,
         hf_split: str,
         hf_subset: str,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
     ) -> None:
         """Index the corpus for retrieval.
 
@@ -363,7 +492,7 @@ class SearchCrossEncoderWrapper:
         hf_split: str,
         hf_subset: str,
         top_k: int,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         top_ranked: TopRankedDocumentsType | None = None,
     ) -> RetrievalOutputType:
         """Search the corpus using the given queries.
@@ -385,6 +514,8 @@ class SearchCrossEncoderWrapper:
             raise ValueError(
                 "CrossEncoder search requires top_ranked documents for reranking."
             )
+        if self.task_corpus is None:
+            raise ValueError("Corpus must be indexed before searching.")
 
         query_id_to_idx = {row["id"]: i for i, row in enumerate(queries)}
         doc_id_to_idx = {doc["id"]: idx for idx, doc in enumerate(self.task_corpus)}
@@ -394,7 +525,8 @@ class SearchCrossEncoderWrapper:
         doc_pairs_ids: list[tuple[str, str]] = []
         for query_id, corpus_ids in top_ranked.items():
             if query_id not in top_ranked:
-                logger.warning(f"No pre-ranked documents found for query {query_id}")
+                msg = f"No pre-ranked documents found for query {query_id}"
+                logger.warning(msg)
                 continue
 
             query_idx = query_id_to_idx[query_id]
@@ -407,13 +539,13 @@ class SearchCrossEncoderWrapper:
             Dataset.from_list(total_queries),
             task_metadata,
             prompt_type=PromptType.document,
-            batch_size=encode_kwargs.get("batch_size", 32),
+            **encode_kwargs,
         )
         corpus_loader = create_dataloader(
             Dataset.from_list(total_docs),
             task_metadata,
             prompt_type=PromptType.document,
-            batch_size=encode_kwargs.get("batch_size", 32),
+            **encode_kwargs,
         )
         predictions = self.model.predict(
             inputs1=queries_loader,
@@ -423,7 +555,7 @@ class SearchCrossEncoderWrapper:
             hf_subset=hf_subset,
         )
 
-        results = {qid: {} for qid in queries["id"]}
+        results: RetrievalOutputType = {qid: {} for qid in queries["id"]}
         for (query_id, corpus_id), score in zip(doc_pairs_ids, predictions):
             results[query_id][corpus_id] = float(score)
 

mteb/models/sentence_transformer_wrapper.py

@@ -1,16 +1,18 @@
 from __future__ import annotations
 
 import logging
+import warnings
 from typing import TYPE_CHECKING, Any
 
 import numpy as np
 import torch
 from packaging.version import Version
 from torch.utils.data import DataLoader
+from typing_extensions import Unpack
 
 from mteb._log_once import LogOnce
 from mteb.models import ModelMeta
-from mteb.types import Array, BatchedInput, PromptType
+from mteb.types import Array, BatchedInput, EncodeKwargs, PromptType
 
 from .abs_encoder import AbsEncoder
 
@@ -25,17 +27,18 @@ SENTENCE_TRANSFORMERS_QUERY_ENCODE_VERSION = "5.0.0"
 
 
 def sentence_transformers_loader(
-    model_name: str, revision: str | None = None, **kwargs
+    model_name: str, revision: str | None = None, device: str | None = None, **kwargs
 ) -> SentenceTransformerEncoderWrapper:
     """Loads a SentenceTransformer model and wraps it in a SentenceTransformerEncoderWrapper.
 
     Args:
         model_name: The name of the SentenceTransformer model to load.
         revision: The revision of the model to load.
+        device: The device used to load the model.
         kwargs: Additional arguments to pass to the SentenceTransformer model.
     """
     return SentenceTransformerEncoderWrapper(
-        model=model_name, revision=revision, **kwargs
+        model=model_name, revision=revision, device=device, **kwargs
     )
 
 
@@ -48,6 +51,7 @@ class SentenceTransformerEncoderWrapper(AbsEncoder):
         self,
         model: str | SentenceTransformer,
         revision: str | None = None,
+        device: str | None = None,
         model_prompts: dict[str, str] | None = None,
         **kwargs,
     ) -> None:
@@ -56,6 +60,7 @@ class SentenceTransformerEncoderWrapper(AbsEncoder):
         Args:
             model: The SentenceTransformer model to use. Can be a string (model name), a SentenceTransformer model, or a CrossEncoder model.
             revision: The revision of the model to use.
+            device: The device used to load the model.
             model_prompts: A dictionary mapping task names to prompt names.
                 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,
@@ -65,22 +70,21 @@ class SentenceTransformerEncoderWrapper(AbsEncoder):
         from sentence_transformers import SentenceTransformer
 
         if isinstance(model, str):
-            self.model = SentenceTransformer(model, revision=revision, **kwargs)
+            self.model = SentenceTransformer(
+                model, revision=revision, device=device, **kwargs
+            )
         else:
             self.model = model
-        from mteb.models.get_model_meta import (
-            _model_meta_from_sentence_transformers,
-        )
 
-        self.mteb_model_meta = _model_meta_from_sentence_transformers(self.model)
+        self.mteb_model_meta = ModelMeta.from_sentence_transformer_model(self.model)
 
         built_in_prompts = getattr(self.model, "prompts", None)
         if built_in_prompts and not model_prompts:
             model_prompts = built_in_prompts
         elif model_prompts and built_in_prompts:
-            logger.warning(
-                f"Model prompts specified, these will overwrite the default model prompts. Current prompts will be:\n {model_prompts}"
-            )
+            msg = f"Model prompts specified, these will overwrite the default model prompts. Current prompts will be:\n {model_prompts}"
+            logger.warning(msg)
+            warnings.warn(msg)
             self.model.prompts = model_prompts
 
         self.model_prompts, invalid_prompts = self.validate_task_to_prompt_name(
@@ -89,9 +93,9 @@ class SentenceTransformerEncoderWrapper(AbsEncoder):
 
         if invalid_prompts:
             invalid_prompts = "\n".join(invalid_prompts)
-            logger.warning(
-                f"Some prompts are not in the expected format and will be ignored. Problems:\n\n{invalid_prompts}"
-            )
+            msg = f"Some prompts are not in the expected format and will be ignored. Problems:\n\n{invalid_prompts}"
+            logger.warning(msg)
+            warnings.warn(msg)
 
         if (
             self.model_prompts
@@ -101,13 +105,15 @@ class SentenceTransformerEncoderWrapper(AbsEncoder):
                 or PromptType.document.value not in self.model_prompts
             )
         ):
-            logger.warning(
-                "SentenceTransformers that use prompts most often need to be configured with at least 'query' and"
-                f" 'document' prompts to ensure optimal performance. Received {self.model_prompts}"
-            )
+            msg = f"SentenceTransformers that use prompts most often need to be configured with at least 'query' and 'document' prompts to ensure optimal performance. Received {self.model_prompts}"
+            logger.warning(msg)
+            warnings.warn(msg)
 
+    def similarity(self, embeddings1: Array, embeddings2: Array) -> Array:
+        """Compute the similarity between two collections of embeddings."""
         if hasattr(self.model, "similarity") and callable(self.model.similarity):
-            self.similarity = self.model.similarity
+            return self.model.similarity(embeddings1, embeddings2)
+        return super().similarity(embeddings1, embeddings2)
 
     def encode(
         self,
@@ -117,7 +123,7 @@ class SentenceTransformerEncoderWrapper(AbsEncoder):
         hf_split: str,
         hf_subset: str,
         prompt_type: PromptType | None = None,
-        **kwargs: Any,
+        **kwargs: Unpack[EncodeKwargs],
     ) -> Array:
         """Encodes the given sentences using the encoder.
 
@@ -153,7 +159,7 @@ class SentenceTransformerEncoderWrapper(AbsEncoder):
         prompt_name = None
         if self.model_prompts is not None:
             prompt_name = self.get_prompt_name(task_metadata, prompt_type)
-            prompt = self.model_prompts.get(prompt_name, None)
+            prompt = self.model_prompts.get(prompt_name, None)  # type: ignore[arg-type]
         if prompt_name:
             prompt_log = f"Using {prompt_name=} for task={task_metadata.name} {prompt_type=} with {prompt=}"
         else:
@@ -196,7 +202,7 @@ class SentenceTransformerMultimodalEncoderWrapper(SentenceTransformerEncoderWrap
         hf_split: str,
         hf_subset: str,
         prompt_type: PromptType | None = None,
-        **kwargs: Any,
+        **kwargs: Unpack[EncodeKwargs],
     ) -> Array:
         """Encodes the given sentences using the encoder.
 
@@ -224,7 +230,7 @@ class SentenceTransformerMultimodalEncoderWrapper(SentenceTransformerEncoderWrap
         prompt_name = None
         if self.model_prompts is not None:
             prompt_name = self.get_prompt_name(task_metadata, prompt_type)
-            prompt = self.model_prompts.get(prompt_name, None)
+            prompt = self.model_prompts.get(prompt_name, None)  # type: ignore[arg-type]
         if prompt_name:
             logger.info(
                 f"Using {prompt_name=} for task={task_metadata.name} {prompt_type=} with {prompt=}"
@@ -237,7 +243,9 @@ class SentenceTransformerMultimodalEncoderWrapper(SentenceTransformerEncoderWrap
         all_embeddings = []
         for batch in inputs:
             batch_column = next(iter(batch.keys()))
-            batched_input = [dict() for _ in range(len(batch[batch_column]))]
+            batched_input: list[dict[str, Any]] = [
+                dict() for _ in range(len(batch[batch_column]))
+            ]
 
             # transform from {"text": [text1, text2], "image": [image1, image2]} to
             # [{"text": text1, "image": image1}, {"text": text2, "image": image2}]
@@ -258,24 +266,36 @@ class SentenceTransformerMultimodalEncoderWrapper(SentenceTransformerEncoderWrap
 
 
 class CrossEncoderWrapper:
-    """Wrapper for CrossEncoder models."""
+    """Wrapper for CrossEncoder models.
+
+    Args:
+        model: The CrossEncoder model to use. Can be a string (model name) or a CrossEncoder model.
+        revision: The revision of the model to use.
+        device: The device used to load the model.
+        query_prefix: A prefix to add to all queries.
+        passage_prefix: A prefix to add to all passages.
+        **kwargs: Additional arguments to pass to the CrossEncoder model.
+    """
 
     def __init__(
         self,
         model: CrossEncoder | str,
         revision: str | None = None,
+        device: str | None = None,
+        query_prefix: str = "",
+        passage_prefix: str = "",
         **kwargs,
     ) -> None:
         from sentence_transformers import CrossEncoder
 
-        from mteb.models.get_model_meta import _model_meta_from_cross_encoder
-
         if isinstance(model, CrossEncoder):
             self.model = model
         elif isinstance(model, str):
-            self.model = CrossEncoder(model, revision=revision, **kwargs)
+            self.model = CrossEncoder(model, revision=revision, device=device, **kwargs)
 
-        self.mteb_model_meta = _model_meta_from_cross_encoder(self.model)
+        self.mteb_model_meta = ModelMeta.from_cross_encoder(self.model)
+        self.query_prefix = query_prefix
+        self.passage_prefix = passage_prefix
 
     def predict(
         self,
@@ -286,7 +306,7 @@ class CrossEncoderWrapper:
         hf_split: str,
         hf_subset: str,
         prompt_type: PromptType | None = None,
-        **kwargs: Any,
+        **kwargs: Unpack[EncodeKwargs],
     ) -> Array:
         """Predicts relevance scores for pairs of inputs. Note that, unlike the encoder, the cross-encoder can compare across inputs.
 
@@ -304,10 +324,10 @@ class CrossEncoderWrapper:
             The predicted relevance scores for each inputs pair.
         """
         all_queries_with_instructions = [
-            text for batch in inputs1 for text in batch["text"]
+            self.query_prefix + text for batch in inputs1 for text in batch["text"]
         ]
         all_corpus_with_instructions = [
-            text for batch in inputs2 for text in batch["text"]
+            self.passage_prefix + text for batch in inputs2 for text in batch["text"]
         ]
 
         return self.model.predict(