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

mteb/abstasks/classification.py

@@ -16,7 +16,7 @@ from sklearn.metrics import (
 
 from mteb._evaluators.sklearn_evaluator import SklearnEvaluator, SklearnModelProtocol
 from mteb.models import EncoderProtocol, MTEBModels
-from mteb.types import HFSubset, ScoresDict
+from mteb.types import EncodeKwargs, HFSubset, ScoresDict
 from mteb.types.statistics import (
     ImageStatistics,
     LabelStatistics,
@@ -98,9 +98,8 @@ class AbsTaskClassification(AbsTask):
             text: str (for text) or PIL.Image (for image). Column name can be changed via `input_column_name` attribute.
             label: int. Column name can be changed via `label_column_name` attribute.
         evaluator_model: The model to use for evaluation. Can be any sklearn compatible model. Default is `LogisticRegression`.
-            Full details of api in [`SklearnModelProtocol`][mteb._evaluators.sklearn_evaluator.SklearnModelProtocol].
-        samples_per_label: Number of samples per label to use for training the evaluator model. Default is 8.
-        n_experiments: Number of experiments to run. Default is 10.
+       samples_per_label: Number of samples per label to use for training the evaluator model. Default is 8.
+       n_experiments: Number of experiments to run. Default is 10.
         train_split: Name of the split to use for training the evaluator model. Default is "train".
         label_column_name: Name of the column containing the labels. Default is "label".
         input_column_name: Name of the column containing the input data. Default is "text".
@@ -126,7 +125,7 @@ class AbsTaskClassification(AbsTask):
         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]:
@@ -143,6 +142,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,19 +177,22 @@ 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],
+        encode_kwargs: EncodeKwargs,
         hf_split: str,
         hf_subset: str,
         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 +242,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 +250,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, EncodeKwargs, 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,15 +153,19 @@ class AbsTaskClustering(AbsTask):
 
     def _evaluate_subset(
         self,
-        model: EncoderProtocol,
+        model: MTEBModels,
         data_split: Dataset,
         *,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         hf_split: str,
         hf_subset: str,
         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,8 +8,8 @@ from scipy.optimize import linear_sum_assignment
 from sklearn import metrics
 
 from mteb._evaluators import ClusteringEvaluator
-from mteb.models import EncoderProtocol
-from mteb.types import ScoresDict
+from mteb.models import EncoderProtocol, MTEBModels
+from mteb.types import EncodeKwargs, ScoresDict
 from mteb.types.statistics import (
     ImageStatistics,
     LabelStatistics,
@@ -80,15 +80,18 @@ class AbsTaskClusteringLegacy(AbsTask):
 
     def _evaluate_subset(
         self,
-        model: EncoderProtocol,
+        model: MTEBModels,
         data_split: Dataset,
         *,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         hf_split: str,
         hf_subset: str,
         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,8 @@ 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 import EncodeKwargs
 from mteb.types.statistics import (
     ImageStatistics,
     SplitDescriptiveStatistics,
@@ -116,15 +117,17 @@ class AbsTaskImageTextPairClassification(AbsTask):
 
     def _evaluate_subset(
         self,
-        model: EncoderProtocol,
+        model: MTEBModels,
         data_split: Dataset,
         *,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         hf_split: str,
         hf_subset: str,
         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 +157,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, EncodeKwargs
 
 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):
@@ -69,25 +70,28 @@ class AbsTaskMultilabelClassification(AbsTaskClassification):
         input_column_name: Name of the column containing the input text.
         label_column_name: Name of the column containing the labels.
         samples_per_label: Number of samples to use pr. label. These samples are embedded and a classifier is fit using the labels and samples.
-        evaluator: Classifier to use for evaluation. Must implement the SklearnModelProtocol.
+        evaluator_model: Classifier to use for evaluation. Must implement the SklearnModelProtocol.
     """
 
-    evaluator: SklearnModelProtocol = KNeighborsClassifier(n_neighbors=5)
+    evaluator_model: SklearnModelProtocol = KNeighborsClassifier(n_neighbors=5)
     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],
+        encode_kwargs: EncodeKwargs,
         hf_split: str,
         hf_subset: str,
         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]
@@ -165,7 +169,7 @@ class AbsTaskMultilabelClassification(AbsTaskClassification):
             y_train = train_split.select(sample_indices)[self.label_column_name]
             y_train = binarizer.transform(y_train)
             y_pred, current_classifier = _evaluate_classifier(
-                X_train, y_train, X_test, self.evaluator
+                X_train, y_train, X_test, self.evaluator_model
             )
             if prediction_folder:
                 all_predictions.append(y_pred.tolist())
@@ -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)

mteb/abstasks/pair_classification.py

@@ -18,8 +18,8 @@ from mteb.abstasks._statistics_calculation import (
 )
 from mteb.abstasks.abstask import AbsTask
 from mteb.models.model_meta import ScoringFunction
-from mteb.models.models_protocols import EncoderProtocol
-from mteb.types import PromptType
+from mteb.models.models_protocols import EncoderProtocol, MTEBModels
+from mteb.types import EncodeKwargs, PromptType
 from mteb.types.statistics import (
     ImageStatistics,
     LabelStatistics,
@@ -44,8 +44,8 @@ class PairClassificationDescriptiveStatistics(SplitDescriptiveStatistics):
     """
 
     num_samples: int
-    number_of_characters: int
-    unique_pairs: int
+    number_of_characters: int | None
+    unique_pairs: int | None
 
     text1_statistics: TextStatistics | None
     image1_statistics: ImageStatistics | None
@@ -79,15 +79,18 @@ class AbsTaskPairClassification(AbsTask):
 
     def _evaluate_subset(
         self,
-        model: EncoderProtocol,
+        model: MTEBModels,
         data_split: Dataset,
         *,
         hf_split: str,
         hf_subset: str,
-        encode_kwargs: dict[str, str],
+        encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
         **kwargs,
     ) -> dict[str, float]:
+        if not isinstance(model, EncoderProtocol):
+            raise TypeError("Expected model to be an instance of EncoderProtocol")
+
         if self.metadata.modalities == ["text"]:
             # for compatibility with v1 version where datasets were stored in a single row
             data_split = data_split[0] if len(data_split) == 1 else data_split
@@ -120,7 +123,7 @@ class AbsTaskPairClassification(AbsTask):
         self, similarity_scores: PairClassificationDistances, labels: list[int]
     ) -> dict[str, float]:
         logger.info("Computing metrics...")
-        labels = np.asarray(labels)
+        np_labels = np.asarray(labels)
         output_scores = {}
         max_scores = defaultdict(list)
         for short_name, scores, reverse in [
@@ -142,7 +145,7 @@ class AbsTaskPairClassification(AbsTask):
             ],
             [ScoringFunction.DOT_PRODUCT.value, similarity_scores["dot_scores"], True],
         ]:
-            metrics = self._compute_metrics_values(scores, labels, reverse)
+            metrics = self._compute_metrics_values(scores, np_labels, reverse)  # type: ignore[arg-type]
             for metric_name, metric_value in metrics.items():
                 output_scores[f"{short_name}_{metric_name}"] = metric_value
                 max_scores[metric_name].append(metric_value)
@@ -237,6 +240,12 @@ class AbsTaskPairClassification(AbsTask):
 
     def _push_dataset_to_hub(self, repo_name: str) -> None:
         # previously pair classification datasets were stored in a single row
+        if self.dataset is None:
+            # overall this shouldn't happen as we check for dataset before pushing to hub
+            # added here for type checking purposes
+            raise RuntimeError(
+                "Dataset not loaded. To load dataset run `task.load_data()`."
+            )
         if self.metadata.is_multilingual:
             for subset in self.dataset:
                 for split in self.dataset[subset]:
@@ -290,13 +299,13 @@ class AbsTaskPairClassification(AbsTask):
         )
 
     def _find_best_acc_and_threshold(
-        self, scores: np.ndarray, labels: np.ndarray, high_score_more_similar: bool
+        self, scores: list[float], labels: np.ndarray, high_score_more_similar: bool
     ) -> tuple[float, float]:
         rows = list(zip(scores, labels))
         rows = sorted(rows, key=lambda x: x[0], reverse=high_score_more_similar)
 
         max_acc = 0
-        best_threshold = -1
+        best_threshold = -1.0
         positive_so_far = 0
         remaining_negatives = sum(np.array(labels) == 0)
 
@@ -323,7 +332,7 @@ class AbsTaskPairClassification(AbsTask):
 
         rows = sorted(rows, key=lambda x: x[0], reverse=high_score_more_similar)
 
-        best_f1 = best_precision = best_recall = 0
+        best_f1 = best_precision = best_recall = 0.0
         threshold = 0
         nextract = 0
         ncorrect = 0

mteb/abstasks/regression.py

@@ -84,10 +84,10 @@ class AbsTaskRegression(AbsTaskClassification):
         n_samples: Number of samples to use for training the regression model. If the dataset has fewer samples than n_samples, all samples are used.
         abstask_prompt: Prompt to use for the task for instruction model if not prompt is provided in TaskMetadata.prompt.
         evaluator_model: The model to use for evaluation. Can be any sklearn compatible model. Default is `LinearRegression`.
-            Full details of api in [`SklearnModelProtocol`][mteb._evaluators.sklearn_evaluator.SklearnModelProtocol].
+
     """
 
-    evaluator: type[SklearnModelProtocol] = SklearnEvaluator
+    evaluator: type[SklearnEvaluator] = SklearnEvaluator
     evaluator_model: SklearnModelProtocol = LinearRegression(n_jobs=-1)
 
     train_split: str = "train"
@@ -113,7 +113,7 @@ class AbsTaskRegression(AbsTaskClassification):
             )["train"]
         return train_split_sampled, []
 
-    def _calculate_scores(
+    def _calculate_scores(  # type: ignore[override]
         self,
         y_test: np.ndarray | list[int],
         y_pred: np.ndarray,
@@ -183,7 +183,7 @@ class AbsTaskRegression(AbsTaskClassification):
 
         return dataset_dict
 
-    def _calculate_descriptive_statistics_from_split(
+    def _calculate_descriptive_statistics_from_split(  # type: ignore[override]
         self, split: str, hf_subset: str | None = None, compute_overall: bool = False
     ) -> RegressionDescriptiveStatistics:
         train_text = []

mteb/abstasks/retrieval.py

@@ -1,7 +1,7 @@
 import json
 import logging
 from collections import defaultdict
-from collections.abc import Callable, Sequence
+from collections.abc import Callable, Mapping, Sequence
 from pathlib import Path
 from time import time
 from typing import Any, Literal
@@ -25,6 +25,7 @@ from mteb.models import (
     SearchProtocol,
 )
 from mteb.types import (
+    EncodeKwargs,
     HFSubset,
     QueryDatasetType,
     RelevantDocumentsType,
@@ -184,17 +185,17 @@ class AbsTaskRetrieval(AbsTask):
             return queries, corpus
 
         if self.metadata.is_multilingual:
-            for subset in self.queries:
-                for split in self.queries[subset]:
-                    queries = self.queries[subset][split]
-                    corpus = self.corpus[subset][split]
+            for subset in self.queries:  # type: ignore[attr-defined]
+                for split in self.queries[subset]:  # type: ignore[attr-defined]
+                    queries = self.queries[subset][split]  # type: ignore[attr-defined]
+                    corpus = self.corpus[subset][split]  # type: ignore[attr-defined]
 
                     (
                         self.dataset[subset][split]["queries"],
                         self.dataset[subset][split]["corpus"],
                     ) = _process_split(queries, corpus)
 
-                    self.dataset[subset][split]["relevant_docs"] = self.relevant_docs[
+                    self.dataset[subset][split]["relevant_docs"] = self.relevant_docs[  # type: ignore[attr-defined]
                         subset
                     ][split]
                     if hasattr(self, "instructions"):
@@ -211,15 +212,15 @@ class AbsTaskRetrieval(AbsTask):
                         ][split]
         else:
             subset = "default"
-            for split in self.queries:
-                queries = self.queries[split]
-                corpus = self.corpus[split]
+            for split in self.queries:  # type: ignore[attr-defined]
+                queries = self.queries[split]  # type: ignore[attr-defined]
+                corpus = self.corpus[split]  # type: ignore[attr-defined]
                 (
                     self.dataset[subset][split]["queries"],
                     self.dataset[subset][split]["corpus"],
                 ) = _process_split(queries, corpus)
 
-                self.dataset[subset][split]["relevant_docs"] = self.relevant_docs[
+                self.dataset[subset][split]["relevant_docs"] = self.relevant_docs[  # type: ignore[attr-defined]
                     split
                 ].copy()
                 if hasattr(self, "instructions"):
@@ -235,9 +236,9 @@ class AbsTaskRetrieval(AbsTask):
                         split
                     ].copy()
 
-        del self.queries
-        del self.corpus
-        del self.relevant_docs
+        del self.queries  # type: ignore[attr-defined]
+        del self.corpus  # type: ignore[attr-defined]
+        del self.relevant_docs  # type: ignore[attr-defined]
         if hasattr(self, "instructions"):
             del self.instructions
         if hasattr(self, "top_ranked"):
@@ -283,10 +284,10 @@ class AbsTaskRetrieval(AbsTask):
         split: str = "test",
         subsets_to_run: list[HFSubset] | None = None,
         *,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
-        **kwargs,
-    ) -> dict[HFSubset, ScoresDict]:
+        **kwargs: Any,
+    ) -> Mapping[HFSubset, ScoresDict]:
         """Evaluate the model on the retrieval task.
 
         Args:
@@ -320,7 +321,7 @@ class AbsTaskRetrieval(AbsTask):
         self,
         model: MTEBModels,
         data_split: RetrievalSplitData,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         hf_split: str,
         hf_subset: str,
         prediction_folder: Path | None = None,
@@ -357,6 +358,8 @@ class AbsTaskRetrieval(AbsTask):
             **kwargs,
         )
 
+        search_model: SearchProtocol
+
         if isinstance(model, EncoderProtocol) and not isinstance(model, SearchProtocol):
             search_model = SearchEncoderWrapper(model)
         elif isinstance(model, CrossEncoderProtocol):
@@ -578,11 +581,12 @@ class AbsTaskRetrieval(AbsTask):
                 if isinstance(data[split][subset_item], Dataset):
                     sections[split] = data[split][subset_item]
                 elif converter is not None:
+                    subset_data = data[split][subset_item]
+                    if subset_data is None:
+                        continue
+
                     sections[split] = Dataset.from_list(
-                        [
-                            converter(idx, item)
-                            for idx, item in data[split][subset_item].items()
-                        ]
+                        [converter(idx, item) for idx, item in subset_data.items()]
                     )
                 else:
                     raise ValueError(
@@ -680,7 +684,7 @@ class AbsTaskRetrieval(AbsTask):
 
                 top_k_sorted = defaultdict(list)
                 for query_id, values in top_ranked.items():
-                    sorted_keys = sorted(values, key=values.get, reverse=True)
+                    sorted_keys = sorted(values, key=lambda k: values[k], reverse=True)
                     top_k_sorted[query_id] = sorted_keys[: self._top_k]
 
                 self.dataset[subset][split]["top_ranked"] = top_k_sorted
@@ -688,10 +692,10 @@ class AbsTaskRetrieval(AbsTask):
 
 
 def _process_relevant_docs(
-    collection: dict[str, dict[str, float]],
+    collection: Mapping[str, Mapping[str, int]],
     hf_subset: str,
     split: str,
-) -> dict[str, dict[str, float]]:
+) -> dict[str, dict[str, int]]:
     """Collections can contain overlapping ids in different splits. Prepend split and subset to avoid this
 
     Returns: