small

app.py

@@ -2,13 +2,42 @@ import gradio as gr
 from datasets import load_dataset as load_dataset_hf
 import os
 
+# Predefined dataset names (configs will be fetched dynamically)
+PREDEFINED_DATASETS = [
+    "aliangdw/rfm",
+    "abraranwar/libero_rfm",
+]
+
 def load_rfm_dataset(dataset_name, config_name):
     """Load the RFM dataset from HuggingFace Hub."""
     try:
+        # Validate inputs
+        if not dataset_name or not config_name:
+            return None, "❌ Please provide both dataset name and configuration"
+        
+        # Try to load the dataset
         dataset = load_dataset_hf(dataset_name, name=config_name, split="train")
+        
+        # Check if dataset has the expected structure
+        expected_features = ["task", "frames", "quality_label", "is_robot", "data_source"]
+        missing_features = [f for f in expected_features if f not in dataset.features]
+        
+        if missing_features:
+            return None, f"⚠️ Dataset loaded but missing expected features: {missing_features}"
+        
+        # Check if dataset has any samples
+        if len(dataset) == 0:
+            return None, f"⚠️ Dataset {dataset_name}/{config_name} is empty"
+        
         return dataset, f"✅ Loaded {len(dataset)} trajectories from {dataset_name}/{config_name}"
     except Exception as e:
-        return None, f"❌ Error loading dataset: {e}"
+        error_msg = str(e)
+        if "not found" in error_msg.lower():
+            return None, f"❌ Dataset or configuration not found: {dataset_name}/{config_name}"
+        elif "authentication" in error_msg.lower():
+            return None, f"❌ Authentication required for {dataset_name}"
+        else:
+            return None, f"❌ Error loading dataset: {error_msg}"
 
 def get_available_configs(dataset_name):
     """Get available configurations for a dataset."""
@@ -20,6 +49,24 @@ def get_available_configs(dataset_name):
         print(f"Error getting configs: {e}")
         return []
 
+def update_config_choices_with_custom(dataset_name):
+    """Update config choices by fetching from the dataset."""
+    if not dataset_name:
+        return gr.update(choices=[], value="")
+    
+    try:
+        # Always try to fetch configs from the dataset
+        configs = get_available_configs(dataset_name)
+        if configs:
+            current_value = configs[0]
+            return gr.update(choices=configs, value=current_value)
+        else:
+            return gr.update(choices=[], value="")
+    except Exception as e:
+        # If fetching fails, allow custom input
+        print(f"Warning: Could not fetch configs for {dataset_name}: {e}")
+        return gr.update(choices=[], value="")
+
 def visualize_trajectory(dataset, index, dataset_name=None):
     """
     Function to retrieve a trajectory and its metadata from the dataset.
@@ -32,7 +79,7 @@ def visualize_trajectory(dataset, index, dataset_name=None):
         
         # Get metadata
         task = item["task"]
-        quality_label = item["quality_label"]
+        optimal = item["optimal"]
         is_robot = item["is_robot"]
         data_source = item["data_source"]
         
@@ -56,7 +103,7 @@ def visualize_trajectory(dataset, index, dataset_name=None):
         
         **Language Task:** {task}
         
-        **Quality:** {quality_label}
+        **Optimal:** {optimal}
         
         **Data Type:** {'Robot' if is_robot else 'Human'}
         
@@ -78,20 +125,27 @@ with gr.Blocks(title="RFM Dataset Visualizer") as demo:
     
     # Dataset selection
     with gr.Row():
-        with gr.Column(scale=1):
-            dataset_name_input = gr.Textbox(
+        with gr.Column(scale=2):
+            dataset_name_input = gr.Dropdown(
+                choices=PREDEFINED_DATASETS,
                 value="aliangdw/rfm",
                 label="Dataset Name",
-                placeholder="username/dataset-name"
+                allow_custom_value=True,
+                placeholder="Select dataset or enter custom"
             )
         
-        with gr.Column(scale=1):
-            config_name_input = gr.Textbox(
-                value="libero_10",
+        with gr.Column(scale=2):
+            config_name_input = gr.Dropdown(
+                choices=[],
+                value="",
                 label="Configuration Name",
-                placeholder="config-name"
+                allow_custom_value=True,
+                placeholder="Select config or enter custom"
             )
         
+        with gr.Column(scale=1):
+            refresh_btn = gr.Button("🔄 Refresh Configs", variant="secondary", size="sm")
+        
         with gr.Column(scale=1):
             load_btn = gr.Button("Load Dataset", variant="primary")
     
@@ -196,7 +250,25 @@ with gr.Blocks(title="RFM Dataset Visualizer") as demo:
                 interactive=False
             )
     
+    def update_config_choices(dataset_name):
+        """Update the config dropdown choices based on selected dataset."""
+        return update_config_choices_with_custom(dataset_name)
+    
     # Connect the components
+    # Update config choices when dataset changes
+    dataset_name_input.change(
+        fn=update_config_choices,
+        inputs=[dataset_name_input],
+        outputs=[config_name_input]
+    )
+    
+    # Refresh configs button for custom datasets
+    refresh_btn.click(
+        fn=update_config_choices_with_custom,
+        inputs=[dataset_name_input],
+        outputs=[config_name_input]
+    )
+    
     load_btn.click(
         fn=load_dataset,
         inputs=[dataset_name_input, config_name_input],
@@ -237,9 +309,17 @@ with gr.Blocks(title="RFM Dataset Visualizer") as demo:
         outputs=slider
     )
     
-    # Load initial dataset
+    # Load initial dataset and configs
     demo.load(
-        fn=lambda: load_dataset("aliangdw/rfm", "libero_10"),
+        fn=lambda: ("aliangdw/rfm", "libero_10"),  # Set initial values
+        outputs=[dataset_name_input, config_name_input]
+    ).then(
+        fn=update_config_choices_with_custom,
+        inputs=[dataset_name_input],
+        outputs=[config_name_input]
+    ).then(
+        fn=lambda dataset_name, config_name: load_dataset(dataset_name, config_name),
+        inputs=[dataset_name_input, config_name_input],
         outputs=[current_dataset, status_output, dataset_info, current_index]
     ).then(
         fn=update_slider_range,