Add features, label mapping in text classification

app.py

@@ -7,6 +7,7 @@ import time
 from pathlib import Path
 
 import json
+import logging
 
 import pandas as pd
 
@@ -64,16 +65,20 @@ def text_classificaiton_match_label_case_unsensative(id2label_mapping, label):
     for model_label in id2label_mapping.keys():
         if model_label.upper() == label.upper():
             return model_label, label
+    return None, label
 
 
 def text_classification_map_model_and_dataset_labels(id2label, dataset_features):
     id2label_mapping = {id2label[k]: None for k in id2label.keys()}
+    dataset_labels = None
     for feature in dataset_features.values():
         if not isinstance(feature, datasets.ClassLabel):
             continue
         if len(feature.names) != len(id2label_mapping.keys()):
             continue
 
+        dataset_labels = feature.names
+
         # Try to match labels
         for label in feature.names:
             if label in id2label_mapping.keys():
@@ -81,9 +86,86 @@ def text_classification_map_model_and_dataset_labels(id2label, dataset_features)
             else:
                 # Try to find case unsensative
                 model_label, label = text_classificaiton_match_label_case_unsensative(id2label_mapping, label)
-            id2label_mapping[model_label] = label
+            if model_label is not None:
+                id2label_mapping[model_label] = label
+
+    return id2label_mapping, dataset_labels
 
-    return id2label_mapping
+
+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
+
+    # Check whether we need to infer the text input column
+    infer_text_input_column = True
+    if "text" in column_mapping.keys():
+        dataset_text_column = column_mapping["text"]
+        if dataset_text_column in dataset_features.keys():
+            infer_text_input_column = False
+        else:
+            logging.warning(f"Provided {dataset_text_column} is not in Dataset columns")
+
+    if infer_text_input_column:
+        # Try to retrieve one
+        candidates = [f for f in dataset_features if dataset_features[f].dtype == "string"]
+        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
+
+    # Load dataset as DataFrame
+    df = ds.to_pandas()
+
+    # Retrieve all labels
+    id2label_mapping = {}
+    id2label = ppl.model.config.id2label
+    label2id = {v: k for k, v in id2label.items()}
+    prediction_result = None
+    try:
+        # Use the first item to test prediction
+        results = ppl({"text": df.head(1).at[0, column_mapping["text"]]}, top_k=None)
+        prediction_result = {
+            f'{result["label"]}({label2id[result["label"]]})': result["score"] for result in results
+        }
+    except Exception:
+        # Pipeline prediction failed, need to provide labels
+        return column_mapping, None, None
+
+    # Infer labels
+    id2label_mapping, dataset_labels = text_classification_map_model_and_dataset_labels(id2label, dataset_features)
+    if "label" in column_mapping.keys():
+        if not isinstance(column_mapping["label"], dict) or set(column_mapping["label"].values()) != set(dataset_labels):
+            logging.warning(f'Provided {column_mapping["label"]} does not match labels in Dataset')
+            return column_mapping, prediction_result, None
+
+        if isinstance(column_mapping["label"], dict):
+            for model_label in id2label_mapping.keys():
+                id2label_mapping[model_label] = column_mapping["label"][str(label2id[model_label])]
+    elif None in id2label_mapping.values():
+        column_mapping["label"] = {
+            i: None for i in id2label.keys()
+        }
+        return column_mapping, prediction_result, None
+
+    id2label_df = pd.DataFrame({
+        "ID": [i for i in id2label.keys()],
+        "Model labels": [id2label[label] for label in id2label.keys()],
+        "Dataset labels": [id2label_mapping[id2label[label]] for label in id2label.keys()],
+    })
+    if "label" not in column_mapping.keys():
+        column_mapping["label"] = {
+            i: id2label_mapping[id2label[i]] for i in id2label.keys()
+        }
+
+    return column_mapping, prediction_result, id2label_df
 
 
 def try_validate(model_id, dataset_id, dataset_config, dataset_split, column_mapping):
@@ -133,7 +215,7 @@ def try_validate(model_id, dataset_id, dataset_config, dataset_split, column_map
         )
 
     # TODO: Validate column mapping by running once
-    prediction_result = {}
+    prediction_result = None
     id2label_df = None
     if isinstance(ppl, TextClassificationPipeline):
         try:
@@ -141,39 +223,32 @@ def try_validate(model_id, dataset_id, dataset_config, dataset_split, column_map
         except Exception:
             column_mapping = {}
 
-        # Retrieve all labels
-        id2label_mapping = {}
-        try:
-            results = ppl({"text": "Test"}, top_k=None)
-            prediction_result = {
-                result["label"]: result["score"] for result in results
-            }
-        except Exception as e:
-            # Pipeline is not executable
-            pass
-
-        # We assume dataset is ok here
-        ds = datasets.load_dataset(d_id, config)[split]
-        try:
-            id2label = ppl.model.config.id2label
-            id2label_mapping = text_classification_map_model_and_dataset_labels(ppl.model.config.id2label, ds.features)
-            id2label_df = pd.DataFrame({
-                "ID": [i for i in id2label.keys()],
-                "Model labels": [id2label[label] for label in id2label.keys()],
-                "Dataset labels": [id2label_mapping[id2label[label]] for label in id2label.keys()],
-            })
-            if "label" not in column_mapping.keys():
-                column_mapping["label"] = {
-                    i: id2label_mapping[id2label[i]] for i in id2label.keys()
-                }
-        except AttributeError:
-            # Dataset does not have features
-            pass
+        column_mapping, prediction_result, id2label_df = \
+            text_classification_fix_column_mapping(column_mapping, ppl, d_id, config, split)
 
         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 (
+            config, split,
+            gr.update(interactive=False),   # Submit button
+            gr.update(visible=False),   # Model prediction preview
+            gr.update(visible=False),   # Label mapping preview
+            gr.update(value=column_mapping, visible=True, interactive=True),    # Column mapping
+        )
+    elif id2label_df is None:
+        gr.Warning('The prediction result does not conform the labels in the dataset. Please provide label mappings in "Advance" settings.')
+        return (
+            config, split,
+            gr.update(interactive=False),   # Submit button
+            gr.update(value=prediction_result, visible=True),   # Model prediction preview
+            gr.update(visible=False),   # Label mapping preview
+            gr.update(value=column_mapping, visible=True, interactive=True),    # Column mapping
+        )
+
     gr.Info("Model and dataset validations passed. Your can submit the evaluation task.")
 
     return (
@@ -248,7 +323,6 @@ with gr.Blocks(theme=theme) as iface:
                 ],
                 value=0,
             )
-            run_local = gr.Checkbox(value=True, label="Run in this Space")
             example_labels = gr.Label(label='Model pipeline test prediction result', visible=False)
 
         with gr.Column():
@@ -278,16 +352,18 @@ with gr.Blocks(theme=theme) as iface:
             id2label_mapping_dataframe = gr.DataFrame(visible=False)
 
     with gr.Row():
-        column_mapping_input = gr.Textbox(
-            value="",
-            lines=5,
-            label="Column mapping",
-            placeholder="Description of mapping of columns in model to dataset, in json format, e.g.:\n"
-                        '{\n'
-                        '   "text": "context",\n'
-                        '   "label": {0: "Positive", 1: "Negative"}\n'
-                        '}',
-        )
+        with gr.Accordion("Advance", open=False):
+            run_local = gr.Checkbox(value=True, label="Run in this Space")
+            column_mapping_input = gr.Textbox(
+                value="",
+                lines=5,
+                label="Column mapping",
+                placeholder="Description of mapping of columns in model to dataset, in json format, e.g.:\n"
+                            '{\n'
+                            '   "text": "context",\n'
+                            '   "label": {0: "Positive", 1: "Negative"}\n'
+                            '}',
+            )
 
     with gr.Row():
         validate_btn = gr.Button("Validate model and dataset", variant="primary")