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

mteb/tasks/clustering/rus/ru_sci_bench_grnti_clustering_p2p.py

@@ -32,7 +32,7 @@ class RuSciBenchGRNTIClusteringP2P(AbsTaskClustering):
         prompt="Identify the category of scientific papers based on the titles and abstracts",
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         self.dataset = self.dataset.rename_columns(
             {"label": "labels", "text": "sentences"}
         )

mteb/tasks/clustering/rus/ru_sci_bench_oecd_clustering_p2p.py

@@ -32,7 +32,7 @@ class RuSciBenchOECDClusteringP2P(AbsTaskClustering):
         prompt="Identify the category of scientific papers based on the titles and abstracts",
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         self.dataset = self.dataset.rename_columns(
             {"label": "labels", "text": "sentences"}
         )

mteb/tasks/clustering/zho/cmteb_clustering.py

@@ -51,7 +51,7 @@ class CLSClusteringFastS2S(AbsTaskClustering):
         adapted_from=["CLSClusteringS2S"],
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         ds = {}
         for split in self.metadata.eval_splits:
             labels = list(itertools.chain.from_iterable(self.dataset[split]["labels"]))
@@ -110,7 +110,7 @@ class CLSClusteringFastP2P(AbsTaskClustering):
         adapted_from=["CLSClusteringP2P"],
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         ds = {}
         for split in self.metadata.eval_splits:
             labels = list(itertools.chain.from_iterable(self.dataset[split]["labels"]))
@@ -239,7 +239,7 @@ class ThuNewsClusteringFastS2S(AbsTaskClustering):
         adapted_from=["ThuNewsClusteringS2S"],
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         ds = {}
         for split in self.metadata.eval_splits:
             labels = list(itertools.chain.from_iterable(self.dataset[split]["labels"]))
@@ -298,7 +298,7 @@ class ThuNewsClusteringFastP2P(AbsTaskClustering):
         adapted_from=["ThuNewsClusteringP2P"],
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         ds = {}
         for split in self.metadata.eval_splits:
             labels = list(itertools.chain.from_iterable(self.dataset[split]["labels"]))

mteb/tasks/image_text_pair_classification/eng/image_co_de.py

@@ -53,7 +53,7 @@ class ImageCoDe(AbsTaskImageTextPairClassification):
 """,
     )
 
-    def load_data(self) -> None:
+    def load_data(self, num_proc: int = 1, **kwargs) -> None:
         if self.data_loaded:
             return
 

mteb/tasks/image_text_pair_classification/eng/sugar_crepe.py

@@ -45,7 +45,7 @@ class SugarCrepe(AbsTaskImageTextPairClassification):
 """,
     )
 
-    def load_data(self) -> None:
+    def load_data(self, num_proc: int = 1, **kwargs) -> None:
         """Load dataset from HuggingFace hub"""
         if self.data_loaded:
             return

mteb/tasks/instruction_reranking/multilingual/m_follow_ir.py

@@ -175,7 +175,7 @@ class mFollowIRCrossLingual(AbsTaskRetrieval):  # noqa: N801
 """,
     )
 
-    def load_data(self) -> None:
+    def load_data(self, num_proc: int = 1, **kwargs) -> None:
         if self.data_loaded:
             return
 
@@ -243,7 +243,7 @@ class mFollowIR(AbsTaskRetrieval):  # noqa: N801
 """,
     )
 
-    def load_data(self) -> None:
+    def load_data(self, num_proc: int = 1, **kwargs) -> None:
         if self.data_loaded:
             return
 

mteb/tasks/multichoice/eng/cv_bench.py

@@ -123,7 +123,7 @@ class CVBenchCount(AbsTaskRetrieval):
 """,
     )
 
-    def load_data(self) -> None:
+    def load_data(self, num_proc: int = 1, **kwargs) -> None:
         self.corpus, self.queries, self.relevant_docs, self.top_ranked = _load_data(
             path=self.metadata.dataset["path"],
             splits=self.metadata.eval_splits,
@@ -165,7 +165,7 @@ class CVBenchRelation(AbsTaskRetrieval):
 """,
     )
 
-    def load_data(self) -> None:
+    def load_data(self, num_proc: int = 1, **kwargs) -> None:
         self.corpus, self.queries, self.relevant_docs, self.top_ranked = _load_data(
             path=self.metadata.dataset["path"],
             splits=self.metadata.eval_splits,
@@ -207,7 +207,7 @@ class CVBenchDepth(AbsTaskRetrieval):
 """,
     )
 
-    def load_data(self) -> None:
+    def load_data(self, num_proc: int = 1, **kwargs) -> None:
         self.corpus, self.queries, self.relevant_docs, self.top_ranked = _load_data(
             path=self.metadata.dataset["path"],
             splits=self.metadata.eval_splits,
@@ -249,7 +249,7 @@ class CVBenchDistance(AbsTaskRetrieval):
 """,
     )
 
-    def load_data(self) -> None:
+    def load_data(self, num_proc: int = 1, **kwargs) -> None:
         self.corpus, self.queries, self.relevant_docs, self.top_ranked = _load_data(
             path=self.metadata.dataset["path"],
             splits=self.metadata.eval_splits,

mteb/tasks/multilabel_classification/ita/emit_classification.py

@@ -39,7 +39,7 @@ class EmitClassification(AbsTaskMultilabelClassification):
 """,
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         self.dataset = self.dataset.rename_columns({"emotion_labels": "label"})
         unused_cols = [
             col

mteb/tasks/multilabel_classification/mlt/maltese_news_classification.py

@@ -40,7 +40,7 @@ Borg, Claudia},
 """,
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         self.dataset = self.dataset.rename_columns({"labels": "label"})
         remove_cols = [
             col

mteb/tasks/multilabel_classification/rus/ru_toixic_multilabelclassification_okmlcup.py

@@ -29,5 +29,5 @@ class RuToxicOKMLCUPMultilabelClassification(AbsTaskMultilabelClassification):
         bibtex_citation="""""",
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         self.dataset = self.dataset.rename_column("labels", "label")

mteb/tasks/multilabel_classification/swe/swedish_patent_cpc_group_classification.py

@@ -42,7 +42,7 @@ class SwedishPatentCPCGroupClassification(AbsTaskMultilabelClassification):
 """,
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         self.dataset = self.stratified_subsampling(
             self.dataset, seed=self.seed, splits=["train"], n_samples=8192
         )

mteb/tasks/multilabel_classification/swe/swedish_patent_cpc_subclass_classification.py

@@ -42,7 +42,7 @@ class SwedishPatentCPCSubclassClassification(AbsTaskMultilabelClassification):
 """,
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         self.dataset = self.stratified_subsampling(
             self.dataset, seed=self.seed, splits=["train"], n_samples=8192
         )

mteb/tasks/pair_classification/ara/ar_entail.py

@@ -39,7 +39,7 @@ class ArEntail(AbsTaskPairClassification):
 """,
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         _dataset = {}
         for split in self.metadata.eval_splits:
             _dataset[split] = [

mteb/tasks/pair_classification/dan/talemaader_pc.py

@@ -38,5 +38,5 @@ class TalemaaderPC(AbsTaskPairClassification):
 """,
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         self.dataset = self.dataset.rename_column("label", "labels")

mteb/tasks/pair_classification/deu/false_friends_de_en_pc.py

@@ -34,7 +34,7 @@ class FalseFriendsDeEnPC(AbsTaskPairClassification):
 """,
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         _dataset = {}
         for split in self.metadata.eval_splits:
             hf_dataset = self.dataset[split]

mteb/tasks/pair_classification/eng/pub_chem_ai_sentence_paraphrase_pc.py

@@ -45,7 +45,7 @@ class PubChemAISentenceParaphrasePC(AbsTaskPairClassification):
 """,
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         _dataset = {}
         for split in self.metadata.eval_splits:
             hf_dataset = self.dataset[split]

mteb/tasks/pair_classification/eng/pub_chem_smilespc.py

@@ -106,7 +106,7 @@ class PubChemSMILESPC(AbsTaskPairClassification):
         self.dataset_transform()
         self.data_loaded = True
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         self.dataset = self.stratified_subsampling(
             self.dataset,
             seed=self.seed,

mteb/tasks/pair_classification/eng/pub_chem_synonym_pc.py

@@ -45,7 +45,7 @@ class PubChemSynonymPC(AbsTaskPairClassification):
 """,
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         _dataset = {}
 
         for split in self.metadata.eval_splits:

mteb/tasks/pair_classification/eng/pub_chem_wiki_paragraphs_pc.py

@@ -45,7 +45,7 @@ class PubChemWikiParagraphsPC(AbsTaskPairClassification):
 """,
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         _dataset = {}
         for split in self.metadata.eval_splits:
             hf_dataset = self.dataset[split]

mteb/tasks/pair_classification/eng/sprint_duplicate_questions_pc.py

@@ -52,6 +52,6 @@ Tsujii, Jun{'}ichi},
 """,
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         self.dataset = self.dataset.rename_column("sent1", "sentence1")
         self.dataset = self.dataset.rename_column("sent2", "sentence2")

mteb/tasks/pair_classification/eng/twitter_sem_eval2015_pc.py

@@ -47,6 +47,6 @@ Jurgens, David},
         prompt="Retrieve tweets that are semantically similar to the given tweet",
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         self.dataset = self.dataset.rename_column("sent1", "sentence1")
         self.dataset = self.dataset.rename_column("sent2", "sentence2")

mteb/tasks/pair_classification/eng/twitter_url_corpus_pc.py

@@ -47,6 +47,6 @@ Riedel, Sebastian},
         prompt="Retrieve tweets that are semantically similar to the given tweet",
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         self.dataset = self.dataset.rename_column("sent1", "sentence1")
         self.dataset = self.dataset.rename_column("sent2", "sentence2")

mteb/tasks/pair_classification/fas/fa_mteb_pair_classification.py

@@ -39,7 +39,7 @@ class CExaPPC(AbsTaskPairClassification):
 """,
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         _dataset = {}
         self.dataset = self.dataset.map(
             lambda example: {"label": 1 if example["label"] == "paraphrase" else 0}
@@ -80,7 +80,7 @@ class SynPerChatbotRAGFAQPC(AbsTaskPairClassification):
         bibtex_citation=""" """,
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         _dataset = {}
         for split in self.metadata.eval_splits:
             _dataset[split] = [
@@ -118,7 +118,7 @@ class FarsiParaphraseDetection(AbsTaskPairClassification):
         bibtex_citation=""" """,
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         _dataset = {}
         for split in self.metadata.eval_splits:
             _dataset[split] = [
@@ -156,7 +156,7 @@ class SynPerTextKeywordsPC(AbsTaskPairClassification):
         bibtex_citation=""" """,
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         _dataset = {}
         for split in self.metadata.eval_splits:
             _dataset[split] = [
@@ -194,7 +194,7 @@ class SynPerQAPC(AbsTaskPairClassification):
         bibtex_citation=""" """,
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         _dataset = {}
         for split in self.metadata.eval_splits:
             _dataset[split] = [

mteb/tasks/pair_classification/fas/fars_tail.py

@@ -38,7 +38,7 @@ class FarsTail(AbsTaskPairClassification):
 """,  # after removing neutral
     )
 
-    def load_data(self) -> None:
+    def load_data(self, num_proc: int = 1, **kwargs) -> None:
         if self.data_loaded:
             return
         path = self.metadata.dataset["path"]
@@ -52,7 +52,7 @@ class FarsTail(AbsTaskPairClassification):
         self.dataset_transform()
         self.data_loaded = True
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         _dataset = {}
         self.dataset = self.dataset.filter(lambda x: x["label"] != "n")
         self.dataset = self.dataset.map(

mteb/tasks/pair_classification/hye/armenian_paraphrase_pc.py

@@ -36,6 +36,6 @@ class ArmenianParaphrasePC(AbsTaskPairClassification):
 """,
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         self.dataset = self.dataset.rename_column("sent1", "sentence1")
         self.dataset = self.dataset.rename_column("sent2", "sentence2")

mteb/tasks/pair_classification/ita/dis_co_tex_pair_classification.py

@@ -37,7 +37,7 @@ class DisCoTexPairClassification(AbsTaskPairClassification):
 """,
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         self.dataset = self.dataset.remove_columns(["id", "source"])
         self.dataset = self.dataset.map(
             lambda x: {

mteb/tasks/pair_classification/kor/klue_nli.py

@@ -37,7 +37,7 @@ class KlueNLI(AbsTaskPairClassification):
 """,  # 3000 - neutral samples
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         _dataset = {}
         for split in self.metadata.eval_splits:
             # keep labels 0=entailment and 2=contradiction, and map them as 1 and 0 for binary classification

mteb/tasks/pair_classification/multilingual/rte3.py

@@ -52,7 +52,7 @@ Dolan, Bill},
         # sum of 4 languages after neutral filtering
     )
 
-    def load_data(self) -> None:
+    def load_data(self, num_proc: int = 1, **kwargs) -> None:
         """Load dataset from HuggingFace hub"""
         if self.data_loaded:
             return
@@ -62,7 +62,7 @@ Dolan, Bill},
         self.dataset_transform()
         self.data_loaded = True
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         _dataset = {}
         for lang in self.hf_subsets:
             _dataset[lang] = {}

mteb/tasks/pair_classification/multilingual/xnli.py

@@ -60,7 +60,7 @@ in Natural Language Processing},
 """,
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         _dataset = {}
         for lang in self.hf_subsets:
             _dataset[lang] = {}

mteb/tasks/pair_classification/pol/polish_pc.py

@@ -57,7 +57,7 @@ Piperidis, Stelios},
 """,
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         self.dataset = self.dataset.rename_column("sent1", "sentence1")
         self.dataset = self.dataset.rename_column("sent2", "sentence2")
 
@@ -105,7 +105,7 @@ class PpcPC(AbsTaskPairClassification):
 """,
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         self.dataset = self.dataset.rename_column("sent1", "sentence1")
         self.dataset = self.dataset.rename_column("sent2", "sentence2")
 
@@ -151,7 +151,7 @@ Kan, Min-Yen},
 """,
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         self.dataset = self.dataset.rename_column("sent1", "sentence1")
         self.dataset = self.dataset.rename_column("sent2", "sentence2")
 
@@ -203,6 +203,6 @@ Piperidis, Stelios},
 """,
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         self.dataset = self.dataset.rename_column("sent1", "sentence1")
         self.dataset = self.dataset.rename_column("sent2", "sentence2")

mteb/tasks/pair_classification/por/assin2_rte.py

@@ -36,7 +36,7 @@ class Assin2RTE(AbsTaskPairClassification):
 """,
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         _dataset = {}
         self.dataset = self.stratified_subsampling(
             self.dataset,

mteb/tasks/pair_classification/por/sick_br_pc.py

@@ -47,7 +47,7 @@ and de Paiva, Valeria},
 """,
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         _dataset = {}
 
         # Do not process the subsets we won't use

mteb/tasks/pair_classification/rus/terra.py

@@ -50,7 +50,7 @@ class TERRa(AbsTaskPairClassification):
         **_terra_metadata,
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         self.dataset = self.dataset.rename_column("sent1", "sentence1")
         self.dataset = self.dataset.rename_column("sent2", "sentence2")
 
@@ -71,6 +71,6 @@ class TERRaV2(AbsTaskPairClassification):
         **_terra_metadata,
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         self.dataset = self.dataset.rename_column("sent1", "sentence1")
         self.dataset = self.dataset.rename_column("sent2", "sentence2")

mteb/tasks/pair_classification/vie/sprint_duplicate_questions_pcvn.py

@@ -37,6 +37,6 @@ class SprintDuplicateQuestionsPCVN(AbsTaskPairClassification):
         adapted_from=["SprintDuplicateQuestions"],
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         self.dataset = self.dataset.rename_column("sent1", "sentence1")
         self.dataset = self.dataset.rename_column("sent2", "sentence2")

mteb/tasks/pair_classification/vie/twitter_sem_eval2015_pcvn.py

@@ -37,6 +37,6 @@ class TwitterSemEval2015PCVN(AbsTaskPairClassification):
         adapted_from=["TwitterSemEval2015"],
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         self.dataset = self.dataset.rename_column("sent1", "sentence1")
         self.dataset = self.dataset.rename_column("sent2", "sentence2")

mteb/tasks/pair_classification/vie/twitter_url_corpus_pcvn.py

@@ -37,6 +37,6 @@ class TwitterURLCorpusPC(AbsTaskPairClassification):
         adapted_from=["TwitterURLCorpus"],
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         self.dataset = self.dataset.rename_column("sent1", "sentence1")
         self.dataset = self.dataset.rename_column("sent2", "sentence2")

mteb/tasks/pair_classification/zho/cmteb_pair_classification.py

@@ -37,7 +37,7 @@ class Ocnli(AbsTaskPairClassification):
         prompt="Retrieve semantically similar text.",
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         self.dataset = self.dataset.rename_column("sent1", "sentence1")
         self.dataset = self.dataset.rename_column("sent2", "sentence2")
 
@@ -112,6 +112,6 @@ Lan, Zhenzhong},
         prompt="Retrieve semantically similar text.",
     )
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         self.dataset = self.dataset.rename_column("sent1", "sentence1")
         self.dataset = self.dataset.rename_column("sent2", "sentence2")

mteb/tasks/retrieval/ara/sadeem_question_retrieval.py

@@ -37,7 +37,7 @@ class SadeemQuestionRetrieval(AbsTaskRetrieval):
 """,
     )
 
-    def load_data(self) -> None:
+    def load_data(self, num_proc: int = 1, **kwargs) -> None:
         if self.data_loaded:
             return
 

mteb/tasks/retrieval/code/code_edit_search_retrieval.py

@@ -53,7 +53,7 @@ class CodeEditSearchRetrieval(AbsTaskRetrieval):
 """,
     )
 
-    def load_data(self) -> None:
+    def load_data(self, num_proc: int = 1, **kwargs) -> None:
         if self.data_loaded:
             return
 

mteb/tasks/retrieval/code/code_rag.py

@@ -51,7 +51,7 @@ class CodeRAGProgrammingSolutionsRetrieval(AbsTaskRetrieval):
         **common_args,
     )
 
-    def load_data(self) -> None:
+    def load_data(self, num_proc: int = 1, **kwargs) -> None:
         """Load dataset from HuggingFace hub"""
         if self.data_loaded:
             return
@@ -108,7 +108,7 @@ class CodeRAGOnlineTutorialsRetrieval(AbsTaskRetrieval):
         **common_args,
     )
 
-    def load_data(self) -> None:
+    def load_data(self, num_proc: int = 1, **kwargs) -> None:
         """Load dataset from HuggingFace hub"""
         if self.data_loaded:
             return
@@ -168,7 +168,7 @@ class CodeRAGLibraryDocumentationSolutionsRetrieval(AbsTaskRetrieval):
         **common_args,
     )
 
-    def load_data(self) -> None:
+    def load_data(self, num_proc: int = 1, **kwargs) -> None:
         """Load dataset from HuggingFace hub"""
         if self.data_loaded:
             return
@@ -225,7 +225,7 @@ class CodeRAGStackoverflowPostsRetrieval(AbsTaskRetrieval):
         **common_args,
     )
 
-    def load_data(self) -> None:
+    def load_data(self, num_proc: int = 1, **kwargs) -> None:
         """Load dataset from HuggingFace hub"""
         if self.data_loaded:
             return

mteb/tasks/retrieval/code/code_search_net_cc_retrieval.py

@@ -99,7 +99,7 @@ class CodeSearchNetCCRetrieval(AbsTaskRetrieval):
 """,
     )
 
-    def load_data(self) -> None:
+    def load_data(self, num_proc: int = 1, **kwargs) -> None:
         if self.data_loaded:
             return
 

mteb/tasks/retrieval/code/coir_code_search_net_retrieval.py

@@ -97,7 +97,7 @@ class COIRCodeSearchNetRetrieval(AbsTaskRetrieval):
 """,
     )
 
-    def load_data(self) -> None:
+    def load_data(self, num_proc: int = 1, **kwargs) -> None:
         if self.data_loaded:
             return
 

mteb/tasks/retrieval/code/ds1000_retrieval.py

@@ -34,7 +34,7 @@ class DS1000Retrieval(AbsTaskRetrieval):
 """,
     )
 
-    def load_data(self) -> None:
+    def load_data(self, num_proc: int = 1, **kwargs) -> None:
         if self.data_loaded:
             return
 

mteb/tasks/retrieval/code/fresh_stack_retrieval.py

@@ -37,7 +37,7 @@ class FreshStackRetrieval(AbsTaskRetrieval):
 """,
     )
 
-    def load_data(self) -> None:
+    def load_data(self, num_proc: int = 1, **kwargs) -> None:
         if self.data_loaded:
             return
 

mteb/tasks/retrieval/code/human_eval_retrieval.py

@@ -34,7 +34,7 @@ class HumanEvalRetrieval(AbsTaskRetrieval):
 }""",
     )
 
-    def load_data(self) -> None:
+    def load_data(self, num_proc: int = 1, **kwargs) -> None:
         if self.data_loaded:
             return
 

mteb/tasks/retrieval/code/mbpp_retrieval.py

@@ -34,7 +34,7 @@ class MBPPRetrieval(AbsTaskRetrieval):
 """,
     )
 
-    def load_data(self) -> None:
+    def load_data(self, num_proc: int = 1, **kwargs) -> None:
         if self.data_loaded:
             return
 

mteb/tasks/retrieval/code/wiki_sql_retrieval.py

@@ -36,7 +36,7 @@ class WikiSQLRetrieval(AbsTaskRetrieval):
 """,
     )
 
-    def load_data(self) -> None:
+    def load_data(self, num_proc: int = 1, **kwargs) -> None:
         if self.data_loaded:
             return