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

mteb/models/model_meta.py

@@ -3,42 +3,52 @@ 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,
+    model_info,
     repo_exists,
 )
 from huggingface_hub.errors import (
     EntryNotFoundError,
     GatedRepoError,
+    HFValidationError,
     NotASafetensorsRepoError,
     RepositoryNotFoundError,
     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__)
@@ -55,6 +65,10 @@ FRAMEWORKS = Literal[
     "PyLate",
     "ColBERT",
     "ColPali",
+    "GGUF",
+    "safetensors",
+    "ONNX",
+    "Transformers",
 ]
 
 MODEL_TYPES = Literal["dense", "cross-encoder", "late-interaction"]
@@ -81,9 +95,6 @@ def _get_loader_name(
     return loader.__name__
 
 
-_SENTENCE_TRANSFORMER_LIB_NAME = "Sentence Transformers"
-
-
 class ModelMeta(BaseModel):
     """The model metadata object.
 
@@ -91,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).
@@ -131,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
@@ -189,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:
@@ -250,7 +274,7 @@ class ModelMeta(BaseModel):
             )
         return v
 
-    def load_model(self, **kwargs: Any) -> MTEBModels:
+    def load_model(self, device: str | None = None, **kwargs: Any) -> MTEBModels:
         """Loads the model using the specified loader function."""
         if self.loader is None:
             raise NotImplementedError(
@@ -262,11 +286,11 @@ class ModelMeta(BaseModel):
         # Allow overwrites
         _kwargs = self.loader_kwargs.copy()
         _kwargs.update(kwargs)
+        if device is not None:
+            _kwargs["device"] = device
 
-        model: EncoderProtocol = self.loader(
-            self.name, revision=self.revision, **_kwargs
-        )
-        model.mteb_model_meta = self  # type: ignore
+        model: MTEBModels = self.loader(self.name, revision=self.revision, **_kwargs)
+        model.mteb_model_meta = self  # type: ignore[misc]
         return model
 
     def model_name_as_path(self) -> str:
@@ -307,7 +331,7 @@ class ModelMeta(BaseModel):
         embedding_dim = None
         max_tokens = None
 
-        if model_name and compute_metadata and repo_exists(model_name):
+        if model_name and compute_metadata and _repo_exists(model_name):
             reference = "https://huggingface.co/" + model_name
             card = ModelCard.load(model_name)
             card_data: ModelCardData = card.data
@@ -318,22 +342,17 @@ class ModelMeta(BaseModel):
                 model_config = None
                 logger.warning(f"Can't get configuration for {model_name}. Error: {e}")
 
-            if (
-                card_data.library_name == _SENTENCE_TRANSFORMER_LIB_NAME
-                or _SENTENCE_TRANSFORMER_LIB_NAME in card_data.tags
-            ):
-                frameworks.append(_SENTENCE_TRANSFORMER_LIB_NAME)
-            else:
-                msg = "Model library not recognized, defaulting to Sentence Transformers loader."
-                logger.warning(msg)
-                warnings.warn(msg)
+            hf_frameworks = (
+                cls._get_frameworks_from_hf_tags(model_name) if model_name else []
+            )
+            frameworks.extend(hf_frameworks)
 
             if revision is None:
                 revisions = _get_repo_commits(model_name, "model")
                 revision = revisions[0].commit_id if revisions else None
 
             release_date = cls.fetch_release_date(model_name)
-            model_license = card_data.license
+            model_license = card_data.license if card_data.license != "other" else None
             n_parameters = cls._calculate_num_parameters_from_hub(model_name)
             memory_usage_mb = cls._calculate_memory_usage_mb(model_name, n_parameters)
             if model_config and hasattr(model_config, "hidden_size"):
@@ -386,8 +405,14 @@ class ModelMeta(BaseModel):
             else model.model_card_data.base_model
         )
         meta = cls._from_hub(name, revision, compute_metadata)
-        if _SENTENCE_TRANSFORMER_LIB_NAME not in meta.framework:
-            meta.framework.append("Sentence Transformers")
+        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()
@@ -413,11 +438,9 @@ class ModelMeta(BaseModel):
             The generated ModelMeta.
         """
         meta = cls._from_hub(model, revision, compute_metadata)
-        if _SENTENCE_TRANSFORMER_LIB_NAME not in meta.framework:
-            meta.framework.append("Sentence Transformers")
         meta.modalities = ["text"]
 
-        if model and compute_metadata and repo_exists(model):
+        if model and compute_metadata and _repo_exists(model):
             # have max_seq_length field
             sbert_config = _get_json_from_hub(
                 model, "sentence_bert_config.json", "model", revision=revision
@@ -435,7 +458,7 @@ class ModelMeta(BaseModel):
                 and config_sbert.get("similarity_fn_name") is not None
             ):
                 meta.similarity_fn_name = ScoringFunction.from_str(
-                    config_sbert.get("similarity_fn_name")
+                    config_sbert["similarity_fn_name"]
                 )
             else:
                 meta.similarity_fn_name = ScoringFunction.COSINE
@@ -461,8 +484,15 @@ class ModelMeta(BaseModel):
         from mteb.models import CrossEncoderWrapper
 
         meta = cls._from_hub(model.model.name_or_path, revision, compute_metadata)
-        if _SENTENCE_TRANSFORMER_LIB_NAME not in meta.framework:
-            meta.framework.append("Sentence Transformers")
+        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
@@ -487,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)
@@ -511,10 +541,12 @@ class ModelMeta(BaseModel):
                 if adapted_training_datasets is not None:
                     training_datasets |= adapted_training_datasets
             except (ValueError, KeyError) as e:
-                logger.warning(f"Could not get source model: {e} in MTEB")
+                msg = f"Could not get source model: {e} in MTEB"
+                logger.warning(msg)
+                warnings.warn(msg)
 
         return_dataset = training_datasets.copy()
-        visited = set()
+        visited: set[str] = set()
 
         for dataset in training_datasets:
             similar_tasks = _collect_similar_tasks(dataset, visited)
@@ -540,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))
@@ -548,6 +580,8 @@ class ModelMeta(BaseModel):
 
     @staticmethod
     def _calculate_num_parameters_from_hub(model_name: str | None = None) -> int | None:
+        if not model_name:
+            return None
         try:
             safetensors_metadata = get_safetensors_metadata(model_name)
             if len(safetensors_metadata.parameter_count) >= 0:
@@ -561,7 +595,7 @@ class ModelMeta(BaseModel):
             logger.warning(
                 f"Can't calculate number of parameters for {model_name}. Got error {e}"
             )
-            return None
+        return None
 
     def calculate_num_parameters_from_hub(self) -> int | None:
         """Calculates the number of parameters in the model.
@@ -624,7 +658,7 @@ class ModelMeta(BaseModel):
         if "API" in self.framework or self.name is None:
             return None
 
-        return self._calculate_memory_usage_mb(self.model_name, self.n_parameters)
+        return self._calculate_memory_usage_mb(self.name, self.n_parameters)
 
     @staticmethod
     def fetch_release_date(model_name: str) -> StrDate | None:
@@ -640,6 +674,43 @@ class ModelMeta(BaseModel):
             return release_date
         return None
 
+    @staticmethod
+    def _get_frameworks_from_hf_tags(model_name: str) -> list[FRAMEWORKS]:
+        """Extract frameworks supported by the model from HuggingFace model tags.
+
+        Args:
+            model_name: HuggingFace model name
+
+        Returns:
+            List of framework names found in tags. Defaults to empty list if no frameworks found.
+        """
+        try:
+            info = model_info(model_name)
+            if not info.tags:
+                return []
+        except Exception as e:
+            logger.warning(
+                f"Failed to fetch frameworks from HuggingFace tags for {model_name}: {e}"
+            )
+            return []
+
+        # Mapping from HuggingFace tags to MTEB framework names
+        tag_to_framework: dict[str, FRAMEWORKS] = {
+            "sentence-transformers": "Sentence Transformers",
+            "transformers": "Transformers",
+            "onnx": "ONNX",
+            "safetensors": "safetensors",
+            "gguf": "GGUF",
+        }
+
+        frameworks: list[FRAMEWORKS] = []
+
+        for framework_tag in tag_to_framework.keys():
+            if framework_tag in info.tags:
+                frameworks.append(tag_to_framework[framework_tag])
+
+        return frameworks
+
     def to_python(self) -> str:
         """Returns a string representation of the model."""
         return _pydantic_instance_to_code(self)
@@ -784,3 +855,19 @@ def _get_file_on_hub(
     except (GatedRepoError, RepositoryNotFoundError, EntryNotFoundError) as e:
         logger.warning(f"Can't get file {file_name} of {repo_id}: {e}")
         return None
+
+
+def _repo_exists(repo_id: str, repo_type: str | None = None) -> bool:
+    """Checks if a repository exists on HuggingFace Hub.
+
+    Repo exists will raise HFValidationError for invalid local paths
+
+    Args:
+        repo_id: The repository ID.
+        repo_type: The type of repository (e.g., "model", "dataset", "space").
+    """
+    try:
+        return repo_exists(repo_id=repo_id, repo_type=repo_type)
+    except HFValidationError as e:
+        logger.warning(f"Can't check existence of {repo_id}: {e}")
+        return False

mteb/models/models_protocols.py

@@ -1,20 +1,23 @@
-from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
-
-from torch.utils.data import DataLoader
+from __future__ import annotations
 
-from mteb.abstasks.task_metadata import TaskMetadata
-from mteb.types import (
-    Array,
-    BatchedInput,
-    CorpusDatasetType,
-    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
@@ -28,7 +31,8 @@ class SearchProtocol(Protocol):
         task_metadata: TaskMetadata,
         hf_split: str,
         hf_subset: str,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
+        num_proc: int,
     ) -> None:
         """Index the corpus for retrieval.
 
@@ -38,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.
         """
         ...
 
@@ -49,8 +54,9 @@ 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,
+        num_proc: int,
     ) -> RetrievalOutputType:
         """Search the corpus using the given queries.
 
@@ -63,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.
@@ -70,7 +77,7 @@ class SearchProtocol(Protocol):
         ...
 
     @property
-    def mteb_model_meta(self) -> "ModelMeta":
+    def mteb_model_meta(self) -> ModelMeta:
         """Metadata of the model"""
         ...
 
@@ -83,12 +90,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 +115,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.
 
@@ -168,7 +182,7 @@ class EncoderProtocol(Protocol):
         ...
 
     @property
-    def mteb_model_meta(self) -> "ModelMeta":
+    def mteb_model_meta(self) -> ModelMeta:
         """Metadata of the model"""
         ...
 
@@ -181,12 +195,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 +221,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.
 
@@ -220,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,13 +1,23 @@
+from __future__ import annotations
+
 import logging
-from collections.abc import Callable
+import warnings
+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__)
 
@@ -32,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
@@ -108,7 +117,7 @@ class FaissSearchIndex:
             ids = ids.tolist()
 
         if issubclass(self.index_type, faiss.IndexFlatL2):
-            similarities = -np.sqrt(np.maximum(similarities, 0))
+            similarities = (-np.sqrt(np.maximum(similarities, 0))).tolist()
 
         return similarities, ids
 
@@ -116,8 +125,8 @@ class FaissSearchIndex:
         self,
         embeddings: Array,
         top_k: int,
-        top_ranked: TopRankedDocumentsType | None = None,
-        query_idx_to_id: dict[int, str] | None = None,
+        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]] = []
@@ -127,15 +136,17 @@ class FaissSearchIndex:
             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}")
+                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
+            d = self.index.d  # type: ignore[union-attr]
             candidate_embs = np.vstack(
-                [self.index.reconstruct(idx) for idx in candidate_indices]
+                [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)