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

mteb/models/search_encoder_index/search_indexes/faiss_search_index.py

@@ -0,0 +1,157 @@
+import logging
+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))
+
+        return similarities, ids
+
+    def _reranking(
+        self,
+        embeddings: Array,
+        top_k: int,
+        top_ranked: TopRankedDocumentsType | None = None,
+        query_idx_to_id: dict[int, str] | None = None,
+    ) -> 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:
+                logger.warning(f"No top-ranked documents for query {query_id}")
+                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
+            candidate_embs = np.vstack(
+                [self.index.reconstruct(idx) for idx in candidate_indices]
+            )
+            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 = []

mteb/models/search_wrappers.py

@@ -21,6 +21,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 +29,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,
@@ -56,6 +63,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,
@@ -90,7 +113,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,27 +128,74 @@ 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
@@ -147,7 +217,7 @@ class SearchEncoderWrapper:
         top_k: int,
         encode_kwargs: dict[str, Any],
     ) -> 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)...")
         itr = range(0, len(self.task_corpus), self.corpus_chunk_size)
 
         result_heaps = {qid: [] for qid in query_idx_to_id.values()}
@@ -165,7 +235,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,
@@ -180,7 +250,7 @@ class SearchEncoderWrapper:
 
             # get top-k values
             cos_scores_top_k_values, cos_scores_top_k_idx = torch.topk(
-                torch.tensor(scores),
+                torch.as_tensor(scores),
                 min(
                     top_k + 1,
                     len(scores[1]) if len(scores) > 1 else len(scores[-1]),
@@ -191,19 +261,46 @@ class SearchEncoderWrapper:
             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))
+            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(
@@ -230,7 +327,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,
@@ -278,14 +375,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: list[tuple[float, str]],
+        query_id: str,
+        ranked_ids: list[str],
+        scores_top_k_idx: torch.Tensor,
+        scores_top_k_values: torch.Tensor,
+    ) -> 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(
@@ -407,13 +524,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,

mteb/models/sentence_transformer_wrapper.py

@@ -68,11 +68,8 @@ class SentenceTransformerEncoderWrapper(AbsEncoder):
             self.model = SentenceTransformer(model, revision=revision, **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:
@@ -268,14 +265,12 @@ class CrossEncoderWrapper:
     ) -> 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.mteb_model_meta = _model_meta_from_cross_encoder(self.model)
+        self.mteb_model_meta = ModelMeta.from_cross_encoder(self.model)
 
     def predict(
         self,