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

mteb/abstasks/abstask.py

@@ -1,10 +1,11 @@
 import json
 import logging
+import warnings
 from abc import ABC, abstractmethod
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
 from copy import copy
 from pathlib import Path
-from typing import Any, cast
+from typing import Any, Literal, cast
 
 import numpy as np
 from datasets import ClassLabel, Dataset, DatasetDict, load_dataset
@@ -78,8 +79,8 @@ class AbsTask(ABC):
     """
 
     metadata: TaskMetadata
-    abstask_prompt: str | None = None
-    _eval_splits: list[str] | None = None
+    abstask_prompt: str
+    _eval_splits: Sequence[str] | None = None
     dataset: dict[HFSubset, DatasetDict] | None = None
     data_loaded: bool = False
     hf_subsets: list[HFSubset]
@@ -102,9 +103,9 @@ class AbsTask(ABC):
     def check_if_dataset_is_superseded(self) -> None:
         """Check if the dataset is superseded by a newer version."""
         if self.superseded_by:
-            logger.warning(
-                f"Dataset '{self.metadata.name}' is superseded by '{self.superseded_by}', you might consider using the newer version of the dataset."
-            )
+            msg = f"Dataset '{self.metadata.name}' is superseded by '{self.superseded_by}'. We recommend using the newer version of the dataset unless you are running a specific benchmark. See `get_task('{self.superseded_by}').metadata.description` to get a description of the task and changes."
+            logger.warning(msg)
+            warnings.warn(msg)
 
     def dataset_transform(self):
         """A transform operations applied to the dataset after loading.
@@ -123,7 +124,7 @@ class AbsTask(ABC):
         encode_kwargs: dict[str, Any],
         prediction_folder: Path | None = None,
         **kwargs: Any,
-    ) -> dict[HFSubset, ScoresDict]:
+    ) -> Mapping[HFSubset, ScoresDict]:
         """Evaluates an MTEB compatible model on the task.
 
         Args:
@@ -195,12 +196,12 @@ class AbsTask(ABC):
     @abstractmethod
     def _evaluate_subset(
         self,
-        model: EncoderProtocol,
+        model: MTEBModels,
         data_split: Dataset,
         *,
-        encode_kwargs: dict[str, Any],
         hf_split: str,
         hf_subset: str,
+        encode_kwargs: dict[str, Any],
         prediction_folder: Path | None = None,
         **kwargs: Any,
     ) -> ScoresDict:
@@ -210,7 +211,7 @@ class AbsTask(ABC):
 
     def _save_task_predictions(
         self,
-        predictions: dict[str, Any] | list[Any],
+        predictions: Mapping[str, Any] | list[Any],
         model: MTEBModels,
         prediction_folder: Path,
         hf_split: str,
@@ -226,7 +227,7 @@ class AbsTask(ABC):
             hf_subset: The subset of the dataset (e.g. "en").
         """
         predictions_path = self._predictions_path(prediction_folder)
-        existing_results = {
+        existing_results: dict[str, Any] = {
             "mteb_model_meta": {
                 "model_name": model.mteb_model_meta.name,
                 "revision": model.mteb_model_meta.revision,
@@ -326,7 +327,7 @@ class AbsTask(ABC):
                     )
         else:
             # some of monolingual datasets explicitly adding the split name to the dataset name
-            self.dataset = load_dataset(**self.metadata.dataset)  # type: ignore
+            self.dataset = load_dataset(**self.metadata.dataset)
         self.dataset_transform()
         self.data_loaded = True
 
@@ -362,15 +363,19 @@ class AbsTask(ABC):
         """
         from mteb.abstasks import AbsTaskClassification
 
-        if self.metadata.descriptive_stat_path.exists() and not overwrite_results:
+        existing_stats = self.metadata.descriptive_stats
+
+        if existing_stats is not None and not overwrite_results:
             logger.info("Loading metadata descriptive statistics from cache.")
-            return self.metadata.descriptive_stats
+            return existing_stats
 
         if not self.data_loaded:
             self.load_data()
 
         descriptive_stats: dict[str, DescriptiveStatistics] = {}
-        hf_subset_stat = "hf_subset_descriptive_stats"
+        hf_subset_stat: Literal["hf_subset_descriptive_stats"] = (
+            "hf_subset_descriptive_stats"
+        )
         eval_splits = self.metadata.eval_splits
         if isinstance(self, AbsTaskClassification):
             eval_splits.append(self.train_split)
@@ -381,7 +386,7 @@ class AbsTask(ABC):
             logger.info(f"Processing metadata for split {split}")
             if self.metadata.is_multilingual:
                 descriptive_stats[split] = (
-                    self._calculate_descriptive_statistics_from_split(
+                    self._calculate_descriptive_statistics_from_split(  # type: ignore[assignment]
                         split, compute_overall=True
                     )
                 )
@@ -400,7 +405,7 @@ class AbsTask(ABC):
                     descriptive_stats[split][hf_subset_stat][hf_subset] = split_details
             else:
                 split_details = self._calculate_descriptive_statistics_from_split(split)
-                descriptive_stats[split] = split_details
+                descriptive_stats[split] = split_details  # type: ignore[assignment]
 
         with self.metadata.descriptive_stat_path.open("w") as f:
             json.dump(descriptive_stats, f, indent=4)
@@ -437,7 +442,7 @@ class AbsTask(ABC):
 
         return self.metadata.languages
 
-    def filter_eval_splits(self, eval_splits: list[str] | None) -> Self:
+    def filter_eval_splits(self, eval_splits: Sequence[str] | None) -> Self:
         """Filter the evaluation splits of the task.
 
         Args:
@@ -451,9 +456,9 @@ class AbsTask(ABC):
 
     def filter_languages(
         self,
-        languages: list[str] | None,
-        script: list[str] | None = None,
-        hf_subsets: list[HFSubset] | None = None,
+        languages: Sequence[str] | None,
+        script: Sequence[str] | None = None,
+        hf_subsets: Sequence[HFSubset] | None = None,
         exclusive_language_filter: bool = False,
     ) -> Self:
         """Filter the languages of the task.
@@ -499,12 +504,14 @@ class AbsTask(ABC):
         self.hf_subsets = subsets_to_keep
         return self
 
-    def _add_main_score(self, scores: dict[HFSubset, ScoresDict]) -> None:
+    def _add_main_score(self, scores: ScoresDict) -> None:
         scores["main_score"] = scores[self.metadata.main_score]
 
     def _upload_dataset_to_hub(
         self, repo_name: str, fields: list[str] | dict[str, str]
     ) -> None:
+        if self.dataset is None:
+            raise ValueError("Dataset not loaded")
         if self.metadata.is_multilingual:
             for config in self.metadata.eval_langs:
                 logger.info(f"Converting {config} of {self.metadata.name}")
@@ -574,7 +581,7 @@ class AbsTask(ABC):
         return False
 
     @property
-    def eval_splits(self) -> list[str]:
+    def eval_splits(self) -> Sequence[str]:
         """Returns the evaluation splits of the task."""
         if self._eval_splits:
             return self._eval_splits
@@ -607,9 +614,8 @@ class AbsTask(ABC):
             self.data_loaded = False
             logger.info(f"Unloaded dataset {self.metadata.name} from memory.")
         else:
-            logger.warning(
-                f"Dataset {self.metadata.name} is not loaded, cannot unload it."
-            )
+            msg = f"Dataset `{self.metadata.name}` is not loaded, cannot unload it."
+            logger.warning(msg)
 
     @property
     def superseded_by(self) -> str | None:

mteb/abstasks/aggregate_task_metadata.py

@@ -5,7 +5,6 @@ from pydantic import ConfigDict, Field, model_validator
 from typing_extensions import Self
 
 from mteb.types import (
-    HFSubset,
     ISOLanguageScript,
     Languages,
     Licenses,
@@ -60,14 +59,7 @@ class AggregateTaskMetadata(TaskMetadata):
     reference: str | None = None
     bibtex_citation: str | None = None
 
-    @property
-    def hf_subsets_to_langscripts(self) -> dict[HFSubset, list[ISOLanguageScript]]:
-        """Return a dictionary mapping huggingface subsets to languages."""
-        if isinstance(self.eval_langs, dict):
-            return self.eval_langs
-        return {"default": self.eval_langs}  # type: ignore
-
-    @model_validator(mode="after")  # type: ignore
+    @model_validator(mode="after")
     def _compute_unfilled_cases(self) -> Self:
         if not self.eval_langs:
             self.eval_langs = self._compute_eval_langs()

mteb/abstasks/aggregated_task.py

@@ -1,10 +1,11 @@
 import logging
+import warnings
+from collections.abc import Mapping
 from pathlib import Path
 from typing import Any
 
 import numpy as np
 from datasets import Dataset, DatasetDict
-from typing_extensions import Self
 
 from mteb.models.models_protocols import MTEBModels
 from mteb.results.task_result import TaskResult
@@ -32,7 +33,7 @@ class AbsTaskAggregate(AbsTask):
 
     def task_results_to_scores(
         self, task_results: list[TaskResult]
-    ) -> dict[str, dict[HFSubset, ScoresDict]]:
+    ) -> dict[str, Mapping[HFSubset, ScoresDict]]:
         """The function that aggregated scores. Can be redefined to allow for custom aggregations.
 
         Args:
@@ -41,7 +42,7 @@ class AbsTaskAggregate(AbsTask):
         Returns:
             A dictionary with the aggregated scores.
         """
-        scores = {}
+        scores: dict[str, Mapping[HFSubset, ScoresDict]] = {}
         subsets = (
             self.metadata.eval_langs.keys()
             if isinstance(self.metadata.eval_langs, dict)
@@ -113,33 +114,13 @@ class AbsTaskAggregate(AbsTask):
         )
         mteb_versions = {tr.mteb_version for tr in task_results}
         if len(mteb_versions) != 1:
-            logger.warning(
-                f"All tasks of {self.metadata.name} is not run using the same version."
-            )
+            msg = f"All tasks of {self.metadata.name} is not run using the same version. different versions found are: {mteb_versions}"
+            logger.warning(msg)
+            warnings.warn(msg)
             task_res.mteb_version = None
         task_res.mteb_version = task_results[0].mteb_version
         return task_res
 
-    def check_if_dataset_is_superseded(self) -> None:
-        """Check if the dataset is superseded by a newer version"""
-        if self.superseded_by:
-            logger.warning(
-                f"Dataset '{self.metadata.name}' is superseded by '{self.superseded_by}', you might consider using the newer version of the dataset."
-            )
-
-    def filter_eval_splits(self, eval_splits: list[str] | None) -> Self:
-        """Filter the evaluation splits of the task.
-
-        Args:
-            eval_splits: List of splits to evaluate on. If None, all splits in metadata
-                are used.
-
-        Returns:
-            The task with filtered evaluation splits.
-        """
-        self._eval_splits = eval_splits
-        return self
-
     def evaluate(
         self,
         model: MTEBModels,

mteb/abstasks/classification.py

@@ -143,6 +143,9 @@ class AbsTaskClassification(AbsTask):
         if not self.data_loaded:
             self.load_data()
 
+        if self.dataset is None:
+            raise RuntimeError("Dataset not loaded.")
+
         if "random_state" in self.evaluator_model.get_params():
             self.evaluator_model = self.evaluator_model.set_params(
                 random_state=self.seed
@@ -175,11 +178,11 @@ class AbsTaskClassification(AbsTask):
             )
             self._add_main_score(scores[hf_subset])
 
-        return scores
+        return scores  # type: ignore[return-value]
 
     def _evaluate_subset(
         self,
-        model: EncoderProtocol,
+        model: MTEBModels,
         data_split: DatasetDict,
         *,
         encode_kwargs: dict[str, Any],
@@ -188,6 +191,9 @@ class AbsTaskClassification(AbsTask):
         prediction_folder: Path | None = None,
         **kwargs: Any,
     ) -> FullClassificationMetrics:
+        if not isinstance(model, EncoderProtocol):
+            raise TypeError("Expected model to be an instance of EncoderProtocol")
+
         train_split = data_split[self.train_split]
         eval_split = data_split[hf_split]
 
@@ -237,7 +243,7 @@ class AbsTaskClassification(AbsTask):
             # ap will be none for non binary classification tasks
             k: (
                 float(np.mean(values))
-                if (values := [s[k] for s in scores if s[k] is not None])
+                if (values := [s[k] for s in scores if s[k] is not None])  # type: ignore[literal-required]
                 else np.nan
             )
             for k in scores[0].keys()
@@ -245,7 +251,7 @@ class AbsTaskClassification(AbsTask):
         logger.info(f"Running {self.metadata.name} - Finished.")
         return FullClassificationMetrics(
             scores_per_experiment=scores,
-            **avg_scores,
+            **avg_scores,  # type: ignore[typeddict-item]
         )
 
     def _calculate_scores(

mteb/abstasks/clustering.py

@@ -3,7 +3,7 @@ import logging
 import random
 from collections import defaultdict
 from pathlib import Path
-from typing import Any
+from typing import Any, cast
 
 import numpy as np
 from datasets import Dataset, DatasetDict
@@ -11,8 +11,8 @@ from sklearn.cluster import MiniBatchKMeans
 from sklearn.metrics.cluster import v_measure_score
 
 from mteb._create_dataloaders import create_dataloader
-from mteb.models import EncoderProtocol
-from mteb.types import HFSubset, ScoresDict
+from mteb.models import EncoderProtocol, MTEBModels
+from mteb.types import Array, HFSubset, ScoresDict
 from mteb.types.statistics import (
     ImageStatistics,
     LabelStatistics,
@@ -34,7 +34,7 @@ MultilingualDataset = dict[HFSubset, DatasetDict]
 
 
 def _evaluate_clustering_bootstrapped(
-    embeddings: np.ndarray,
+    embeddings: Array,
     labels: list[list[str]],
     n_clusters: int,
     cluster_size: int,
@@ -61,21 +61,21 @@ def _evaluate_clustering_bootstrapped(
         max_depth = max(map(len, labels))
     # Evaluate on each level til max depth
     for i_level in range(max_depth):
-        level_labels = []
+        level_labels: list[str | int] = []
         # Assign -1 to gold label if the level is not there
         for label in labels:
             if len(label) > i_level:
                 level_labels.append(label[i_level])
             else:
                 level_labels.append(-1)
-        level_labels = np.array(level_labels)
+        np_level_labels = np.array(level_labels)
         valid_idx = np.array(
-            [level_label != -1 for level_label in level_labels]
+            [level_label != -1 for level_label in np_level_labels]
         )  # Could be level_labels != -1 but fails with FutureWarning: elementwise comparison failed
-        level_labels = level_labels[valid_idx]
+        np_level_labels = np_level_labels[valid_idx]
         level_embeddings = embeddings[valid_idx]
         clustering_model = MiniBatchKMeans(
-            n_clusters=np.unique(level_labels).size,
+            n_clusters=np.unique(np_level_labels).size,
             batch_size=kmean_batch_size,
             init="k-means++",
             n_init=1,  # default when kmeans++ is used
@@ -87,7 +87,7 @@ def _evaluate_clustering_bootstrapped(
             cluster_indices = rng_state.choices(range(n_embeddings), k=cluster_size)
 
             _embeddings = level_embeddings[cluster_indices]
-            _labels = level_labels[cluster_indices]
+            _labels = np_level_labels[cluster_indices]
             cluster_assignment = clustering_model.fit_predict(_embeddings)
             v_measure = v_measure_score(_labels, cluster_assignment)
             v_measures[f"Level {i_level}"].append(v_measure)
@@ -153,7 +153,7 @@ class AbsTaskClustering(AbsTask):
 
     def _evaluate_subset(
         self,
-        model: EncoderProtocol,
+        model: MTEBModels,
         data_split: Dataset,
         *,
         encode_kwargs: dict[str, Any],
@@ -162,6 +162,10 @@ class AbsTaskClustering(AbsTask):
         prediction_folder: Path | None = None,
         **kwargs: Any,
     ) -> ScoresDict:
+        if not isinstance(model, EncoderProtocol):
+            raise TypeError(
+                "Expected encoder model to be an instance of EncoderProtocol."
+            )
         if (
             self.max_document_to_embed is not None
             and self.max_fraction_of_documents_to_embed is not None
@@ -182,13 +186,13 @@ class AbsTaskClustering(AbsTask):
                     self.max_fraction_of_documents_to_embed * len(data_split)
                 )
             else:
-                max_documents_to_embed = self.max_document_to_embed
+                max_documents_to_embed = cast(int, self.max_document_to_embed)
 
-            max_documents_to_embed = min(len(data_split), max_documents_to_embed)  # type: ignore
+            max_documents_to_embed = min(len(data_split), max_documents_to_embed)
             example_indices = self.rng_state.sample(
                 range(len(data_split)), k=max_documents_to_embed
             )
-            downsampled_dataset = data_split.select(example_indices)  # type: ignore
+            downsampled_dataset = data_split.select(example_indices)
 
         downsampled_dataset = downsampled_dataset.select_columns(
             [self.input_column_name, self.label_column_name]

mteb/abstasks/clustering_legacy.py

@@ -8,7 +8,7 @@ from scipy.optimize import linear_sum_assignment
 from sklearn import metrics
 
 from mteb._evaluators import ClusteringEvaluator
-from mteb.models import EncoderProtocol
+from mteb.models import EncoderProtocol, MTEBModels
 from mteb.types import ScoresDict
 from mteb.types.statistics import (
     ImageStatistics,
@@ -80,7 +80,7 @@ class AbsTaskClusteringLegacy(AbsTask):
 
     def _evaluate_subset(
         self,
-        model: EncoderProtocol,
+        model: MTEBModels,
         data_split: Dataset,
         *,
         encode_kwargs: dict[str, Any],
@@ -89,6 +89,9 @@ class AbsTaskClusteringLegacy(AbsTask):
         prediction_folder: Path | None = None,
         **kwargs: Any,
     ) -> ScoresDict:
+        if not isinstance(model, EncoderProtocol):
+            raise TypeError("Expected model to be an instance of EncoderProtocol")
+
         data_split = data_split.select_columns(
             [self.input_column_name, self.label_column_name]
         )
@@ -139,9 +142,6 @@ class AbsTaskClusteringLegacy(AbsTask):
             }
             return scores
 
-        data_split = data_split.select_columns(
-            [self.input_column_name, self.label_column_name]
-        )
         evaluator = self.evaluator(
             data_split,
             input_column_name=self.input_column_name,
@@ -151,10 +151,10 @@ class AbsTaskClusteringLegacy(AbsTask):
             hf_subset=hf_subset,
             **kwargs,
         )
-        clusters = evaluator(model, encode_kwargs=encode_kwargs)
+        evaluate_clusters = evaluator(model, encode_kwargs=encode_kwargs)
         if prediction_folder:
             self._save_task_predictions(
-                clusters,
+                evaluate_clusters,
                 model,
                 prediction_folder,
                 hf_subset=hf_subset,
@@ -163,7 +163,7 @@ class AbsTaskClusteringLegacy(AbsTask):
 
         return self._compute_metrics(
             data_split[self.label_column_name],
-            clusters,
+            evaluate_clusters,
         )
 
     def _compute_metrics(

mteb/abstasks/image/image_text_pair_classification.py

@@ -12,7 +12,7 @@ from mteb.abstasks._statistics_calculation import (
     calculate_text_statistics,
 )
 from mteb.abstasks.abstask import AbsTask
-from mteb.models.models_protocols import EncoderProtocol
+from mteb.models.models_protocols import EncoderProtocol, MTEBModels
 from mteb.types.statistics import (
     ImageStatistics,
     SplitDescriptiveStatistics,
@@ -116,7 +116,7 @@ class AbsTaskImageTextPairClassification(AbsTask):
 
     def _evaluate_subset(
         self,
-        model: EncoderProtocol,
+        model: MTEBModels,
         data_split: Dataset,
         *,
         encode_kwargs: dict[str, Any],
@@ -125,6 +125,8 @@ class AbsTaskImageTextPairClassification(AbsTask):
         prediction_folder: Path | None = None,
         **kwargs: Any,
     ) -> ImageTextPairClassificationMetrics:
+        if not isinstance(model, EncoderProtocol):
+            raise TypeError("Expected model to be an instance of EncoderProtocol")
         select_columns = []
         for columns in (self.images_column_names, self.texts_column_names):
             if isinstance(columns, str):
@@ -154,7 +156,7 @@ class AbsTaskImageTextPairClassification(AbsTask):
             hf_subset=hf_subset,
             **kwargs,
         )
-        scores = evaluator(model, encode_kwargs=encode_kwargs)
+        scores: list[torch.Tensor] = evaluator(model, encode_kwargs=encode_kwargs)  # type: ignore[assignment]
         if prediction_folder:
             self._save_task_predictions(
                 [score.tolist() for score in scores],

mteb/abstasks/multilabel_classification.py

@@ -16,7 +16,8 @@ from typing_extensions import override
 from mteb._create_dataloaders import create_dataloader
 from mteb._evaluators.classification_metrics import hamming_score
 from mteb._evaluators.sklearn_evaluator import SklearnModelProtocol
-from mteb.models import EncoderProtocol
+from mteb.models import EncoderProtocol, MTEBModels
+from mteb.types import Array
 
 from .classification import AbsTaskClassification
 
@@ -24,14 +25,14 @@ logger = logging.getLogger(__name__)
 
 
 def _evaluate_classifier(
-    embeddings_train: np.ndarray,
+    embeddings_train: Array,
     y_train: np.ndarray,
-    embeddings_test: np.ndarray,
+    embeddings_test: Array,
     classifier: SklearnModelProtocol,
 ) -> tuple[np.ndarray, SklearnModelProtocol]:
-    classifier: SklearnModelProtocol = clone(classifier)
-    classifier.fit(embeddings_train, y_train)
-    return classifier.predict(embeddings_test), classifier
+    classifier_copy: SklearnModelProtocol = clone(classifier)
+    classifier_copy.fit(embeddings_train, y_train)
+    return classifier_copy.predict(embeddings_test), classifier_copy
 
 
 class MultilabelClassificationMetrics(TypedDict):
@@ -72,14 +73,14 @@ class AbsTaskMultilabelClassification(AbsTaskClassification):
         evaluator: Classifier to use for evaluation. Must implement the SklearnModelProtocol.
     """
 
-    evaluator: SklearnModelProtocol = KNeighborsClassifier(n_neighbors=5)
+    evaluator: SklearnModelProtocol = KNeighborsClassifier(n_neighbors=5)  # type: ignore[assignment]
     input_column_name: str = "text"
     label_column_name: str = "label"
 
     @override
-    def _evaluate_subset(
+    def _evaluate_subset(  # type: ignore[override]
         self,
-        model: EncoderProtocol,
+        model: MTEBModels,
         data_split: DatasetDict,
         *,
         encode_kwargs: dict[str, Any],
@@ -88,6 +89,9 @@ class AbsTaskMultilabelClassification(AbsTaskClassification):
         prediction_folder: Path | None = None,
         **kwargs: Any,
     ) -> FullMultilabelClassificationMetrics:
+        if not isinstance(model, EncoderProtocol):
+            raise TypeError("Expected model to be an instance of EncoderProtocol")
+
         if isinstance(data_split, DatasetDict):
             data_split = data_split.select_columns(
                 [self.input_column_name, self.label_column_name]
@@ -185,19 +189,20 @@ class AbsTaskMultilabelClassification(AbsTaskClassification):
             )
 
         avg_scores: dict[str, Any] = {
-            k: np.mean([s[k] for s in scores]) for k in scores[0].keys()
+            k: np.mean([s[k] for s in scores])  # type: ignore[literal-required]
+            for k in scores[0].keys()
         }
         logger.info("Running multilabel classification - Finished.")
         return FullMultilabelClassificationMetrics(
             scores_per_experiment=scores,
-            **avg_scores,
+            **avg_scores,  # type: ignore[typeddict-item]
         )
 
-    def _calculate_scores(
+    def _calculate_scores(  # type: ignore[override]
         self,
         y_test: np.ndarray,
         y_pred: np.ndarray,
-        x_test_embedding: np.ndarray,
+        x_test_embedding: Array,
         current_classifier: SklearnModelProtocol,
     ) -> MultilabelClassificationMetrics:
         accuracy = current_classifier.score(x_test_embedding, y_test)
@@ -232,10 +237,9 @@ class AbsTaskMultilabelClassification(AbsTaskClassification):
         """
         sample_indices = []
         if idxs is None:
-            idxs = np.arange(len(y))
+            idxs = list(np.arange(len(y)))
         self.np_rng.shuffle(idxs)
-        idxs = idxs.tolist()
-        label_counter = defaultdict(int)
+        label_counter: dict[int, int] = defaultdict(int)
         for i in idxs:
             if any((label_counter[label] < samples_per_label) for label in y[i]):
                 sample_indices.append(i)