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

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

@@ -1,3 +1,5 @@
+from typing import Any
+
 import datasets
 
 from mteb.abstasks.pair_classification import AbsTaskPairClassification
@@ -74,8 +76,7 @@ class PubChemSMILESPC(AbsTaskPairClassification):
 """,
     )
 
-    def load_data(self):
-        """Load dataset from HuggingFace hub"""
+    def load_data(self, num_proc: int = 1, **kwargs: Any) -> None:
         if self.data_loaded:
             return
 
@@ -106,7 +107,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/reranking/multilingual/wikipedia_reranking_multilingual.py

@@ -44,7 +44,7 @@ class WikipediaRerankingMultilingual(AbsTaskRetrieval):
         dialect=[],
         sample_creation="LM-generated and verified",
         bibtex_citation=r"""
-@online{wikidump,
+@online{wikidump2024,
   author = {Wikimedia Foundation},
   title = {Wikimedia Downloads},
   url = {https://dumps.wikimedia.org},

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

@@ -48,14 +48,14 @@ class CodeRAGProgrammingSolutionsRetrieval(AbsTaskRetrieval):
             "path": "code-rag-bench/programming-solutions",
             "revision": "1064f7bba54d5400d4836f5831fe4c2332a566a6",
         },
-        **common_args,  # type: ignore
+        **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
-        self.dataset = datasets.load_dataset(**self.metadata.dataset)  # type: ignore
+        self.dataset = datasets.load_dataset(**self.metadata.dataset)
         self.dataset_transform()
         self.data_loaded = True
 
@@ -71,7 +71,7 @@ class CodeRAGProgrammingSolutionsRetrieval(AbsTaskRetrieval):
         self.queries = {}
 
         split = self.metadata.eval_splits[0]
-        ds: datasets.Dataset = self.dataset[split]  # type: ignore
+        ds: datasets.Dataset = self.dataset[split]
         ds = ds.shuffle(seed=42)
 
         self.queries[split] = {}
@@ -105,14 +105,14 @@ class CodeRAGOnlineTutorialsRetrieval(AbsTaskRetrieval):
             "path": "code-rag-bench/online-tutorials",
             "revision": "095bb77130082e4690d6c3a031997b03487bf6e2",
         },
-        **common_args,  # type: ignore
+        **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
-        self.dataset = datasets.load_dataset(**self.metadata.dataset)  # type: ignore
+        self.dataset = datasets.load_dataset(**self.metadata.dataset)
         self.dataset_transform()
         self.data_loaded = True
 
@@ -128,7 +128,7 @@ class CodeRAGOnlineTutorialsRetrieval(AbsTaskRetrieval):
         self.queries = {}
 
         split = self.metadata.eval_splits[0]
-        ds: datasets.Dataset = self.dataset[split]  # type: ignore
+        ds: datasets.Dataset = self.dataset[split]
         ds = ds.shuffle(seed=42)
 
         self.queries[split] = {}
@@ -165,14 +165,14 @@ class CodeRAGLibraryDocumentationSolutionsRetrieval(AbsTaskRetrieval):
             "path": "code-rag-bench/library-documentation",
             "revision": "b530d3b5a25087d2074e731b76232db85b9e9107",
         },
-        **common_args,  # type: ignore
+        **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
-        self.dataset = datasets.load_dataset(**self.metadata.dataset)  # type: ignore
+        self.dataset = datasets.load_dataset(**self.metadata.dataset)
         self.dataset_transform()
         self.data_loaded = True
 
@@ -188,7 +188,7 @@ class CodeRAGLibraryDocumentationSolutionsRetrieval(AbsTaskRetrieval):
         self.queries = {}
 
         split = self.metadata.eval_splits[0]
-        ds: datasets.Dataset = self.dataset[split]  # type: ignore
+        ds: datasets.Dataset = self.dataset[split]
         ds = ds.shuffle(seed=42)
 
         self.queries[split] = {}
@@ -222,14 +222,14 @@ class CodeRAGStackoverflowPostsRetrieval(AbsTaskRetrieval):
             "path": "code-rag-bench/stackoverflow-posts",
             "revision": "04e05d86cb0ac467b29a5d87f4c56eac99dfc0a4",
         },
-        **common_args,  # type: ignore
+        **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
-        self.dataset = datasets.load_dataset(**self.metadata.dataset)  # type: ignore
+        self.dataset = datasets.load_dataset(**self.metadata.dataset)
         self.dataset_transform()
         self.data_loaded = True
 
@@ -245,7 +245,7 @@ class CodeRAGStackoverflowPostsRetrieval(AbsTaskRetrieval):
         self.queries = {}
 
         split = self.metadata.eval_splits[0]
-        ds: datasets.Dataset = self.dataset[split]  # type: ignore
+        ds: datasets.Dataset = self.dataset[split]
         ds = ds.shuffle(seed=42)
 
         self.queries[split] = {}

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