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

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, Mapping, 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,25 +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 (
-    EncodeKwargs,
-    HFSubset,
-    QueryDatasetType,
-    RelevantDocumentsType,
-    RetrievalOutputType,
-    ScoresDict,
-)
 from mteb.types.statistics import (
-    ImageStatistics,
-    RelevantDocsStatistics,
     SplitDescriptiveStatistics,
-    TextStatistics,
-    TopRankedStatistics,
 )
 
 from ._statistics_calculation import (
@@ -53,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__)
 
 
@@ -137,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"):
@@ -204,6 +215,7 @@ class AbsTaskRetrieval(AbsTask):
                             _combine_queries_with_instructions_datasets(
                                 self.dataset[subset][split]["queries"],
                                 instructions,
+                                num_proc,
                             )
                         )
                     if hasattr(self, "top_ranked"):
@@ -229,9 +241,10 @@ 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()
@@ -244,13 +257,13 @@ 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
 
         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"]
 
@@ -266,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(
@@ -286,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.
@@ -297,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,
@@ -314,6 +330,7 @@ class AbsTaskRetrieval(AbsTask):
             subsets_to_run,
             encode_kwargs=encode_kwargs,
             prediction_folder=prediction_folder,
+            num_proc=num_proc,
             **kwargs,
         )
 
@@ -325,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.
@@ -336,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:
@@ -375,6 +394,7 @@ class AbsTaskRetrieval(AbsTask):
         results = retriever(
             search_model,
             encode_kwargs=encode_kwargs,
+            num_proc=num_proc,
         )
         end_time = time()
         logger.debug(
@@ -449,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"]
@@ -556,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],
@@ -597,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:
@@ -630,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

@@ -1,5 +1,7 @@
+from __future__ import annotations
+
 import logging
-from typing import TypedDict
+from typing import TYPE_CHECKING, TypedDict
 
 from datasets import (
     Dataset,
@@ -11,13 +13,14 @@ from datasets import (
     load_dataset,
 )
 
-from mteb.types import (
-    CorpusDatasetType,
-    InstructionDatasetType,
-    QueryDatasetType,
-    RelevantDocumentsType,
-    TopRankedDocumentsType,
-)
+if TYPE_CHECKING:
+    from mteb.types import (
+        CorpusDatasetType,
+        InstructionDatasetType,
+        QueryDatasetType,
+        RelevantDocumentsType,
+        TopRankedDocumentsType,
+    )
 
 logger = logging.getLogger(__name__)
 
@@ -73,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,
@@ -117,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:
-        logger.info("Loading Corpus...")
-
+    def _load_corpus(self, num_proc: int) -> CorpusDatasetType:
         config = f"{self.config}-corpus" if self.config is not None else "corpus"
-        corpus_ds = self._load_dataset_split(config)
+        logger.info("Loading corpus subset: %s", 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"
@@ -139,13 +151,13 @@ class RetrievalDatasetLoader:
         logger.debug("Doc Example: %s", corpus_ds[0])
         return corpus_ds
 
-    def _load_queries(self) -> QueryDatasetType:
-        logger.info("Loading Queries...")
-
+    def _load_queries(self, num_proc: int) -> QueryDatasetType:
         config = f"{self.config}-queries" if self.config is not None else "queries"
+        logger.info("Loading queries subset: %s", config)
+
         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"
@@ -156,10 +168,10 @@ class RetrievalDatasetLoader:
 
         return queries_ds
 
-    def _load_qrels(self) -> RelevantDocumentsType:
-        logger.info("Loading qrels...")
-
+    def _load_qrels(self, num_proc: int) -> RelevantDocumentsType:
         config = f"{self.config}-qrels" if self.config is not None else "default"
+
+        logger.info("Loading qrels subset: %s", config)
         if config == "default" and config not in self.dataset_configs:
             if "qrels" in self.dataset_configs:
                 config = "qrels"
@@ -168,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(
@@ -191,13 +203,12 @@ class RetrievalDatasetLoader:
         logger.info("Loaded %d %s qrels.", len(qrels_dict), self.split.upper())
         return qrels_dict
 
-    def _load_top_ranked(self) -> TopRankedDocumentsType:
-        logger.info("Loading Top Ranked")
-
+    def _load_top_ranked(self, num_proc: int) -> TopRankedDocumentsType:
         config = (
             f"{self.config}-top_ranked" if self.config is not None else "top_ranked"
         )
-        top_ranked_ds = self._load_dataset_split(config)
+        logger.info("Loading top ranked subset: %s", config)
+        top_ranked_ds = self._load_dataset_split(config, num_proc)
         top_ranked_ds = top_ranked_ds.cast(
             Features(
                 {
@@ -215,13 +226,12 @@ class RetrievalDatasetLoader:
         logger.info(f"Top ranked loaded: {len(top_ranked_ds)}")
         return top_ranked_dict
 
-    def _load_instructions(self) -> InstructionDatasetType:
-        logger.info("Loading Instructions")
-
+    def _load_instructions(self, num_proc: int) -> InstructionDatasetType:
         config = (
             f"{self.config}-instruction" if self.config is not None else "instruction"
         )
-        instructions_ds = self._load_dataset_split(config)
+        logger.info("Loading instruction subset: %s", config)
+        instructions_ds = self._load_dataset_split(config, num_proc)
         instructions_ds = instructions_ds.cast(
             Features(
                 {
@@ -236,6 +246,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 = {
@@ -248,4 +259,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

@@ -1,19 +1,14 @@
+from __future__ import annotations
+
 import logging
-from pathlib import Path
-from typing import Any, TypedDict, cast
+from typing import TYPE_CHECKING, Any, TypedDict, cast
 
-from datasets import Dataset
 from scipy.stats import pearsonr, spearmanr
 
 from mteb._evaluators import AnySTSEvaluator
-from mteb._evaluators.any_sts_evaluator import STSEvaluatorScores
-from mteb.models import EncoderProtocol, MTEBModels
-from mteb.types import EncodeKwargs, PromptType
+from mteb.models import EncoderProtocol
 from mteb.types.statistics import (
-    ImageStatistics,
-    ScoreStatistics,
     SplitDescriptiveStatistics,
-    TextStatistics,
 )
 
 from ._statistics_calculation import (
@@ -23,6 +18,20 @@ from ._statistics_calculation import (
 )
 from .abstask import AbsTask
 
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from datasets import Dataset
+
+    from mteb._evaluators.any_sts_evaluator import STSEvaluatorScores
+    from mteb.models import MTEBModels
+    from mteb.types import EncodeKwargs, PromptType
+    from mteb.types.statistics import (
+        ImageStatistics,
+        ScoreStatistics,
+        TextStatistics,
+    )
+
 logger = logging.getLogger(__name__)
 
 
@@ -109,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):
@@ -127,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(
@@ -182,7 +196,7 @@ class AbsTaskSTS(AbsTask):
         self, split: str, hf_subset: str | None = None, compute_overall: bool = False
     ) -> AnySTSDescriptiveStatistics:
         first_column, second_column = self.column_names
-        self.dataset = cast(dict[str, dict[str, Dataset]], self.dataset)
+        self.dataset = cast("dict[str, dict[str, Dataset]]", self.dataset)
 
         if hf_subset:
             sentence1 = self.dataset[hf_subset][split][first_column]
@@ -236,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/task_metadata.py

@@ -1,11 +1,12 @@
+from __future__ import annotations
+
 import json
 import logging
 from collections.abc import Sequence
 from pathlib import Path
-from typing import Any, Literal, cast
+from typing import TYPE_CHECKING, Any, Literal, cast
 
 from huggingface_hub import (
-    CardData,
     DatasetCard,
     DatasetCardData,
     constants,
@@ -17,13 +18,11 @@ from pydantic import (
     ConfigDict,
     field_validator,
 )
-from typing_extensions import Required, TypedDict
+from typing_extensions import Required, TypedDict  # noqa: TC002
 
 import mteb
 from mteb.languages import check_language_code
 from mteb.types import (
-    HFSubset,
-    ISOLanguageScript,
     Languages,
     Licenses,
     Modalities,
@@ -31,7 +30,17 @@ from mteb.types import (
     StrDate,
     StrURL,
 )
-from mteb.types.statistics import DescriptiveStatistics
+
+if TYPE_CHECKING:
+    from huggingface_hub import (
+        CardData,
+    )
+
+    from mteb.types import (
+        HFSubset,
+        ISOLanguageScript,
+    )
+    from mteb.types.statistics import DescriptiveStatistics
 
 logger = logging.getLogger(__name__)
 
@@ -368,7 +377,7 @@ class TaskMetadata(BaseModel):
         """Return a dictionary mapping huggingface subsets to languages."""
         if isinstance(self.eval_langs, dict):
             return self.eval_langs
-        return {"default": cast(list[str], self.eval_langs)}
+        return {"default": cast("list[str]", self.eval_langs)}
 
     @property
     def intext_citation(self, include_cite: bool = True) -> str:
@@ -697,7 +706,7 @@ class TaskMetadata(BaseModel):
             for val in self.eval_langs.values():
                 languages.extend(val)
         else:
-            languages = cast(list[str], 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 = []

mteb/abstasks/text/bitext_mining.py

@@ -1,7 +1,8 @@
+from __future__ import annotations
+
 import logging
 from collections import defaultdict
-from pathlib import Path
-from typing import Any, ClassVar, TypedDict, cast
+from typing import TYPE_CHECKING, Any, ClassVar, TypedDict, cast
 
 from datasets import Dataset, DatasetDict
 from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score
@@ -9,9 +10,15 @@ from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_sc
 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 EncodeKwargs, HFSubset, ScoresDict
-from mteb.types.statistics import SplitDescriptiveStatistics, TextStatistics
+from mteb.models import EncoderProtocol
+from mteb.types.statistics import SplitDescriptiveStatistics
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from mteb.models import MTEBModels
+    from mteb.types import EncodeKwargs, HFSubset, ScoresDict
+    from mteb.types.statistics import TextStatistics
 
 logger = logging.getLogger(__name__)
 
@@ -75,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"""
@@ -82,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
 
@@ -90,7 +98,7 @@ class AbsTaskBitextMining(AbsTask):
         if subsets_to_run is not None:
             hf_subsets = [s for s in hf_subsets if s in subsets_to_run]
 
-        encoder_model = cast(EncoderProtocol, model)
+        encoder_model = cast("EncoderProtocol", model)
 
         if self.dataset is None:
             raise ValueError("Dataset is not loaded.")
@@ -105,6 +113,7 @@ class AbsTaskBitextMining(AbsTask):
                 hf_subset="parallel",
                 encode_kwargs=encode_kwargs,
                 prediction_folder=prediction_folder,
+                num_proc=num_proc,
                 **kwargs,
             )
         else:
@@ -124,10 +133,11 @@ class AbsTaskBitextMining(AbsTask):
                     hf_subset=hf_subset,
                     encode_kwargs=encode_kwargs,
                     prediction_folder=prediction_folder,
+                    num_proc=num_proc,
                     **kwargs,
                 )
 
-        return cast(dict[HFSubset, ScoresDict], scores)
+        return cast("dict[HFSubset, ScoresDict]", scores)
 
     def _get_pairs(self, parallel: bool) -> list[tuple[str, str]]:
         pairs = self._DEFAULT_PAIR
@@ -145,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)
@@ -164,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(
@@ -257,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.")
 
@@ -280,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:
@@ -292,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)