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

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:

mteb/abstasks/retrieval_dataset_loaders.py

@@ -136,7 +136,7 @@ class RetrievalDatasetLoader:
                 "_id", "id"
             )
         logger.info("Loaded %d %s Documents.", len(corpus_ds), self.split.upper())
-        logger.info("Doc Example: %s", corpus_ds[0])
+        logger.debug("Doc Example: %s", corpus_ds[0])
         return corpus_ds
 
     def _load_queries(self) -> QueryDatasetType:
@@ -152,7 +152,7 @@ class RetrievalDatasetLoader:
             )
 
         logger.info("Loaded %d %s queries.", len(queries_ds), self.split.upper())
-        logger.info("Query Example: %s", queries_ds[0])
+        logger.debug("Query Example: %s", queries_ds[0])
 
         return queries_ds
 

mteb/abstasks/sts.py

@@ -7,7 +7,8 @@ from scipy.stats import pearsonr, spearmanr
 
 from mteb._evaluators import AnySTSEvaluator
 from mteb._evaluators.any_sts_evaluator import STSEvaluatorScores
-from mteb.models import EncoderProtocol
+from mteb.models import EncoderProtocol, MTEBModels
+from mteb.types import EncodeKwargs, PromptType
 from mteb.types.statistics import (
     ImageStatistics,
     ScoreStatistics,
@@ -89,23 +90,30 @@ class AbsTaskSTS(AbsTask):
         min_score: Minimum possible score in the dataset.
         max_score: Maximum possible score in the dataset.
         abstask_prompt: Prompt to use for the task for instruction model if not prompt is provided in TaskMetadata.prompt.
+        input1_prompt_type: Type of prompt of first input. Used for asymmetric tasks.
+        input2_prompt_type: Type of prompt of second input. Used for asymmetric tasks.
     """
 
     abstask_prompt = "Retrieve semantically similar text."
     column_names: tuple[str, str] = ("sentence1", "sentence2")
     min_score: int = 0
     max_score: int = 5
+    input1_prompt_type: PromptType | None = None
+    input2_prompt_type: PromptType | None = None
 
     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,
     ) -> STSMetrics:
+        if not isinstance(model, EncoderProtocol):
+            raise TypeError("Expected model to be an instance of EncoderProtocol")
+
         normalized_scores = list(map(self._normalize, data_split["score"]))
         data_split = data_split.select_columns(list(self.column_names))
 
@@ -115,6 +123,8 @@ class AbsTaskSTS(AbsTask):
             task_metadata=self.metadata,
             hf_split=hf_split,
             hf_subset=hf_subset,
+            input1_prompt_type=self.input1_prompt_type,
+            input2_prompt_type=self.input2_prompt_type,
             **kwargs,
         )
         scores = evaluator(model, encode_kwargs=encode_kwargs)
@@ -135,7 +145,7 @@ class AbsTaskSTS(AbsTask):
     ) -> STSMetrics:
         def compute_corr(x: list[float], y: list[float]) -> tuple[float, float]:
             """Return (pearson, spearman) correlations between x and y."""
-            return pearsonr(x, y)[0], spearmanr(x, y)[0]
+            return float(pearsonr(x, y)[0]), float(spearmanr(x, y)[0])
 
         cosine_pearson, cosine_spearman = compute_corr(
             normalized_scores, scores["cosine_scores"]

mteb/abstasks/task_metadata.py

@@ -2,9 +2,10 @@ import json
 import logging
 from collections.abc import Sequence
 from pathlib import Path
-from typing import Any, Literal
+from typing import Any, Literal, cast
 
 from huggingface_hub import (
+    CardData,
     DatasetCard,
     DatasetCardData,
     constants,
@@ -107,6 +108,7 @@ The domains follow the categories used in the [Universal Dependencies project](h
 SampleCreationMethod = Literal[
     "found",
     "created",
+    "created and machine-translated",
     "human-translated and localized",
     "human-translated",
     "machine-translated",
@@ -149,7 +151,7 @@ _TASK_TYPE = (
     "InstructionReranking",
 ) + MIEB_TASK_TYPE
 
-TaskType = Literal[_TASK_TYPE]
+TaskType = Literal[_TASK_TYPE]  # type: ignore[valid-type]
 """The type of the task. E.g. includes "Classification", "Retrieval" and "Clustering"."""
 
 
@@ -191,8 +193,10 @@ AnnotatorType = Literal[
 """The type of the annotators. Is often important for understanding the quality of a dataset."""
 
 
-PromptDict = TypedDict(
-    "PromptDict", {prompt_type.value: str for prompt_type in PromptType}, total=False
+PromptDict = TypedDict(  # type: ignore[misc]
+    "PromptDict",
+    {prompt_type.value: str for prompt_type in PromptType},
+    total=False,
 )
 """A dictionary containing the prompt used for the task.
 
@@ -364,7 +368,7 @@ class TaskMetadata(BaseModel):
         """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
+        return {"default": cast(list[str], self.eval_langs)}
 
     @property
     def intext_citation(self, include_cite: bool = True) -> str:
@@ -375,9 +379,8 @@ class TaskMetadata(BaseModel):
         if include_cite and cite:
             # check for whitespace in the citation
             if " " in cite:
-                logger.warning(
-                    "Citation contains whitespace. Please ensure that the citation is correctly formatted."
-                )
+                msg = "Citation contains whitespace. Please ensure that the citation is correctly formatted."
+                logger.warning(msg)
             return f"\\cite{{{cite}}}"
         return cite
 
@@ -413,7 +416,7 @@ class TaskMetadata(BaseModel):
         for subset, subset_value in stats.items():
             if subset == "hf_subset_descriptive_stats":
                 continue
-            n_samples[subset] = subset_value["num_samples"]  # type: ignore
+            n_samples[subset] = subset_value["num_samples"]
         return n_samples
 
     @property
@@ -446,7 +449,7 @@ class TaskMetadata(BaseModel):
         Raises:
             ValueError: If the prompt type is not recognized.
         """
-        if prompt_type is None:
+        if prompt_type is None or self.category is None:
             return self.modalities
         query_modalities, doc_modalities = self.category.split("2")
         category_to_modality: dict[str, Modalities] = {
@@ -466,7 +469,7 @@ class TaskMetadata(BaseModel):
 
     def _create_dataset_card_data(
         self,
-        existing_dataset_card_data: DatasetCardData | None = None,
+        existing_dataset_card_data: CardData | None = None,
     ) -> tuple[DatasetCardData, dict[str, Any]]:
         """Create a DatasetCardData object from the task metadata.
 
@@ -482,7 +485,6 @@ class TaskMetadata(BaseModel):
         dataset_type = [
             *self._hf_task_type(),
             *self._hf_task_category(),
-            *self._hf_subtypes(),
         ]
         languages = self._hf_languages()
 
@@ -501,12 +503,13 @@ class TaskMetadata(BaseModel):
 
         tags = ["mteb"] + self.modalities
 
-        descriptive_stats = self.descriptive_stats
-        if descriptive_stats is not None:
-            for split, split_stat in descriptive_stats.items():
+        descriptive_stats = ""
+        if self.descriptive_stats is not None:
+            descriptive_stats_ = self.descriptive_stats
+            for split, split_stat in descriptive_stats_.items():
                 if len(split_stat.get("hf_subset_descriptive_stats", {})) > 10:
                     split_stat.pop("hf_subset_descriptive_stats", {})
-            descriptive_stats = json.dumps(descriptive_stats, indent=4)
+            descriptive_stats = json.dumps(descriptive_stats_, indent=4)
 
         dataset_card_data_params = existing_dataset_card_data.to_dict()
         # override the existing values
@@ -583,10 +586,8 @@ class TaskMetadata(BaseModel):
 
     def _hf_subtypes(self) -> list[str]:
         # to get full list of available task_ids execute
-        # requests.post("https://huggingface.co/api/validate-yaml", json={
-        #   "content": "---\ntask_ids: 'test'\n---",
-        #   "repoType": "dataset"
-        # })
+        # https://huggingface.co/api/datasets-tags-by-type?type=task_ids
+        # ref https://huggingface-openapi.hf.space/#tag/datasets/GET/api/datasets-tags-by-type
         mteb_to_hf_subtype = {
             "Article retrieval": ["document-retrieval"],
             "Conversational retrieval": ["conversational", "utterance-retrieval"],
@@ -608,7 +609,7 @@ class TaskMetadata(BaseModel):
                 "hate-speech-detection",
             ],
             "Thematic clustering": [],
-            "Scientific Reranking": [],
+            "Scientific Reranking": ["text-scoring"],
             "Claim verification": ["fact-checking", "fact-checking-retrieval"],
             "Topic classification": ["topic-classification"],
             "Code retrieval": [],
@@ -616,21 +617,21 @@ class TaskMetadata(BaseModel):
             "Cross-Lingual Semantic Discrimination": [],
             "Textual Entailment": ["natural-language-inference"],
             "Counterfactual Detection": [],
-            "Emotion classification": [],
+            "Emotion classification": ["sentiment-classification"],
             "Reasoning as Retrieval": [],
             "Rendered Texts Understanding": [],
             "Image Text Retrieval": [],
             "Object recognition": [],
             "Scene recognition": [],
             "Caption Pairing": ["image-captioning"],
-            "Emotion recognition": [],
+            "Emotion recognition": ["sentiment-scoring"],
             "Textures recognition": [],
             "Activity recognition": [],
             "Tumor detection": [],
             "Duplicate Detection": [],
             "Rendered semantic textual similarity": [
                 "semantic-similarity-scoring",
-                "rendered semantic textual similarity",
+                "semantic-similarity-classification",
             ],
             "Intent classification": [
                 "intent-classification",
@@ -644,10 +645,8 @@ class TaskMetadata(BaseModel):
 
     def _hf_task_type(self) -> list[str]:
         # to get full list of task_types execute:
-        # requests.post("https://huggingface.co/api/validate-yaml", json={
-        #     "content": "---\ntask_categories: ['test']\n---", "repoType": "dataset"
-        # }).json()
-        # or look at https://huggingface.co/tasks
+        # https://huggingface.co/api/datasets-tags-by-type?type=task_categories
+        # ref https://huggingface-openapi.hf.space/#tag/datasets/GET/api/datasets-tags-by-type
         mteb_task_type_to_datasets = {
             # Text
             "BitextMining": ["translation"],
@@ -666,7 +665,7 @@ class TaskMetadata(BaseModel):
             "Any2AnyRetrieval": ["visual-document-retrieval"],
             "Any2AnyMultilingualRetrieval": ["visual-document-retrieval"],
             "VisionCentricQA": ["visual-question-answering"],
-            "ImageClustering": ["image-clustering"],
+            "ImageClustering": ["image-feature-extraction"],
             "ImageClassification": ["image-classification"],
             "ImageMultilabelClassification": ["image-classification"],
             "DocumentUnderstanding": ["visual-document-retrieval"],
@@ -694,11 +693,11 @@ class TaskMetadata(BaseModel):
 
     def _hf_languages(self) -> list[str]:
         languages: list[str] = []
-        if self.is_multilingual:
-            for val in list(self.eval_langs.values()):
+        if self.is_multilingual and isinstance(self.eval_langs, dict):
+            for val in self.eval_langs.values():
                 languages.extend(val)
         else:
-            languages = self.eval_langs
+            languages = cast(list[str], self.eval_langs)
         # value "python" is not valid. It must be an ISO 639-1, 639-2 or 639-3 code (two/three letters),
         # or a special value like "code", "multilingual".
         readme_langs = []
@@ -710,7 +709,7 @@ class TaskMetadata(BaseModel):
                 readme_langs.append(lang_name)
         return sorted(set(readme_langs))
 
-    def _hf_license(self) -> str:
+    def _hf_license(self) -> str | None:
         dataset_license = self.license
         if dataset_license:
             license_mapping = {

mteb/abstasks/text/bitext_mining.py

@@ -1,7 +1,7 @@
 import logging
 from collections import defaultdict
 from pathlib import Path
-from typing import Any, ClassVar, TypedDict
+from typing import Any, ClassVar, TypedDict, cast
 
 from datasets import Dataset, DatasetDict
 from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score
@@ -10,7 +10,7 @@ from mteb._evaluators import BitextMiningEvaluator
 from mteb.abstasks._statistics_calculation import calculate_text_statistics
 from mteb.abstasks.abstask import AbsTask
 from mteb.models import EncoderProtocol, MTEBModels
-from mteb.types import HFSubset, ScoresDict
+from mteb.types import EncodeKwargs, HFSubset, ScoresDict
 from mteb.types.statistics import SplitDescriptiveStatistics, TextStatistics
 
 logger = logging.getLogger(__name__)
@@ -73,11 +73,14 @@ class AbsTaskBitextMining(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]:
         """Added load for "parallel" datasets"""
+        if not isinstance(model, EncoderProtocol):
+            raise TypeError("Expected model to be an instance of EncoderProtocol")
+
         if not self.data_loaded:
             self.load_data()
 
@@ -87,11 +90,16 @@ class AbsTaskBitextMining(AbsTask):
         if subsets_to_run is not None:
             hf_subsets = [s for s in hf_subsets if s in subsets_to_run]
 
-        scores = {}
+        encoder_model = cast(EncoderProtocol, model)
+
+        if self.dataset is None:
+            raise ValueError("Dataset is not loaded.")
+
+        scores: dict[str, BitextMiningMetrics] = {}
         if self.parallel_subsets:
-            scores = self._evaluate_subset(
-                model,
-                self.dataset[split],  # type: ignore
+            scores = self._evaluate_subset(  # type: ignore[assignment]
+                encoder_model,
+                self.dataset[split],
                 parallel=True,
                 hf_split=split,
                 hf_subset="parallel",
@@ -109,8 +117,8 @@ class AbsTaskBitextMining(AbsTask):
                     data_split = self.dataset[split]
                 else:
                     data_split = self.dataset[hf_subset][split]
-                scores[hf_subset] = self._evaluate_subset(
-                    model,
+                scores[hf_subset] = self._evaluate_subset(  # type: ignore[assignment]
+                    encoder_model,
                     data_split,
                     hf_split=split,
                     hf_subset=hf_subset,
@@ -119,32 +127,32 @@ class AbsTaskBitextMining(AbsTask):
                     **kwargs,
                 )
 
-        return scores
+        return cast(dict[HFSubset, ScoresDict], scores)
 
     def _get_pairs(self, parallel: bool) -> list[tuple[str, str]]:
         pairs = self._DEFAULT_PAIR
         if parallel:
-            pairs = [langpair.split("-") for langpair in self.hf_subsets]
+            pairs = [langpair.split("-") for langpair in self.hf_subsets]  # type: ignore[misc]
         return pairs
 
-    def _evaluate_subset(
+    def _evaluate_subset(  # type: ignore[override]
         self,
         model: EncoderProtocol,
         data_split: Dataset,
         *,
         hf_split: str,
         hf_subset: str,
-        parallel: bool = False,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
+        parallel: bool = False,
         **kwargs,
-    ) -> ScoresDict:
+    ) -> BitextMiningMetrics | dict[str, BitextMiningMetrics]:
         pairs = self._get_pairs(parallel)
 
         evaluator = BitextMiningEvaluator(
             data_split,
             task_metadata=self.metadata,
-            pair_columns=pairs,  # type: ignore
+            pair_columns=pairs,
             hf_split=hf_split,
             hf_subset=hf_subset,
             **kwargs,
@@ -168,16 +176,16 @@ class AbsTaskBitextMining(AbsTask):
             )
 
         if parallel:
-            metrics = {}
+            parallel_metrics = {}
             for keys, nearest_neighbors in neighbours.items():
-                metrics[keys] = self._compute_metrics(nearest_neighbors, gold)
+                parallel_metrics[keys] = self._compute_metrics(nearest_neighbors, gold)
 
-            for v in metrics.values():
+            for v in parallel_metrics.values():
                 self._add_main_score(v)
-        else:
-            def_pair_str = "-".join(self._DEFAULT_PAIR[0])
-            metrics = self._compute_metrics(neighbours[def_pair_str], gold)
-            self._add_main_score(metrics)
+            return parallel_metrics
+        def_pair_str = "-".join(self._DEFAULT_PAIR[0])
+        metrics = self._compute_metrics(neighbours[def_pair_str], gold)
+        self._add_main_score(metrics)
         return metrics
 
     def _compute_metrics(
@@ -250,8 +258,11 @@ class AbsTaskBitextMining(AbsTask):
         )
 
     def _push_dataset_to_hub(self, repo_name: str) -> None:
+        if self.dataset is None:
+            raise ValueError("Dataset is not loaded.")
+
         if self.metadata.is_multilingual:
-            dataset = defaultdict(dict)
+            dataset: dict[str, dict[str, list[str]]] = defaultdict(dict)
             for config in self.metadata.eval_langs:
                 logger.info(f"Converting {config} of {self.metadata.name}")
 
@@ -266,10 +277,10 @@ class AbsTaskBitextMining(AbsTask):
                     for split in self.dataset[config]:
                         dataset[split][lang_1] = self.dataset[config][split][sent_1]
                         dataset[split][lang_2] = self.dataset[config][split][sent_2]
-            for split in dataset:
-                dataset[split] = Dataset.from_dict(dataset[split])
-            dataset = DatasetDict(dataset)
-            dataset.push_to_hub(repo_name)
+            dataset_dict = DatasetDict(
+                {split: Dataset.from_dict(dataset[split]) for split in dataset}
+            )
+            dataset_dict.push_to_hub(repo_name)
         else:
             sentences = {}
             for split in self.dataset:

mteb/abstasks/text/reranking.py

@@ -16,7 +16,7 @@ else:
 
 logger = logging.getLogger(__name__)
 
-OLD_FORMAT_RERANKING_TASKS = []
+OLD_FORMAT_RERANKING_TASKS: list[str] = []
 
 
 @deprecated(
@@ -100,12 +100,14 @@ class AbsTaskReranking(AbsTaskRetrieval):
         if self.metadata.name not in OLD_FORMAT_RERANKING_TASKS:
             return
 
-        logging.info(
+        logger.info(
             f"Transforming old format to standard format for {self.metadata.name}"
         )
 
         given_dataset = copy(given_dataset)
-        self.dataset = defaultdict(lambda: defaultdict(dict))
+        self.dataset: dict[str, dict[str, RetrievalSplitData]] = defaultdict(
+            lambda: defaultdict(dict)  # type: ignore[arg-type]
+        )
 
         hf_subsets = self.hf_subsets
 
@@ -115,19 +117,19 @@ class AbsTaskReranking(AbsTaskRetrieval):
                 if hf_subset in cur_dataset:
                     cur_dataset = cur_dataset[hf_subset]
             elif "name" in self.metadata.dataset:
-                cur_dataset = datasets.load_dataset(**self.metadata.dataset)  # type: ignore
+                cur_dataset = datasets.load_dataset(**self.metadata.dataset)
                 assert hf_subset == "default", (
                     f"Only default subset is supported for {self.metadata.name} since `name` is given in the metadata."
                 )
             else:
                 cur_dataset = datasets.load_dataset(
                     **self.metadata.dataset, name=hf_subset
-                )  # type: ignore
+                )
 
             for split in cur_dataset:
                 corpus = []
                 queries = []
-                relevant_docs = defaultdict(dict)
+                relevant_docs: dict[str, dict[str, int]] = defaultdict(dict)
                 top_ranked = defaultdict(list)
 
                 # Create an enumerated dataset to pass indices

mteb/abstasks/text/summarization.py

@@ -1,6 +1,5 @@
 import logging
 from pathlib import Path
-from typing import Any
 
 import numpy as np
 from datasets import Dataset
@@ -12,7 +11,8 @@ from mteb.abstasks._statistics_calculation import (
     calculate_text_statistics,
 )
 from mteb.abstasks.abstask import AbsTask
-from mteb.models import EncoderProtocol
+from mteb.models import EncoderProtocol, MTEBModels
+from mteb.types import EncodeKwargs
 from mteb.types.statistics import (
     ScoreStatistics,
     SplitDescriptiveStatistics,
@@ -77,17 +77,22 @@ class AbsTaskSummarization(AbsTask):
 
     def _evaluate_subset(
         self,
-        model: EncoderProtocol,
+        model: MTEBModels,
         data_split: Dataset,
         *,
         hf_split: str,
         hf_subset: str,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
         **kwargs,
     ) -> SummarizationMetrics:
+        if not isinstance(model, EncoderProtocol):
+            raise TypeError("Expected model to be an instance of EncoderProtocol")
+
         normalized_scores = [
-            (np.array(x) - self.min_score) / (self.max_score - self.min_score)
+            (
+                (np.array(x) - self.min_score) / (self.max_score - self.min_score)
+            ).tolist()
             for x in data_split[self.relevancy_column_name]
         ]
         evaluator = self.evaluator(