add load dataset button

app.py

@@ -198,6 +198,75 @@ def validate_request(request: GenerationRequest, oauth_token: Optional[Union[gr.
     return request
 
 
+def load_dataset_info(dataset_name, dataset_token=None):
+    """Load dataset information and return choices for dropdowns"""
+    if not dataset_name.strip():
+        return (
+            gr.update(choices=[], value=None),  # config
+            gr.update(choices=[], value=None),  # split
+            gr.update(choices=[], value=None),  # prompt_column
+            gr.update(value="", interactive=True),  # output_dataset_name
+            gr.update(interactive=False),  # num_output_samples
+            "Please enter a dataset name first."
+        )
+
+    try:
+        # Get dataset info
+        dataset_infos = get_dataset_infos(dataset_name, token=dataset_token)
+
+        if not dataset_infos:
+            raise Exception("No configs found for this dataset")
+
+        # Get available configs
+        config_choices = list(dataset_infos.keys())
+        default_config = config_choices[0] if config_choices else None
+
+        # Get splits and features for the default config
+        if default_config:
+            config_info = dataset_infos[default_config]
+            split_choices = list(config_info.splits.keys())
+            default_split = split_choices[0] if split_choices else None
+
+            # Get column choices (features)
+            column_choices = list(config_info.features.keys())
+            default_column = None
+            # Try to find a likely prompt column
+            for col in column_choices:
+                if any(keyword in col.lower() for keyword in ['prompt', 'text', 'question', 'input']):
+                    default_column = col
+                    break
+            if not default_column and column_choices:
+                default_column = column_choices[0]
+        else:
+            split_choices = []
+            column_choices = []
+            default_split = None
+            default_column = None
+
+        # Generate a suggested output dataset name
+        dataset_base_name = dataset_name.split('/')[-1] if '/' in dataset_name else dataset_name
+        suggested_output_name = f"{dataset_base_name}-synthetic"
+
+        return (
+            gr.update(choices=config_choices, value=default_config, interactive=True),  # config
+            gr.update(choices=split_choices, value=default_split, interactive=True),  # split
+            gr.update(choices=column_choices, value=default_column, interactive=True),  # prompt_column
+            gr.update(value=suggested_output_name, interactive=True),  # output_dataset_name
+            gr.update(interactive=True),  # num_output_samples
+            f"✅ Dataset info loaded successfully! Found {len(config_choices)} config(s), {len(split_choices)} split(s), and {len(column_choices)} column(s)."
+        )
+
+    except Exception as e:
+        return (
+            gr.update(choices=[], value=None, interactive=False),  # config
+            gr.update(choices=[], value=None, interactive=False),  # split
+            gr.update(choices=[], value=None, interactive=False),  # prompt_column
+            gr.update(value="", interactive=False),  # output_dataset_name
+            gr.update(interactive=False),  # num_output_samples
+            f"❌ Error loading dataset info: {str(e)}"
+        )
+
+
 def add_request_to_db(request: GenerationRequest):
     url: str = os.getenv("SUPABASE_URL")
     key: str = os.getenv("SUPABASE_KEY")
@@ -342,8 +411,9 @@ def main():
                         with gr.Row():
                             with gr.Column():
                                 input_dataset_name = gr.Textbox(label="Input Dataset Name", placeholder="e.g., simplescaling/s1K-1.1")
-                                prompt_column = gr.Textbox(label="Prompt Column", placeholder="e.g., text, prompt, question", value=None, interactive=False, info="Click Load Info to populate")
-                                
+                                load_info_btn = gr.Button("📊 Load Dataset Info", size="sm", variant="secondary")
+                                load_info_status = gr.Markdown("", visible=True)
+
                             with gr.Column():
                                 output_dataset_name = gr.Textbox(label="Output Dataset Name", placeholder="e.g., my-generated-dataset, must be unique. Will be created under the org 'synthetic-data-universe'", value=None, interactive=False, info="Click Load Info to populate")
                         
@@ -519,6 +589,13 @@ def main():
                 except Exception as e:
                     return f"Error: {str(e)}"
             
+            # Wire up the Load Dataset Info button
+            load_info_btn.click(
+                load_dataset_info,
+                inputs=[input_dataset_name],
+                outputs=[input_dataset_config, input_dataset_split, prompt_column, output_dataset_name, num_output_samples, load_info_status]
+            )
+
             submit_btn.click(
                 submit_request,
                 inputs=[input_dataset_name, input_dataset_split, input_dataset_config, output_dataset_name, prompt_column, model_name_or_path,