add conditions to extract labels from dataset

text_classification.py

@@ -15,7 +15,15 @@ def get_labels_and_features_from_dataset(dataset_id, dataset_config, split):
     try:
         ds = datasets.load_dataset(dataset_id, dataset_config)[split]
         dataset_features = ds.features
-        labels = dataset_features["label"].names
+        label_keys = [i for i in dataset_features.keys() if i.startswith('label')]
+        if len(label_keys) == 0:
+            raise ValueError("Dataset does not have label column")
+        if not isinstance(dataset_features[label_keys[0]], datasets.ClassLabel):
+            if hasattr(dataset_features[label_keys[0]], 'feature'):
+                label_feat = dataset_features[label_keys[0]].feature
+                labels = label_feat.names
+        else:
+            labels = [dataset_features[label_keys[0]].names]
         features = [f for f in dataset_features.keys() if f != "label"]
         return labels, features
     except Exception as e: