clean code

text_classification.py

@@ -22,16 +22,14 @@ class HuggingFaceInferenceAPIResponse:
 def get_labels_and_features_from_dataset(ds):
     try:
         dataset_features = ds.features
-        label_keys = [i for i in dataset_features.keys() if i == 'label']
+        label_keys = [i for i in dataset_features.keys() if i.startswith("label")]
         features = [f for f in dataset_features.keys() if not f.startswith("label")]
-        if len(label_keys) == 0: # no labels found            
-            label_keys = [i for i in dataset_features.keys() if i.startswith('label')]
 
         if len(label_keys) == 0: # no labels found
             # return everything for post processing
             return list(dataset_features.keys()), list(dataset_features.keys()), None
         if not isinstance(dataset_features[label_keys[0]], datasets.ClassLabel):
-            if hasattr(dataset_features[label_keys[0]], 'feature'):
+            if hasattr(dataset_features[label_keys[0]], "feature"):
                 label_feat = dataset_features[label_keys[0]].feature
                 labels = label_feat.names
         else: