return label keys and add in mapping before submission

text_classification.py

@@ -26,23 +26,22 @@ def get_labels_and_features_from_dataset(ds):
         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())
+            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'):
                 label_feat = dataset_features[label_keys[0]].feature
                 labels = label_feat.names
         else:
             labels = dataset_features[label_keys[0]].names
-        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

@@ -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,7 +390,7 @@ 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(f"""Label mapping corrupted: {CONFIRM_MAPPING_DETAILS_FAIL_RAW}. 
@@ -408,9 +408,9 @@ def construct_label_and_feature_mapping(all_mappings, ds_labels, ds_features):
         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"})
+    if len(label_keys) > 0:
+      feature_mapping.update({"label": label_keys[0]})
     return label_mapping, feature_mapping
 
 def show_hf_token_info(token):
@@ -426,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(