[v2] Select required columns to run (#3190)

* select required columns to run

* fix tasks

Author: Samoed

mteb/_evaluators/clustering_evaluator.py

@@ -6,10 +6,9 @@
 from datasets import Dataset
 from scipy.optimize import linear_sum_assignment
 from sklearn import cluster, metrics
-from torch.utils.data import DataLoader
 
 from mteb.abstasks.task_metadata import TaskMetadata
-from mteb.create_dataloaders import create_image_dataloader
+from mteb.create_dataloaders import create_dataloader
 from mteb.models import Encoder
 
 from .evaluator import Evaluator
@@ -38,29 +37,19 @@ def __init__(
         self.hf_split = hf_split
         self.hf_subset = hf_subset
 
-    def create_dataloader(self, batch_size: int) -> DataLoader:
-        if self.task_metadata.modalities == ["image"]:
-            return create_image_dataloader(
-                self.dataset,
-                image_column_name=self.input_column_name,
-                batch_size=batch_size,
-            )
-        elif self.task_metadata.modalities == ["text"]:
-            return DataLoader(self.dataset)
-        else:
-            raise ValueError(
-                f"Unsupported modality {self.task_metadata.modalities}. "
-                "Currently only 'image' modality is supported."
-            )
-
     def __call__(
         self,
         model: Encoder,
         *,
         encode_kwargs: dict[str, Any],
         v_measure_only: bool = False,
     ):
-        data_loader = self.create_dataloader(batch_size=encode_kwargs["batch_size"])
+        data_loader = create_dataloader(
+            self.dataset,
+            self.task_metadata,
+            input_column=self.input_column_name,
+            batch_size=encode_kwargs["batch_size"],
+        )
 
         embeddings = model.encode(
             data_loader,

mteb/abstasks/AbsTaskAnyClassification.py

@@ -109,6 +109,9 @@ def evaluate(
                 ds = self.dataset
             else:
                 ds = self.dataset[hf_subset]
+
+            if isinstance(ds, (Dataset, DatasetDict)):
+                ds = ds.select_columns([self.label_column_name, self.input_column_name])
             scores[hf_subset] = self._evaluate_subset(
                 model,
                 ds,

mteb/abstasks/AbsTaskAnyClustering.py

@@ -76,8 +76,8 @@ def _evaluate_subset(
         ):
             v_measures = []
             for cluster_set in tqdm.tqdm(dataset, desc="Clustering"):
-                clustering_dataset = Dataset.from_dict(cluster_set).rename_column(
-                    original_column_name="sentences", new_column_name="text"
+                clustering_dataset = Dataset.from_dict(cluster_set).select_columns(
+                    [self.input_column_name, self.label_column_name]
                 )
                 evaluator = self.evaluator(
                     clustering_dataset,
@@ -103,6 +103,9 @@ def _evaluate_subset(
             self._add_main_score(scores)
             return scores
 
+        dataset = dataset.select_columns(
+            [self.input_column_name, self.label_column_name]
+        )
         evaluator = self.evaluator(
             dataset,
             input_column_name=self.input_column_name,

mteb/abstasks/AbsTaskAnySTS.py

@@ -79,6 +79,8 @@ def _evaluate_subset(
         **kwargs: Any,
     ) -> ScoresDict:
         normalized_scores = list(map(self.normalize, data_split["score"]))
+        data_split = data_split.select_columns(list(self.column_names))
+
         evaluator = AnySTSEvaluator(
             data_split,
             self.column_names,

mteb/abstasks/AbsTaskAnyZeroShotClassification.py

@@ -60,12 +60,6 @@ class AbsTaskAnyZeroShotClassification(AbsTask):
     input_column_name: str = "image"
     label_column_name: str = "label"
 
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-
-    def _add_main_score(self, scores) -> None:
-        scores["main_score"] = scores[self.metadata.main_score]
-
     def _calculate_descriptive_statistics_from_split(
         self, split: str, hf_subset: str | None = None, compute_overall: bool = False
     ) -> ZeroShotClassificationDescriptiveStatistics:
@@ -114,6 +108,9 @@ def _evaluate_subset(
         **kwargs,
     ) -> ScoresDict:
         candidate_labels = self.get_candidate_labels()
+        dataset = dataset.select_columns(
+            [self.input_column_name, self.label_column_name]
+        )
         evaluator = ZeroShotClassificationEvaluator(
             dataset,
             self.input_column_name,

mteb/abstasks/AbsTaskClusteringFast.py

@@ -10,12 +10,12 @@
 from datasets import Dataset, DatasetDict
 from sklearn.cluster import MiniBatchKMeans
 from sklearn.metrics.cluster import v_measure_score
-from torch.utils.data import DataLoader
 
 from mteb.models import Encoder
 from mteb.types import HFSubset
 from mteb.types.statistics import DescriptiveStatistics, LabelStatistics, TextStatistics
 
+from ..create_dataloaders import create_dataloader
 from ._statistics_calculation import (
     calculate_label_statistics,
     calculate_text_statistics,
@@ -126,6 +126,8 @@ class AbsTaskClusteringFast(AbsTask):
     k_mean_batch_size: int = 512
     max_depth = None
     abstask_prompt = "Identify categories in user passages."
+    input_column_name: str = "sentences"
+    label_column_name: str = "labels"
 
     def _evaluate_subset(
         self,
@@ -164,19 +166,24 @@ def _evaluate_subset(
             )
             downsampled_dataset = dataset.select(example_indices)  # type: ignore
 
-        downsampled_dataset = downsampled_dataset.rename_column(
-            original_column_name="sentences", new_column_name="text"
+        downsampled_dataset = downsampled_dataset.select_columns(
+            [self.input_column_name, self.label_column_name]
         )
         embeddings = model.encode(
-            DataLoader(downsampled_dataset),
+            create_dataloader(
+                downsampled_dataset,
+                self.metadata,
+                input_column=self.input_column_name,
+                batch_size=encode_kwargs["batch_size"],
+            ),
             task_metadata=self.metadata,
             hf_subset=hf_subset,
             hf_split=hf_split,
             **encode_kwargs,
         )
 
         labels = []
-        for label in downsampled_dataset["labels"]:
+        for label in downsampled_dataset[self.label_column_name]:
             if not isinstance(label, list):
                 label = [label]
             labels.append(label)
@@ -194,29 +201,27 @@ def _evaluate_subset(
 
         mean_v_measure = np.mean(v_measures)
         v_std = np.std(v_measures)
-        scores = {
+        return {
             "v_measures": all_v_scores,
             "v_measure": float(mean_v_measure),
             "v_measure_std": v_std,
         }
-        self._add_main_score(scores)
-        return scores
 
     def _calculate_descriptive_statistics_from_split(
         self, split: str, hf_subset: str | None = None, compute_overall: bool = False
     ) -> ClusteringFastDescriptiveStatistics:
         if hf_subset:
-            sentences = self.dataset[hf_subset][split]["sentences"]
-            labels = self.dataset[hf_subset][split]["labels"]
+            sentences = self.dataset[hf_subset][split][self.input_column_name]
+            labels = self.dataset[hf_subset][split][self.label_column_name]
         elif compute_overall:
             sentences = []
             labels = []
             for hf_subset in self.metadata.eval_langs:
-                sentences.extend(self.dataset[hf_subset][split]["sentences"])
-                labels.extend(self.dataset[hf_subset][split]["labels"])
+                sentences.extend(self.dataset[hf_subset][split][self.input_column_name])
+                labels.extend(self.dataset[hf_subset][split][self.label_column_name])
         else:
-            sentences = self.dataset[split]["sentences"]
-            labels = self.dataset[split]["labels"]
+            sentences = self.dataset[split][self.input_column_name]
+            labels = self.dataset[split][self.label_column_name]
 
         return ClusteringFastDescriptiveStatistics(
             num_samples=len(sentences),
@@ -225,11 +230,17 @@ def _calculate_descriptive_statistics_from_split(
         )
 
     def _push_dataset_to_hub(self, repo_name: str) -> None:
-        self._upload_dataset_to_hub(repo_name, ["sentences", "labels"])
+        self._upload_dataset_to_hub(
+            repo_name, [self.input_column_name, self.label_column_name]
+        )
 
 
 def convert_to_fast(
-    dataset: DatasetDict, seed: int, max_size: int = 100_000
+    dataset: DatasetDict,
+    input_column_name: str,
+    label_column_name: str,
+    seed: int,
+    max_size: int = 100_000,
 ) -> DatasetDict:
     """Converts a clustering dataset to a fast version. This concats the cluster into two columns, sentences and labels.
     It additionally downsamples the dataset to max_size.
@@ -242,10 +253,12 @@ def convert_to_fast(
         labels = []
         sentences = []
         n_clusters = len(dataset[split])
-        all_labels_set = set(itertools.chain.from_iterable(dataset[split]["labels"]))
+        all_labels_set = set(
+            itertools.chain.from_iterable(dataset[split][label_column_name])
+        )
         for i in range(n_clusters):
-            lab = dataset[split]["labels"][i]
-            sents = dataset[split]["sentences"][i]
+            lab = dataset[split][label_column_name][i]
+            sents = dataset[split][input_column_name][i]
 
             # check that it is the same distribution
             row_label_set = set(lab)
@@ -259,7 +272,9 @@ def convert_to_fast(
                     sentences.append(s)
                     sent_set.add(s)  # ensuring no duplicates
 
-        ds[split] = Dataset.from_dict({"sentences": sentences, "labels": labels})
+        ds[split] = Dataset.from_dict(
+            {input_column_name: sentences, label_column_name: labels}
+        )
 
         if len(ds[split]) > max_size:
             idxs = rng_state.sample(range(len(ds[split])), max_size)
@@ -268,17 +283,20 @@ def convert_to_fast(
     return DatasetDict(ds)
 
 
-def check_label_distribution(ds: DatasetDict) -> None:
+def check_label_distribution(
+    ds: DatasetDict,
+    label_column_name: str = "labels",
+) -> None:
     """For older clustering dataset versions.
     ds is a DatasetDict at the split level
     """
     n_clusters = len(ds)
     if n_clusters > 50:
         return
-    all_labels_set = set(itertools.chain.from_iterable(ds["labels"]))
+    all_labels_set = set(itertools.chain.from_iterable(ds[label_column_name]))
 
     for i in range(n_clusters):
-        lab = ds["labels"][i]
+        lab = ds[label_column_name][i]
 
         # check that it is the same distribution
         row_label_set = set(lab)

mteb/abstasks/AbsTaskMultilabelClassification.py

@@ -77,6 +77,10 @@ def _evaluate_subset(
         encode_kwargs: dict[str, Any],
         **kwargs: Any,
     ) -> ScoresDict:
+        if isinstance(dataset, (Dataset, DatasetDict)):
+            dataset = dataset.select_columns(
+                [self.input_column_name, self.label_column_name]
+            )
         train_split = dataset[self.train_split]
         eval_split = dataset[hf_split]
 

mteb/abstasks/AbsTaskPairClassification.py

@@ -53,6 +53,9 @@ class AbsTaskPairClassification(AbsTask):
     """
 
     abstask_prompt = "Retrieve text that are semantically similar to the given text."
+    sentence1_column_name: str = "sentence1"
+    sentence2_column_name: str = "sentence2"
+    label_column_name: str = "labels"
 
     def _evaluate_subset(
         self,
@@ -69,9 +72,9 @@ def _evaluate_subset(
             "sentence_transformers.evaluation.PairClassificationEvaluator"
         ).setLevel(logging.WARN)
         evaluator = PairClassificationEvaluator(
-            data_split["sentence1"],
-            data_split["sentence2"],
-            data_split["labels"],
+            data_split[self.sentence1_column_name],
+            data_split[self.sentence2_column_name],
+            data_split[self.label_column_name],
             task_metadata=self.metadata,
             hf_split=hf_split,
             hf_subset=hf_subset,
@@ -102,17 +105,19 @@ def _calculate_descriptive_statistics_from_split(
             dataset = dataset[0]
 
         sentence1 = (
-            dataset["sentence1"][0]
-            if len(dataset["sentence1"]) == 1
-            else dataset["sentence1"]
+            dataset[self.sentence1_column_name][0]
+            if len(dataset[self.sentence1_column_name]) == 1
+            else dataset[self.sentence1_column_name]
         )
         sentence2 = (
-            dataset["sentence2"][0]
-            if len(dataset["sentence2"]) == 1
-            else dataset["sentence2"]
+            dataset[self.sentence2_column_name][0]
+            if len(dataset[self.sentence2_column_name]) == 1
+            else dataset[self.sentence2_column_name]
         )
         labels = (
-            dataset["labels"][0] if len(dataset["labels"]) == 1 else dataset["labels"]
+            dataset[self.label_column_name][0]
+            if len(dataset[self.label_column_name]) == 1
+            else dataset[self.label_column_name]
         )
 
         text1_statistics = calculate_text_statistics(sentence1)
@@ -140,4 +145,11 @@ def _push_dataset_to_hub(self, repo_name: str) -> None:
             for split in self.dataset:
                 if len(self.dataset[split]) == 1:
                     self.dataset[split] = self.dataset[split][0]
-        self._upload_dataset_to_hub(repo_name, ["sentence1", "sentence2", "labels"])
+        self._upload_dataset_to_hub(
+            repo_name,
+            [
+                self.sentence1_column_name,
+                self.sentence2_column_name,
+                self.label_column_name,
+            ],
+        )