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

mteb/_evaluators/retrieval_evaluator.py

@@ -1,23 +1,29 @@
-import logging
-from collections.abc import Sequence
-from typing import Any
+from __future__ import annotations
 
-from mteb.abstasks.task_metadata import TaskMetadata
-from mteb.models import SearchProtocol
-from mteb.types import (
-    CorpusDatasetType,
-    QueryDatasetType,
-    RelevantDocumentsType,
-    RetrievalEvaluationResult,
-    RetrievalOutputType,
-    TopRankedDocumentsType,
-)
+import logging
+from typing import TYPE_CHECKING
 
 from .evaluator import Evaluator
 from .retrieval_metrics import (
     calculate_retrieval_scores,
 )
 
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.models import SearchProtocol
+    from mteb.types import (
+        CorpusDatasetType,
+        EncodeKwargs,
+        QueryDatasetType,
+        RelevantDocumentsType,
+        RetrievalEvaluationResult,
+        RetrievalOutputType,
+        TopRankedDocumentsType,
+    )
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -48,7 +54,8 @@ class RetrievalEvaluator(Evaluator):
     def __call__(  # type: ignore[override]
         self,
         search_model: SearchProtocol,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
+        num_proc: int = 1,
     ) -> RetrievalOutputType:
         logger.info("Running retrieval task - Indexing corpus...")
         search_model.index(
@@ -57,6 +64,7 @@ class RetrievalEvaluator(Evaluator):
             hf_split=self.hf_split,
             hf_subset=self.hf_subset,
             encode_kwargs=encode_kwargs,
+            num_proc=num_proc,
         )
         logger.info("Running retrieval task - Searching queries...")
         return search_model.search(
@@ -67,6 +75,7 @@ class RetrievalEvaluator(Evaluator):
             hf_subset=self.hf_subset,
             encode_kwargs=encode_kwargs,
             top_ranked=self.top_ranked,
+            num_proc=num_proc,
         )
 
     def evaluate(

mteb/_evaluators/retrieval_metrics.py

@@ -1,6 +1,8 @@
+from __future__ import annotations
+
 import logging
 from collections import defaultdict
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 import pandas as pd
@@ -8,14 +10,19 @@ import pytrec_eval
 from packaging.version import Version
 from sklearn.metrics import auc
 
-from mteb.types import RelevantDocumentsType, RetrievalEvaluationResult
+from mteb.types import RetrievalEvaluationResult
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+    from mteb.types import RelevantDocumentsType
 
 logger = logging.getLogger(__name__)
 
 
 def mrr(
     qrels: RelevantDocumentsType,
-    results: dict[str, dict[str, float]],
+    results: Mapping[str, Mapping[str, float]],
     k_values: list[int],
 ) -> dict[str, list[float]]:
     mrr_metrics = defaultdict(list)
@@ -32,7 +39,7 @@ def mrr(
             doc_id for doc_id in qrels[query_id] if qrels[query_id][doc_id] > 0
         }
         for k in k_values:
-            rr = 0
+            rr = 0.0
             for rank, hit in enumerate(top_hits[query_id][0:k]):
                 if hit[0] in query_relevant_docs:
                     rr = 1.0 / (rank + 1)
@@ -45,8 +52,8 @@ def recall_cap(
     qrels: RelevantDocumentsType,
     results: dict[str, dict[str, float]],
     k_values: list[int],
-) -> dict[str, list[float]]:
-    capped_recall = defaultdict(list)
+) -> dict[str, list[float | None]]:
+    capped_recall: dict[str, list[float | None]] = defaultdict(list)
 
     k_max = max(k_values)
 
@@ -139,7 +146,7 @@ def calculate_pmrr(original_run, new_run, changed_qrels):
     changes = []
     for qid in changed_qrels.keys():
         if qid + "-og" not in original_run or qid + "-changed" not in new_run:
-            logging.warning(f"Query {qid} not found in the runs for calculating p-MRR")
+            logger.warning(f"Query {qid} not found in the runs for calculating p-MRR")
             continue
         original_qid_run = original_run[qid + "-og"]
         new_qid_run = new_run[qid + "-changed"]
@@ -188,7 +195,7 @@ def evaluate_p_mrr_change(
     Returns:
         A dictionary with the scores, including "p-MRR", "og" and "changed" keys.
     """
-    followir_scores = defaultdict(dict)
+    followir_scores: dict[str, float | dict[str, float]] = defaultdict(dict)
 
     qrels_sep = {
         "og": {k: v for k, v in qrels.items() if k.endswith("-og")},
@@ -227,7 +234,7 @@ def evaluate_p_mrr_change(
             ndcg, _map, recall, precision, naucs, avg_mrr, naucs_mrr, cv_recall, {}
         )
         for key, value in scores_dict.items():
-            followir_scores[name][key] = value
+            followir_scores[name][key] = value  # type: ignore[index]
 
     return followir_scores
 
@@ -254,8 +261,8 @@ def confidence_scores(sim_scores: list[float]) -> dict[str, float]:
     sim_scores_sorted = sorted(sim_scores)[::-1]
 
     cs_max = sim_scores_sorted[0]
-    cs_std = np.std(sim_scores)
-    cs_diff1 = None
+    cs_std = float(np.std(sim_scores))
+    cs_diff1 = 0.0
     if len(sim_scores) > 1:
         cs_diff1 = sim_scores_sorted[0] - sim_scores_sorted[1]
     elif len(sim_scores) == 1:
@@ -410,7 +417,7 @@ def make_score_dict(
     cv_recall: dict[str, float],
     task_scores: dict[str, float],
     previous_results_model_meta: dict[str, Any] | None = None,
-) -> dict[str, float]:
+) -> dict[str, Any]:
     return {
         **{f"ndcg_at_{k.split('@')[1]}": v for (k, v) in ndcg.items()},
         **{f"map_at_{k.split('@')[1]}": v for (k, v) in _map.items()},
@@ -528,7 +535,7 @@ def max_over_subqueries(
 
 
 def calculate_retrieval_scores(
-    results: dict[str, dict[str, float]],
+    results: Mapping[str, Mapping[str, float]],
     qrels: RelevantDocumentsType,
     k_values: list[int],
     skip_first_result: bool = False,
@@ -576,7 +583,7 @@ def calculate_retrieval_scores(
 
 
 def evaluate_abstention(
-    results: dict[str, dict[str, float]],
+    results: Mapping[str, Mapping[str, float]],
     metric_scores: dict[str, list[float]],
 ) -> dict[str, float]:
     """Computes normalized Area Under the Curve on a set of evaluated instances as presented in the paper https://arxiv.org/abs/2402.12997
@@ -591,21 +598,21 @@ def evaluate_abstention(
     all_sim_scores = [list(results[qid].values()) for qid in list(results.keys())]
     all_conf_scores = [confidence_scores(sim_scores) for sim_scores in all_sim_scores]
     conf_fcts = list(all_conf_scores[0].keys())
-    all_conf_scores = {
+    all_conf_scores_ = {
         fct: np.array([x[fct] for x in all_conf_scores]) for fct in conf_fcts
     }
-    metric_scores = {k: np.array(v) for k, v in metric_scores.items()}
+    metric_scores_ = {k: np.array(v) for k, v in metric_scores.items()}
     naucs = {}
 
-    for metric_name, scores in metric_scores.items():
-        for fct, conf_scores in all_conf_scores.items():
+    for metric_name, scores in metric_scores_.items():
+        for fct, conf_scores in all_conf_scores_.items():
             naucs[f"nAUC_{metric_name}_{fct}"] = nauc(conf_scores, scores)
 
     return naucs
 
 
 def calculate_cv_recall(
-    results: dict[str, dict[str, float]],
+    results: Mapping[str, Mapping[str, float]],
     qrels: RelevantDocumentsType,
     k_values: list[int],
     skip_first_result: bool = False,

mteb/_evaluators/sklearn_evaluator.py

@@ -1,27 +1,31 @@
-import logging
-from typing import Any, Protocol
+from __future__ import annotations
 
-import numpy as np
-from datasets import Dataset
-from torch.utils.data import DataLoader
-from typing_extensions import Self
+import logging
+from typing import TYPE_CHECKING, Any, Protocol, cast
 
 from mteb._create_dataloaders import create_dataloader
-from mteb.abstasks.task_metadata import TaskMetadata
-from mteb.models import EncoderProtocol
-from mteb.types import BatchedInput
 
 from .evaluator import Evaluator
 
+if TYPE_CHECKING:
+    import numpy as np
+    from datasets import Dataset
+    from torch.utils.data import DataLoader
+    from typing_extensions import Self
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.models import EncoderProtocol
+    from mteb.types import Array, BatchedInput, EncodeKwargs
+
 logger = logging.getLogger(__name__)
 
 
 class SklearnModelProtocol(Protocol):
-    def fit(self, X: np.ndarray, y: np.ndarray | list[int]) -> None: ...  # noqa: N803
-    def predict(self, X: np.ndarray) -> np.ndarray: ...  # noqa: N803
+    def fit(self, X: Array, y: np.ndarray | list[int]) -> None: ...  # noqa: N803
+    def predict(self, X: Array) -> np.ndarray: ...  # noqa: N803
     def get_params(self) -> dict[str, Any]: ...
-    def set_params(self, **kwargs: dict[str, Any]) -> Self: ...
-    def score(self, X: np.ndarray, y: np.ndarray | list[int]) -> float: ...  # noqa: N803
+    def set_params(self, random_state: int, **kwargs: dict[str, Any]) -> Self: ...
+    def score(self, X: Array, y: np.ndarray | list[int]) -> float: ...  # noqa: N803
 
 
 class SklearnEvaluator(Evaluator):
@@ -50,18 +54,20 @@ class SklearnEvaluator(Evaluator):
         self.evaluator_model = evaluator_model
 
     def create_dataloaders(
-        self, encode_kwargs: dict[str, Any]
+        self, encode_kwargs: EncodeKwargs, num_proc: int
     ) -> tuple[DataLoader[BatchedInput], DataLoader[BatchedInput]]:
         dataloader_train = create_dataloader(
             self.train_dataset,
             self.task_metadata,
             input_column=self.values_column_name,
+            num_proc=num_proc,
             **encode_kwargs,
         )
         dataloader_test = create_dataloader(
             self.eval_dataset,
             self.task_metadata,
             input_column=self.values_column_name,
+            num_proc=num_proc,
             **encode_kwargs,
         )
         return dataloader_train, dataloader_test
@@ -70,15 +76,17 @@ class SklearnEvaluator(Evaluator):
         self,
         model: EncoderProtocol,
         *,
-        encode_kwargs: dict[str, Any],
-        test_cache: np.ndarray | None = None,
-    ) -> tuple[np.ndarray, np.ndarray]:
+        encode_kwargs: EncodeKwargs,
+        test_cache: Array | None = None,
+        num_proc: int = 1,
+    ) -> tuple[np.ndarray, Array]:
         """Classification evaluation by training a sklearn classifier on the embeddings of the training set and evaluating on the embeddings of the test set.
 
         Args:
             model: Encoder
             encode_kwargs: encode kwargs
             test_cache: embeddings of the test set, if already computed
+            num_proc: number of processes to use
 
         Returns:
             Tuple of test predictions and embeddings
@@ -86,6 +94,7 @@ class SklearnEvaluator(Evaluator):
         """
         dataloader_train, dataloader_test = self.create_dataloaders(
             encode_kwargs=encode_kwargs,
+            num_proc=num_proc,
         )
 
         logger.info("Running - Encoding samples...")
@@ -104,6 +113,7 @@ class SklearnEvaluator(Evaluator):
                 hf_subset=self.hf_subset,
                 **encode_kwargs,
             )
+            test_cache = cast("Array", test_cache)
 
         logger.info("Running - Fitting classifier...")
         y_train = self.train_dataset[self.label_column_name]

mteb/_evaluators/text/bitext_mining_evaluator.py

@@ -1,15 +1,19 @@
+from __future__ import annotations
+
 import logging
-from typing import Any
+from typing import TYPE_CHECKING
 
-import numpy as np
 import torch
 from datasets import Dataset
 from tqdm.auto import tqdm
 
 from mteb._create_dataloaders import _create_dataloader_from_texts
 from mteb._evaluators.evaluator import Evaluator
-from mteb.abstasks.task_metadata import TaskMetadata
-from mteb.models import EncoderProtocol
+
+if TYPE_CHECKING:
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.models import EncoderProtocol
+    from mteb.types import Array, EncodeKwargs
 
 logger = logging.getLogger(__name__)
 
@@ -33,7 +37,11 @@ class BitextMiningEvaluator(Evaluator):
         self.task_metadata = task_metadata
 
     def __call__(
-        self, model: EncoderProtocol, *, encode_kwargs: dict[str, Any]
+        self,
+        model: EncoderProtocol,
+        *,
+        encode_kwargs: EncodeKwargs,
+        num_proc: int = 1,
     ) -> dict[str, list[dict[str, float]]]:
         pair_elements = {p for pair in self.pairs for p in pair}
         if isinstance(self.sentences, Dataset):
@@ -48,6 +56,7 @@ class BitextMiningEvaluator(Evaluator):
         for sub in tqdm(subsets):
             dataloader = _create_dataloader_from_texts(
                 self.sentences[sub],
+                num_proc=num_proc,
                 **encode_kwargs,
             )
             embeddings[sub] = model.encode(
@@ -69,11 +78,11 @@ class BitextMiningEvaluator(Evaluator):
 
     def _similarity_search(
         self,
-        query_embeddings: np.ndarray,
-        corpus_embeddings: np.ndarray,
+        query_embeddings: Array,
+        corpus_embeddings: Array,
         model: EncoderProtocol,
         query_chunk_size: int = 100,
-        corpus_chunk_size: int = 500000,
+        corpus_chunk_size: int = 500_000,
     ) -> list[dict[str, float]]:
         """This function performs a cosine similarity search between a list of query embeddings and a list of corpus embeddings.
 
@@ -104,13 +113,15 @@ class BitextMiningEvaluator(Evaluator):
         ):
             query_embeddings = query_embeddings.to(corpus_embeddings.device)
 
-        queries_result_list = [[] for _ in range(len(query_embeddings))]
+        queries_result_list: list[list[dict[str, float]]] = [
+            [] for _ in range(len(query_embeddings))
+        ]
 
         for query_start_idx in range(0, len(query_embeddings), query_chunk_size):
             # Iterate over chunks of the corpus
             for corpus_start_idx in range(0, len(corpus_embeddings), corpus_chunk_size):
                 # Compute cosine similarities
-                similarity_scores = model.similarity(  # type: ignore
+                similarity_scores = model.similarity(
                     query_embeddings[
                         query_start_idx : query_start_idx + query_chunk_size
                     ],
@@ -120,15 +131,17 @@ class BitextMiningEvaluator(Evaluator):
                 )
 
                 # Get top-k scores
-                cos_scores_top_k_values, cos_scores_top_k_idx = torch.topk(
-                    torch.tensor(similarity_scores),
-                    1,
-                    dim=1,
-                    largest=True,
-                    sorted=False,
+                cos_scores_top_k_values_tensor, cos_scores_top_k_idx_tensor = (
+                    torch.topk(
+                        torch.tensor(similarity_scores),
+                        1,
+                        dim=1,
+                        largest=True,
+                        sorted=False,
+                    )
                 )
-                cos_scores_top_k_values = cos_scores_top_k_values.cpu().tolist()
-                cos_scores_top_k_idx = cos_scores_top_k_idx.cpu().tolist()
+                cos_scores_top_k_values = cos_scores_top_k_values_tensor.cpu().tolist()
+                cos_scores_top_k_idx = cos_scores_top_k_idx_tensor.cpu().tolist()
 
                 for query_itr in range(len(similarity_scores)):
                     for sub_corpus_id, score in zip(
@@ -141,11 +154,14 @@ class BitextMiningEvaluator(Evaluator):
                             {"corpus_id": corpus_id, "score": score}
                         )
 
+        result_queries_list: list[dict[str, float]] = [
+            {} for _ in range(len(query_embeddings))
+        ]
         # Sort and strip to top_k results
         for idx in range(len(queries_result_list)):
             queries_result_list[idx] = sorted(
                 queries_result_list[idx], key=lambda x: x["score"], reverse=True
             )
-            queries_result_list[idx] = queries_result_list[idx][0]
+            result_queries_list[idx] = queries_result_list[idx][0]
 
-        return queries_result_list
+        return result_queries_list

mteb/_evaluators/text/summarization_evaluator.py

@@ -1,6 +1,8 @@
+from __future__ import annotations
+
 import logging
 import sys
-from typing import Any, TypedDict
+from typing import TYPE_CHECKING, TypedDict
 
 import numpy as np
 import torch
@@ -9,10 +11,13 @@ from tqdm.auto import tqdm
 
 from mteb._create_dataloaders import _create_dataloader_from_texts
 from mteb._evaluators.evaluator import Evaluator
-from mteb.abstasks.task_metadata import TaskMetadata
-from mteb.models import EncoderProtocol
 from mteb.similarity_functions import cos_sim, dot_score
 
+if TYPE_CHECKING:
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.models import EncoderProtocol
+    from mteb.types import EncodeKwargs
+
 # if later than python 3.13 use typing module
 if sys.version_info >= (3, 13):
     from warnings import deprecated
@@ -94,7 +99,8 @@ class SummarizationEvaluator(Evaluator):
         self,
         model: EncoderProtocol,
         *,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
+        num_proc: int = 1,
     ) -> SummarizationDistances:
         # Get the human & machine summaries for the text in one go for all
         human_lens = [len(human_summaries) for human_summaries in self.human_summaries]
@@ -110,6 +116,7 @@ class SummarizationEvaluator(Evaluator):
                     for human_summaries in self.human_summaries
                     for summary in human_summaries
                 ],
+                num_proc=num_proc,
                 **encode_kwargs,
             ),
             task_metadata=self.task_metadata,
@@ -135,10 +142,10 @@ class SummarizationEvaluator(Evaluator):
         )
 
         # Split the embeddings into the original human & machine summaries
-        embs_human_summaries_all = np.split(
+        embs_human_summaries_all_split = np.split(
             embs_human_summaries_all, np.cumsum(human_lens)[:-1]
         )
-        embs_machine_summaries_all = np.split(
+        embs_machine_summaries_all_split = np.split(
             embs_machine_summaries_all, np.cumsum(machine_lens)[:-1]
         )
 
@@ -148,7 +155,9 @@ class SummarizationEvaluator(Evaluator):
         all_human_scores = []
 
         for i, (embs_human_summaries, embs_machine_summaries) in tqdm(
-            enumerate(zip(embs_human_summaries_all, embs_machine_summaries_all)),
+            enumerate(
+                zip(embs_human_summaries_all_split, embs_machine_summaries_all_split)
+            ),
             desc="Scoring",
             total=len(self.human_summaries),
         ):
@@ -164,7 +173,7 @@ class SummarizationEvaluator(Evaluator):
                 dot_scores = dot_score(emb_machine_summary, embs_human_summaries)
 
                 _sim_score = [
-                    float(model.similarity(emb_machine_summary, emb_human_summary))  # type: ignore
+                    float(model.similarity(emb_machine_summary, emb_human_summary))
                     for emb_human_summary in embs_human_summaries
                 ]
                 sim_score = torch.tensor(_sim_score)
@@ -216,17 +225,19 @@ class SummarizationEvaluator(Evaluator):
             strict=True,
         ):
             cosine_spearman_scores.append(
-                spearmanr(human_scores, cosine_pred_scores).statistic
+                float(spearmanr(human_scores, cosine_pred_scores).statistic)
             )
             cosine_pearson_scores.append(
-                pearsonr(human_scores, cosine_pred_scores).statistic
+                float(pearsonr(human_scores, cosine_pred_scores).statistic)
             )
             dot_spearman_scores.append(
-                spearmanr(human_scores, dot_pred_scores).statistic
+                float(spearmanr(human_scores, dot_pred_scores).statistic)
             )
-            dot_pearson_scores.append(pearsonr(human_scores, dot_pred_scores).statistic)
-            spearman_scores.append(spearmanr(human_scores, sim_scores).statistic)
-            pearson_scores.append(pearsonr(human_scores, sim_scores).statistic)
+            dot_pearson_scores.append(
+                float(pearsonr(human_scores, dot_pred_scores).statistic)
+            )
+            spearman_scores.append(float(spearmanr(human_scores, sim_scores).statistic))
+            pearson_scores.append(float(pearsonr(human_scores, sim_scores).statistic))
 
         return SummarizationMetrics(
             pearson=float(np.mean(pearson_scores)),
@@ -273,10 +284,10 @@ class DeprecatedSummarizationEvaluator(SummarizationEvaluator):
             pearson_scores.append(pearsonr(human_scores, sim_scores))
 
         return SummarizationMetrics(
-            pearson=float(np.mean(pearson_scores)),
-            spearman=float(np.mean(spearman_scores)),
-            cosine_spearman=float(np.mean(cosine_spearman_scores)),
-            cosine_pearson=float(np.mean(cosine_pearson_scores)),
-            dot_pearson=float(np.mean(dot_pearson_scores)),
-            dot_spearman=float(np.mean(dot_spearman_scores)),
+            pearson=float(np.mean(pearson_scores)),  # type: ignore[arg-type]
+            spearman=float(np.mean(spearman_scores)),  # type: ignore[arg-type]
+            cosine_spearman=float(np.mean(cosine_spearman_scores)),  # type: ignore[arg-type]
+            cosine_pearson=float(np.mean(cosine_pearson_scores)),  # type: ignore[arg-type]
+            dot_pearson=float(np.mean(dot_pearson_scores)),  # type: ignore[arg-type]
+            dot_spearman=float(np.mean(dot_spearman_scores)),  # type: ignore[arg-type]
         )

mteb/_evaluators/zeroshot_classification_evaluator.py

@@ -1,5 +1,7 @@
+from __future__ import annotations
+
 import logging
-from typing import Any
+from typing import TYPE_CHECKING
 
 from datasets import Dataset
 
@@ -7,13 +9,17 @@ from mteb._create_dataloaders import (
     _create_dataloader_from_texts,
     create_dataloader,
 )
-from mteb.abstasks.task_metadata import TaskMetadata
-from mteb.models import EncoderProtocol
 from mteb.similarity_functions import similarity
-from mteb.types import Array
 
 from .evaluator import Evaluator
 
+if TYPE_CHECKING:
+    from datasets import Dataset
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.models import EncoderProtocol
+    from mteb.types import Array, EncodeKwargs
+
 logger = logging.getLogger(__name__)
 
 
@@ -38,12 +44,17 @@ class ZeroShotClassificationEvaluator(Evaluator):
         self.hf_subset = hf_subset
 
     def __call__(
-        self, model: EncoderProtocol, *, encode_kwargs: dict[str, Any]
+        self,
+        model: EncoderProtocol,
+        *,
+        encode_kwargs: EncodeKwargs,
+        num_proc: int = 1,
     ) -> Array:
         dataloader = create_dataloader(
             self.dataset,
             input_column=self.input_column_name,
             task_metadata=self.task_metadata,
+            num_proc=num_proc,
             **encode_kwargs,
         )
 

mteb/_helpful_enum.py

@@ -1,6 +1,10 @@
+from __future__ import annotations
+
 from enum import Enum
+from typing import TYPE_CHECKING
 
-from typing_extensions import Self
+if TYPE_CHECKING:
+    from typing_extensions import Self
 
 
 class HelpfulStrEnum(str, Enum):

mteb/abstasks/_data_filter/filters.py

@@ -1,12 +1,18 @@
 """Simplified version of https://gist.github.com/AlexeyVatolin/ea3adc21aa7a767603ff393b22085adc from https://github.com/embeddings-benchmark/mteb/pull/2900"""
 
+from __future__ import annotations
+
 import logging
+from typing import TYPE_CHECKING
 
 import datasets
 import pandas as pd
-from datasets import Dataset, DatasetDict
+from datasets import DatasetDict
+
+if TYPE_CHECKING:
+    from datasets import Dataset
 
-from mteb import TaskMetadata
+    from mteb import TaskMetadata
 
 logger = logging.getLogger(__name__)
 
@@ -61,7 +67,7 @@ def filter_unclear_label(
         for text, label in zip(ds[input_column], ds[label_column]):
             key = text.strip().lower()
             normalized.setdefault(key, set()).add(
-                label if isinstance(label, (str, int, float)) else tuple(label)
+                label if isinstance(label, (str, int, float)) else tuple(label)  # type: ignore[arg-type]
             )
 
     bad_texts = {t for t, labels in normalized.items() if len(labels) > 1}

mteb/abstasks/_data_filter/task_pipelines.py

@@ -1,9 +1,10 @@
+from __future__ import annotations
+
 import logging
+from typing import TYPE_CHECKING
 
 from datasets import DatasetDict
 
-from mteb import TaskMetadata
-from mteb.abstasks import AbsTaskClassification
 from mteb.abstasks._data_filter.filters import (
     deduplicate,
     filter_empty,
@@ -13,6 +14,10 @@ from mteb.abstasks._data_filter.filters import (
     split_train_test,
 )
 
+if TYPE_CHECKING:
+    from mteb import TaskMetadata
+    from mteb.abstasks import AbsTaskClassification
+
 logger = logging.getLogger(__name__)
 
 
@@ -89,6 +94,9 @@ def process_classification(
             subset=None,
         )
 
+    if task.dataset is None:
+        raise ValueError("Task dataset is None.")
+
     new_ds = {}
     for subset in task.dataset:
         new_ds[subset] = clean_dataset(