mteb 2.2.2__py3-none-any.whl → 2.3.1__py3-none-any.whl

mteb/models/model_implementations/nvidia_llama_nemoretriever_colemb.py

@@ -117,19 +117,18 @@ class LlamaNemoretrieverColembed(AbsEncoder):
 
 TRAINING_DATA = {
     # from https://huggingface.co/datasets/vidore/colpali_train_set
-    "DocVQA",
-    "InfoVQA",
-    "TATDQA",
-    "arXivQA",
-    "hotpotqa",
-    "miracl",
+    "VidoreDocVQARetrieval",
+    "VidoreInfoVQARetrieval",
+    "VidoreTatdqaRetrieval",
+    "VidoreArxivQARetrieval",
+    "HotpotQA",
+    "MIRACLRetrieval",
     "NQ",
-    "stackexchange",
+    "StackExchangeClustering",
     "SQuAD",
     "WebInstructSub",
     "docmatix-ir",
-    "vdr-multilingual-train",
-    "colpali_train_set",  # as it contains PDFs
+    "VDRMultilingualRetrieval",
     "VisRAG-Ret-Train-Synthetic-data",
     "VisRAG-Ret-Train-In-domain-data",
     "wiki-ss-nq",
@@ -146,7 +145,7 @@ llama_nemoretriever_colembed_1b_v1 = ModelMeta(
     release_date="2025-06-27",
     modalities=["image", "text"],
     n_parameters=2_418_000_000,
-    memory_usage_mb=9224,
+    memory_usage_mb=4610,
     max_tokens=8192,
     embed_dim=2048,
     license="https://huggingface.co/nvidia/llama-nemoretriever-colembed-1b-v1/blob/main/LICENSE",
@@ -172,7 +171,7 @@ llama_nemoretriever_colembed_3b_v1 = ModelMeta(
     release_date="2025-06-27",
     modalities=["image", "text"],
     n_parameters=4_407_000_000,
-    memory_usage_mb=16811,
+    memory_usage_mb=8403,
     max_tokens=8192,
     embed_dim=3072,
     license="https://huggingface.co/nvidia/llama-nemoretriever-colembed-1b-v1/blob/main/LICENSE",

mteb/models/model_implementations/nvidia_models.py

@@ -146,7 +146,7 @@ NV_embed_v1 = ModelMeta(
     revision="570834afd5fef5bf3a3c2311a2b6e0a66f6f4f2c",
     release_date="2024-09-13",  # initial commit of hf model.
     n_parameters=7_850_000_000,
-    memory_usage_mb=29945,
+    memory_usage_mb=14975,
     embed_dim=4096,
     license="cc-by-nc-4.0",
     max_tokens=32768,

mteb/models/model_implementations/ops_moa_models.py

@@ -27,7 +27,7 @@ ops_moa_conan_embedding = ModelMeta(
     languages=["zho-Hans"],
     loader=OPSWrapper,
     n_parameters=int(343 * 1e6),
-    memory_usage_mb=2e3,
+    memory_usage_mb=1308,
     max_tokens=512,
     embed_dim=1536,
     license="cc-by-nc-4.0",
@@ -58,7 +58,7 @@ ops_moa_yuan_embedding = ModelMeta(
     languages=["zho-Hans"],
     loader=OPSWrapper,
     n_parameters=int(343 * 1e6),
-    memory_usage_mb=2e3,
+    memory_usage_mb=1242,
     max_tokens=512,
     embed_dim=1536,
     license="cc-by-nc-4.0",

mteb/models/model_implementations/promptriever_models.py

@@ -80,7 +80,7 @@ promptriever_llama2 = ModelMeta(
     revision="01c7f73d771dfac7d292323805ebc428287df4f9-30b14e3813c0fa45facfd01a594580c3fe5ecf23",  # base-peft revision
     release_date="2024-09-15",
     n_parameters=7_000_000_000,
-    memory_usage_mb=27,
+    memory_usage_mb=26703,
     max_tokens=4096,
     embed_dim=4096,
     license="apache-2.0",
@@ -115,7 +115,7 @@ promptriever_llama3 = ModelMeta(
     },
     release_date="2024-09-15",
     n_parameters=8_000_000_000,
-    memory_usage_mb=31,
+    memory_usage_mb=30518,
     max_tokens=8192,
     embed_dim=4096,
     license="apache-2.0",
@@ -143,7 +143,7 @@ promptriever_llama3_instruct = ModelMeta(
     revision="5206a32e0bd3067aef1ce90f5528ade7d866253f-8b677258615625122c2eb7329292b8c402612c21",  # base-peft revision
     release_date="2024-09-15",
     n_parameters=8_000_000_000,
-    memory_usage_mb=31,
+    memory_usage_mb=30518,
     max_tokens=8192,
     embed_dim=4096,
     training_datasets={
@@ -175,7 +175,7 @@ promptriever_mistral_v1 = ModelMeta(
     revision="7231864981174d9bee8c7687c24c8344414eae6b-876d63e49b6115ecb6839893a56298fadee7e8f5",  # base-peft revision
     release_date="2024-09-15",
     n_parameters=7_000_000_000,
-    memory_usage_mb=27,
+    memory_usage_mb=26703,
     training_datasets={
         # "samaya-ai/msmarco-w-instructions",
         "mMARCO-NL",  # translation not trained on

mteb/models/model_implementations/qwen3_models.py

@@ -139,7 +139,7 @@ Qwen3_Embedding_0B6 = ModelMeta(
     revision="b22da495047858cce924d27d76261e96be6febc0",  # Commit of @tomaarsen
     release_date="2025-06-05",
     n_parameters=595776512,
-    memory_usage_mb=2272,
+    memory_usage_mb=1136,
     embed_dim=1024,
     max_tokens=32768,
     license="apache-2.0",
@@ -161,7 +161,7 @@ Qwen3_Embedding_4B = ModelMeta(
     revision="636cd9bf47d976946cdbb2b0c3ca0cb2f8eea5ff",  # Commit of @tomaarsen
     release_date="2025-06-05",
     n_parameters=4021774336,
-    memory_usage_mb=15341,
+    memory_usage_mb=7671,
     embed_dim=2560,
     max_tokens=32768,
     license="apache-2.0",
@@ -183,7 +183,7 @@ Qwen3_Embedding_8B = ModelMeta(
     revision="4e423935c619ae4df87b646a3ce949610c66241c",  # Commit of @tomaarsen
     release_date="2025-06-05",
     n_parameters=7567295488,
-    memory_usage_mb=28866,
+    memory_usage_mb=14433,
     embed_dim=4096,
     max_tokens=32768,
     license="apache-2.0",

mteb/models/model_implementations/qzhou_models.py

@@ -63,7 +63,7 @@ QZhou_Embedding = ModelMeta(
     revision="f1e6c03ee3882e7b9fa5cec91217715272e433b8",
     release_date="2025-08-24",
     n_parameters=7_070_619_136,
-    memory_usage_mb=29070,
+    memory_usage_mb=14436,
     embed_dim=3584,
     license="apache-2.0",
     max_tokens=8192,

mteb/models/model_implementations/random_baseline.py

@@ -8,6 +8,10 @@ from torch.utils.data import DataLoader
 
 from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.models.model_meta import ModelMeta
+from mteb.similarity_functions import (
+    select_pairwise_similarity,
+    select_similarity,
+)
 from mteb.types._encoder_io import Array, BatchedInput, PromptType
 
 
@@ -155,15 +159,9 @@ class RandomEncoderBaseline:
         Returns:
             Cosine similarity matrix between the two sets of embeddings
         """
-        norm1 = np.linalg.norm(
-            embeddings1.reshape(-1, self.embedding_dim), axis=1, keepdims=True
-        )
-        norm2 = np.linalg.norm(
-            embeddings2.reshape(-1, self.embedding_dim), axis=1, keepdims=True
+        return select_similarity(
+            embeddings1, embeddings2, self.mteb_model_meta.similarity_fn_name
         )
-        normalized1 = embeddings1 / (norm1 + 1e-10)
-        normalized2 = embeddings2 / (norm2 + 1e-10)
-        return np.dot(normalized1, normalized2.T)
 
     def similarity_pairwise(
         self,
@@ -179,17 +177,9 @@ class RandomEncoderBaseline:
         Returns:
             Cosine similarity for each pair of embeddings
         """
-        norm1 = np.linalg.norm(
-            embeddings1.reshape(-1, self.embedding_dim), axis=1, keepdims=True
-        )
-        norm2 = np.linalg.norm(
-            embeddings2.reshape(-1, self.embedding_dim), axis=1, keepdims=True
+        return select_pairwise_similarity(
+            embeddings1, embeddings2, self.mteb_model_meta.similarity_fn_name
         )
-        normalized1 = embeddings1 / (norm1 + 1e-10)
-        normalized2 = embeddings2 / (norm2 + 1e-10)
-        normalized1 = np.asarray(normalized1)
-        normalized2 = np.asarray(normalized2)
-        return np.sum(normalized1 * normalized2, axis=1)
 
 
 random_encoder_baseline = ModelMeta(

mteb/models/model_implementations/vdr_models.py

@@ -43,5 +43,6 @@ vdr_2b_multi_v1 = ModelMeta(
     public_training_data="https://huggingface.co/datasets/llamaindex/vdr-multilingual-train",
     training_datasets=set(
         # llamaindex/vdr-multilingual-train
+        "VDRMultilingualRetrieval",
     ),
 )

mteb/models/model_implementations/yuan_models_en.py

@@ -0,0 +1,57 @@
+from mteb.models.instruct_wrapper import InstructSentenceTransformerModel
+from mteb.models.model_meta import ModelMeta
+from mteb.models.models_protocols import PromptType
+
+
+def instruction_template(
+    instruction: str, prompt_type: PromptType | None = None
+) -> str:
+    if not instruction or prompt_type == PromptType.document:
+        return ""
+    if isinstance(instruction, dict):
+        if prompt_type is None:
+            instruction = next(iter(instruction.values()))  # TODO
+        else:
+            instruction = instruction[prompt_type]
+    return f"Instruct: {instruction}\nQuery:"
+
+
+training_data = {
+    "T2Retrieval",
+    "DuRetrieval",
+    "MMarcoReranking",
+    "CMedQAv2-reranking",
+    "NQ",
+    "MSMARCO",
+    "HotpotQA",
+    "MrTidyRetrieval",
+    "MIRACLRetrieval",
+    "CodeSearchNet",
+}
+
+
+yuan_embedding_2_en = ModelMeta(
+    loader=InstructSentenceTransformerModel,
+    loader_kwargs=dict(
+        instruction_template=instruction_template,
+        apply_instruction_to_passages=False,
+    ),
+    name="IEITYuan/Yuan-embedding-2.0-en",
+    languages=["eng-Latn"],
+    open_weights=True,
+    revision="b2fd15da3bcae3473c8529593825c15068f09fce",
+    release_date="2025-11-27",
+    n_parameters=595776512,
+    memory_usage_mb=2272,
+    embed_dim=1024,
+    max_tokens=2048,
+    license="apache-2.0",
+    reference="https://huggingface.co/IEITYuan/Yuan-embedding-2.0-en",
+    similarity_fn_name="cosine",
+    framework=["Sentence Transformers", "PyTorch"],
+    use_instructions=True,
+    public_training_code=None,
+    public_training_data=None,
+    training_datasets=training_data,
+    adapted_from="Qwen/Qwen3-Embedding-0.6B",
+)

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,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 = []