Fix feature mapping selection bug (#21)

- fix feature mapping selection bug (3025552358963d437caf5108d98dfa6653caffa3)
- add warning component (104fec1752a8f6d84d7d0a27e1fa4d43f0669364)
- do not show when dataset has problem (0160f4b08cf662a83c76f2fce2cd5d377fa17429)
- Show check config warning (b4afb862a243e1dfc725d284d2d58d4d9e10a777)
- add leaderboard arg and inf token (5927800bda60c3d05cb37dc1637687888aa68008)
- fix inf token ui (ff775723f239f6ccc60834a2b43a5644614ae0ce)
- fix config file io bug (8a5336fb86880b4f03caafb2a6e78c34b7067f58)


Co-authored-by: zcy <ZeroCommand@users.noreply.huggingface.co>

app.py

@@ -1,14 +1,10 @@
 import atexit
-import threading
-
 import gradio as gr
 
 from app_leaderboard import get_demo as get_demo_leaderboard
 from app_text_classification import get_demo as get_demo_text_classification
 from run_jobs import start_process_run_job, stop_thread
 
-if threading.current_thread() is not threading.main_thread():
-    t = threading.current_thread()
 try:
     with gr.Blocks(theme=gr.themes.Soft(primary_hue="green")) as demo:
         with gr.Tab("Text Classification"):
@@ -22,6 +18,6 @@ try:
     demo.launch(share=False)
     atexit.register(stop_thread)
 
-except Exception:
-    print("stop background thread")
+except Exception as e:
+    print("stop background thread: ", e)
     stop_thread()

app_text_classification.py

@@ -4,8 +4,8 @@ from io_utils import (
     read_scanners,
     write_scanners,
     read_inference_type,
-    write_inference_type,
     get_logs_file,
+    write_inference_type,
 )
 from wordings import INTRODUCTION_MD, CONFIRM_MAPPING_DETAILS_MD
 from text_classification_ui_helpers import (
@@ -14,6 +14,8 @@ from text_classification_ui_helpers import (
     check_dataset_and_get_split,
     check_model_and_show_prediction,
     write_column_mapping_to_config,
+    select_run_mode,
+    deselect_run_inference,
 )
 
 MAX_LABELS = 20
@@ -70,6 +72,7 @@ def get_demo(demo):
         run_local = gr.Checkbox(value=True, label="Run in this Space")
         use_inference = read_inference_type(uid) == "hf_inference_api"
         run_inference = gr.Checkbox(value=use_inference, label="Run with Inference API")
+        inference_token = gr.Textbox(value="", label="HF Token for Inference API", visible=False, interactive=True)
 
     with gr.Accordion(label="Scanner Advance Config (optional)", open=False):
         selected = read_scanners(uid)
@@ -94,7 +97,8 @@ def get_demo(demo):
         demo.load(get_logs_file, uid_label, logs, every=0.5)
 
     dataset_id_input.change(
-        check_dataset_and_get_config, inputs=[dataset_id_input, uid_label], outputs=[dataset_config_input]
+        check_dataset_and_get_config, 
+        inputs=[dataset_id_input, uid_label], outputs=[dataset_config_input]
     )
 
     dataset_config_input.change(
@@ -105,8 +109,20 @@ def get_demo(demo):
 
     scanners.change(write_scanners, inputs=[scanners, uid_label])
 
-    run_inference.change(write_inference_type, inputs=[run_inference, uid_label])
-
+    run_inference.change(
+        select_run_mode, 
+        inputs=[run_inference, inference_token, uid_label], 
+        outputs=[inference_token, run_local])
+    
+    run_local.change(
+        deselect_run_inference, 
+        inputs=[run_local], 
+        outputs=[inference_token, run_inference])
+    
+    inference_token.change(
+        write_inference_type, 
+        inputs=[run_inference, inference_token, uid_label])
+    
     gr.on(
         triggers=[label.change for label in column_mappings],
         fn=write_column_mapping_to_config,
@@ -119,6 +135,19 @@ def get_demo(demo):
         ],
     )
 
+    # label.change sometimes does not pass the changed value
+    gr.on(
+        triggers=[label.input for label in column_mappings],
+        fn=write_column_mapping_to_config,
+        inputs=[
+            dataset_id_input,
+            dataset_config_input,
+            dataset_split_input,
+            uid_label,
+            *column_mappings,
+        ],
+    )
+
     gr.on(
         triggers=[
             model_id_input.change,

io_utils.py

@@ -1,6 +1,6 @@
 import os
 import subprocess
-
+import gradio as gr
 import yaml
 
 import pipe
@@ -28,17 +28,21 @@ def read_scanners(uid):
     with open(get_yaml_path(uid), "r") as f:
         config = yaml.load(f, Loader=yaml.FullLoader)
         scanners = config.get("detectors", [])
+    f.close()
     return scanners
 
 
 # convert a list of scanners to yaml file
 def write_scanners(scanners, uid):
-    with open(get_yaml_path(uid), "r+") as f:
+    with open(get_yaml_path(uid), "r") as f:
         config = yaml.load(f, Loader=yaml.FullLoader)
         if config:
             config["detectors"] = scanners
-            # save scanners to detectors in yaml
-            yaml.dump(config, f, Dumper=Dumper)
+    f.close()
+    # save scanners to detectors in yaml
+    with open(get_yaml_path(uid), "w") as f:
+        yaml.dump(config, f, Dumper=Dumper)
+    f.close()
 
 
 # read model_type from yaml file
@@ -47,19 +51,25 @@ def read_inference_type(uid):
     with open(get_yaml_path(uid), "r") as f:
         config = yaml.load(f, Loader=yaml.FullLoader)
         inference_type = config.get("inference_type", "")
+    f.close()
     return inference_type
 
 
 # write model_type to yaml file
-def write_inference_type(use_inference, uid):
-    with open(get_yaml_path(uid), "r+") as f:
+def write_inference_type(use_inference, inference_token, uid):
+    with open(get_yaml_path(uid), "r") as f:
         config = yaml.load(f, Loader=yaml.FullLoader)
         if use_inference:
             config["inference_type"] = "hf_inference_api"
+            config["inference_token"] = inference_token
         else:
             config["inference_type"] = "hf_pipeline"
-        # save inference_type to inference_type in yaml
+    f.close()
+    # save inference_type to inference_type in yaml
+    with open(get_yaml_path(uid), "w") as f:
         yaml.dump(config, f, Dumper=Dumper)
+    f.close()
+
 
 
 # read column mapping from yaml file
@@ -69,6 +79,7 @@ def read_column_mapping(uid):
         config = yaml.load(f, Loader=yaml.FullLoader)
         if config:
             column_mapping = config.get("column_mapping", dict())
+    f.close()
     return column_mapping
 
 
@@ -76,15 +87,18 @@ def read_column_mapping(uid):
 def write_column_mapping(mapping, uid):
     with open(get_yaml_path(uid), "r") as f:
         config = yaml.load(f, Loader=yaml.FullLoader)
+    f.close()
+
     if config is None:
         return
     if mapping is None and "column_mapping" in config.keys():
         del config["column_mapping"]
     else:
         config["column_mapping"] = mapping
+
     with open(get_yaml_path(uid), "w") as f:
-        # save column_mapping to column_mapping in yaml
         yaml.dump(config, f, Dumper=Dumper)
+    f.close()
 
 
 # convert column mapping dataframe to json
@@ -107,6 +121,7 @@ def get_logs_file(uid):
 def write_log_to_user_file(id, log):
     with open(f"./tmp/{id}_log", "a") as f:
         f.write(log)
+    f.close()
 
 
 def save_job_to_pipe(id, job, lock):

run_jobs.py

@@ -1,6 +1,6 @@
 import threading
 import time
-
+import logging
 import pipe
 from io_utils import pop_job_from_pipe
 
@@ -9,7 +9,7 @@ is_running = False
 
 def start_process_run_job():
     try:
-        print("Running jobs in thread")
+        logging.debug("Running jobs in thread")
         global thread, is_running
         thread = threading.Thread(target=run_job)
         thread.daemon = True
@@ -22,7 +22,7 @@ def start_process_run_job():
 
 
 def stop_thread():
-    print("Stop thread")
+    logging.debug("Stop thread")
     global is_running
     is_running = False
 
@@ -34,6 +34,6 @@ def run_job():
             pop_job_from_pipe()
             time.sleep(10)
         except KeyboardInterrupt:
-            print("KeyboardInterrupt stop background thread")
+            logging.debug("KeyboardInterrupt stop background thread")
             is_running = False
             break

text_classification_ui_helpers.py

@@ -9,10 +9,11 @@ import gradio as gr
 from transformers.pipelines import TextClassificationPipeline
 
 from io_utils import (get_yaml_path, read_column_mapping, save_job_to_pipe,
-                      write_column_mapping, write_log_to_user_file)
+                      write_column_mapping, write_log_to_user_file,
+                      write_inference_type)
 from text_classification import (check_model, get_example_prediction,
                                  get_labels_and_features_from_dataset)
-from wordings import CONFIRM_MAPPING_DETAILS_FAIL_RAW
+from wordings import CONFIRM_MAPPING_DETAILS_FAIL_RAW, MAPPING_STYLED_ERROR_WARNING, CHECK_CONFIG_OR_SPLIT_RAW
 
 MAX_LABELS = 20
 MAX_FEATURES = 20
@@ -24,7 +25,7 @@ HF_WRITE_TOKEN = "HF_WRITE_TOKEN"
 
 def check_dataset_and_get_config(dataset_id, uid):
     try:
-        write_column_mapping(None, uid)  # reset column mapping
+        # write_column_mapping(None, uid)  # reset column mapping
         configs = datasets.get_dataset_config_names(dataset_id)
         return gr.Dropdown(configs, value=configs[0], visible=True)
     except Exception:
@@ -41,7 +42,30 @@ def check_dataset_and_get_split(dataset_id, dataset_config):
         # gr.Warning(f"Failed to load dataset {dataset_id} with config {dataset_config}: {e}")
         pass
 
+def select_run_mode(run_inf, inf_token, uid):
+    if run_inf:
+        if len(inf_token) > 0:
+            write_inference_type(run_inf, inf_token, uid)
+        return (
+            gr.update(visible=True),
+            gr.update(value=False))
+    else:
+        return (
+            gr.update(visible=False),
+            gr.update(value=True))
 
+def deselect_run_inference(run_local):
+    if run_local:
+        return (
+            gr.update(visible=False),
+            gr.update(value=False)
+        )
+    else:
+        return (
+            gr.update(visible=True),
+            gr.update(value=True)
+        )
+        
 def write_column_mapping_to_config(
     dataset_id, dataset_config, dataset_split, uid, *labels
 ):
@@ -52,20 +76,17 @@ def write_column_mapping_to_config(
     )
     if labels is None:
         return
-    labels = [*labels]
-    all_mappings = read_column_mapping(uid)
 
-    if all_mappings is None:
-        all_mappings = dict()
+    all_mappings = dict()
 
     if "labels" not in all_mappings.keys():
         all_mappings["labels"] = dict()
     for i, label in enumerate(labels[:MAX_LABELS]):
         if label:
-            all_mappings["labels"][label] = ds_labels[i]
+            all_mappings["labels"][label] = ds_labels[i%len(ds_labels)]
     if "features" not in all_mappings.keys():
         all_mappings["features"] = dict()
-    for i, feat in enumerate(labels[MAX_LABELS : (MAX_LABELS + MAX_FEATURES)]):
+    for _, feat in enumerate(labels[MAX_LABELS : (MAX_LABELS + MAX_FEATURES)]):
         if feat:
             # TODO: Substitute 'text' with more features for zero-shot
             all_mappings["features"]["text"] = feat
@@ -134,7 +155,7 @@ def check_model_and_show_prediction(
 
     # when dataset does not have labels or features
     if not isinstance(ds_labels, list) or not isinstance(ds_features, list):
-        # gr.Warning(CONFIRM_MAPPING_DETAILS_FAIL_RAW)
+        gr.Warning(CHECK_CONFIG_OR_SPLIT_RAW)
         return (
             gr.update(visible=False),
             gr.update(visible=False),
@@ -154,9 +175,8 @@ def check_model_and_show_prediction(
         collections.Counter(model_id2label.values()) != collections.Counter(ds_labels)
         or ds_features[0] != "text"
     ):
-        gr.Warning(CONFIRM_MAPPING_DETAILS_FAIL_RAW)
         return (
-            gr.update(visible=False),
+            gr.update(value=MAPPING_STYLED_ERROR_WARNING, visible=True),
             gr.update(visible=False),
             gr.update(visible=True, open=True),
             *column_mappings,
@@ -192,6 +212,10 @@ def try_submit(m_id, d_id, config, split, local, uid):
         return (gr.update(interactive=True), gr.update(visible=False))
     feature_mapping = all_mappings["features"]
 
+    leaderboard_dataset = None
+    if os.environ.get("SPACE_ID") == "giskardai/giskard-evaluator":
+        leaderboard_dataset = "ZeroCommand/test-giskard-report"
+
     # TODO: Set column mapping for some dataset such as `amazon_polarity`
     if local:
         command = [
@@ -220,6 +244,8 @@ def try_submit(m_id, d_id, config, split, local, uid):
             json.dumps(label_mapping),
             "--scan_config",
             get_yaml_path(uid),
+            "--leaderboard_dataset",
+            leaderboard_dataset,
         ]
 
         eval_str = f"[{m_id}]<{d_id}({config}, {split} set)>"

wordings.py

@@ -21,9 +21,19 @@ CONFIRM_MAPPING_DETAILS_FAIL_RAW = """
                             Sorry, we cannot align the input/output of your dataset with the model. Pleaser double check your model and dataset.
                             """
 
+CHECK_CONFIG_OR_SPLIT_RAW = """
+                            Please check your dataset config or split.
+                            """
+
 PREDICTION_SAMPLE_MD = """
                         <h1 style="text-align: center;">
                         Model Prediction Sample
                         </h1>
                         Here is a sample prediction from your model based on your dataset.
                         """
+
+MAPPING_STYLED_ERROR_WARNING = """
+                        <h3 style="text-align: center;color: coral; background-color: #fff0f3; border-radius: 8px; padding: 10px; ">
+                        Sorry, we cannot auto-align the labels/features of your dataset and model. Please double check.
+                        </h3>
+                        """