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

mteb/models/model_implementations/youtu_models.py

@@ -121,6 +121,7 @@ Youtu_Embedding_V1 = ModelMeta(
     release_date="2025-09-28",
     open_weights=True,
     n_parameters=2672957440,
+    n_embedding_parameters=None,
     memory_usage_mb=None,
     embed_dim=2048,
     license="apache-2.0",

mteb/models/model_implementations/yuan_models.py

@@ -20,6 +20,7 @@ yuan_embedding_2_zh = ModelMeta(
     revision="b5ebcace6f4fc6e5a4d1852557eb2dc2d1040cee",
     release_date="2025-11-24",
     n_parameters=326000000,
+    n_embedding_parameters=21_635_072,
     memory_usage_mb=1242,
     embed_dim=1792,
     license="apache-2.0",

mteb/models/model_implementations/yuan_models_en.py

@@ -1,6 +1,6 @@
 from mteb.models.instruct_wrapper import InstructSentenceTransformerModel
 from mteb.models.model_meta import ModelMeta
-from mteb.models.models_protocols import PromptType
+from mteb.types import PromptType
 
 
 def instruction_template(
@@ -43,6 +43,7 @@ yuan_embedding_2_en = ModelMeta(
     revision="b2fd15da3bcae3473c8529593825c15068f09fce",
     release_date="2025-11-27",
     n_parameters=595776512,
+    n_embedding_parameters=None,
     memory_usage_mb=2272,
     embed_dim=1024,
     max_tokens=2048,

mteb/models/model_meta.py

@@ -3,17 +3,16 @@ from __future__ import annotations
 import json
 import logging
 import warnings
-from collections.abc import Callable, Sequence
+from collections.abc import Callable
 from dataclasses import field
 from enum import Enum
 from functools import partial
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Literal, cast
 
+import numpy as np
 from huggingface_hub import (
-    GitCommitInfo,
     ModelCard,
-    ModelCardData,
     get_safetensors_metadata,
     hf_hub_download,
     list_repo_commits,
@@ -29,18 +28,27 @@ from huggingface_hub.errors import (
     SafetensorsParsingError,
 )
 from pydantic import BaseModel, ConfigDict, field_validator, model_validator
+from sentence_transformers.models import Transformer
+from torch import nn
 from transformers import AutoConfig
-from typing_extensions import Self
 
 from mteb._helpful_enum import HelpfulStrEnum
 from mteb.languages import check_language_code
-from mteb.models.models_protocols import EncoderProtocol, MTEBModels
+from mteb.models.models_protocols import MTEBModels
 from mteb.types import ISOLanguageScript, Licenses, Modalities, StrDate, StrURL
 
 if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from huggingface_hub import (
+        GitCommitInfo,
+        ModelCardData,
+    )
     from sentence_transformers import CrossEncoder, SentenceTransformer
+    from typing_extensions import Self
 
     from mteb.abstasks import AbsTask
+    from mteb.models.models_protocols import EncoderProtocol
 
 
 logger = logging.getLogger(__name__)
@@ -94,8 +102,9 @@ class ModelMeta(BaseModel):
         loader: The function that loads the model. If None it assumes that the model is not implemented.
         loader_kwargs: The keyword arguments to pass to the loader function.
         name: The name of the model, ideally the name on huggingface. It should be in the format "organization/model_name".
-        n_parameters: The number of parameters in the model, e.g. 7_000_000 for a 7M parameter model. Can be None if the number of parameters is not known (e.g. for proprietary models) or
-            if the loader returns a SentenceTransformer model from which it can be derived.
+        n_parameters: The total number of parameters in the model, e.g. `7_000_000` for a 7M parameter model. Can be none in case the number of parameters is unknown.
+        n_embedding_parameters: The number of parameters used for the embedding layer. Can be None if the number of embedding parameters is not known (e.g. for proprietary models).
+        n_active_parameters_override: The number of active parameters used bu model. Should be used **only** for Mixture of Experts models.
         memory_usage_mb: The memory usage of the model in MB. Can be None if the memory usage is not known (e.g. for proprietary models). To calculate it use the `calculate_memory_usage_mb` method.
         max_tokens: The maximum number of tokens the model can handle. Can be None if the maximum number of tokens is not known (e.g. for proprietary
             models).
@@ -134,6 +143,8 @@ class ModelMeta(BaseModel):
     release_date: StrDate | None
     languages: list[ISOLanguageScript] | None
     n_parameters: int | None
+    n_active_parameters_override: int | None = None
+    n_embedding_parameters: int | None = None
     memory_usage_mb: float | None
     max_tokens: float | None
     embed_dim: int | None
@@ -192,6 +203,16 @@ class ModelMeta(BaseModel):
         """
         return "cross-encoder" in self.model_type
 
+    @property
+    def n_active_parameters(self):
+        """Number of active parameters. Assumed to be `n_parameters - n_embedding_parameters`. Can be overwritten using `n_active_parameters_override` e.g. for MoE models."""
+        if self.n_active_parameters_override is not None:
+            return self.n_active_parameters_override
+
+        if self.n_parameters is not None and self.n_embedding_parameters is not None:
+            return self.n_parameters - self.n_embedding_parameters
+        return None
+
     @field_validator("similarity_fn_name", mode="before")
     @classmethod
     def _validate_similarity_fn_name(cls, value: str) -> ScoringFunction | None:
@@ -384,6 +405,14 @@ class ModelMeta(BaseModel):
             else model.model_card_data.base_model
         )
         meta = cls._from_hub(name, revision, compute_metadata)
+        try:
+            first = model[0]
+
+            if isinstance(first, Transformer):
+                emb = first.auto_model.get_input_embeddings()
+                meta.n_embedding_parameters = int(np.prod(emb.weight.shape))
+        except Exception as e:
+            logger.warning(f"Could not calculate embedding parameters for {name}: {e}")
         meta.revision = model.model_card_data.base_model_revision or meta.revision
         meta.max_tokens = model.max_seq_length
         meta.embed_dim = model.get_sentence_embedding_dimension()
@@ -455,6 +484,15 @@ class ModelMeta(BaseModel):
         from mteb.models import CrossEncoderWrapper
 
         meta = cls._from_hub(model.model.name_or_path, revision, compute_metadata)
+        try:
+            emb = model.model.get_input_embeddings()
+
+            if isinstance(emb, nn.Embedding):
+                meta.n_embedding_parameters = int(np.prod(emb.weight.shape))
+        except Exception as e:
+            logger.warning(
+                f"Could not calculate embedding parameters for {model.model.name_or_path}: {e}"
+            )
         meta.revision = model.config._commit_hash or meta.revision
         meta.loader = CrossEncoderWrapper
         meta.embed_dim = None
@@ -479,7 +517,7 @@ class ModelMeta(BaseModel):
         if isinstance(tasks[0], str):
             benchmark_datasets = set(tasks)
         else:
-            tasks = cast(Sequence["AbsTask"], tasks)
+            tasks = cast("Sequence[AbsTask]", tasks)
             benchmark_datasets = set()
             for task in tasks:
                 benchmark_datasets.add(task.metadata.name)
@@ -534,7 +572,7 @@ class ModelMeta(BaseModel):
         if isinstance(tasks[0], str):
             benchmark_datasets = set(tasks)
         else:
-            tasks = cast(Sequence["AbsTask"], tasks)
+            tasks = cast("Sequence[AbsTask]", tasks)
             benchmark_datasets = {task.metadata.name for task in tasks}
         overlap = training_datasets & benchmark_datasets
         perc_overlap = 100 * (len(overlap) / len(benchmark_datasets))

mteb/models/models_protocols.py

@@ -1,22 +1,23 @@
-from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
+from __future__ import annotations
 
-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,
-    TopRankedDocumentsType,
-)
+from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
 
 if TYPE_CHECKING:
+    from torch.utils.data import DataLoader
+    from typing_extensions import Unpack
+
+    from mteb.abstasks.task_metadata import TaskMetadata
     from mteb.models.model_meta import ModelMeta
+    from mteb.types import (
+        Array,
+        BatchedInput,
+        CorpusDatasetType,
+        EncodeKwargs,
+        PromptType,
+        QueryDatasetType,
+        RetrievalOutputType,
+        TopRankedDocumentsType,
+    )
 
 
 @runtime_checkable
@@ -31,6 +32,7 @@ class SearchProtocol(Protocol):
         hf_split: str,
         hf_subset: str,
         encode_kwargs: EncodeKwargs,
+        num_proc: int,
     ) -> None:
         """Index the corpus for retrieval.
 
@@ -40,6 +42,7 @@ class SearchProtocol(Protocol):
             hf_split: Split of current task, allows to know some additional information about current split.
             hf_subset: Subset of current task. Similar to `hf_split` to get more information
             encode_kwargs: Additional arguments to pass to the encoder during indexing.
+            num_proc: Number of processes to use for dataloading.
         """
         ...
 
@@ -53,6 +56,7 @@ class SearchProtocol(Protocol):
         top_k: int,
         encode_kwargs: EncodeKwargs,
         top_ranked: TopRankedDocumentsType | None = None,
+        num_proc: int,
     ) -> RetrievalOutputType:
         """Search the corpus using the given queries.
 
@@ -65,6 +69,7 @@ class SearchProtocol(Protocol):
                 Passed only from Reranking tasks.
             top_k: Number of top documents to return for each query.
             encode_kwargs: Additional arguments to pass to the encoder during indexing.
+            num_proc: Number of processes to use for dataloading.
 
         Returns:
             Dictionary with query IDs as keys with dict as values, where each value is a mapping of document IDs to their relevance scores.
@@ -72,7 +77,7 @@ class SearchProtocol(Protocol):
         ...
 
     @property
-    def mteb_model_meta(self) -> "ModelMeta":
+    def mteb_model_meta(self) -> ModelMeta:
         """Metadata of the model"""
         ...
 
@@ -177,7 +182,7 @@ class EncoderProtocol(Protocol):
         ...
 
     @property
-    def mteb_model_meta(self) -> "ModelMeta":
+    def mteb_model_meta(self) -> ModelMeta:
         """Metadata of the model"""
         ...
 
@@ -236,7 +241,7 @@ class CrossEncoderProtocol(Protocol):
         ...
 
     @property
-    def mteb_model_meta(self) -> "ModelMeta":
+    def mteb_model_meta(self) -> ModelMeta:
         """Metadata of the model"""
         ...
 

mteb/models/search_encoder_index/search_backend_protocol.py

@@ -1,7 +1,11 @@
-from collections.abc import Callable
-from typing import Protocol
+from __future__ import annotations
 
-from mteb.types import Array, TopRankedDocumentsType
+from typing import TYPE_CHECKING, Protocol
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from mteb.types import Array, TopRankedDocumentsType
 
 
 class IndexEncoderSearchProtocol(Protocol):

mteb/models/search_encoder_index/search_indexes/faiss_search_index.py

@@ -1,14 +1,23 @@
+from __future__ import annotations
+
 import logging
 import warnings
-from collections.abc import Callable
+from typing import TYPE_CHECKING
 
 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
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    import faiss
+
+    from mteb.models.models_protocols import EncoderProtocol
+    from mteb.types import Array, TopRankedDocumentsType
+
 
 logger = logging.getLogger(__name__)
 
@@ -33,7 +42,6 @@ class FaissSearchIndex:
             install_instruction="pip install mteb[faiss-cpu]",
         )
 
-        import faiss
         from faiss import IndexFlatIP, IndexFlatL2
 
         # https://github.com/facebookresearch/faiss/wiki/Faiss-indexes

mteb/models/search_wrappers.py

@@ -1,28 +1,35 @@
+from __future__ import annotations
+
 import heapq
 import logging
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import torch
 from datasets import Dataset
-from torch.utils.data import DataLoader
 
 from mteb._create_dataloaders import (
     create_dataloader,
 )
-from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.types import (
-    Array,
-    BatchedInput,
-    CorpusDatasetType,
-    EncodeKwargs,
     PromptType,
-    QueryDatasetType,
-    RetrievalOutputType,
-    TopRankedDocumentsType,
 )
 
-from .models_protocols import CrossEncoderProtocol, EncoderProtocol
-from .search_encoder_index.search_backend_protocol import IndexEncoderSearchProtocol
+if TYPE_CHECKING:
+    from torch.utils.data import DataLoader
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.types import (
+        Array,
+        BatchedInput,
+        CorpusDatasetType,
+        EncodeKwargs,
+        QueryDatasetType,
+        RetrievalOutputType,
+        TopRankedDocumentsType,
+    )
+
+    from .models_protocols import CrossEncoderProtocol, EncoderProtocol
+    from .search_encoder_index.search_backend_protocol import IndexEncoderSearchProtocol
 
 logger = logging.getLogger(__name__)
 
@@ -52,6 +59,7 @@ class SearchEncoderWrapper:
         hf_split: str,
         hf_subset: str,
         encode_kwargs: EncodeKwargs,
+        num_proc: int = 1,
     ) -> None:
         """Index the corpus for retrieval.
 
@@ -61,6 +69,7 @@ class SearchEncoderWrapper:
             hf_split: Split of current task, allows to know some additional information about current split.
             hf_subset: Subset of current task. Similar to `hf_split` to get more information
             encode_kwargs: Additional arguments to pass to the encoder during indexing.
+            num_proc: Number of processes to use for dataloading.
         """
         # Always retain corpus for potential reranking or fallback flows
         self.task_corpus = corpus
@@ -70,6 +79,7 @@ class SearchEncoderWrapper:
                     corpus,
                     task_metadata,
                     prompt_type=PromptType.document,
+                    num_proc=num_proc,
                     **encode_kwargs,
                 ),
                 task_metadata=task_metadata,
@@ -91,6 +101,7 @@ class SearchEncoderWrapper:
         top_k: int,
         encode_kwargs: EncodeKwargs,
         top_ranked: TopRankedDocumentsType | None = None,
+        num_proc: int = 1,
     ) -> RetrievalOutputType:
         """Search the corpus for the given queries.
 
@@ -103,6 +114,7 @@ class SearchEncoderWrapper:
                 Passed only from Reranking tasks.
             top_k: Number of top documents to return for each query.
             encode_kwargs: Additional arguments to pass to the encoder during indexing.
+            num_proc: Number of processes to use for dataloading.
 
         Returns:
             Dictionary with query IDs as keys with dict as values, where each value is a mapping of document IDs to their relevance scores.
@@ -114,6 +126,7 @@ class SearchEncoderWrapper:
             queries,
             task_metadata,
             prompt_type=PromptType.query,
+            num_proc=num_proc,
             **encode_kwargs,
         )
 
@@ -472,6 +485,7 @@ class SearchCrossEncoderWrapper:
         hf_split: str,
         hf_subset: str,
         encode_kwargs: EncodeKwargs,
+        num_proc: int = 1,
     ) -> None:
         """Index the corpus for retrieval.
 
@@ -481,6 +495,7 @@ class SearchCrossEncoderWrapper:
             hf_split: Split of current task, allows to know some additional information about current split.
             hf_subset: Subset of current task. Similar to `hf_split` to get more information
             encode_kwargs: Additional arguments to pass to the encoder during indexing.
+            num_proc: Number of processes to use.
         """
         self.task_corpus = corpus
 
@@ -494,6 +509,7 @@ class SearchCrossEncoderWrapper:
         top_k: int,
         encode_kwargs: EncodeKwargs,
         top_ranked: TopRankedDocumentsType | None = None,
+        num_proc: int = 1,
     ) -> RetrievalOutputType:
         """Search the corpus using the given queries.
 
@@ -506,6 +522,7 @@ class SearchCrossEncoderWrapper:
                 Passed only from Reranking tasks.
             top_k: Number of top documents to return for each query.
             encode_kwargs: Additional arguments to pass to the encoder during indexing.
+            num_proc: Number of processes to use.
 
         Returns:
             Dictionary with query IDs as keys with dict as values, where each value is a mapping of document IDs to their relevance scores.
@@ -539,12 +556,14 @@ class SearchCrossEncoderWrapper:
             Dataset.from_list(total_queries),
             task_metadata,
             prompt_type=PromptType.document,
+            num_proc=num_proc,
             **encode_kwargs,
         )
         corpus_loader = create_dataloader(
             Dataset.from_list(total_docs),
             task_metadata,
             prompt_type=PromptType.document,
+            num_proc=num_proc,
             **encode_kwargs,
         )
         predictions = self.model.predict(

mteb/models/sentence_transformer_wrapper.py

@@ -7,19 +7,20 @@ 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, EncodeKwargs, PromptType
+from mteb.types import PromptType
 
 from .abs_encoder import AbsEncoder
 
 if TYPE_CHECKING:
     from sentence_transformers import CrossEncoder, SentenceTransformer
+    from torch.utils.data import DataLoader
+    from typing_extensions import Unpack
 
     from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.types import Array, BatchedInput, EncodeKwargs
 
 logger = logging.getLogger(__name__)
 

mteb/models/vllm_wrapper.py

@@ -4,23 +4,25 @@ import atexit
 import gc
 import logging
 import os
-from collections.abc import Callable
 from typing import TYPE_CHECKING, Any, Literal
 
 import numpy as np
 import torch
-from torch.utils.data import DataLoader
 
 from mteb._requires_package import requires_package
-from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.models import ModelMeta
 from mteb.models.abs_encoder import AbsEncoder
-from mteb.types import Array, BatchedInput, PromptType
+from mteb.types import PromptType
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from torch.utils.data import DataLoader
     from vllm.config import PoolerConfig  # type: ignore[import-not-found]
-else:
-    PoolerConfig = dict[str, Any]
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.types import Array, BatchedInput
+
 
 logger = logging.getLogger(__name__)
 

mteb/results/benchmark_results.py

@@ -4,34 +4,39 @@ import functools
 import json
 import logging
 import warnings
-from collections.abc import Callable, Iterable, Iterator
 from pathlib import Path
-from typing import Any, Literal, cast
+from typing import TYPE_CHECKING, Any, Literal, cast
 
 import pandas as pd
 from packaging.version import InvalidVersion, Version
 from pydantic import BaseModel, ConfigDict
-from typing_extensions import Self
 
-from mteb.abstasks.abstask import AbsTask
-from mteb.abstasks.task_metadata import (
-    TaskDomain,
-    TaskType,
-)
 from mteb.benchmarks.benchmark import Benchmark
 from mteb.models import ModelMeta
 from mteb.models.get_model_meta import get_model_metas
-from mteb.types import (
-    ISOLanguage,
-    ISOLanguageScript,
-    Modalities,
-    Score,
-    ScoresDict,
-    SplitName,
-)
 
 from .model_result import ModelResult, _aggregate_and_pivot
 
+if TYPE_CHECKING:
+    from collections.abc import Callable, Iterable, Iterator
+
+    from typing_extensions import Self
+
+    from mteb.abstasks.abstask import AbsTask
+    from mteb.abstasks.task_metadata import (
+        TaskDomain,
+        TaskType,
+    )
+    from mteb.types import (
+        ISOLanguage,
+        ISOLanguageScript,
+        Modalities,
+        Score,
+        ScoresDict,
+        SplitName,
+    )
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -144,7 +149,7 @@ class BenchmarkResults(BaseModel):
                     raise ValueError("name in ModelMeta is None. It must be a string.")
                 name_rev[name.name] = name.revision
             else:
-                name_ = cast(str, name)
+                name_ = cast("str", name)
                 name_rev[name_] = revision
 
         for model_res in self.model_results:

mteb/results/model_result.py

@@ -2,30 +2,36 @@ from __future__ import annotations
 
 import logging
 import warnings
-from collections.abc import Callable, Iterable
-from typing import Any, Literal, cast
+from typing import TYPE_CHECKING, Any, Literal, cast
 
 import numpy as np
 import pandas as pd
 from pydantic import BaseModel, ConfigDict, Field
 from typing_extensions import overload
 
-from mteb.abstasks.abstask import AbsTask
-from mteb.abstasks.task_metadata import (
-    TaskDomain,
-    TaskType,
-)
 from mteb.types import (
-    ISOLanguage,
-    ISOLanguageScript,
     Modalities,
-    Score,
-    ScoresDict,
-    SplitName,
 )
 
 from .task_result import TaskError, TaskResult
 
+if TYPE_CHECKING:
+    from collections.abc import Callable, Iterable
+
+    from mteb.abstasks.abstask import AbsTask
+    from mteb.abstasks.task_metadata import (
+        TaskDomain,
+        TaskType,
+    )
+    from mteb.types import (
+        ISOLanguage,
+        ISOLanguageScript,
+        Score,
+        ScoresDict,
+        SplitName,
+    )
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -83,7 +89,7 @@ class ModelResult(BaseModel):
     model_revision: str | None
     task_results: list[TaskResult]
     default_modalities: list[Modalities] = Field(
-        default_factory=lambda: [cast(Modalities, "text")], alias="modalities"
+        default_factory=lambda: [cast("Modalities", "text")], alias="modalities"
     )
     model_config = (
         ConfigDict(  # to free up the name model_* which is otherwise protected
@@ -202,8 +208,8 @@ class ModelResult(BaseModel):
             aggregation = aggregation if aggregation is not None else np.mean
         else:
             use_fast = True
-        aggregation = cast(Callable[[list[Score]], Any], aggregation)
-        getter = cast(Callable[[ScoresDict], Score], getter)
+        aggregation = cast("Callable[[list[Score]], Any]", aggregation)
+        getter = cast("Callable[[ScoresDict], Score]", getter)
 
         if format == "wide":
             scores = {}