add pre-check for column mapping values

app.py

@@ -10,7 +10,7 @@ import json
 
 from transformers.pipelines import TextClassificationPipeline
 
-from text_classification import text_classification_fix_column_mapping
+from text_classification import check_column_mapping_keys_validity, text_classification_fix_column_mapping
 
 
 HF_REPO_ID = 'HF_REPO_ID'
@@ -233,15 +233,23 @@ with gr.Blocks(theme=theme) as iface:
             column_mapping = '{}'
             if id2label_mapping_dataframe is not None:
                 column_mapping = id2label_mapping_dataframe.to_json(orient="split")
-                print(column_mapping)   
-            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)
-            else:
+            if check_column_mapping_keys_validity(column_mapping) is False:
+                gr.Warning('Label mapping table has invalid contents. Please check again.')
                 return (gr.update(interactive=False), 
-                        gr.update(visible=True),
-                        gr.update(visible=False),
-                        gr.update(visible=False),
-                        gr.update(visible=False))
+                        gr.update(),
+                        gr.update(),
+                        gr.update(),
+                        gr.update(),
+                        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)
+                else:
+                    return (gr.update(interactive=False), 
+                            gr.update(visible=True),
+                            gr.update(visible=False),
+                            gr.update(visible=False),
+                            gr.update(visible=False))
         with gr.Row():
             gr.Markdown('''
                 <h1 style="text-align: center;">

text_classification.py

@@ -1,7 +1,6 @@
 import datasets
-
 import logging
-
+import json
 import pandas as pd
 
 
@@ -35,6 +34,16 @@ def text_classification_map_model_and_dataset_labels(id2label, dataset_features)
 
     return id2label_mapping, dataset_labels
 
+def check_column_mapping_keys_validity(column_mapping):
+    # 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"]])
+    
+    return user_labels == model_labels
+
 
 def text_classification_fix_column_mapping(column_mapping, ppl, d_id, config, split):
     # We assume dataset is ok here