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

mteb/models/models_protocols.py

@@ -1,12 +1,14 @@
 from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
 
 from torch.utils.data import DataLoader
+from typing_extensions import Unpack
 
 from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.types import (
     Array,
     BatchedInput,
     CorpusDatasetType,
+    EncodeKwargs,
     PromptType,
     QueryDatasetType,
     RetrievalOutputType,
@@ -28,7 +30,7 @@ class SearchProtocol(Protocol):
         task_metadata: TaskMetadata,
         hf_split: str,
         hf_subset: str,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
     ) -> None:
         """Index the corpus for retrieval.
 
@@ -49,7 +51,7 @@ class SearchProtocol(Protocol):
         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.
@@ -83,12 +85,19 @@ class EncoderProtocol(Protocol):
     In general the interface is kept aligned with sentence-transformers interface. In cases where exceptions occurs these are handled within MTEB.
     """
 
-    def __init__(self, model_name: str, revision: str | None, **kwargs: Any) -> None:
+    def __init__(
+        self,
+        model_name: str,
+        revision: str | None,
+        device: str | None = None,
+        **kwargs: Any,
+    ) -> None:
         """The initialization function for the encoder. Used when calling it from the mteb run CLI.
 
         Args:
             model_name: Name of the model
             revision: revision of the model
+            device: Device used to load the model
             kwargs: Any additional kwargs
         """
         ...
@@ -101,7 +110,7 @@ class EncoderProtocol(Protocol):
         hf_split: str,
         hf_subset: str,
         prompt_type: PromptType | None = None,
-        **kwargs: Any,
+        **kwargs: Unpack[EncodeKwargs],
     ) -> Array:
         """Encodes the given sentences using the encoder.
 
@@ -181,12 +190,19 @@ class CrossEncoderProtocol(Protocol):
     In general the interface is kept aligned with sentence-transformers interface. In cases where exceptions occurs these are handled within MTEB.
     """
 
-    def __init__(self, model_name: str, revision: str | None, **kwargs: Any) -> None:
+    def __init__(
+        self,
+        model_name: str,
+        revision: str | None,
+        device: str | None = None,
+        **kwargs: Any,
+    ) -> None:
         """The initialization function for the encoder. Used when calling it from the mteb run CLI.
 
         Args:
             model_name: Name of the model
             revision: revision of the model
+            device: Device used to load the model
             kwargs: Any additional kwargs
         """
         ...
@@ -200,7 +216,7 @@ class CrossEncoderProtocol(Protocol):
         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.
 

mteb/models/search_encoder_index/__init__.py

@@ -0,0 +1,7 @@
+from .search_backend_protocol import IndexEncoderSearchProtocol
+from .search_indexes import FaissSearchIndex
+
+__all__ = [
+    "FaissSearchIndex",
+    "IndexEncoderSearchProtocol",
+]

mteb/models/search_encoder_index/search_backend_protocol.py

@@ -0,0 +1,50 @@
+from collections.abc import Callable
+from typing import Protocol
+
+from mteb.types import Array, TopRankedDocumentsType
+
+
+class IndexEncoderSearchProtocol(Protocol):
+    """Protocol for search backends used in encoder-based retrieval."""
+
+    def add_documents(
+        self,
+        embeddings: Array,
+        idxs: list[str],
+    ) -> None:
+        """Add documents to the search backend.
+
+        Args:
+            embeddings: Embeddings of the documents to add.
+            idxs: IDs of the documents to add.
+        """
+
+    def search(
+        self,
+        embeddings: Array,
+        top_k: int,
+        similarity_fn: Callable[[Array, Array], Array],
+        top_ranked: TopRankedDocumentsType | None = None,
+        query_idx_to_id: dict[int, str] | None = None,
+    ) -> tuple[list[list[float]], list[list[int]]]:
+        """Search through added corpus embeddings or rerank top-ranked documents.
+
+        Supports both full-corpus and reranking search modes:
+            - Full-corpus mode: `top_ranked=None`, uses added corpus embeddings.
+            - Reranking mode:  `top_ranked` contains mapping {query_id: [doc_ids]}.
+
+        Args:
+            embeddings: Query embeddings, shape (num_queries, dim).
+            top_k: Number of top results to return.
+            similarity_fn: Function to compute similarity between query and corpus.
+            top_ranked: Mapping of query_id -> list of candidate doc_ids. Used for reranking.
+            query_idx_to_id: Mapping of query index -> query_id. Used for reranking.
+
+        Returns:
+            A tuple (top_k_values, top_k_indices), for each query:
+                - top_k_values: List of top-k similarity scores.
+                - top_k_indices: List of indices of the top-k documents in the added corpus.
+        """
+
+    def clear(self) -> None:
+        """Clear all stored documents and embeddings from the backend."""

mteb/models/search_encoder_index/search_indexes/__init__.py

@@ -0,0 +1,5 @@
+from .faiss_search_index import FaissSearchIndex
+
+__all__ = [
+    "FaissSearchIndex",
+]

mteb/models/search_encoder_index/search_indexes/faiss_search_index.py

@@ -0,0 +1,160 @@
+import logging
+import warnings
+from collections.abc import Callable
+
+import numpy as np
+import torch
+
+from mteb._requires_package import requires_package
+from mteb.models.model_meta import ScoringFunction
+from mteb.models.models_protocols import EncoderProtocol
+from mteb.types import Array, TopRankedDocumentsType
+
+logger = logging.getLogger(__name__)
+
+
+class FaissSearchIndex:
+    """FAISS-based backend for encoder-based search.
+
+    Supports both full-corpus retrieval and reranking (via `top_ranked`).
+
+    Notes:
+        - Stores *all* embeddings in memory (IndexFlatIP or IndexFlatL2).
+        - Expects embeddings to be normalized if cosine similarity is desired.
+    """
+
+    _normalize: bool = False
+
+    def __init__(self, model: EncoderProtocol) -> None:
+        requires_package(
+            self,
+            "faiss",
+            "FAISS-based search",
+            install_instruction="pip install mteb[faiss-cpu]",
+        )
+
+        import faiss
+        from faiss import IndexFlatIP, IndexFlatL2
+
+        # https://github.com/facebookresearch/faiss/wiki/Faiss-indexes
+        if model.mteb_model_meta.similarity_fn_name is ScoringFunction.DOT_PRODUCT:
+            self.index_type = IndexFlatIP
+        elif model.mteb_model_meta.similarity_fn_name is ScoringFunction.COSINE:
+            self.index_type = IndexFlatIP
+            self._normalize = True
+        elif model.mteb_model_meta.similarity_fn_name is ScoringFunction.EUCLIDEAN:
+            self.index_type = IndexFlatL2
+        else:
+            raise ValueError(
+                f"FAISS backend does not support similarity function {model.mteb_model_meta.similarity_fn_name}. "
+                f"Available: {ScoringFunction.DOT_PRODUCT}, {ScoringFunction.COSINE}."
+            )
+
+        self.idxs: list[str] = []
+        self.index: faiss.Index | None = None
+
+    def add_documents(self, embeddings: Array, idxs: list[str]) -> None:
+        """Add all document embeddings and their IDs to FAISS index."""
+        import faiss
+
+        if isinstance(embeddings, torch.Tensor):
+            embeddings = embeddings.detach().cpu().numpy()
+
+        embeddings = embeddings.astype(np.float32)
+        self.idxs.extend(idxs)
+
+        if self._normalize:
+            faiss.normalize_L2(embeddings)
+
+        dim = embeddings.shape[1]
+        if self.index is None:
+            self.index = self.index_type(dim)
+
+        self.index.add(embeddings)
+        logger.info(f"FAISS index built with {len(idxs)} vectors of dim {dim}.")
+
+    def search(
+        self,
+        embeddings: Array,
+        top_k: int,
+        similarity_fn: Callable[[Array, Array], Array],
+        top_ranked: TopRankedDocumentsType | None = None,
+        query_idx_to_id: dict[int, str] | None = None,
+    ) -> tuple[list[list[float]], list[list[int]]]:
+        """Search using FAISS."""
+        import faiss
+
+        if self.index is None:
+            raise ValueError("No index built. Call add_document() first.")
+
+        if isinstance(embeddings, torch.Tensor):
+            embeddings = embeddings.detach().cpu().numpy()
+
+        if self._normalize:
+            faiss.normalize_L2(embeddings)
+
+        if top_ranked is not None:
+            if query_idx_to_id is None:
+                raise ValueError("query_idx_to_id must be provided when reranking.")
+
+            similarities, ids = self._reranking(
+                embeddings,
+                top_k,
+                top_ranked=top_ranked,
+                query_idx_to_id=query_idx_to_id,
+            )
+        else:
+            similarities, ids = self.index.search(embeddings.astype(np.float32), top_k)
+            similarities = similarities.tolist()
+            ids = ids.tolist()
+
+        if issubclass(self.index_type, faiss.IndexFlatL2):
+            similarities = (-np.sqrt(np.maximum(similarities, 0))).tolist()
+
+        return similarities, ids
+
+    def _reranking(
+        self,
+        embeddings: Array,
+        top_k: int,
+        top_ranked: TopRankedDocumentsType,
+        query_idx_to_id: dict[int, str],
+    ) -> tuple[list[list[float]], list[list[int]]]:
+        doc_id_to_idx = {doc_id: i for i, doc_id in enumerate(self.idxs)}
+        scores_all: list[list[float]] = []
+        idxs_all: list[list[int]] = []
+
+        for query_idx, query_emb in enumerate(embeddings):
+            query_id = query_idx_to_id[query_idx]
+            ranked_ids = top_ranked.get(query_id)
+            if not ranked_ids:
+                msg = f"No top-ranked documents for query {query_id}"
+                logger.warning(msg)
+                warnings.warn(msg)
+                scores_all.append([])
+                idxs_all.append([])
+                continue
+
+            candidate_indices = [doc_id_to_idx[doc_id] for doc_id in ranked_ids]
+            d = self.index.d  # type: ignore[union-attr]
+            candidate_embs = np.vstack(
+                [self.index.reconstruct(idx) for idx in candidate_indices]  # type: ignore[union-attr]
+            )
+            sub_reranking_index = self.index_type(d)
+            sub_reranking_index.add(candidate_embs)
+
+            # Search returns scores and indices in one call
+            scores, local_indices = sub_reranking_index.search(
+                query_emb.reshape(1, -1).astype(np.float32),
+                min(top_k, len(candidate_indices)),
+            )
+            # faiss will output 2d arrays even for single query
+            scores_all.append(scores[0].tolist())
+            idxs_all.append(local_indices[0].tolist())
+
+        return scores_all, idxs_all
+
+    def clear(self) -> None:
+        """Clear all stored documents and embeddings from the backend."""
+        self.index = None
+        self.idxs = []