Add files

app.py

@@ -57,7 +57,7 @@ def search_datasets_fn(query: str) -> List[str]:
         return []
 
 
-def load_episodes_for_dataset(repo_id: str):
+def load_episodes_for_dataset(repo_id: str, progress=gr.Progress()):
     """Download dataset and list available episodes"""
     if not repo_id:
         return ""
@@ -66,7 +66,10 @@ def load_episodes_for_dataset(repo_id: str):
     temp_dir = tempfile.mkdtemp(prefix="episode_delete_")
     
     try:
+        progress(0, desc="Downloading dataset...")
         download_dataset(repo_id, temp_dir, hf_token=token)
+        
+        progress(0.7, desc="Listing episodes...")
         episodes = list_episodes(temp_dir)
         
         # Cleanup temp directory
@@ -78,8 +81,16 @@ def load_episodes_for_dataset(repo_id: str):
         # Return info about available episodes
         return f"Found {len(episodes)} episodes: {', '.join(map(str, episodes))}"
     except Exception as e:
-        print(f"Error loading episodes: {e}")
-        return f"Error: {e}"
+        import traceback
+        error_msg = f"Error: {str(e)}\n{traceback.format_exc()}"
+        print(error_msg)
+        # Cleanup on error
+        try:
+            if temp_dir and os.path.exists(temp_dir):
+                shutil.rmtree(temp_dir, ignore_errors=True)
+        except Exception:
+            pass
+        return error_msg
 
 
 def delete_episodes_stream(repo_id: str, episode_indexes_str: str, dest_repo_id: str):
@@ -184,6 +195,14 @@ with gr.Blocks(title="LeRobot Episode Deleter") as demo:
     # Load initial datasets
     _initial_choices = search_datasets_fn("griffinlabs-cortex")
     
+    with gr.Row():
+        org_input = gr.Textbox(
+            label="Organization or keyword",
+            value="griffinlabs-cortex",
+            placeholder="e.g., lerobot, griffinlabs-cortex"
+        )
+        load_btn = gr.Button("Load Datasets")
+    
     dataset_dropdown = gr.Dropdown(
         label="Select dataset",
         choices=_initial_choices,
@@ -211,6 +230,16 @@ with gr.Blocks(title="LeRobot Episode Deleter") as demo:
     progress_log = gr.Textbox(label="Progress log", lines=20)
     
     # Event handlers
+    def load_datasets_from_org(org_name):
+        results = search_datasets_fn(org_name)
+        return gr.update(choices=results, value=None)
+    
+    load_btn.click(
+        load_datasets_from_org,
+        inputs=org_input,
+        outputs=dataset_dropdown,
+    )
+    
     dataset_dropdown.change(
         load_episodes_for_dataset,
         inputs=dataset_dropdown,