update inference api arg in yaml; update error handling

app.py

@@ -11,7 +11,7 @@ import json
 from transformers.pipelines import TextClassificationPipeline
 
 from text_classification import check_column_mapping_keys_validity, text_classification_fix_column_mapping
-from utils import read_scanners, write_scanners, convert_column_mapping_to_json
+from utils import read_scanners, write_scanners, read_model_type, write_model_type, convert_column_mapping_to_json
 
 HF_REPO_ID = 'HF_REPO_ID'
 HF_SPACE_ID = 'SPACE_ID'
@@ -160,10 +160,14 @@ def try_validate(m_id, ppl, dataset_id, dataset_config, dataset_split, column_ma
     )
 
 
-def try_submit(m_id, d_id, config, split, column_mappings, local):
+def try_submit(m_id, d_id, config, split, id2label_mapping_dataframe, feature_mapping_dataframe, local):
     label_mapping = {}
-    for i, label in column_mappings["Model Prediction Labels"].items():
+    for i, label in id2label_mapping_dataframe["Model Prediction Labels"].items():
         label_mapping.update({str(i): label})
+    
+    feature_mapping = {}
+    for i, feature in feature_mapping_dataframe["Dataset Features"].items():
+        feature_mapping.update({feature_mapping_dataframe["Model Input Features"][i]: feature})
 
     # TODO: Set column mapping for some dataset such as `amazon_polarity`
 
@@ -180,9 +184,9 @@ def try_submit(m_id, d_id, config, split, column_mappings, local):
             "--discussion_repo", os.environ.get(HF_REPO_ID) or os.environ.get(HF_SPACE_ID),
             "--output_format", "markdown",
             "--output_portal", "huggingface",
-            # TODO: "--feature_mapping", json.dumps(column_mapping),
+            "--feature_mapping", json.dumps(feature_mapping),
             "--label_mapping", json.dumps(label_mapping),
-            "--scan_config", "./scan_config.yaml",
+            "--scan_config", "./config.yaml",
         ]
 
         eval_str = f"[{m_id}]<{d_id}({config}, {split} set)>"
@@ -227,11 +231,12 @@ with gr.Blocks(theme=theme) as iface:
         def gate_validate_btn(model_id, dataset_id, dataset_config, dataset_split, id2label_mapping_dataframe=None, feature_mapping_dataframe=None):
             column_mapping = '{}'
             _, ppl = check_model(model_id=model_id)
+
             if id2label_mapping_dataframe is not None:
                 labels = convert_column_mapping_to_json(id2label_mapping_dataframe.value, label="data")
                 features = convert_column_mapping_to_json(feature_mapping_dataframe.value, label="text")
                 column_mapping = json.dumps({**labels, **features}, indent=2)
-                print('229 >>>>> ', column_mapping)
+
             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), 
@@ -261,10 +266,11 @@ with gr.Blocks(theme=theme) as iface:
                 ''')
         with gr.Row():
             run_local = gr.Checkbox(value=True, label="Run in this Space")
-            run_inference = gr.Checkbox(value=False, label="Run with Inference API")
+            use_inference = read_model_type('./config.yaml')[0] == 'hf_inference_api'
+            run_inference = gr.Checkbox(value=use_inference, label="Run with Inference API")
         
         with gr.Row() as advanced_row:
-            selected = read_scanners('./scan_config.yaml')
+            selected = read_scanners('./config.yaml')
             scan_config = selected + ['data_leakage']
             scanners = gr.CheckboxGroup(choices=scan_config, value=selected, label='Scan Settings', visible=True)
 
@@ -282,8 +288,10 @@ with gr.Blocks(theme=theme) as iface:
             dataset_config_input = gr.Dropdown(['default'], value='default', label='Dataset Config', visible=False)
             dataset_split_input = gr.Dropdown(['default'], value='default', label='Dataset Split', visible=False)
 
-            dataset_id_input.change(check_dataset_and_get_config, dataset_id_input, dataset_config_input)
-            dataset_config_input.change(
+            dataset_id_input.blur(check_dataset_and_get_config, dataset_id_input, dataset_config_input)
+            dataset_id_input.submit(check_dataset_and_get_config, dataset_id_input, dataset_config_input)
+
+            dataset_config_input.blur(
                 check_dataset_and_get_split, 
                 inputs=[dataset_config_input, dataset_id_input], 
                 outputs=[dataset_split_input])
@@ -319,16 +327,16 @@ with gr.Blocks(theme=theme) as iface:
             size="lg",
         )
 
-        model_id_input.change(gate_validate_btn, 
+        model_id_input.blur(gate_validate_btn, 
                                 inputs=[model_id_input, dataset_id_input, dataset_config_input, dataset_split_input], 
                                 outputs=[run_btn, loading_row, preview_row, example_input, example_labels, id2label_mapping_dataframe, feature_mapping_dataframe])
-        dataset_id_input.change(gate_validate_btn, 
+        dataset_id_input.blur(gate_validate_btn, 
                                 inputs=[model_id_input, dataset_id_input, dataset_config_input, dataset_split_input], 
                                 outputs=[run_btn, loading_row, preview_row, example_input,  example_labels, id2label_mapping_dataframe, feature_mapping_dataframe])
-        dataset_config_input.change(gate_validate_btn,
+        dataset_config_input.input(gate_validate_btn,
                                 inputs=[model_id_input, dataset_id_input, dataset_config_input, dataset_split_input],
                                 outputs=[run_btn, loading_row, preview_row, example_input, example_labels, id2label_mapping_dataframe, feature_mapping_dataframe])
-        dataset_split_input.change(gate_validate_btn,
+        dataset_split_input.input(gate_validate_btn,
                                 inputs=[model_id_input, dataset_id_input, dataset_config_input, dataset_split_input],
                                 outputs=[run_btn, loading_row, preview_row, example_input, example_labels, id2label_mapping_dataframe, feature_mapping_dataframe])
         id2label_mapping_dataframe.input(gate_validate_btn,
@@ -338,6 +346,10 @@ with gr.Blocks(theme=theme) as iface:
                                 inputs=[model_id_input, dataset_id_input, dataset_config_input, dataset_split_input, id2label_mapping_dataframe, feature_mapping_dataframe],
                                 outputs=[run_btn, loading_row, preview_row, example_input, example_labels, id2label_mapping_dataframe, feature_mapping_dataframe])
         scanners.change(write_scanners, inputs=scanners)
+        run_inference.change(
+            write_model_type,
+            inputs=[run_inference]
+        )
 
         run_btn.click(
             try_submit,
@@ -347,6 +359,7 @@ with gr.Blocks(theme=theme) as iface:
                 dataset_config_input,
                 dataset_split_input,
                 id2label_mapping_dataframe,
+                feature_mapping_dataframe,
                 run_local,
             ],
             outputs=[

scan_config.yaml → config.yaml

@@ -5,4 +5,6 @@ detectors:
   - performance
   - underconfidence
   - overconfidence
-  - spurious_correlation
+  - spurious_correlation
+model_type:
+  - hf_inference_api

text_classification.py

@@ -35,7 +35,18 @@ def text_classification_map_model_and_dataset_labels(id2label, dataset_features)
 
     return id2label_mapping, dataset_labels
 
-
+'''
+params:
+    column_mapping: dict
+    example: {
+        "text": "sentences",
+        "label": {
+            "label0": "LABEL_0",
+            "label1": "LABEL_1"
+        }
+    }
+    ppl: pipeline
+'''
 def check_column_mapping_keys_validity(column_mapping, ppl):
     # get the element in all the list elements
     column_mapping = json.loads(column_mapping)
@@ -49,16 +60,7 @@ def check_column_mapping_keys_validity(column_mapping, ppl):
     
     return user_labels == model_labels == original_labels
 
-
-def text_classification_fix_column_mapping(column_mapping, ppl, d_id, config, split):
-    # We assume dataset is ok here
-    ds = datasets.load_dataset(d_id, config)[split]
-    try:
-        dataset_features = ds.features
-    except AttributeError:
-        # Dataset does not have features, need to provide everything
-        return None, None, None, None, None
-    
+def infer_text_input_column(column_mapping, dataset_features):
     # Check whether we need to infer the text input column
     infer_text_input_column = True
     feature_map_df = None
@@ -79,9 +81,19 @@ def text_classification_fix_column_mapping(column_mapping, ppl, d_id, config, sp
         if len(candidates) > 0:
             logging.debug(f"Candidates are {candidates}")
             column_mapping["text"] = candidates[0]
-        else:
-            # Not found a text feature
-            return column_mapping, None, None, feature_map_df 
+    
+    return column_mapping, feature_map_df 
+
+def text_classification_fix_column_mapping(column_mapping, ppl, d_id, config, split):
+    # We assume dataset is ok here
+    ds = datasets.load_dataset(d_id, config)[split]
+    try:
+        dataset_features = ds.features
+    except AttributeError:
+        # Dataset does not have features, need to provide everything
+        return None, None, None, None, None
+
+    column_mapping, feature_map_df = infer_text_input_column(column_mapping, dataset_features)
 
     # Load dataset as DataFrame
     df = ds.to_pandas()
@@ -97,7 +109,6 @@ def text_classification_fix_column_mapping(column_mapping, ppl, d_id, config, sp
         v: k for k, v in id2label_mapping.items()
     }
 
-    # TODO: convert dataframe column mapping to json properly
     if "data" in column_mapping.keys():
         if isinstance(column_mapping["data"], list):
             # Use the column mapping passed by user
@@ -114,6 +125,7 @@ def text_classification_fix_column_mapping(column_mapping, ppl, d_id, config, sp
         "Model Prediction Labels": [id2label_mapping_dataset_model[label] for label in dataset_labels],
     })
     
+    # get a sample prediction from the model on the dataset
     prediction_input = None
     prediction_result = None
     try:

utils.py

@@ -1,5 +1,11 @@
 import yaml
-import sys
+
+YAML_PATH = "./config.yaml"
+
+class Dumper(yaml.Dumper):
+    def increase_indent(self, flow=False, *args, **kwargs):
+        return super().increase_indent(flow=flow, indentless=False)
+    
 # read scanners from yaml file
 # return a list of scanners
 def read_scanners(path):
@@ -11,9 +17,33 @@ def read_scanners(path):
 
 # convert a list of scanners to yaml file
 def write_scanners(scanners):
-    with open("./scan_config.yaml", "w") as f:
+    with open(YAML_PATH, "r") as f:
+        config = yaml.load(f, Loader=yaml.FullLoader)
+
+    config["detectors"] = scanners
+    with open(YAML_PATH, "w") as f:
         # save scanners to detectors in yaml
-        yaml.dump({"detectors": scanners}, f)
+        yaml.dump(config, f, Dumper=Dumper)
+
+# read model_type from yaml file
+def read_model_type(path):
+    model_type = ""
+    with open(path, "r") as f:
+        config = yaml.load(f, Loader=yaml.FullLoader)
+        model_type = config.get("model_type", None)
+    return model_type
+
+# write model_type to yaml file
+def write_model_type(use_inference):
+    with open(YAML_PATH, "r") as f:
+        config = yaml.load(f, Loader=yaml.FullLoader)
+    if use_inference:
+        config["model_type"] = ['hf_inference_api']
+    else:
+        config["model_type"] = ['hf_pipeline']
+    with open(YAML_PATH, "w") as f:
+        # save model_type to model_type in yaml
+        yaml.dump(config, f, Dumper=Dumper)
 
 # convert column mapping dataframe to json
 def convert_column_mapping_to_json(df, label=""):