Unify dataset checking, show header dataset

app_text_classification.py

@@ -2,14 +2,16 @@ import uuid
 
 import gradio as gr
 
-from io_utils import (get_logs_file, read_scanners, write_scanners)
-from text_classification_ui_helpers import (check_dataset_and_get_config,
-                                            check_dataset_and_get_split,
-                                            align_columns_and_show_prediction,
-                                            deselect_run_inference,
-                                            select_run_mode, try_submit,
-                                            write_column_mapping_to_config,
-                                            precheck_model_ds_enable_example_btn)
+from io_utils import get_logs_file, read_scanners, write_scanners
+from text_classification_ui_helpers import (
+    align_columns_and_show_prediction,
+    check_dataset,
+    deselect_run_inference,
+    precheck_model_ds_enable_example_btn,
+    select_run_mode,
+    try_submit,
+    write_column_mapping_to_config,
+)
 from wordings import CONFIRM_MAPPING_DETAILS_MD, INTRODUCTION_MD
 
 MAX_LABELS = 40
@@ -38,16 +40,19 @@ def get_demo():
         )
 
     with gr.Row():
-        dataset_config_input = gr.Dropdown(label="Dataset Config", visible=False)
-        dataset_split_input = gr.Dropdown(label="Dataset Split", visible=False)
+        dataset_config_input = gr.Dropdown(label="Dataset Config", visible=False, allow_custom_value=True)
+        dataset_split_input = gr.Dropdown(label="Dataset Split", visible=False, allow_custom_value=True)
 
     with gr.Row():
         example_btn = gr.Button(
-            "Auto-align Columns & Get Sample Prediction", 
-            visible=True, 
+            "Auto-align Columns & Get Sample Prediction",
+            visible=True,
             variant="primary",
-            interactive=False)
+            interactive=False,
+        )
 
+    with gr.Row():
+        first_line_ds = gr.DataFrame(label="Dataset preview", visible=False)
     with gr.Row():
         example_input = gr.HTML(visible=False)
     with gr.Row():
@@ -103,18 +108,29 @@ def get_demo():
         )
 
     with gr.Row():
-        logs = gr.Textbox(value=get_logs_file, label="Giskard Bot Evaluation Log:", visible=False, every=0.5)
+        logs = gr.Textbox(
+            value=get_logs_file,
+            label="Giskard Bot Evaluation Log:",
+            visible=False,
+            every=0.5,
+        )
 
     dataset_id_input.change(
-        check_dataset_and_get_config,
+        check_dataset,
         inputs=[dataset_id_input],
-        outputs=[dataset_config_input],
+        outputs=[dataset_config_input, dataset_split_input, first_line_ds],
     )
 
     dataset_config_input.change(
-        check_dataset_and_get_split,
+        check_dataset,
         inputs=[dataset_id_input, dataset_config_input],
-        outputs=[dataset_split_input],
+        outputs=[dataset_config_input, dataset_split_input, first_line_ds],
+    )
+
+    dataset_split_input.change(
+        check_dataset,
+        inputs=[dataset_id_input, dataset_config_input, dataset_split_input],
+        outputs=[dataset_config_input, dataset_split_input, first_line_ds],
     )
 
     scanners.change(write_scanners, inputs=[scanners, uid_label])
@@ -155,15 +171,17 @@ def get_demo():
             model_id_input.change,
             dataset_id_input.change,
             dataset_config_input.change,
-            dataset_split_input.change], 
+            dataset_split_input.change,
+        ],
         fn=precheck_model_ds_enable_example_btn,
         inputs=[
             model_id_input,
             dataset_id_input,
             dataset_config_input,
             dataset_split_input,
-        ], 
-        outputs=[example_btn])
+        ],
+        outputs=[example_btn],
+    )
 
     gr.on(
         triggers=[
@@ -222,6 +240,6 @@ def get_demo():
     gr.on(
         triggers=[label.input for label in column_mappings],
         fn=enable_run_btn,
-        inputs=None,    # FIXME
+        inputs=None,  # FIXME
         outputs=[run_btn],
     )

text_classification_ui_helpers.py

@@ -7,15 +7,27 @@ import uuid
 
 import datasets
 import gradio as gr
+import pandas as pd
 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)
-from text_classification import (check_model, get_example_prediction,
-                                 get_labels_and_features_from_dataset)
-from wordings import (CHECK_CONFIG_OR_SPLIT_RAW,
-                      CONFIRM_MAPPING_DETAILS_FAIL_RAW,
-                      MAPPING_STYLED_ERROR_WARNING, get_styled_input)
+from io_utils import (
+    get_yaml_path,
+    read_column_mapping,
+    save_job_to_pipe,
+    write_column_mapping,
+    write_log_to_user_file,
+)
+from text_classification import (
+    check_model,
+    get_example_prediction,
+    get_labels_and_features_from_dataset,
+)
+from wordings import (
+    CHECK_CONFIG_OR_SPLIT_RAW,
+    CONFIRM_MAPPING_DETAILS_FAIL_RAW,
+    MAPPING_STYLED_ERROR_WARNING,
+    get_styled_input,
+)
 
 MAX_LABELS = 40
 MAX_FEATURES = 20
@@ -32,24 +44,50 @@ HF_GSK_HUB_UNLOCK_TOKEN = "GSK_HUB_UNLOCK_TOKEN"
 LEADERBOARD = "giskard-bot/evaluator-leaderboard"
 
 
-def check_dataset_and_get_config(dataset_id):
-    try:
-        # 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:
-        # Dataset may not exist
-        pass
+logger = logging.getLogger(__file__)
 
 
-def check_dataset_and_get_split(dataset_id, dataset_config):
+def check_dataset(dataset_id, dataset_config=None, dataset_split=None):
+    configs = ["default"]
+    splits = ["default"]
+    logger.info(f"Loading {dataset_id}, {dataset_config}, {dataset_split}")
     try:
-        splits = list(datasets.load_dataset(dataset_id, dataset_config).keys())
-        return gr.Dropdown(splits, value=splits[0], visible=True)
-    except Exception:
+        configs = datasets.get_dataset_config_names(dataset_id)
+        splits = list(
+            datasets.load_dataset(
+                dataset_id, configs[0] if not dataset_config else dataset_config
+            ).keys()
+        )
+        if dataset_config == None:
+            dataset_config = configs[0]
+            dataset_split = splits[0]
+        elif dataset_split == None:
+            dataset_split = splits[0]
+    except Exception as e:
         # Dataset may not exist
-        # gr.Warning(f"Failed to load dataset {dataset_id} with config {dataset_config}: {e}")
-        pass
+        logger.warn(
+            f"Failed to load dataset {dataset_id} with config {dataset_config}: {e}"
+        )
+        if dataset_config == None:
+            return (
+                gr.Dropdown(configs, value=configs[0], visible=True),
+                gr.Dropdown(splits, value=splits[0], visible=True),
+                gr.DataFrame(pd.DataFrame(), visible=False),
+            )
+        elif dataset_split == None:
+            return (
+                gr.Dropdown(configs, value=dataset_config, visible=True),
+                gr.Dropdown(splits, value=splits[0], visible=True),
+                gr.DataFrame(pd.DataFrame(), visible=False),
+            )
+
+    dataset_dict = datasets.load_dataset(dataset_id, dataset_config)
+    dataframe: pd.DataFrame = dataset_dict[dataset_split].to_pandas().head(5)
+    return (
+        gr.Dropdown(configs, value=dataset_config, visible=True),
+        gr.Dropdown(splits, value=dataset_split, visible=True),
+        gr.DataFrame(dataframe, visible=True),
+    )
 
 
 def select_run_mode(run_inf):