Add dataset probing and validation

app.py

@@ -1,10 +1,50 @@
 import gradio as gr
+import datasets
 
 
 theme = gr.themes.Soft(
     primary_hue="green",
 )
 
+
+def check_dataset(dataset_id, dataset_config="default", dataset_split="test"):
+    try:
+        configs = datasets.get_dataset_config_names(dataset_id)
+    except Exception:
+        # Dataset may not exist
+        return None, dataset_config, dataset_split
+
+    if dataset_config not in configs:
+        # Need to choose dataset subset (config)
+        return dataset_id, configs, dataset_split
+
+    ds = datasets.load_dataset(dataset_id, dataset_config)
+
+    if isinstance(ds, datasets.DatasetDict):
+        # Need to choose dataset split
+        if dataset_split not in ds.keys():
+            return dataset_id, None, list(ds.keys())
+    elif not isinstance(ds, datasets.Dataset):
+        # Unknown type
+        return dataset_id, None, None
+    return dataset_id, dataset_config, dataset_split
+
+
+def try_submit(dataset_id, dataset_config, dataset_split):
+    # Validate dataset
+    d_id, config, split = check_dataset(dataset_id=dataset_id, dataset_config=dataset_config, dataset_split=dataset_split)
+
+    if d_id is None:
+        gr.Warning(f'Dataset "{dataset_id}" is not accessible. Please set your HF_TOKEN if it is a private dataset.')
+    if isinstance(config, list):
+        gr.Warning(f'Dataset "{dataset_id}" does have "{dataset_config}" config. Please choose a valid config.')
+        config = gr.Dropdown.update(choices=config, value=config[0])
+    if isinstance(split, list):
+        gr.Warning(f'Dataset "{dataset_id}" does have "{dataset_split}" split. Please choose a valid split.')
+        split = gr.Dropdown.update(choices=split, value=split[0])
+
+    return config, split
+
 with gr.Blocks(theme=theme) as iface:
     with gr.Row():
         with gr.Column():
@@ -28,13 +68,38 @@ with gr.Blocks(theme=theme) as iface:
                 placeholder="tweet_eval",
             )
 
-            gr.Dropdown(
+            dataset_config_input = gr.Dropdown(
                 label="Hugging Face dataset subset",
+                choices=[
+                    "default",
+                ],
+                allow_custom_value=True,
+                value="default",
             )
 
-            gr.Dropdown(
+            dataset_split_input = gr.Dropdown(
                 label="Hugging Face dataset split",
+                choices=[
+                    "test",
+                ],
+                allow_custom_value=True,
+                value="test",
             )
 
+    with gr.Row():
+        run_btn = gr.Button("Validate and submit", variant="primary")
+        run_btn.click(
+            try_submit,
+            inputs=[
+                dataset_id_input,
+                dataset_config_input,
+                dataset_split_input
+            ],
+            outputs=[
+                dataset_config_input,
+                dataset_split_input
+            ],
+        )
+
 iface.queue(max_size=20)
 iface.launch()