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

mteb/abstasks/pair_classification.py

@@ -1,16 +1,15 @@
+from __future__ import annotations
+
 import hashlib
 import logging
 from collections import defaultdict
-from pathlib import Path
+from typing import TYPE_CHECKING
 
 import numpy as np
 from datasets import Dataset
 from sklearn.metrics import average_precision_score
 
 from mteb._evaluators import PairClassificationEvaluator
-from mteb._evaluators.pair_classification_evaluator import (
-    PairClassificationDistances,
-)
 from mteb.abstasks._statistics_calculation import (
     calculate_image_statistics,
     calculate_label_statistics,
@@ -19,14 +18,25 @@ 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.types.statistics import (
-    ImageStatistics,
-    LabelStatistics,
     SplitDescriptiveStatistics,
-    TextStatistics,
 )
 
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from mteb._evaluators.pair_classification_evaluator import (
+        PairClassificationDistances,
+    )
+    from mteb.models.models_protocols import MTEBModels
+    from mteb.types import EncodeKwargs, PromptType
+    from mteb.types.statistics import (
+        ImageStatistics,
+        LabelStatistics,
+        TextStatistics,
+    )
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -44,8 +54,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 +89,19 @@ 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,
+        num_proc: int = 1,
         **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
@@ -102,7 +116,11 @@ class AbsTaskPairClassification(AbsTask):
             input2_prompt_type=self.input2_prompt_type,
             **kwargs,
         )
-        similarity_scores = evaluator(model, encode_kwargs=encode_kwargs)
+        similarity_scores = evaluator(
+            model,
+            encode_kwargs=encode_kwargs,
+            num_proc=num_proc,
+        )
 
         if prediction_folder:
             self._save_task_predictions(
@@ -120,7 +138,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 +160,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)
@@ -235,8 +253,14 @@ class AbsTaskPairClassification(AbsTask):
             labels_statistics=calculate_label_statistics(labels),
         )
 
-    def _push_dataset_to_hub(self, repo_name: str) -> None:
+    def _push_dataset_to_hub(self, repo_name: str, num_proc: int = 1) -> 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]:
@@ -253,6 +277,7 @@ class AbsTaskPairClassification(AbsTask):
                 self.input2_column_name,
                 self.label_column_name,
             ],
+            num_proc=num_proc,
         )
 
     def _compute_metrics_values(
@@ -290,13 +315,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 +348,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

@@ -1,29 +1,37 @@
+from __future__ import annotations
+
 import logging
-from typing import TypedDict
+from typing import TYPE_CHECKING, TypedDict
 
 import datasets
 import numpy as np
 import pandas as pd
-from datasets import Dataset
 from scipy.stats import kendalltau
 from sklearn.linear_model import LinearRegression
 from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
 
-from mteb._evaluators.sklearn_evaluator import SklearnEvaluator, SklearnModelProtocol
+from mteb._evaluators.sklearn_evaluator import SklearnEvaluator
 from mteb.abstasks._statistics_calculation import (
     calculate_image_statistics,
     calculate_score_statistics,
     calculate_text_statistics,
 )
 from mteb.types.statistics import (
-    ImageStatistics,
-    ScoreStatistics,
     SplitDescriptiveStatistics,
-    TextStatistics,
 )
 
 from .classification import AbsTaskClassification
 
+if TYPE_CHECKING:
+    from datasets import Dataset
+
+    from mteb._evaluators.sklearn_evaluator import SklearnModelProtocol
+    from mteb.types.statistics import (
+        ImageStatistics,
+        ScoreStatistics,
+        TextStatistics,
+    )
+
 logger = logging.getLogger(__name__)
 
 
@@ -84,10 +92,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 +121,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 +191,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,13 +1,13 @@
+from __future__ import annotations
+
 import json
 import logging
 from collections import defaultdict
-from collections.abc import Callable, Sequence
 from pathlib import Path
 from time import time
-from typing import Any, Literal
+from typing import TYPE_CHECKING, Any, Literal
 
 from datasets import Dataset, DatasetDict, concatenate_datasets
-from typing_extensions import Self
 
 from mteb._create_dataloaders import (
     _combine_queries_with_instruction_text,
@@ -19,24 +19,12 @@ from mteb._evaluators.retrieval_metrics import make_score_dict
 from mteb.models import (
     CrossEncoderProtocol,
     EncoderProtocol,
-    MTEBModels,
     SearchCrossEncoderWrapper,
     SearchEncoderWrapper,
     SearchProtocol,
 )
-from mteb.types import (
-    HFSubset,
-    QueryDatasetType,
-    RelevantDocumentsType,
-    RetrievalOutputType,
-    ScoresDict,
-)
 from mteb.types.statistics import (
-    ImageStatistics,
-    RelevantDocsStatistics,
     SplitDescriptiveStatistics,
-    TextStatistics,
-    TopRankedStatistics,
 )
 
 from ._statistics_calculation import (
@@ -52,6 +40,30 @@ from .retrieval_dataset_loaders import (
     _combine_queries_with_instructions_datasets,
 )
 
+if TYPE_CHECKING:
+    from collections.abc import Callable, Mapping, Sequence
+
+    from typing_extensions import Self
+
+    from mteb.models import (
+        MTEBModels,
+    )
+    from mteb.types import (
+        EncodeKwargs,
+        HFSubset,
+        QueryDatasetType,
+        RelevantDocumentsType,
+        RetrievalOutputType,
+        ScoresDict,
+    )
+    from mteb.types.statistics import (
+        ImageStatistics,
+        RelevantDocsStatistics,
+        TextStatistics,
+        TopRankedStatistics,
+    )
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -136,7 +148,7 @@ class AbsTaskRetrieval(AbsTask):
             )
         )
 
-    def convert_v1_dataset_format_to_v2(self):
+    def convert_v1_dataset_format_to_v2(self, num_proc: int) -> None:
         """Convert dataset from v1 (from `self.queries`, `self.document`) format to v2 format (`self.dotaset`)."""
         # check if dataset is `v1` version
         if not hasattr(self, "queries"):
@@ -184,17 +196,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"):
@@ -203,6 +215,7 @@ class AbsTaskRetrieval(AbsTask):
                             _combine_queries_with_instructions_datasets(
                                 self.dataset[subset][split]["queries"],
                                 instructions,
+                                num_proc,
                             )
                         )
                     if hasattr(self, "top_ranked"):
@@ -211,15 +224,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"):
@@ -228,28 +241,29 @@ class AbsTaskRetrieval(AbsTask):
                         _combine_queries_with_instructions_datasets(
                             self.dataset[subset][split]["queries"],
                             instructions,
+                            num_proc,
                         )
                     )
-                if hasattr(self, "top_ranked"):
+                if hasattr(self, "top_ranked") and self.top_ranked:
                     self.dataset[subset][split]["top_ranked"] = self.top_ranked[
                         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"):
             del self.top_ranked
 
-    def load_data(self) -> None:
+    def load_data(self, num_proc: int = 1, **kwargs) -> None:
         """Load the dataset for the retrieval task."""
         if self.data_loaded:
             return
 
         dataset_path = self.metadata.dataset["path"]
-        eval_splits = self.metadata.eval_splits
+        eval_splits = self.eval_splits
         trust_remote_code = self.metadata.dataset.get("trust_remote_code", False)
         revision = self.metadata.dataset["revision"]
 
@@ -265,16 +279,18 @@ class AbsTaskRetrieval(AbsTask):
                 trust_remote_code=trust_remote_code,
                 split=split,
                 config=hf_subset,
-            ).load()
+            ).load(
+                num_proc=num_proc,
+            )
 
         if self.metadata.is_multilingual:
-            for lang in self.metadata.eval_langs:
+            for lang in self.hf_subsets:
                 for split in eval_splits:
                     _process_data(split, lang)
         else:
             for split in eval_splits:
                 _process_data(split)
-        self.dataset_transform()
+        self.dataset_transform(num_proc=num_proc)
         self.data_loaded = True
 
     def evaluate(
@@ -283,10 +299,11 @@ 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]:
+        num_proc: int = 1,
+        **kwargs: Any,
+    ) -> Mapping[HFSubset, ScoresDict]:
         """Evaluate the model on the retrieval task.
 
         Args:
@@ -296,16 +313,16 @@ class AbsTaskRetrieval(AbsTask):
             subsets_to_run: Optional list of subsets to evaluate on
             encode_kwargs: Keyword arguments passed to the encoder
             prediction_folder: Folder to save model predictions
+            num_proc: Number of processes to use
             **kwargs: Additional keyword arguments passed to the evaluator
 
-
         Returns:
             Dictionary mapping subsets to their evaluation scores
         """
         if not self.data_loaded:
-            self.load_data()
+            self.load_data(num_proc=num_proc)
         # TODO: convert all tasks directly https://github.com/embeddings-benchmark/mteb/issues/2030
-        self.convert_v1_dataset_format_to_v2()
+        self.convert_v1_dataset_format_to_v2(num_proc=num_proc)
 
         return super().evaluate(
             model,
@@ -313,6 +330,7 @@ class AbsTaskRetrieval(AbsTask):
             subsets_to_run,
             encode_kwargs=encode_kwargs,
             prediction_folder=prediction_folder,
+            num_proc=num_proc,
             **kwargs,
         )
 
@@ -320,10 +338,11 @@ 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,
+        num_proc: int = 1,
         **kwargs,
     ) -> ScoresDict:
         """Evaluate a model on a specific subset of the data.
@@ -335,6 +354,7 @@ class AbsTaskRetrieval(AbsTask):
             hf_split: Split to evaluate on
             hf_subset: Subset to evaluate on
             prediction_folder: Folder with results prediction
+            num_proc: Number of processes to use
             **kwargs: Additional keyword arguments passed to the evaluator
 
         Returns:
@@ -357,6 +377,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):
@@ -372,6 +394,7 @@ class AbsTaskRetrieval(AbsTask):
         results = retriever(
             search_model,
             encode_kwargs=encode_kwargs,
+            num_proc=num_proc,
         )
         end_time = time()
         logger.debug(
@@ -446,9 +469,13 @@ class AbsTaskRetrieval(AbsTask):
         return {}
 
     def _calculate_descriptive_statistics_from_split(
-        self, split: str, hf_subset: str | None = None, compute_overall: bool = False
+        self,
+        split: str,
+        hf_subset: str | None = None,
+        compute_overall: bool = False,
+        num_proc: int = 1,
     ) -> RetrievalDescriptiveStatistics:
-        self.convert_v1_dataset_format_to_v2()
+        self.convert_v1_dataset_format_to_v2(num_proc)
         if hf_subset and hf_subset in self.dataset:
             split_data = self.dataset[hf_subset][split]
             queries = split_data["queries"]
@@ -553,8 +580,8 @@ class AbsTaskRetrieval(AbsTask):
             top_ranked_statistics=top_ranked_statistics,
         )
 
-    def _push_dataset_to_hub(self, repo_name: str) -> None:
-        self.convert_v1_dataset_format_to_v2()
+    def _push_dataset_to_hub(self, repo_name: str, num_proc: int = 1) -> None:
+        self.convert_v1_dataset_format_to_v2(num_proc)
 
         def _push_section(
             data: dict[str, RetrievalSplitData],
@@ -578,11 +605,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(
@@ -593,6 +621,7 @@ class AbsTaskRetrieval(AbsTask):
                     repo_name,
                     hf_subset_name,
                     commit_message=f"Add {hf_subset_name}-{subset_item}",
+                    num_proc=num_proc,
                 )
 
         for subset in self.dataset:
@@ -626,6 +655,7 @@ class AbsTaskRetrieval(AbsTask):
                 repo_name,
                 f"{subset}-qrels" if subset != "default" else "qrels",
                 commit_message=f"Add {subset}-qrels",
+                num_proc=num_proc,
             )
 
             _push_section(
@@ -680,7 +710,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 +718,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: