fix feature mapping by adding multi labels case

text_classification.py

@@ -22,7 +22,12 @@ 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.startswith('label')]
+        label_keys = [i for i in dataset_features.keys() if i == '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')]
+            features += label_keys
+
         if len(label_keys) == 0: # no labels found
             # return everything for post processing
             return list(dataset_features.keys()), list(dataset_features.keys())
@@ -32,7 +37,6 @@ def get_labels_and_features_from_dataset(ds):
                 labels = label_feat.names
         else:
             labels = dataset_features[label_keys[0]].names
-        features = [f for f in dataset_features.keys() if not f.startswith("label")]
         return labels, features
     except Exception as e:
         logging.warning(

text_classification_ui_helpers.py

@@ -138,7 +138,7 @@ def list_labels_and_features_from_dataset(ds_labels, ds_features, model_labels,
         ds_labels = list(shared_labels)
     if len(ds_labels) > MAX_LABELS:
         ds_labels = ds_labels[:MAX_LABELS]
-        gr.Warning(f"The number of labels is truncated to length {MAX_LABELS}")
+        gr.Warning(f"Too many labels to display for this spcae. We do not support more than {MAX_LABELS} in this space. You can use cli tool at https://github.com/Giskard-AI/cicd.")
 
     # sort labels to make sure the order is consistent
     # prediction gives the order based on probability
@@ -393,10 +393,12 @@ def enable_run_btn(uid, run_inference, inference_token, model_id, dataset_id, da
 def construct_label_and_feature_mapping(all_mappings, ds_labels, ds_features):
     label_mapping = {}
     if len(all_mappings["labels"].keys()) != len(ds_labels):
-        logger.warn("Label mapping corrupted: " + CONFIRM_MAPPING_DETAILS_FAIL_RAW)
+        logger.warn(f"""Label mapping corrupted: {CONFIRM_MAPPING_DETAILS_FAIL_RAW}. 
+                    \nall_mappings: {all_mappings}\nds_labels: {ds_labels}""")
       
     if len(all_mappings["features"].keys()) != len(ds_features):
-        logger.warn("Feature mapping corrupted: " + CONFIRM_MAPPING_DETAILS_FAIL_RAW)
+        logger.warn(f"""Feature mapping corrupted: {CONFIRM_MAPPING_DETAILS_FAIL_RAW}. 
+                    \nall_mappings: {all_mappings}\nds_features: {ds_features}""")
 
     for i, label in zip(range(len(ds_labels)),  ds_labels):
         # align the saved labels with dataset labels order
@@ -405,7 +407,10 @@ def construct_label_and_feature_mapping(all_mappings, ds_labels, ds_features):
     if "features" not in all_mappings.keys():
         logger.warning("features not in all_mappings")
         gr.Warning(CONFIRM_MAPPING_DETAILS_FAIL_RAW)
+
+    special_classlabel = [feature for feature in ds_features if feature.startswith("label")]
     feature_mapping = all_mappings["features"]
+    feature_mapping.update({"label": special_classlabel[0] if special_classlabel else "label"})
     return label_mapping, feature_mapping
 
 def show_hf_token_info(token):