mteb 2.5.1__py3-none-any.whl → 2.5.3__py3-none-any.whl

mteb/models/model_meta.py

@@ -26,7 +26,7 @@ from huggingface_hub.errors import (
     RepositoryNotFoundError,
     SafetensorsParsingError,
 )
-from pydantic import BaseModel, ConfigDict, field_validator
+from pydantic import BaseModel, ConfigDict, field_validator, model_validator
 from transformers import AutoConfig
 from typing_extensions import Self
 
@@ -57,6 +57,8 @@ FRAMEWORKS = Literal[
     "ColPali",
 ]
 
+MODEL_TYPES = Literal["dense", "cross-encoder", "late-interaction"]
+
 
 class ScoringFunction(HelpfulStrEnum):
     """The scoring function used by the models."""
@@ -114,7 +116,7 @@ class ModelMeta(BaseModel):
             a benchmark as well as mark dataset contaminations.
         adapted_from: Name of the model from which this model is adapted. For quantizations, fine-tunes, long doc extensions, etc.
         superseded_by: Name of the model that supersedes this model, e.g., nvidia/NV-Embed-v2 supersedes v1.
-        is_cross_encoder: Whether the model can act as a cross-encoder or not.
+        model_type: A list of strings representing the type of model.
         modalities: A list of strings representing the modalities the model supports. Default is ["text"].
         contacts: The people to contact in case of a problem in the model, preferably a GitHub handle.
     """
@@ -144,10 +146,49 @@ class ModelMeta(BaseModel):
     adapted_from: str | None = None
     superseded_by: str | None = None
     modalities: list[Modalities] = ["text"]
-    is_cross_encoder: bool | None = None
+    model_type: list[MODEL_TYPES] = ["dense"]
     citation: str | None = None
     contacts: list[str] | None = None
 
+    @model_validator(mode="before")
+    @classmethod
+    def handle_legacy_is_cross_encoder(cls, data: Any) -> Any:
+        """Handle legacy is_cross_encoder field by converting it to model_type.
+
+        This validator handles backward compatibility for the deprecated is_cross_encoder field.
+        If is_cross_encoder=True is provided, it adds "cross_encoder" to model_type.
+        """
+        if isinstance(data, dict) and "is_cross_encoder" in data:
+            is_cross_encoder_value = data.pop("is_cross_encoder")
+
+            if is_cross_encoder_value is not None:
+                warnings.warn(
+                    "is_cross_encoder is deprecated and will be removed in a future version. "
+                    "Use model_type=['cross-encoder'] instead.",
+                    DeprecationWarning,
+                    stacklevel=2,
+                )
+
+                model_type = data.get("model_type", ["dense"])
+
+                if is_cross_encoder_value:
+                    if "cross-encoder" not in model_type:
+                        data["model_type"] = ["cross-encoder"]
+                else:
+                    if "cross-encoder" in model_type:
+                        model_type = [t for t in model_type if t != "cross-encoder"]
+                        data["model_type"] = model_type if model_type else ["dense"]
+
+        return data
+
+    @property
+    def is_cross_encoder(self) -> bool:
+        """Returns True if the model is a cross-encoder.
+
+        Derived from model_type field. A model is considered a cross-encoder if "cross-encoder" is in its model_type list.
+        """
+        return "cross-encoder" in self.model_type
+
     @field_validator("similarity_fn_name", mode="before")
     @classmethod
     def _validate_similarity_fn_name(cls, value: str) -> ScoringFunction | None:
@@ -183,6 +224,7 @@ class ModelMeta(BaseModel):
             else dict_repr["training_datasets"]
         )
         dict_repr["loader"] = _get_loader_name(loader)
+        dict_repr["is_cross_encoder"] = self.is_cross_encoder
         return dict_repr
 
     @field_validator("languages")
@@ -425,6 +467,7 @@ class ModelMeta(BaseModel):
         meta.loader = CrossEncoderWrapper
         meta.embed_dim = None
         meta.modalities = ["text"]
+        meta.model_type = ["cross-encoder"]
         return meta
 
     def is_zero_shot_on(self, tasks: Sequence[AbsTask] | Sequence[str]) -> bool | None:
@@ -468,7 +511,9 @@ 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()

mteb/models/search_encoder_index/search_indexes/faiss_search_index.py

@@ -1,4 +1,5 @@
 import logging
+import warnings
 from collections.abc import Callable
 
 import numpy as np
@@ -127,7 +128,9 @@ 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

mteb/models/search_wrappers.py

@@ -340,7 +340,8 @@ class SearchEncoderWrapper:
         for query_idx, query_embedding in enumerate(query_embeddings):
             query_id = query_idx_to_id[query_idx]
             if query_id not in top_ranked:
-                logger.warning(f"No pre-ranked documents found for query {query_id}")
+                msg = f"No pre-ranked documents found for query {query_id}"
+                logger.warning(msg)
                 continue
 
             ranked_ids = top_ranked[query_id]
@@ -511,7 +512,8 @@ class SearchCrossEncoderWrapper:
         doc_pairs_ids: list[tuple[str, str]] = []
         for query_id, corpus_ids in top_ranked.items():
             if query_id not in top_ranked:
-                logger.warning(f"No pre-ranked documents found for query {query_id}")
+                msg = f"No pre-ranked documents found for query {query_id}"
+                logger.warning(msg)
                 continue
 
             query_idx = query_id_to_idx[query_id]

mteb/models/sentence_transformer_wrapper.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+import warnings
 from typing import TYPE_CHECKING, Any
 
 import numpy as np
@@ -75,9 +76,9 @@ class SentenceTransformerEncoderWrapper(AbsEncoder):
         if built_in_prompts and not model_prompts:
             model_prompts = built_in_prompts
         elif model_prompts and built_in_prompts:
-            logger.warning(
-                f"Model prompts specified, these will overwrite the default model prompts. Current prompts will be:\n {model_prompts}"
-            )
+            msg = f"Model prompts specified, these will overwrite the default model prompts. Current prompts will be:\n {model_prompts}"
+            logger.warning(msg)
+            warnings.warn(msg)
             self.model.prompts = model_prompts
 
         self.model_prompts, invalid_prompts = self.validate_task_to_prompt_name(
@@ -86,9 +87,9 @@ class SentenceTransformerEncoderWrapper(AbsEncoder):
 
         if invalid_prompts:
             invalid_prompts = "\n".join(invalid_prompts)
-            logger.warning(
-                f"Some prompts are not in the expected format and will be ignored. Problems:\n\n{invalid_prompts}"
-            )
+            msg = f"Some prompts are not in the expected format and will be ignored. Problems:\n\n{invalid_prompts}"
+            logger.warning(msg)
+            warnings.warn(msg)
 
         if (
             self.model_prompts
@@ -98,10 +99,9 @@ class SentenceTransformerEncoderWrapper(AbsEncoder):
                 or PromptType.document.value not in self.model_prompts
             )
         ):
-            logger.warning(
-                "SentenceTransformers that use prompts most often need to be configured with at least 'query' and"
-                f" 'document' prompts to ensure optimal performance. Received {self.model_prompts}"
-            )
+            msg = f"SentenceTransformers that use prompts most often need to be configured with at least 'query' and 'document' prompts to ensure optimal performance. Received {self.model_prompts}"
+            logger.warning(msg)
+            warnings.warn(msg)
 
         if hasattr(self.model, "similarity") and callable(self.model.similarity):
             self.similarity = self.model.similarity

mteb/results/benchmark_results.py

@@ -1,3 +1,4 @@
+import functools
 import json
 import logging
 import warnings
@@ -32,6 +33,24 @@ from .model_result import ModelResult, _aggregate_and_pivot
 logger = logging.getLogger(__name__)
 
 
+# Global cache for model metas and version parsing
+@functools.lru_cache
+def _get_cached_model_metas() -> dict[str, str | None]:
+    """Cache model metas to avoid repeated calls."""
+    return {meta.name: meta.revision for meta in get_model_metas()}
+
+
+@functools.lru_cache(maxsize=10000)
+def _parse_version_cached(version_str: str | None) -> Version | None:
+    """Cache version parsing to avoid repeated parsing."""
+    if version_str is None:
+        return None
+    try:
+        return Version(version_str)
+    except (InvalidVersion, TypeError):
+        return None
+
+
 class BenchmarkResults(BaseModel):
     """Data class to hold the benchmark results of a model.
 
@@ -174,40 +193,6 @@ class BenchmarkResults(BaseModel):
         Returns:
             A new BenchmarkResults object with the revisions joined.
         """
-
-        def parse_version(version_str: str) -> Version | None:
-            try:
-                return Version(version_str)
-            except (InvalidVersion, TypeError):
-                return None
-
-        def keep_best(group: pd.DataFrame) -> pd.DataFrame:
-            # Filtering out task_results where no scores are present
-            group = group[group["has_scores"]]
-            is_main_revision = group["revision"] == group["main_revision"]
-            # If the main revision is present we select that
-            if is_main_revision.sum() > 0:
-                return group[is_main_revision].head(n=1)
-            unique_revisions = group["revision"].unique()
-
-            # ensure None/NA/"external" revisions is filtered out
-            group.loc[group["revision"].isna(), "revision"] = "no_revision_available"
-            group.loc[group["revision"] == "external", "revision"] = (
-                "no_revision_available"
-            )
-
-            # Filtering out no_revision_available if other revisions are present
-            if (len(unique_revisions) > 1) and (
-                "no_revision_available" in unique_revisions
-            ):
-                group = group[group["revision"] != "no_revision_available"]
-            # If there are any not-NA mteb versions, we select the latest one
-            if group["mteb_version"].notna().any():
-                group = group.dropna(subset=["mteb_version"])
-                group = group.sort_values("mteb_version", ascending=False)
-                return group.head(n=1)
-            return group.head(n=1)
-
         records = []
         for model_result in self:
             for task_result in model_result.task_results:
@@ -224,17 +209,54 @@ class BenchmarkResults(BaseModel):
         if not records:
             return BenchmarkResults.model_construct(model_results=[])
         task_df = pd.DataFrame.from_records(records)
-        model_to_main_revision = {
-            meta.name: meta.revision for meta in get_model_metas()
-        }
+
+        # Use cached model metas
+        model_to_main_revision = _get_cached_model_metas()
         task_df["main_revision"] = task_df["model"].map(model_to_main_revision)  # type: ignore
-        task_df["mteb_version"] = task_df["mteb_version"].map(parse_version)  # type: ignore
-        task_df = (
-            task_df.groupby(["model", "task_name"])
-            .apply(keep_best)
-            .reset_index(drop=True)
+
+        # Use cached version parsing
+        task_df["mteb_version"] = task_df["mteb_version"].map(_parse_version_cached)  # type: ignore
+
+        # Filter out rows without scores first
+        task_df = task_df[task_df["has_scores"]]
+
+        # Optimize groupby with vectorized operations
+        # Sort by priority: main_revision match, then mteb_version (descending), then revision
+        task_df["is_main_revision"] = task_df["revision"] == task_df["main_revision"]
+
+        # Handle None/NA/external revisions
+        task_df["revision_clean"] = task_df["revision"].copy()
+        task_df.loc[task_df["revision"].isna(), "revision_clean"] = (
+            "no_revision_available"
+        )
+        task_df.loc[task_df["revision"] == "external", "revision_clean"] = (
+            "no_revision_available"
         )
+
+        # Create a priority column for sorting
+        # Higher priority = better to keep
+        # Priority: main_revision (1000), has valid mteb_version (100), has valid revision (10)
+        task_df["priority"] = 0
+        task_df.loc[task_df["is_main_revision"], "priority"] += 1000
+        task_df.loc[task_df["mteb_version"].notna(), "priority"] += 100
+        task_df.loc[
+            task_df["revision_clean"] != "no_revision_available", "priority"
+        ] += 10
+
+        # Sort by priority (desc), mteb_version (desc), and take first per group
+        task_df = task_df.sort_values(
+            ["model", "task_name", "priority", "mteb_version"],
+            ascending=[True, True, False, False],
+            na_position="last",
+        )
+
+        task_df = task_df.groupby(["model", "task_name"], as_index=False).first()
+
+        # Reconstruct model results
         model_results = []
+        # Group by original revision to maintain deterministic behavior
+        # After the first() selection above, each (model, task_name) is unique,
+        # so grouping by original revision ensures consistent ModelResult creation
         for (model, model_revision), group in task_df.groupby(["model", "revision"]):
             model_result = ModelResult.model_construct(
                 model_name=model,
@@ -342,7 +364,9 @@ class BenchmarkResults(BaseModel):
             scores_data.extend(model_result._get_score_for_table())
 
         if not scores_data:
-            logger.warning("No scores data available. Returning empty DataFrame.")
+            msg = "No scores data available. Returning empty DataFrame."
+            logger.warning(msg)
+            warnings.warn(msg)
             return pd.DataFrame()
 
         # Create DataFrame

mteb/results/model_result.py

@@ -292,7 +292,9 @@ class ModelResult(BaseModel):
         scores_data = self._get_score_for_table()
 
         if not scores_data:
-            logger.warning("No scores data available. Returning empty DataFrame.")
+            msg = "No scores data available. Returning empty DataFrame."
+            logger.warning(msg)
+            warnings.warn(msg)
             return pd.DataFrame()
 
         # Create DataFrame

mteb/results/task_result.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import json
 import logging
+import warnings
 from argparse import Namespace
 from collections import defaultdict
 from collections.abc import Callable, Iterable
@@ -462,7 +463,9 @@ class TaskResult(BaseModel):
                     if main_score in hf_subset_scores:
                         hf_subset_scores["main_score"] = hf_subset_scores[main_score]
                     else:
-                        logger.warning(f"Main score {main_score} not found in scores")
+                        msg = f"Main score {main_score} not found in scores"
+                        logger.warning(msg)
+                        warnings.warn(msg)
                         hf_subset_scores["main_score"] = None
 
         # specific fixes:
@@ -633,21 +636,23 @@ class TaskResult(BaseModel):
             task = get_task(self.task_name)
 
         splits = task.eval_splits
-        hf_subsets = task.hf_subsets
-        hf_subsets = set(hf_subsets)
+        hf_subsets = set(task.hf_subsets)  # Convert to set once
 
         new_scores = {}
         seen_splits = set()
         for split in self.scores:
             if split not in splits:
                 continue
-            new_scores[split] = []
             seen_subsets = set()
-            for _scores in self.scores[split]:
-                if _scores["hf_subset"] not in hf_subsets:
-                    continue
-                new_scores[split].append(_scores)
+            # Use list comprehension for better performance
+            new_scores[split] = [
+                _scores
+                for _scores in self.scores[split]
+                if _scores["hf_subset"] in hf_subsets
+            ]
+            for _scores in new_scores[split]:
                 seen_subsets.add(_scores["hf_subset"])
+
             if seen_subsets != hf_subsets:
                 missing_subsets = hf_subsets - seen_subsets
                 if len(missing_subsets) > 2:
@@ -656,17 +661,17 @@ class TaskResult(BaseModel):
                 else:
                     missing_subsets_str = str(missing_subsets)
 
-                logger.warning(
-                    f"{task.metadata.name}: Missing subsets {missing_subsets_str} for split {split}"
-                )
+                msg = f"{task.metadata.name}: Missing subsets {missing_subsets_str} for split {split}"
+                logger.warning(msg)
+                warnings.warn(msg)
             seen_splits.add(split)
         if seen_splits != set(splits):
-            logger.warning(
-                f"{task.metadata.name}: Missing splits {set(splits) - seen_splits}"
-            )
-        new_res = {**self.to_dict(), "scores": new_scores}
-        new_res = TaskResult.from_validated(**new_res)
-        return new_res
+            msg = f"{task.metadata.name}: Missing splits {set(splits) - seen_splits}"
+            logger.warning(msg)
+            warnings.warn(msg)
+        data = self.model_dump()
+        data["scores"] = new_scores
+        return type(self).model_construct(**data)
 
     def is_mergeable(
         self,

{mteb-2.5.1.dist-info → mteb-2.5.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mteb
-Version: 2.5.1
+Version: 2.5.3
 Summary: Massive Text Embedding Benchmark
 Author-email: MTEB Contributors <niklas@huggingface.co>, Kenneth Enevoldsen <kenneth.enevoldsen@cas.au.dk>, Nouamane Tazi <nouamane@huggingface.co>, Nils Reimers <info@nils-reimers.de>
 Maintainer-email: Kenneth Enevoldsen <kenneth.enevoldsen@cas.au.dk>, Roman Solomatin <risolomatin@gmail.com>, Isaac Chung <chungisaac1217@gmail.com>