mteb 2.7.4__py3-none-any.whl → 2.7.6__py3-none-any.whl

mteb/abstasks/classification.py

@@ -136,6 +136,7 @@ class AbsTaskClassification(AbsTask):
         *,
         encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
+        num_proc: int = 1,
         **kwargs: Any,
     ) -> dict[HFSubset, ScoresDict]:
         """Evaluate a model on the classification task.
@@ -149,7 +150,7 @@ class AbsTaskClassification(AbsTask):
             )
 
         if not self.data_loaded:
-            self.load_data()
+            self.load_data(num_proc=num_proc)
 
         if self.dataset is None:
             raise RuntimeError("Dataset not loaded.")
@@ -182,6 +183,7 @@ class AbsTaskClassification(AbsTask):
                 hf_subset=hf_subset,
                 encode_kwargs=encode_kwargs,
                 prediction_folder=prediction_folder,
+                num_proc=num_proc,
                 **kwargs,
             )
             self._add_main_score(scores[hf_subset])
@@ -197,6 +199,7 @@ class AbsTaskClassification(AbsTask):
         hf_split: str,
         hf_subset: str,
         prediction_folder: Path | None = None,
+        num_proc: int = 1,
         **kwargs: Any,
     ) -> FullClassificationMetrics:
         if not isinstance(model, EncoderProtocol):
@@ -230,7 +233,10 @@ class AbsTaskClassification(AbsTask):
                 evaluator_model=self.evaluator_model,
             )
             y_pred, test_cache = evaluator(
-                model, encode_kwargs=encode_kwargs, test_cache=test_cache
+                model,
+                encode_kwargs=encode_kwargs,
+                test_cache=test_cache,
+                num_proc=num_proc,
             )
             if prediction_folder:
                 all_predictions.append(y_pred.tolist())
@@ -372,11 +378,12 @@ class AbsTaskClassification(AbsTask):
             label_statistics=label_statistics,
         )
 
-    def _push_dataset_to_hub(self, repo_name: str) -> None:
+    def _push_dataset_to_hub(self, repo_name: str, num_proc: int = 1) -> None:
         self._upload_dataset_to_hub(
             repo_name,
             [
                 self.input_column_name,
                 self.label_column_name,
             ],
+            num_proc=num_proc,
         )

mteb/abstasks/clustering.py

@@ -169,6 +169,7 @@ class AbsTaskClustering(AbsTask):
         hf_split: str,
         hf_subset: str,
         prediction_folder: Path | None = None,
+        num_proc: int = 1,
         **kwargs: Any,
     ) -> ScoresDict:
         if not isinstance(model, EncoderProtocol):
@@ -213,6 +214,7 @@ class AbsTaskClustering(AbsTask):
                 downsampled_dataset,
                 self.metadata,
                 input_column=self.input_column_name,
+                num_proc=num_proc,
                 **encode_kwargs,
             ),
             task_metadata=self.metadata,
@@ -296,9 +298,11 @@ class AbsTaskClustering(AbsTask):
             labels_statistics=label_statistics,
         )
 
-    def _push_dataset_to_hub(self, repo_name: str) -> None:
+    def _push_dataset_to_hub(self, repo_name: str, num_proc: int = 1) -> None:
         self._upload_dataset_to_hub(
-            repo_name, [self.input_column_name, self.label_column_name]
+            repo_name,
+            [self.input_column_name, self.label_column_name],
+            num_proc=num_proc,
         )
 
 

mteb/abstasks/clustering_legacy.py

@@ -95,6 +95,7 @@ class AbsTaskClusteringLegacy(AbsTask):
         hf_split: str,
         hf_subset: str,
         prediction_folder: Path | None = None,
+        num_proc: int = 1,
         **kwargs: Any,
     ) -> ScoresDict:
         if not isinstance(model, EncoderProtocol):
@@ -159,7 +160,11 @@ class AbsTaskClusteringLegacy(AbsTask):
             hf_subset=hf_subset,
             **kwargs,
         )
-        evaluate_clusters = evaluator(model, encode_kwargs=encode_kwargs)
+        evaluate_clusters = evaluator(
+            model,
+            encode_kwargs=encode_kwargs,
+            num_proc=num_proc,
+        )
         if prediction_folder:
             self._save_task_predictions(
                 evaluate_clusters,
@@ -238,11 +243,12 @@ class AbsTaskClusteringLegacy(AbsTask):
             label_statistics=label_statistics,
         )
 
-    def _push_dataset_to_hub(self, repo_name: str) -> None:
+    def _push_dataset_to_hub(self, repo_name: str, num_proc: int = 1) -> None:
         self._upload_dataset_to_hub(
             repo_name,
             [
                 self.input_column_name,
                 self.label_column_name,
             ],
+            num_proc=num_proc,
         )

mteb/abstasks/image/image_text_pair_classification.py

@@ -134,6 +134,7 @@ class AbsTaskImageTextPairClassification(AbsTask):
         hf_split: str,
         hf_subset: str,
         prediction_folder: Path | None = None,
+        num_proc: int = 1,
         **kwargs: Any,
     ) -> ImageTextPairClassificationMetrics:
         if not isinstance(model, EncoderProtocol):
@@ -167,7 +168,9 @@ class AbsTaskImageTextPairClassification(AbsTask):
             hf_subset=hf_subset,
             **kwargs,
         )
-        scores: list[torch.Tensor] = evaluator(model, encode_kwargs=encode_kwargs)  # type: ignore[assignment]
+        scores: list[torch.Tensor] = evaluator(
+            model, encode_kwargs=encode_kwargs, num_proc=num_proc
+        )  # type: ignore[assignment]
         if prediction_folder:
             self._save_task_predictions(
                 [score.tolist() for score in scores],
@@ -215,7 +218,7 @@ class AbsTaskImageTextPairClassification(AbsTask):
             accuracy=torch.Tensor(all_correct_scores).float().mean().item(),
         )
 
-    def _push_dataset_to_hub(self, repo_name: str) -> None:
+    def _push_dataset_to_hub(self, repo_name: str, num_proc: int = 1) -> None:
         text_columns = (
             [self.texts_column_names]
             if isinstance(self.texts_column_names, str)
@@ -230,4 +233,5 @@ class AbsTaskImageTextPairClassification(AbsTask):
         self._upload_dataset_to_hub(
             repo_name,
             [*text_columns, *image_columns],
+            num_proc=num_proc,
         )

mteb/abstasks/multilabel_classification.py

@@ -93,6 +93,7 @@ class AbsTaskMultilabelClassification(AbsTaskClassification):
         hf_split: str,
         hf_subset: str,
         prediction_folder: Path | None = None,
+        num_proc: int = 1,
         **kwargs: Any,
     ) -> FullMultilabelClassificationMetrics:
         if not isinstance(model, EncoderProtocol):
@@ -125,6 +126,7 @@ class AbsTaskMultilabelClassification(AbsTaskClassification):
             unique_train_dataset,
             self.metadata,
             input_column=self.input_column_name,
+            num_proc=num_proc,
             **encode_kwargs,
         )
 

mteb/abstasks/pair_classification.py

@@ -96,6 +96,7 @@ class AbsTaskPairClassification(AbsTask):
         hf_subset: str,
         encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
+        num_proc: int = 1,
         **kwargs,
     ) -> dict[str, float]:
         if not isinstance(model, EncoderProtocol):
@@ -115,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(
@@ -248,7 +253,7 @@ 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
@@ -272,6 +277,7 @@ class AbsTaskPairClassification(AbsTask):
                 self.input2_column_name,
                 self.label_column_name,
             ],
+            num_proc=num_proc,
         )
 
     def _compute_metrics_values(

mteb/abstasks/retrieval.py

@@ -148,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"):
@@ -215,6 +215,7 @@ class AbsTaskRetrieval(AbsTask):
                             _combine_queries_with_instructions_datasets(
                                 self.dataset[subset][split]["queries"],
                                 instructions,
+                                num_proc,
                             )
                         )
                     if hasattr(self, "top_ranked"):
@@ -240,6 +241,7 @@ class AbsTaskRetrieval(AbsTask):
                         _combine_queries_with_instructions_datasets(
                             self.dataset[subset][split]["queries"],
                             instructions,
+                            num_proc,
                         )
                     )
                 if hasattr(self, "top_ranked") and self.top_ranked:
@@ -255,7 +257,7 @@ class AbsTaskRetrieval(AbsTask):
         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
@@ -277,7 +279,9 @@ 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:
@@ -286,7 +290,7 @@ class AbsTaskRetrieval(AbsTask):
         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(
@@ -297,6 +301,7 @@ class AbsTaskRetrieval(AbsTask):
         *,
         encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
+        num_proc: int = 1,
         **kwargs: Any,
     ) -> Mapping[HFSubset, ScoresDict]:
         """Evaluate the model on the retrieval task.
@@ -308,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,
@@ -325,6 +330,7 @@ class AbsTaskRetrieval(AbsTask):
             subsets_to_run,
             encode_kwargs=encode_kwargs,
             prediction_folder=prediction_folder,
+            num_proc=num_proc,
             **kwargs,
         )
 
@@ -336,6 +342,7 @@ class AbsTaskRetrieval(AbsTask):
         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.
@@ -347,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:
@@ -386,6 +394,7 @@ class AbsTaskRetrieval(AbsTask):
         results = retriever(
             search_model,
             encode_kwargs=encode_kwargs,
+            num_proc=num_proc,
         )
         end_time = time()
         logger.debug(
@@ -460,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"]
@@ -567,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],
@@ -608,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:
@@ -641,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(

mteb/abstasks/retrieval_dataset_loaders.py

@@ -76,28 +76,36 @@ class RetrievalDatasetLoader:
         self.config = config if config != "default" else None
         self.dataset_configs = get_dataset_config_names(self.hf_repo, self.revision)
 
-    def load(self) -> RetrievalSplitData:
+    def load(
+        self,
+        num_proc: int = 1,
+    ) -> RetrievalSplitData:
         """Loads the dataset split for the specified configuration.
 
+        Args:
+            num_proc: The number of processes to use.
+
         Returns:
             A dictionary containing the corpus, queries, relevant documents, instructions (if applicable), and top-ranked documents (if applicable).
         """
         top_ranked = None
 
-        qrels = self._load_qrels()
-        corpus = self._load_corpus()
-        queries = self._load_queries()
+        qrels = self._load_qrels(num_proc)
+        corpus = self._load_corpus(num_proc)
+        queries = self._load_queries(num_proc)
 
         queries = queries.filter(
             lambda x: x["id"] in qrels.keys(), desc="Filtering queries by qrels"
         )
 
         if any(c.endswith("top_ranked") for c in self.dataset_configs):
-            top_ranked = self._load_top_ranked()
+            top_ranked = self._load_top_ranked(num_proc)
 
         if any(c.endswith("instruction") for c in self.dataset_configs):
-            instructions = self._load_instructions()
-            queries = _combine_queries_with_instructions_datasets(queries, instructions)
+            instructions = self._load_instructions(num_proc)
+            queries = _combine_queries_with_instructions_datasets(
+                queries, instructions, num_proc
+            )
 
         return RetrievalSplitData(
             corpus=corpus,
@@ -120,20 +128,21 @@ class RetrievalDatasetLoader:
             f"Split {self.split} not found in {splits}. Please specify a valid split."
         )
 
-    def _load_dataset_split(self, config: str) -> Dataset:
+    def _load_dataset_split(self, config: str, num_proc: int) -> Dataset:
         return load_dataset(
             self.hf_repo,
             config,
             split=self._get_split(config),
             trust_remote_code=self.trust_remote_code,
             revision=self.revision,
+            num_proc=num_proc,
         )
 
-    def _load_corpus(self) -> CorpusDatasetType:
+    def _load_corpus(self, num_proc: int) -> CorpusDatasetType:
         logger.info("Loading Corpus...")
 
         config = f"{self.config}-corpus" if self.config is not None else "corpus"
-        corpus_ds = self._load_dataset_split(config)
+        corpus_ds = self._load_dataset_split(config, num_proc)
         if "_id" in corpus_ds.column_names:
             corpus_ds = corpus_ds.cast_column("_id", Value("string")).rename_column(
                 "_id", "id"
@@ -142,13 +151,13 @@ class RetrievalDatasetLoader:
         logger.debug("Doc Example: %s", corpus_ds[0])
         return corpus_ds
 
-    def _load_queries(self) -> QueryDatasetType:
+    def _load_queries(self, num_proc: int) -> QueryDatasetType:
         logger.info("Loading Queries...")
 
         config = f"{self.config}-queries" if self.config is not None else "queries"
         if "query" in self.dataset_configs:
             config = "query"
-        queries_ds = self._load_dataset_split(config)
+        queries_ds = self._load_dataset_split(config, num_proc)
         if "_id" in queries_ds.column_names:
             queries_ds = queries_ds.cast_column("_id", Value("string")).rename_column(
                 "_id", "id"
@@ -159,7 +168,7 @@ class RetrievalDatasetLoader:
 
         return queries_ds
 
-    def _load_qrels(self) -> RelevantDocumentsType:
+    def _load_qrels(self, num_proc: int) -> RelevantDocumentsType:
         logger.info("Loading qrels...")
 
         config = f"{self.config}-qrels" if self.config is not None else "default"
@@ -171,7 +180,7 @@ class RetrievalDatasetLoader:
                     "No qrels or default config found. Please specify a valid config or ensure the dataset has qrels."
                 )
 
-        qrels_ds = self._load_dataset_split(config)
+        qrels_ds = self._load_dataset_split(config, num_proc)
         qrels_ds = qrels_ds.select_columns(["query-id", "corpus-id", "score"])
 
         qrels_ds = qrels_ds.cast(
@@ -194,13 +203,13 @@ class RetrievalDatasetLoader:
         logger.info("Loaded %d %s qrels.", len(qrels_dict), self.split.upper())
         return qrels_dict
 
-    def _load_top_ranked(self) -> TopRankedDocumentsType:
+    def _load_top_ranked(self, num_proc: int) -> TopRankedDocumentsType:
         logger.info("Loading Top Ranked")
 
         config = (
             f"{self.config}-top_ranked" if self.config is not None else "top_ranked"
         )
-        top_ranked_ds = self._load_dataset_split(config)
+        top_ranked_ds = self._load_dataset_split(config, num_proc)
         top_ranked_ds = top_ranked_ds.cast(
             Features(
                 {
@@ -218,13 +227,13 @@ class RetrievalDatasetLoader:
         logger.info(f"Top ranked loaded: {len(top_ranked_ds)}")
         return top_ranked_dict
 
-    def _load_instructions(self) -> InstructionDatasetType:
+    def _load_instructions(self, num_proc: int) -> InstructionDatasetType:
         logger.info("Loading Instructions")
 
         config = (
             f"{self.config}-instruction" if self.config is not None else "instruction"
         )
-        instructions_ds = self._load_dataset_split(config)
+        instructions_ds = self._load_dataset_split(config, num_proc)
         instructions_ds = instructions_ds.cast(
             Features(
                 {
@@ -239,6 +248,7 @@ class RetrievalDatasetLoader:
 def _combine_queries_with_instructions_datasets(
     queries_dataset: QueryDatasetType,
     instruction_dataset: InstructionDatasetType | dict[str, str],
+    num_proc: int,
 ) -> Dataset:
     if isinstance(instruction_dataset, Dataset):
         instruction_to_query_idx = {
@@ -251,4 +261,4 @@ def _combine_queries_with_instructions_datasets(
         row["instruction"] = instruction_to_query_idx[row["id"]]
         return row
 
-    return queries_dataset.map(_add_instruction_to_query)
+    return queries_dataset.map(_add_instruction_to_query, num_proc=num_proc)

mteb/abstasks/sts.py

@@ -118,6 +118,7 @@ class AbsTaskSTS(AbsTask):
         hf_split: str,
         hf_subset: str,
         prediction_folder: Path | None = None,
+        num_proc: int = 1,
         **kwargs: Any,
     ) -> STSMetrics:
         if not isinstance(model, EncoderProtocol):
@@ -136,7 +137,11 @@ class AbsTaskSTS(AbsTask):
             input2_prompt_type=self.input2_prompt_type,
             **kwargs,
         )
-        scores = evaluator(model, encode_kwargs=encode_kwargs)
+        scores = evaluator(
+            model,
+            encode_kwargs=encode_kwargs,
+            num_proc=num_proc,
+        )
 
         if prediction_folder:
             self._save_task_predictions(
@@ -245,9 +250,11 @@ class AbsTaskSTS(AbsTask):
             label_statistics=labels_statistics,
         )
 
-    def _push_dataset_to_hub(self, repo_name: str) -> None:
+    def _push_dataset_to_hub(self, repo_name: str, num_proc: int = 1) -> None:
         self._upload_dataset_to_hub(
-            repo_name, [self.column_names[0], self.column_names[1], "score"]
+            repo_name,
+            [self.column_names[0], self.column_names[1], "score"],
+            num_proc=num_proc,
         )
 
     def _normalize(self, x: float) -> float:

mteb/abstasks/text/bitext_mining.py

@@ -82,6 +82,7 @@ class AbsTaskBitextMining(AbsTask):
         *,
         encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
+        num_proc: int = 1,
         **kwargs: Any,
     ) -> dict[HFSubset, ScoresDict]:
         """Added load for "parallel" datasets"""
@@ -89,7 +90,7 @@ class AbsTaskBitextMining(AbsTask):
             raise TypeError("Expected model to be an instance of EncoderProtocol")
 
         if not self.data_loaded:
-            self.load_data()
+            self.load_data(num_proc=num_proc)
 
         hf_subsets = self.hf_subsets
 
@@ -112,6 +113,7 @@ class AbsTaskBitextMining(AbsTask):
                 hf_subset="parallel",
                 encode_kwargs=encode_kwargs,
                 prediction_folder=prediction_folder,
+                num_proc=num_proc,
                 **kwargs,
             )
         else:
@@ -131,6 +133,7 @@ class AbsTaskBitextMining(AbsTask):
                     hf_subset=hf_subset,
                     encode_kwargs=encode_kwargs,
                     prediction_folder=prediction_folder,
+                    num_proc=num_proc,
                     **kwargs,
                 )
 
@@ -152,6 +155,7 @@ class AbsTaskBitextMining(AbsTask):
         encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
         parallel: bool = False,
+        num_proc: int = 1,
         **kwargs,
     ) -> BitextMiningMetrics | dict[str, BitextMiningMetrics]:
         pairs = self._get_pairs(parallel)
@@ -171,7 +175,7 @@ class AbsTaskBitextMining(AbsTask):
             else data_split["gold"]
         )
 
-        neighbours = evaluator(model, encode_kwargs=encode_kwargs)
+        neighbours = evaluator(model, encode_kwargs=encode_kwargs, num_proc=num_proc)
 
         if prediction_folder:
             self._save_task_predictions(
@@ -264,7 +268,7 @@ class AbsTaskBitextMining(AbsTask):
             sentence2_statistics=text2_statistics,
         )
 
-    def _push_dataset_to_hub(self, repo_name: str) -> None:
+    def _push_dataset_to_hub(self, repo_name: str, num_proc: int = 1) -> None:
         if self.dataset is None:
             raise ValueError("Dataset is not loaded.")
 
@@ -287,7 +291,7 @@ class AbsTaskBitextMining(AbsTask):
             dataset_dict = DatasetDict(
                 {split: Dataset.from_dict(dataset[split]) for split in dataset}
             )
-            dataset_dict.push_to_hub(repo_name)
+            dataset_dict.push_to_hub(repo_name, num_proc=num_proc)
         else:
             sentences = {}
             for split in self.dataset:
@@ -299,4 +303,4 @@ class AbsTaskBitextMining(AbsTask):
                     }
                 )
             sentences = DatasetDict(sentences)
-            sentences.push_to_hub(repo_name)
+            sentences.push_to_hub(repo_name, num_proc=num_proc)

mteb/abstasks/text/reranking.py

@@ -34,7 +34,7 @@ class AbsTaskReranking(AbsTaskRetrieval):
         For dataformat and other information, see [AbsTaskRetrieval][mteb.abstasks.retrieval.AbsTaskRetrieval].
     """
 
-    def load_data(self) -> None:
+    def load_data(self, num_proc: int = 1, **kwargs) -> None:
         """Load the dataset."""
         if self.data_loaded:
             return
@@ -43,7 +43,7 @@ class AbsTaskReranking(AbsTaskRetrieval):
             self.transform_old_dataset_format()
         else:
             # use AbsTaskRetrieval default to load the data
-            return super().load_data()
+            return super().load_data(num_proc=num_proc)
 
     def _process_example(self, example: dict, split: str, query_idx: int) -> dict:
         """Process a single example from the dataset.

mteb/abstasks/text/summarization.py

@@ -94,6 +94,7 @@ class AbsTaskSummarization(AbsTask):
         hf_subset: str,
         encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
+        num_proc: int = 1,
         **kwargs,
     ) -> SummarizationMetrics:
         if not isinstance(model, EncoderProtocol):
@@ -115,7 +116,7 @@ class AbsTaskSummarization(AbsTask):
             hf_subset=hf_subset,
             **kwargs,
         )
-        scores = evaluator(model, encode_kwargs=encode_kwargs)
+        scores = evaluator(model, encode_kwargs=encode_kwargs, num_proc=num_proc)
         if prediction_folder:
             self._save_task_predictions(
                 scores,