small

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from datasets import load_dataset as load_dataset_hf
+from datasets import load_dataset as load_dataset_hf, get_dataset_config_names
 import os
 
 # Predefined dataset names (configs will be fetched dynamically)
@@ -42,11 +42,11 @@ def load_rfm_dataset(dataset_name, config_name):
 def get_available_configs(dataset_name):
     """Get available configurations for a dataset."""
     try:
-        dataset_info = load_dataset_hf(dataset_name, trust_remote_code=True)
-        configs = list(dataset_info.keys())
+        # Use the dedicated function to get config names
+        configs = get_dataset_config_names(dataset_name)
         return configs
     except Exception as e:
-        print(f"Error getting configs: {e}")
+        print(f"Error getting configs for {dataset_name}: {e}")
         return []
 
 def update_config_choices_with_custom(dataset_name):
@@ -79,7 +79,7 @@ def visualize_trajectory(dataset, index, dataset_name=None):
         
         # Get metadata
         task = item["task"]
-        optimal = item["optimal"]
+        quality_label = item["quality_label"]
         is_robot = item["is_robot"]
         data_source = item["data_source"]
         
@@ -103,7 +103,7 @@ def visualize_trajectory(dataset, index, dataset_name=None):
         
         **Language Task:** {task}
         
-        **Optimal:** {optimal}
+        **Quality Label:** {quality_label}
         
         **Data Type:** {'Robot' if is_robot else 'Human'}