polish up and add more information

app.py

@@ -59,11 +59,10 @@ def check_dataset(dataset_id, dataset_config="default", dataset_split="test"):
         return dataset_id, None, None
     return dataset_id, dataset_config, dataset_split
 
-def try_validate(model_id, dataset_id, dataset_config, dataset_split, column_mapping='{}'):
+def try_validate(m_id, ppl, dataset_id, dataset_config, dataset_split, column_mapping='{}'):
     # Validate model
-    m_id, ppl = check_model(model_id=model_id)
     if m_id is None:
-        gr.Warning(f'Model "{model_id}" is not accessible. Please set your HF_TOKEN if it is a private model.')
+        gr.Warning('Model is not accessible. Please set your HF_TOKEN if it is a private model.')
         return (
             gr.update(interactive=False),   # Submit button
             gr.update(visible=True),       # Loading row
@@ -73,7 +72,7 @@ def try_validate(model_id, dataset_id, dataset_config, dataset_split, column_map
             gr.update(visible=False),       # Label mapping preview
         )
     if isinstance(ppl, Exception):
-        gr.Warning(f'Failed to load "{model_id} model": {ppl}')
+        gr.Warning(f'Failed to load model": {ppl}')
         return (
             gr.update(interactive=False),   # Submit button
             gr.update(visible=True),       # Loading row
@@ -124,8 +123,6 @@ def try_validate(model_id, dataset_id, dataset_config, dataset_split, column_map
 
         column_mapping = json.dumps(column_mapping, indent=2)
 
-    del ppl
-
     if prediction_result is None:
         gr.Warning('The model failed to predict with the first row in the dataset. Please provide column mappings in "Advance" settings.')
         return (
@@ -212,7 +209,6 @@ with gr.Blocks(theme=theme) as iface:
         def check_dataset_and_get_config(dataset_id):
             try:
                 configs = datasets.get_dataset_config_names(dataset_id)
-                print(configs)
                 return gr.Dropdown(configs, value=configs[0], visible=True)
             except Exception:
                 # Dataset may not exist
@@ -221,19 +217,19 @@ with gr.Blocks(theme=theme) as iface:
         def check_dataset_and_get_split(dataset_config, dataset_id):
             try:
                 splits = list(datasets.load_dataset(dataset_id, dataset_config).keys())
-                print('splits: ',splits)
                 return gr.Dropdown(splits, value=splits[0], visible=True)
             except Exception as e:
                 # Dataset may not exist
-                print(e)
+                gr.Warning(f"Failed to load dataset {dataset_id} with config {dataset_config}: {e}")
                 pass
         
         def gate_validate_btn(model_id, dataset_id, dataset_config, dataset_split, id2label_mapping_dataframe=None):
-            print('model_id: ',model_id)
             column_mapping = '{}'
+            m_id, ppl = check_model(model_id=model_id)
+
             if id2label_mapping_dataframe is not None:
                 column_mapping = id2label_mapping_dataframe.to_json(orient="split")
-            if check_column_mapping_keys_validity(column_mapping) is False:
+            if check_column_mapping_keys_validity(column_mapping, ppl) is False:
                 gr.Warning('Label mapping table has invalid contents. Please check again.')
                 return (gr.update(interactive=False), 
                         gr.update(),
@@ -243,12 +239,15 @@ with gr.Blocks(theme=theme) as iface:
                         gr.update())
             else:
                 if model_id and dataset_id and dataset_config and dataset_split:
-                    return try_validate(model_id, dataset_id, dataset_config, dataset_split, column_mapping)
+                    return try_validate(m_id, ppl, dataset_id, dataset_config, dataset_split, column_mapping)
                 else:
+                    del ppl
+
                     return (gr.update(interactive=False), 
                             gr.update(visible=True),
                             gr.update(visible=False),
                             gr.update(visible=False),
+                            gr.update(visible=False),
                             gr.update(visible=False))
         with gr.Row():
             gr.Markdown('''

text_classification.py

@@ -34,15 +34,19 @@ def text_classification_map_model_and_dataset_labels(id2label, dataset_features)
 
     return id2label_mapping, dataset_labels
 
-def check_column_mapping_keys_validity(column_mapping):
+
+def check_column_mapping_keys_validity(column_mapping, ppl):
     # get the element in all the list elements
     column_mapping = json.loads(column_mapping)
     if "data" not in column_mapping.keys():
         return True
     user_labels = set([pair[0] for pair in column_mapping["data"]])
     model_labels = set([pair[1] for pair in column_mapping["data"]])
+
+    id2label = ppl.model.config.id2label
+    original_labels = set(id2label.values())
     
-    return user_labels == model_labels
+    return user_labels == model_labels == original_labels
 
 
 def text_classification_fix_column_mapping(column_mapping, ppl, d_id, config, split):
@@ -100,7 +104,6 @@ def text_classification_fix_column_mapping(column_mapping, ppl, d_id, config, sp
         if isinstance(column_mapping["data"], list):
             # Use the column mapping passed by user
             for user_label, model_label in column_mapping["data"]:
-                print(user_label, model_label)
                 id2label_mapping[model_label] = user_label
     elif None in id2label_mapping.values():
         column_mapping["label"] = {
@@ -108,7 +111,9 @@ def text_classification_fix_column_mapping(column_mapping, ppl, d_id, config, sp
         }
         return column_mapping, prediction_result, None
 
-    print(id2label_mapping)
+    prediction_result = {
+        f'[{label2id[result["label"]]}]{result["label"]}(original) - {id2label_mapping[result["label"]]}(mapped)': result["score"] for result in results
+    }
     id2label_df = pd.DataFrame({
         "Dataset Labels": dataset_labels,
         "Model Prediction Labels": [id2label_mapping[label] for label in dataset_labels],