fix: all dataset transform

mteb/abstasks/abstask.py

@@ -116,14 +116,15 @@ def check_if_dataset_is_superseded(self) -> None:
             logger.warning(msg)
             warnings.warn(msg)
 
-    def dataset_transform(self, num_proc: int = 1):
+    def dataset_transform(self, num_proc: int = 1, **kwargs: Any) -> None:
         """A transform operations applied to the dataset after loading.
 
         This method is useful when the dataset from Huggingface is not in an `mteb` compatible format.
         Override this method if your dataset requires additional transformation.
 
         Args:
             num_proc: Number of processes to use for the transformation.
+            kwargs: Additional keyword arguments passed to the load_dataset function. Keep for forward compatibility.
         """
         pass
 

mteb/tasks/bitext_mining/multilingual/bible_nlp_bitext_mining.py

@@ -914,7 +914,7 @@ def load_data(self, **kwargs: Any) -> None:
         self.dataset_transform()
         self.data_loaded = True
 
-    def dataset_transform(self) -> None:
+    def dataset_transform(self, num_proc: int = 1, **kwargs) -> None:
         # Convert to standard format
         for lang in self.hf_subsets:
             l1, l2 = (l.split("_")[0] for l in lang.split("-"))

mteb/tasks/classification/ben/bengali_document_classification.py

@@ -43,7 +43,7 @@ class BengaliDocumentClassification(AbsTaskClassification):
         superseded_by="BengaliDocumentClassification.v2",
     )
 
-    def dataset_transform(self) -> None:
+    def dataset_transform(self, num_proc: int = 1, **kwargs) -> None:
         self.dataset = self.dataset.rename_columns(
             {"article": "text", "category": "label"}
         )
@@ -92,7 +92,7 @@ class BengaliDocumentClassificationV2(AbsTaskClassification):
 """,
     )
 
-    def dataset_transform(self) -> None:
+    def dataset_transform(self, num_proc: int = 1, **kwargs) -> None:
         self.dataset = self.stratified_subsampling(
             self.dataset, seed=self.seed, splits=["test"]
         )

mteb/tasks/classification/ces/czech_product_review_sentiment_classification.py

@@ -46,7 +46,7 @@ class CzechProductReviewSentimentClassification(AbsTaskClassification):
     )
     samples_per_label = 16
 
-    def dataset_transform(self) -> None:
+    def dataset_transform(self, num_proc: int = 1, **kwargs) -> None:
         self.dataset = self.dataset.rename_columns(
             {"comment": "text", "rating_str": "label"}
         )
@@ -99,7 +99,7 @@ class CzechProductReviewSentimentClassificationV2(AbsTaskClassification):
     )
     samples_per_label = 16
 
-    def dataset_transform(self) -> None:
+    def dataset_transform(self, num_proc: int = 1, **kwargs) -> None:
         self.dataset = self.stratified_subsampling(
             self.dataset, seed=self.seed, splits=["test"]
         )

mteb/tasks/classification/ces/czech_so_me_sentiment_classification.py

@@ -46,7 +46,7 @@ class CzechSoMeSentimentClassification(AbsTaskClassification):
     )
     samples_per_label = 16
 
-    def dataset_transform(self) -> None:
+    def dataset_transform(self, num_proc: int = 1, **kwargs) -> None:
         self.dataset = self.dataset.rename_columns(
             {"comment": "text", "sentiment_int": "label"}
         )

mteb/tasks/classification/multilingual/hin_dialect_classification.py

@@ -60,7 +60,7 @@ class HinDialectClassification(AbsTaskClassification):
 """,
     )
 
-    def dataset_transform(self) -> None:
+    def dataset_transform(self, num_proc: int = 1, **kwargs) -> None:
         self.dataset = self.dataset.rename_columns(
             {"folksong": "text", "language": "label"}
         )

mteb/tasks/classification/multilingual/indic_lang_classification.py

@@ -137,6 +137,6 @@ def load_data(self, **kwargs: Any) -> None:
         self.dataset_transform()
         self.data_loaded = True
 
-    def dataset_transform(self) -> None:
+    def dataset_transform(self, num_proc: int = 1, **kwargs) -> None:
         self.dataset = self.dataset.remove_columns(["language", "script"])
         self.dataset = self.dataset.rename_columns({"native sentence": "text"})

mteb/tasks/classification/multilingual/indic_sentiment_classification.py

@@ -52,7 +52,7 @@ class IndicSentimentClassification(AbsTaskClassification):
 """,
     )
 
-    def dataset_transform(self) -> None:
+    def dataset_transform(self, num_proc: int = 1, **kwargs) -> None:
         label_map = {"Negative": 0, "Positive": 1}
         # Convert to standard format
         for lang in self.hf_subsets:

mteb/tasks/classification/multilingual/language_classification.py

@@ -66,7 +66,7 @@ class LanguageClassification(AbsTaskClassification):
 """,
     )
 
-    def dataset_transform(self) -> None:
+    def dataset_transform(self, num_proc: int = 1, **kwargs) -> None:
         self.dataset = self.dataset.rename_columns({"labels": "label"})
         self.dataset = self.stratified_subsampling(
             self.dataset, seed=self.seed, splits=["test"]

mteb/tasks/classification/multilingual/south_african_lang_classification.py

@@ -49,7 +49,7 @@ class SouthAfricanLangClassification(AbsTaskClassification):
 """,
     )
 
-    def dataset_transform(self) -> None:
+    def dataset_transform(self, num_proc: int = 1, **kwargs) -> None:
         self.dataset = self.dataset.rename_columns(
             {" text": "text", "lang_id": "label"}
         )

mteb/tasks/classification/slk/slovak_movie_review_sentiment_classification.py

@@ -35,7 +35,7 @@ class SlovakMovieReviewSentimentClassification(AbsTaskClassification):
         superseded_by="SlovakMovieReviewSentimentClassification.v2",
     )
 
-    def dataset_transform(self) -> None:
+    def dataset_transform(self, num_proc: int = 1, **kwargs) -> None:
         self.dataset = self.dataset.rename_columns({"comment": "text"})
 
         self.dataset = self.stratified_subsampling(
@@ -76,7 +76,7 @@ class SlovakMovieReviewSentimentClassificationV2(AbsTaskClassification):
         adapted_from=["SlovakMovieReviewSentimentClassification"],
     )
 
-    def dataset_transform(self) -> None:
+    def dataset_transform(self, num_proc: int = 1, **kwargs) -> None:
         self.dataset = self.stratified_subsampling(
             self.dataset, seed=self.seed, splits=["test"]
         )

mteb/tasks/classification/swa/swahili_news_classification.py

@@ -37,7 +37,7 @@ class SwahiliNewsClassification(AbsTaskClassification):
         superseded_by="SwahiliNewsClassification.v2",
     )
 
-    def dataset_transform(self) -> None:
+    def dataset_transform(self, num_proc: int = 1, **kwargs) -> None:
         self.dataset = self.dataset.rename_columns(
             {"content": "text", "category": "label"}
         )
@@ -81,7 +81,7 @@ class SwahiliNewsClassificationV2(AbsTaskClassification):
         adapted_from=["SwahiliNewsClassification"],
     )
 
-    def dataset_transform(self) -> None:
+    def dataset_transform(self, num_proc: int = 1, **kwargs) -> None:
         self.dataset = self.stratified_subsampling(
             self.dataset, seed=self.seed, splits=["train"]
         )

mteb/tasks/clustering/deu/ten_k_gnad_clustering_p2p.py

@@ -63,7 +63,7 @@ class TenKGnadClusteringP2PFast(AbsTaskClustering):
         adapted_from=["TenKGnadClusteringP2P"],
     )
 
-    def dataset_transform(self) -> None:
+    def dataset_transform(self, num_proc: int = 1, **kwargs) -> None:
         ds = _convert_to_fast(
             self.dataset, self.input_column_name, self.label_column_name, self.seed
         )

mteb/tasks/clustering/deu/ten_k_gnad_clustering_s2s.py

@@ -63,7 +63,7 @@ class TenKGnadClusteringS2SFast(AbsTaskClustering):
         adapted_from=["TenKGnadClusteringS2S"],
     )
 
-    def dataset_transform(self) -> None:
+    def dataset_transform(self, num_proc: int = 1, **kwargs) -> None:
         ds = _convert_to_fast(
             self.dataset, self.input_column_name, self.label_column_name, self.seed
         )

mteb/tasks/clustering/nob/vg_hierarchical_clustering.py

@@ -45,7 +45,7 @@ class VGHierarchicalClusteringP2P(AbsTaskClustering):
         prompt="Identify the categories (e.g. sports) of given articles in Norwegian",
     )
 
-    def dataset_transform(self) -> None:
+    def dataset_transform(self, num_proc: int = 1, **kwargs) -> None:
         self.dataset = self.dataset.rename_columns(
             {"article": "sentences", "classes": "labels"}
         )
@@ -92,7 +92,7 @@ class VGHierarchicalClusteringS2S(AbsTaskClustering):
         prompt="Identify the categories (e.g. sports) of given articles in Norwegian",
     )
 
-    def dataset_transform(self) -> None:
+    def dataset_transform(self, num_proc: int = 1, **kwargs) -> None:
         self.dataset = self.dataset.rename_columns(
             {"ingress": "sentences", "classes": "labels"}
         )

mteb/tasks/multilabel_classification/nld/covid_disinformation_nl_multi_label_classification.py

@@ -66,7 +66,7 @@ class CovidDisinformationNLMultiLabelClassification(AbsTaskMultilabelClassificat
         },
     )
 
-    def dataset_transform(self) -> None:
+    def dataset_transform(self, num_proc: int = 1, **kwargs) -> None:
         labels = [
             "q2_label",
             "q3_label",

mteb/tasks/pair_classification/multilingual/pub_chem_wiki_pair_classification.py

@@ -60,7 +60,7 @@ class PubChemWikiPairClassification(AbsTaskPairClassification):
 """,
     )
 
-    def dataset_transform(self, num_proc: int = 1) -> None:
+    def dataset_transform(self, num_proc: int = 1, **kwargs) -> None:
         _dataset = {}
         for lang in self.hf_subsets:
             _dataset[lang] = {}

mteb/tasks/retrieval/code/code_rag.py

@@ -59,7 +59,7 @@ def load_data(self, num_proc: int = 1, **kwargs) -> None:
         self.dataset_transform()
         self.data_loaded = True
 
-    def dataset_transform(self) -> None:
+    def dataset_transform(self, num_proc: int = 1, **kwargs) -> None:
         """And transform to a retrieval dataset, which have the following attributes
 
         self.corpus = Dict[doc_id, Dict[str, str]] #id => dict with document data like title and text
@@ -116,7 +116,7 @@ def load_data(self, num_proc: int = 1, **kwargs) -> None:
         self.dataset_transform()
         self.data_loaded = True
 
-    def dataset_transform(self) -> None:
+    def dataset_transform(self, num_proc: int = 1, **kwargs) -> None:
         """And transform to a retrieval dataset, which have the following attributes
 
         self.corpus = Dict[doc_id, Dict[str, str]] #id => dict with document data like title and text
@@ -176,7 +176,7 @@ def load_data(self, num_proc: int = 1, **kwargs) -> None:
         self.dataset_transform()
         self.data_loaded = True
 
-    def dataset_transform(self) -> None:
+    def dataset_transform(self, num_proc: int = 1, **kwargs) -> None:
         """And transform to a retrieval dataset, which have the following attributes
 
         self.corpus = Dict[doc_id, Dict[str, str]] #id => dict with document data like title and text
@@ -233,7 +233,7 @@ def load_data(self, num_proc: int = 1, **kwargs) -> None:
         self.dataset_transform()
         self.data_loaded = True
 
-    def dataset_transform(self) -> None:
+    def dataset_transform(self, num_proc: int = 1, **kwargs) -> None:
         """And transform to a retrieval dataset, which have the following attributes
 
         self.corpus = Dict[doc_id, Dict[str, str]] #id => dict with document data like title and text

mteb/tasks/retrieval/dan/dan_fever_retrieval.py

@@ -55,7 +55,7 @@ def load_data(self, num_proc: int = 1, **kwargs) -> None:
         self.dataset_transform()
         self.data_loaded = True
 
-    def dataset_transform(self) -> None:
+    def dataset_transform(self, num_proc: int = 1, **kwargs) -> None:
         """And transform to a retrieval dataset, which have the following attributes
 
         self.corpus = dict[doc_id, dict[str, str]] #id => dict with document data like title and text

mteb/tasks/retrieval/dan/tv2_nordretrieval.py

@@ -68,7 +68,7 @@ def load_data(self, num_proc: int = 1, **kwargs) -> None:
         self.dataset_transform()
         self.data_loaded = True
 
-    def dataset_transform(self) -> None:
+    def dataset_transform(self, num_proc: int = 1, **kwargs) -> None:
         """And transform to a retrieval dataset, which have the following attributes
 
         self.corpus = dict[doc_id, dict[str, str]] #id => dict with document data like title and text

mteb/tasks/retrieval/dan/twitter_hjerne_retrieval.py

@@ -44,7 +44,7 @@ def load_data(self, num_proc: int = 1, **kwargs) -> None:
         self.dataset_transform()
         self.data_loaded = True
 
-    def dataset_transform(self) -> None:
+    def dataset_transform(self, num_proc: int = 1, **kwargs) -> None:
         """And transform to a retrieval dataset, which have the following attributes
 
         self.corpus = dict[doc_id, dict[str, str]] #id => dict with document data like title and text

mteb/tasks/retrieval/nob/norquad.py

@@ -58,7 +58,7 @@ def load_data(self, num_proc: int = 1, **kwargs) -> None:
         self.dataset_transform()
         self.data_loaded = True
 
-    def dataset_transform(self) -> None:
+    def dataset_transform(self, num_proc: int = 1, **kwargs) -> None:
         """And transform to a retrieval dataset, which have the following attributes
 
         self.corpus = dict[doc_id, dict[str, str]] #id => dict with document data like title and text

mteb/tasks/retrieval/nob/snl_retrieval.py

@@ -45,7 +45,7 @@ def load_data(self, num_proc: int = 1, **kwargs) -> None:
         self.dataset_transform()
         self.data_loaded = True
 
-    def dataset_transform(self) -> None:
+    def dataset_transform(self, num_proc: int = 1, **kwargs) -> None:
         """And transform to a retrieval dataset, which have the following attributes
 
         self.corpus = dict[doc_id, dict[str, str]] #id => dict with document data like title and text

mteb/tasks/sts/multilingual/sem_rel24_sts.py

@@ -66,6 +66,6 @@ class SemRel24STS(AbsTaskSTS):
     min_score = 0
     max_score = 1
 
-    def dataset_transform(self) -> None:
+    def dataset_transform(self, num_proc: int = 1, **kwargs) -> None:
         for lang, subset in self.dataset.items():
             self.dataset[lang] = subset.rename_column("label", "score")

mteb/tasks/sts/multilingual/sts_benchmark_multilingual_sts.py

@@ -56,6 +56,6 @@ class STSBenchmarkMultilingualSTS(AbsTaskSTS):
     min_score = 0
     max_score = 5
 
-    def dataset_transform(self) -> None:
+    def dataset_transform(self, num_proc: int = 1, **kwargs) -> None:
         for lang, subset in self.dataset.items():
             self.dataset[lang] = subset.rename_column("similarity_score", "score")

mteb/tasks/sts/por/assin2_sts.py

@@ -39,7 +39,7 @@ class Assin2STS(AbsTaskSTS):
     min_score = 1
     max_score = 5
 
-    def dataset_transform(self) -> None:
+    def dataset_transform(self, num_proc: int = 1, **kwargs) -> None:
         self.dataset = self.dataset.rename_columns(
             {
                 "premise": "sentence1",