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

mteb/abstasks/_statistics_calculation.py

@@ -2,9 +2,8 @@ from __future__ import annotations
 
 import hashlib
 from collections import Counter
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 
-from mteb.types import TopRankedDocumentsType
 from mteb.types.statistics import (
     ImageStatistics,
     LabelStatistics,
@@ -15,8 +14,12 @@ from mteb.types.statistics import (
 )
 
 if TYPE_CHECKING:
+    from collections.abc import Mapping
+
     from PIL import Image
 
+    from mteb.types import TopRankedDocumentsType
+
 
 def calculate_text_statistics(texts: list[str]) -> TextStatistics:
     """Calculate descriptive statistics for a list of texts.
@@ -52,7 +55,7 @@ def calculate_image_statistics(images: list[Image.Image]) -> ImageStatistics:
     seen_hashes: set[str] = set()
 
     for img in images:
-        width, height = img.size  # type: ignore
+        width, height = img.size
         img_heights.append(height)
         img_widths.append(width)
 
@@ -82,17 +85,24 @@ def calculate_label_statistics(labels: list[int | list[int]]) -> LabelStatistics
         LabelStatistics: A dictionary containing the descriptive statistics.
 
     """
+    total_labels: list[int | None] = []
+
     if not isinstance(labels[0], list):
-        label_len = [1] * len(labels)
-        total_label_len = len(labels)
-        total_labels = labels
+        # single label classification
+        single_label = cast("list[int]", labels)
+        label_len = [1] * len(single_label)
+        total_label_len = len(single_label)
+        total_labels.extend(single_label)
     elif isinstance(labels[0], list):
         # multilabel classification
-        label_len = [len(l) for l in labels]
+        multilabel_labels = cast("list[list[int]]", labels)
+        label_len = [len(l) for l in multilabel_labels]
         total_label_len = sum(label_len)
-        total_labels = []
-        for l in labels:
-            total_labels.extend(l if len(l) > 0 else [None])
+        for l in multilabel_labels:
+            if l and len(l) > 0:
+                total_labels.extend(l)
+            else:
+                total_labels.append(None)
     else:
         raise ValueError(
             "Labels must be a list of integers or a list of lists of integers."
@@ -159,7 +169,7 @@ def calculate_top_ranked_statistics(
 
 
 def calculate_relevant_docs_statistics(
-    relevant_docs: dict[str, dict[str, float]],
+    relevant_docs: Mapping[str, Mapping[str, int]],
 ) -> RelevantDocsStatistics:
     qrels_lengths = [len(relevant_docs[qid]) for qid in relevant_docs]
     unique_qrels = len({doc for qid in relevant_docs for doc in relevant_docs[qid]})

mteb/abstasks/_stratification.py

@@ -39,6 +39,7 @@ Bibtex:
 """
 
 import itertools
+from typing import Any
 
 import numpy as np
 import scipy.sparse as sp
@@ -119,8 +120,10 @@ def _get_most_desired_combination(samples_with_combination: dict):
         if support_size == 0:
             continue
         if currently_chosen is None or (
-            best_number_of_combinations < number_of_combinations  # type: ignore
-            and best_support_size > support_size  # type: ignore
+            best_number_of_combinations is not None
+            and best_support_size is not None
+            and best_number_of_combinations < number_of_combinations
+            and best_support_size > support_size
         ):
             currently_chosen = combination
             best_number_of_combinations, best_support_size = (
@@ -162,7 +165,7 @@ class IterativeStratification(_BaseKFold):
         self._rng_state = check_random_state(random_state)
         need_shuffle = shuffle or random_state is not None
         self.order = order
-        super().__init__(  # type: ignore
+        super().__init__(
             n_splits,
             shuffle=need_shuffle,
             random_state=self._rng_state if need_shuffle else None,
@@ -172,8 +175,7 @@ class IterativeStratification(_BaseKFold):
             self.percentage_per_fold = sample_distribution_per_fold
         else:
             self.percentage_per_fold = [
-                1 / float(self.n_splits)
-                for _ in range(self.n_splits)  # type: ignore
+                1 / float(self.n_splits) for _ in range(self.n_splits)
             ]
 
     def _prepare_stratification(
@@ -182,9 +184,9 @@ class IterativeStratification(_BaseKFold):
         list[list[int]],
         dict[int, bool],
         list[list[int]],
-        list[list[list[int]]],
-        dict[tuple[int, ...], list[int]],
-        list[list[int]],
+        list[list[Any]],
+        dict[str, list[Any]],
+        list[list[Any]],
     ]:
         """Prepares variables for performing stratification
 
@@ -206,14 +208,14 @@ class IterativeStratification(_BaseKFold):
         """
         self.n_samples, self.n_labels = y.shape
         self.desired_samples_per_fold = np.array(
-            [self.percentage_per_fold[i] * self.n_samples for i in range(self.n_splits)]  # type: ignore
+            [self.percentage_per_fold[i] * self.n_samples for i in range(self.n_splits)]
         )
         rows = sp.lil_matrix(y).rows
         rows_used = dict.fromkeys(range(self.n_samples), False)
         all_combinations = []
-        per_row_combinations = [[] for i in range(self.n_samples)]
-        samples_with_combination = {}
-        folds = [[] for _ in range(self.n_splits)]  # type: ignore
+        per_row_combinations: list[list[Any]] = [[] for i in range(self.n_samples)]
+        samples_with_combination: dict[str, list[Any]] = {}
+        folds: list[list[int]] = [[] for _ in range(self.n_splits)]
 
         # for every row
         for sample_index, label_assignment in enumerate(rows):
@@ -229,21 +231,19 @@ class IterativeStratification(_BaseKFold):
                 all_combinations.append(combination)
                 per_row_combinations[sample_index].append(combination)
 
-        all_combinations = [list(x) for x in set(all_combinations)]
-
         self.desired_samples_per_combination_per_fold = {
             combination: np.array(
                 [
                     len(evidence_for_combination) * self.percentage_per_fold[j]
-                    for j in range(self.n_splits)  # type: ignore
+                    for j in range(self.n_splits)
                 ]
             )
             for combination, evidence_for_combination in samples_with_combination.items()
         }
         return (
-            rows,
+            rows.tolist(),
             rows_used,
-            all_combinations,
+            [list(x) for x in set(all_combinations)],
             per_row_combinations,
             samples_with_combination,
             folds,
@@ -328,7 +328,7 @@ class IterativeStratification(_BaseKFold):
             per_row_combinations,
             samples_with_combination,
             folds,
-        ) = self._prepare_stratification(y)  # type: ignore
+        ) = self._prepare_stratification(y)
 
         self._distribute_positive_evidence(
             rows_used, folds, samples_with_combination, per_row_combinations

mteb/abstasks/abstask.py

@@ -1,28 +1,38 @@
+from __future__ import annotations
+
 import json
 import logging
+import warnings
 from abc import ABC, abstractmethod
 from collections.abc import Sequence
 from copy import copy
 from pathlib import Path
-from typing import Any, cast
+from typing import TYPE_CHECKING, Any, Literal, cast
 
 import numpy as np
 from datasets import ClassLabel, Dataset, DatasetDict, load_dataset
 from sklearn.preprocessing import MultiLabelBinarizer
 from tqdm.auto import tqdm
-from typing_extensions import Self
 
 from mteb._set_seed import _set_seed
-from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.languages import LanguageScripts
 from mteb.models import (
     CrossEncoderProtocol,
     EncoderProtocol,
-    MTEBModels,
     SearchProtocol,
 )
-from mteb.types import HFSubset, Modalities, ScoresDict
-from mteb.types.statistics import DescriptiveStatistics, SplitDescriptiveStatistics
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+    from typing_extensions import Self
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.models import (
+        MTEBModels,
+    )
+    from mteb.types import EncodeKwargs, HFSubset, Modalities, ScoresDict
+    from mteb.types.statistics import DescriptiveStatistics, SplitDescriptiveStatistics
 
 logger = logging.getLogger(__name__)
 
@@ -78,8 +88,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,15 +112,18 @@ 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):
+    def dataset_transform(self, num_proc: int = 1):
         """A transform operations applied to the dataset after loading.
 
         This method is useful when the dataset from Huggingface is not in an `mteb` compatible format.
         Override this method if your dataset requires additional transformation.
+
+        Args:
+            num_proc: Number of processes to use for the transformation.
         """
         pass
 
@@ -120,10 +133,11 @@ class AbsTask(ABC):
         split: str = "test",
         subsets_to_run: list[HFSubset] | None = None,
         *,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
+        num_proc: int = 1,
         **kwargs: Any,
-    ) -> dict[HFSubset, ScoresDict]:
+    ) -> Mapping[HFSubset, ScoresDict]:
         """Evaluates an MTEB compatible model on the task.
 
         Args:
@@ -132,6 +146,7 @@ class AbsTask(ABC):
             subsets_to_run: List of huggingface subsets (HFSubsets) to evaluate. If None, all subsets are evaluated.
             encode_kwargs: Additional keyword arguments that are passed to the model's `encode` method.
             prediction_folder: Folder to save model predictions
+            num_proc: Number of processes to use for loading the dataset or processing.
             kwargs: Additional keyword arguments that are passed to the _evaluate_subset method.
 
         Returns:
@@ -161,7 +176,7 @@ class AbsTask(ABC):
         if not self.data_loaded:
             self.load_data()
 
-        self.dataset = cast(dict[HFSubset, DatasetDict], self.dataset)
+        self.dataset = cast("dict[HFSubset, DatasetDict]", self.dataset)
 
         scores = {}
         if self.hf_subsets is None:
@@ -187,6 +202,7 @@ class AbsTask(ABC):
                 hf_subset=hf_subset,
                 encode_kwargs=encode_kwargs,
                 prediction_folder=prediction_folder,
+                num_proc=num_proc,
                 **kwargs,
             )
             self._add_main_score(scores[hf_subset])
@@ -195,13 +211,14 @@ 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: EncodeKwargs,
         prediction_folder: Path | None = None,
+        num_proc: int = 1,
         **kwargs: Any,
     ) -> ScoresDict:
         raise NotImplementedError(
@@ -210,7 +227,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 +243,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,
@@ -306,11 +323,15 @@ class AbsTask(ABC):
             )  # only take the specified test split.
         return dataset_dict
 
-    def load_data(self) -> None:
+    def load_data(self, num_proc: int = 1, **kwargs: Any) -> None:
         """Loads dataset from HuggingFace hub
 
         This is the main loading function for Task. Do not overwrite this, instead we recommend using `dataset_transform`, which is called after the
         dataset is loaded using `datasets.load_dataset`.
+
+        Args:
+            num_proc: Number of processes to use for loading the dataset.
+            kwargs: Additional keyword arguments passed to the load_dataset function. Keep for forward compatibility.
         """
         if self.data_loaded:
             return
@@ -323,11 +344,12 @@ class AbsTask(ABC):
                     self.dataset[hf_subset] = load_dataset(
                         name=hf_subset,
                         **self.metadata.dataset,
+                        num_proc=num_proc,
                     )
         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_transform()
+            self.dataset = load_dataset(**self.metadata.dataset, num_proc=num_proc)
+        self.dataset_transform(num_proc=num_proc)
         self.data_loaded = True
 
     def fast_load(self) -> None:
@@ -350,27 +372,32 @@ class AbsTask(ABC):
             self.dataset[lang] = DatasetDict(subset)
 
     def calculate_descriptive_statistics(
-        self, overwrite_results: bool = False
+        self, overwrite_results: bool = False, num_proc: int = 1
     ) -> dict[str, DescriptiveStatistics]:
         """Calculates descriptive statistics from the dataset.
 
         Args:
             overwrite_results: Whether to overwrite existing results. If False and results already exist, the existing results will be loaded from cache.
+            num_proc: Number of processes to use for loading the dataset.
 
         Returns:
             A dictionary containing descriptive statistics for each split.
         """
         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()
+            self.load_data(num_proc=num_proc)
 
         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 +408,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 +427,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 +464,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 +478,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 +526,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]
+        self, repo_name: str, fields: list[str] | dict[str, str], num_proc: int = 1
     ) -> 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}")
@@ -526,7 +555,10 @@ class AbsTask(ABC):
                         )
                 sentences = DatasetDict(sentences)
                 sentences.push_to_hub(
-                    repo_name, config, commit_message=f"Add {config} dataset"
+                    repo_name,
+                    config,
+                    commit_message=f"Add {config} dataset",
+                    num_proc=num_proc,
                 )
         else:
             sentences = {}
@@ -543,16 +575,19 @@ class AbsTask(ABC):
                         {field: self.dataset[split][field] for field in fields}
                     )
             sentences = DatasetDict(sentences)
-            sentences.push_to_hub(repo_name, commit_message="Add dataset")
+            sentences.push_to_hub(
+                repo_name, commit_message="Add dataset", num_proc=num_proc
+            )
 
-    def _push_dataset_to_hub(self, repo_name: str) -> None:
+    def _push_dataset_to_hub(self, repo_name: str, num_proc: int = 1) -> None:
         raise NotImplementedError
 
-    def push_dataset_to_hub(self, repo_name: str) -> None:
+    def push_dataset_to_hub(self, repo_name: str, num_proc: int = 1) -> None:
         """Push the dataset to the HuggingFace Hub.
 
         Args:
             repo_name: The name of the repository to push the dataset to.
+            num_proc: Number of processes to use for loading the dataset.
 
         Examples:
             >>> import mteb
@@ -564,7 +599,7 @@ class AbsTask(ABC):
         if not self.data_loaded:
             self.load_data()
 
-        self._push_dataset_to_hub(repo_name)
+        self._push_dataset_to_hub(repo_name, num_proc)
         # dataset repo not creating when pushing card
         self.metadata.push_dataset_card_to_hub(repo_name)
 
@@ -574,7 +609,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 +642,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

@@ -1,29 +1,39 @@
+from __future__ import annotations
+
 import logging
 from datetime import datetime
+from typing import TYPE_CHECKING
 
 from pydantic import ConfigDict, Field, model_validator
-from typing_extensions import Self
 
 from mteb.types import (
-    HFSubset,
-    ISOLanguageScript,
     Languages,
-    Licenses,
-    Modalities,
-    StrDate,
 )
 
 from .abstask import AbsTask
 from .task_metadata import (
-    AnnotatorType,
     MetadataDatasetDict,
-    SampleCreationMethod,
-    TaskDomain,
     TaskMetadata,
-    TaskSubtype,
     TaskType,
 )
 
+if TYPE_CHECKING:
+    from typing_extensions import Self
+
+    from mteb.types import (
+        ISOLanguageScript,
+        Licenses,
+        Modalities,
+        StrDate,
+    )
+
+    from .task_metadata import (
+        AnnotatorType,
+        SampleCreationMethod,
+        TaskDomain,
+        TaskSubtype,
+    )
+
 logger = logging.getLogger(__name__)
 
 
@@ -60,14 +70,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,18 +1,26 @@
+from __future__ import annotations
+
 import logging
-from pathlib import Path
-from typing import Any
+import warnings
+from typing import TYPE_CHECKING, 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
-from mteb.types import HFSubset, ScoresDict
-from mteb.types.statistics import DescriptiveStatistics
 
 from .abstask import AbsTask
-from .aggregate_task_metadata import AggregateTaskMetadata
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+    from pathlib import Path
+
+    from datasets import Dataset, DatasetDict
+
+    from mteb.models.models_protocols import MTEBModels
+    from mteb.types import EncodeKwargs, HFSubset, ScoresDict
+    from mteb.types.statistics import DescriptiveStatistics
+
+    from .aggregate_task_metadata import AggregateTaskMetadata
 
 logger = logging.getLogger(__name__)
 
@@ -32,7 +40,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 +49,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,40 +121,20 @@ 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,
         split: str = "test",
         subsets_to_run: list[HFSubset] | None = None,
         *,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
         **kwargs: Any,
     ) -> dict[HFSubset, ScoresDict]:
@@ -160,7 +148,7 @@ class AbsTaskAggregate(AbsTask):
         self,
         model: MTEBModels,
         data_split: DatasetDict | Dataset,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         **kwargs: Any,
     ) -> ScoresDict:
         raise NotImplementedError(