Fix-feature-mapping-for-multi-labels

text_classification.py

@@ -22,23 +22,24 @@ 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.startswith("label")]
+        features = [f for f in dataset_features.keys() if not f.startswith("label")]
+
         if len(label_keys) == 0: # no labels found
             # return everything for post processing
-            return list(dataset_features.keys()), list(dataset_features.keys())
+            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:
             labels = dataset_features[label_keys[0]].names
-        features = [f for f in dataset_features.keys() if not f.startswith("label")]
-        return labels, features
+        return labels, features, label_keys
     except Exception as e:
         logging.warning(
             f"Get Labels/Features Failed for dataset: {e}"
         )
-        return None, None
+        return None, None, None
 
 def check_model_task(model_id):
     # check if model is valid on huggingface

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
@@ -198,7 +198,7 @@ def precheck_model_ds_enable_example_btn(
     try:
         ds = datasets.load_dataset(dataset_id, dataset_config, trust_remote_code=True)
         df: pd.DataFrame = ds[dataset_split].to_pandas().head(5)
-        ds_labels, ds_features = get_labels_and_features_from_dataset(ds[dataset_split])
+        ds_labels, ds_features, _ = get_labels_and_features_from_dataset(ds[dataset_split])
 
         if model_task is None or model_task != "text-classification":
           gr.Warning(NOT_TEXT_CLASSIFICATION_MODEL_RAW)
@@ -300,7 +300,7 @@ def align_columns_and_show_prediction(
     model_labels = list(prediction_response.keys())
     
     ds = datasets.load_dataset(dataset_id, dataset_config, split=dataset_split, trust_remote_code=True)
-    ds_labels, ds_features = get_labels_and_features_from_dataset(ds)
+    ds_labels, ds_features, _ = get_labels_and_features_from_dataset(ds)
 
     # when dataset does not have labels or features
     if not isinstance(ds_labels, list) or not isinstance(ds_features, list):
@@ -390,13 +390,15 @@ def enable_run_btn(uid, run_inference, inference_token, model_id, dataset_id, da
         return gr.update(interactive=False)
     return gr.update(interactive=True)
 
-def construct_label_and_feature_mapping(all_mappings, ds_labels, ds_features):
+def construct_label_and_feature_mapping(all_mappings, ds_labels, ds_features, label_keys=None):
     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)
+
     feature_mapping = all_mappings["features"]
+    if len(label_keys) > 0:
+      feature_mapping.update({"label": label_keys[0]})
     return label_mapping, feature_mapping
 
 def show_hf_token_info(token):
@@ -421,8 +426,8 @@ def try_submit(m_id, d_id, config, split, inference, inference_token, uid):
 
     # get ds labels and features again for alignment
     ds = datasets.load_dataset(d_id, config, split=split, trust_remote_code=True)
-    ds_labels, ds_features = get_labels_and_features_from_dataset(ds)
-    label_mapping, feature_mapping = construct_label_and_feature_mapping(all_mappings, ds_labels, ds_features)
+    ds_labels, ds_features, label_keys = get_labels_and_features_from_dataset(ds)
+    label_mapping, feature_mapping = construct_label_and_feature_mapping(all_mappings, ds_labels, ds_features, label_keys)
 
     eval_str = f"[{m_id}]<{d_id}({config}, {split} set)>"
     save_job_to_pipe(