Upload folder using huggingface_hub

example_training_dataset.csv

@@ -0,0 +1,4 @@
+Anchor,Positive,Negative
+MY_FAVORITE_NEWS,Denial of service and source code exposure in React Server Components,An SVG is all you need
+MY_FAVORITE_NEWS,The highest quality codebase,Litestream VFS
+

src/config.py

@@ -45,9 +45,6 @@ class AppConfig:
     # Anchor text used for contrastive learning dataset generation
     QUERY_ANCHOR: Final[str] = "MY_FAVORITE_NEWS"
 
-    # Number of titles shown for user selection in the Gradio interface
-    TOP_TITLES_COUNT: Final[int] = 10
-
     # Default export path for the dataset CSV
     DATASET_EXPORT_FILENAME: Final[Path] = ARTIFACTS_DIR.joinpath("training_dataset.csv")
 

src/session_manager.py

@@ -45,8 +45,6 @@ class HackerNewsFineTuner:
         self.model: Optional[SentenceTransformer] = None
         self.vibe_checker: Optional[VibeChecker] = None
         self.titles: List[str] = [] 
-        self.target_titles: List[str] = [] 
-        self.number_list: List[int] = [] 
         self.last_hn_dataset: List[List[str]] = [] 
         self.imported_dataset: List[List[str]] = [] 
 
@@ -66,9 +64,12 @@ class HackerNewsFineTuner:
 
     ## Data and Model Management ##
 
-    def refresh_data_and_model(self) -> Tuple[gr.update, gr.update]:
+    def refresh_data_and_model(self) -> Tuple[List[str], str]:
         """
         Reloads model and fetches data.
+        Returns:
+            - List of titles (for the UI)
+            - Status message string
         """
         print(f"[{self.session_id}] Reloading model and data...")
 
@@ -84,34 +85,23 @@ class HackerNewsFineTuner:
             print(error_msg)
             self.model = None
             self._update_vibe_checker()
-            return (
-                gr.update(choices=[], label="Model Load Failed"),
-                gr.update(value=error_msg)
-            )
+            return [], error_msg
 
         # 2. Fetch fresh news data
         news_feed, status_msg = read_hacker_news_rss(self.config)
-        titles_out, target_titles_out = [], []
+        titles_out = []
         status_value: str = f"Ready. Session ID: {self.session_id[:8]}... | Status: {status_msg}"
 
         if news_feed is not None and news_feed.entries:
-            titles_out = [item.title for item in news_feed.entries[:self.config.TOP_TITLES_COUNT]]
-            target_titles_out = [item.title for item in news_feed.entries[self.config.TOP_TITLES_COUNT:]]
+            titles_out = [item.title for item in news_feed.entries]
         else:
-            titles_out = ["Error fetching news.", "Check console."]
+            titles_out = ["Error fetching news."]
             gr.Warning(f"Data reload failed. {status_msg}")
 
         self.titles = titles_out
-        self.target_titles = target_titles_out
-        self.number_list = list(range(len(self.titles)))
-
-        return (
-            gr.update(
-                choices=self.titles,
-                label=f"Hacker News Top {len(self.titles)} (Select your favorites)"
-            ),
-            gr.update(value=status_value)
-        )
+
+        # Return raw list of titles + status text
+        return self.titles, status_value
 
     # --- Import Dataset/Export ---
     def import_additional_dataset(self, file_path: str) -> str:
@@ -123,6 +113,7 @@ class HackerNewsFineTuner:
                 reader = csv.reader(f)
                 try:
                     header = next(reader)
+                    # Simple heuristic to detect if header exists
                     if not (header and header[0].lower().strip() == 'anchor'):
                         f.seek(0)
                 except StopIteration:
@@ -188,54 +179,75 @@ class HackerNewsFineTuner:
 
 
     ## Training Logic ##
-    def _create_hn_dataset(self, selected_ids: List[int]) -> Tuple[List[List[str]], str, str]:
-        total_ids, selected_ids = set(self.number_list), set(selected_ids)
-        non_selected_ids = total_ids - selected_ids
-        is_minority = len(selected_ids) < (len(total_ids) / 2)
-
-        anchor_ids, pool_ids = (non_selected_ids, list(selected_ids)) if is_minority else (selected_ids, list(non_selected_ids))
-
-        def get_titles(anchor_id, pool_id):
-            return (self.titles[pool_id], self.titles[anchor_id]) if is_minority else (self.titles[anchor_id], self.titles[pool_id])
+    def _create_hn_dataset(self, pos_ids: List[int], neg_ids: List[int]) -> List[List[str]]:
+        """
+        Creates triplets (Anchor, Positive, Negative) from the selected indices.
+        Uses cycling to balance the dataset if the number of positives != negatives.
+        """
+        if not pos_ids or not neg_ids:
+            return []
 
-        if not pool_ids or not anchor_ids:
-             return [], "", "" 
+        # Convert indices to actual title strings
+        pos_titles = [self.titles[i] for i in pos_ids]
+        neg_titles = [self.titles[i] for i in neg_ids]
 
-        fav_idx = pool_ids[0] if is_minority else list(anchor_ids)[0]
-        non_fav_idx = list(anchor_ids)[0] if is_minority else pool_ids[0]
+        dataset = []
 
-        hn_dataset = []
-        pool_cycler = cycle(pool_ids)
-        for anchor_id in sorted(list(anchor_ids)):
-            fav, non_fav = get_titles(anchor_id, next(pool_cycler))
-            hn_dataset.append([self.config.QUERY_ANCHOR, fav, non_fav])
+        # We need to pair every Positive with a Negative.
+        # Strategy: Iterate over the longer list and cycle through the shorter list
+        # to ensure every selected item is used at least once and the dataset is balanced.
+        
+        if len(pos_titles) >= len(neg_titles):
+            # More positives than negatives: Iterate positives, reuse negatives
+            neg_cycle = cycle(neg_titles)
+            for p_title in pos_titles:
+                dataset.append([self.config.QUERY_ANCHOR, p_title, next(neg_cycle)])
+        else:
+            # More negatives than positives: Iterate negatives, reuse positives
+            pos_cycle = cycle(pos_titles)
+            for n_title in neg_titles:
+                dataset.append([self.config.QUERY_ANCHOR, next(pos_cycle), n_title])
 
-        return hn_dataset, self.titles[fav_idx], self.titles[non_fav_idx]
+        return dataset
 
-    def training(self, selected_ids: List[int]) -> str:
+    def training(self, pos_ids: List[int], neg_ids: List[int]) -> str:
+        """
+        Main training entry point.
+        Args:
+            pos_ids: Indices of stories marked as "Favorite"
+            neg_ids: Indices of stories marked as "Dislike"
+        """
         if self.model is None:
              raise gr.Error("Model not loaded.")
-        if not selected_ids:
-            raise gr.Error("Select at least one title.")
-        if len(selected_ids) == len(self.number_list):
-            raise gr.Error("Cannot select all titles.")
-
-        hn_dataset, _, _ = self._create_hn_dataset(selected_ids)
-        self.last_hn_dataset = hn_dataset
-        final_dataset = self.last_hn_dataset + self.imported_dataset
         
-        if not final_dataset:
-            raise gr.Error("Dataset is empty.")
+        # Validation
+        if not pos_ids:
+            raise gr.Error("Please select at least one 'Favorite' story.")
+        if not neg_ids:
+            raise gr.Error("Please select at least one 'Dislike' story.")
+        
+        # Generate Dataset
+        hn_dataset = self._create_hn_dataset(pos_ids, neg_ids)
+        
+        # Merge with imported dataset if it exists
+        if self.imported_dataset:
+            # If we have both, combine them
+            self.last_hn_dataset = hn_dataset + self.imported_dataset
+        else:
+            self.last_hn_dataset = hn_dataset
+                    
+        if not self.last_hn_dataset:
+            raise gr.Error("Dataset generation failed (Empty dataset).")
 
         def semantic_search_fn() -> str:
-            return get_top_hits(model=self.model, target_titles=self.target_titles, task_name=self.config.TASK_NAME, query=self.config.QUERY_ANCHOR)
+            return get_top_hits(model=self.model, target_titles=self.titles, task_name=self.config.TASK_NAME, query=self.config.QUERY_ANCHOR)
 
         result = "### Search (Before):\n" + f"{semantic_search_fn()}\n\n"
-        print(f"[{self.session_id}] Starting Training...")
+        print(f"[{self.session_id}] Starting Training with {len(self.last_hn_dataset)} examples...")
         
         train_with_dataset(
             model=self.model, 
-            dataset=final_dataset, 
+            dataset=self.last_hn_dataset, 
             output_dir=self.output_dir, 
             task_name=self.config.TASK_NAME, 
             search_fn=semantic_search_fn
@@ -247,6 +259,9 @@ class HackerNewsFineTuner:
         result += "### Search (After):\n" + f"{semantic_search_fn()}"
         return result
 
+    def is_model_tuned(self) -> bool:
+        return True if self.last_hn_dataset else False
+
     ## Vibe Check Logic ##
     def get_vibe_check(self, news_text: str) -> Tuple[str, str, gr.update]:
         info_text = f"**Session:** {self.session_id[:6]}<br>**Model:** `{self.config.MODEL_NAME}`{' - Fine-tuned' if self.last_hn_dataset else ''}"
@@ -303,4 +318,4 @@ class HackerNewsFineTuner:
                    f"| [{item['title']}]({item['link']}) "
                    f"| [Comments]({item['comments']}) "
                    f"| {item['published']} |\n")
-        return md
+        return md

src/ui.py

@@ -1,26 +1,64 @@
 import gradio as gr
-from typing import Optional
+from typing import Optional, Dict, List
 from datetime import datetime
 
 from .config import AppConfig
 from .session_manager import HackerNewsFineTuner
 
+# --- Constants for Labels ---
+LABEL_FAV = "👍"
+LABEL_NEU = "😐"
+LABEL_DIS = "👎"
+
 # --- Session Wrappers ---
 
 def refresh_wrapper(app):
     """
     Initializes the session if it's not already created, then runs the refresh.
-    Returns the app instance to update the State.
+    Returns:
+    1. App instance
+    2. Stories List (List[str])
+    3. Empty Labels Dict (Dict)
+    4. Log text
     """
     if app is None or callable(app) or isinstance(app, type):
         print("Initializing new HackerNewsFineTuner session...")
         app = HackerNewsFineTuner(AppConfig)
     
     # Run the refresh logic
-    update1, update2 = app.refresh_data_and_model()
+    # choices_list is a simple list of strings: ["Title 1", "Title 2", ...]
+    choices_list, log_update = app.refresh_data_and_model()
+    
+    # Reset user labels
+    empty_labels = {}
+    
+    return app, choices_list, empty_labels, log_update
+
+def on_app_load(app, profile: Optional[gr.OAuthProfile] = None):
+    """
+    Combined wrapper for initial load:
+    1. Initializes/Refreshes App Session
+    2. Checks OAuth Profile to enable/disable Hub features
+    """
+    # 1. Reuse the logic from refresh_wrapper
+    app, stories, labels, text_update = refresh_wrapper(app)
     
-    # Return 3 items: The App Instance (for State), Choice Update, Text Update
-    return app, update1, update2
+    # 2. Check Login Status
+    is_logged_in = profile is not None
+    username = profile.username if is_logged_in else None
+    
+    hub_interactive = gr.update(interactive=is_logged_in)
+    
+    # Return items matching the output signature of demo.load
+    return app, stories, labels, text_update, hub_interactive, hub_interactive, username
+
+def update_repo_preview(username, repo_name):
+    """Updates the markdown preview to show 'username/repo_name'."""
+    if not username:
+        return "⚠️ Sign in to see the target repository path."
+    
+    clean_repo = repo_name.strip() if repo_name else "..."
+    return f"Target Repository: **`{username}/{clean_repo}`**"
 
 def import_wrapper(app, file):
     return app.import_additional_dataset(file)
@@ -32,19 +70,31 @@ def download_model_wrapper(app):
     return app.download_model()
 
 def push_to_hub_wrapper(app, repo_name, oauth_token: Optional[gr.OAuthToken]):
-    """
-    Wrapper for pushing the model to the Hugging Face Hub.
-    Gradio automatically injects 'oauth_token' if the user is logged in via LoginButton.
-    """
     if oauth_token is None:
         return "⚠️ You must be logged in to push to the Hub. Please sign in above."
-    
-    # Extract the token string from the OAuthToken object
     token_str = oauth_token.token
     return app.upload_model(repo_name, token_str)
 
-def training_wrapper(app, selected_ids):
-    return app.training(selected_ids)
+def training_wrapper(app, stories: List[str], labels: Dict[int, str]):
+    """
+    Parses the Stories and Labels to extract Positive and Negative indices.
+    stories: List of titles
+    labels: Dictionary of {index: LABEL_FAV | LABEL_DIS | LABEL_NEU}
+    """
+    pos_ids = []
+    neg_ids = []
+    
+    # Iterate through all available stories by index
+    for i in range(len(stories)):
+        # Get label for this index, default to Neutral if not set
+        label = labels.get(i, LABEL_NEU)
+        
+        if label == LABEL_FAV:
+            pos_ids.append(i)
+        elif label == LABEL_DIS:
+            neg_ids.append(i)
+            
+    return app.training(pos_ids, neg_ids)
 
 def vibe_check_wrapper(app, text):
     return app.get_vibe_check(text)
@@ -57,124 +107,248 @@ def mood_feed_wrapper(app):
 
 def build_interface() -> gr.Blocks:
     with gr.Blocks(title="EmbeddingGemma Modkit") as demo:
-        # Initialize state as None. It will be populated by refresh_wrapper on load.
         session_state = gr.State()
+        username_state = gr.State()
+        
+        # State variables for the Feed List and User Choices
+        stories_state = gr.State([]) 
+        labels_state = gr.State({})
+        reset_counter = gr.State(0)
 
         with gr.Column():
             gr.Markdown("# 🤖 EmbeddingGemma Modkit: Fine-Tuning and Mood Reader")
             gr.Markdown("This project provides a set of tools to fine-tune [EmbeddingGemma](https://huggingface.co/google/embeddinggemma-300m) to understand your personal taste in Hacker News titles and then use it to score and rank new articles based on their \"vibe\". The core idea is to measure the \"vibe\" of a news title by calculating the semantic similarity between its embedding and the embedding of a fixed anchor phrase, **`MY_FAVORITE_NEWS`**.<br>See [README](https://huggingface.co/spaces/google/embeddinggemma-modkit/blob/main/README.md) for more details.")
-            gr.LoginButton(value="(Optional) Sign in to Hugging Face, if you want to push fine-tuned model to your repo.")
         
         with gr.Tab("🚀 Fine-Tuning & Evaluation"):
-            with gr.Column():
-                gr.Markdown("## Fine-Tuning & Semantic Search\nSelect titles to fine-tune the model towards making them more similar to **`MY_FAVORITE_NEWS`**.")
-                with gr.Row():
-                    favorite_list = gr.CheckboxGroup(choices=[], type="index", label="Hacker News Top Stories", show_select_all=True)
-                    output = gr.Textbox(lines=14, label="Training and Search Results", value="Loading data...")
-                
+            
+            # --- Model Indicator ---
+            gr.Dropdown(
+                choices=[f"{AppConfig.MODEL_NAME}"], 
+                value=f"{AppConfig.MODEL_NAME}", 
+                label="Base Model for Fine-tuning", 
+                interactive=False
+            )
+        
+            # --- Step 0: Login ---
+            with gr.Accordion("0️⃣ Step 0: Sign In (Optional)", open=True):
+                gr.Markdown("Sign in to Hugging Face if you plan to push your fine-tuned model to the Hub later (Step 3).")
                 with gr.Row():
-                    clear_reload_btn = gr.Button("Clear & Reload")
-                    run_training_btn = gr.Button("🚀 Run Fine-Tuning", variant="primary")
+                    gr.LoginButton(value="Sign in with Hugging Face")
+                    with gr.Column(scale=3):
+                        gr.Markdown("")
+            
+            # --- Step 1: Data Selection ---
+            with gr.Accordion("1️⃣ Step 1: Select Data Source", open=True):
+                gr.Markdown("Select titles from the live Hacker News feed **OR** upload your own CSV dataset to prepare your training data.")
                 
-                gr.Markdown("--- \n ## Dataset & Model Management")
-                gr.Markdown("To train on your own data, upload a CSV file with the following columns (no header required, or header ignored if present):\n1. **Anchor**: A fixed anchor phrase, `MY_FAVORITE_NEWS`.\n2. **Positive**: A title or contents that you like.\n3. **Negative**: A title or contents that you don't like.\n\nExample CSV Row:\n```\nMY_FAVORITE_NEWS,What is machine learning?,How to write a compiler from scratch.\n```")
-                import_file = gr.File(label="Upload Additional Dataset (.csv)", file_types=[".csv"], height=50)
+                with gr.Column():
+                    # Option A: Live Feed (Radio List)
+                    with gr.Accordion("Option A: Live Hacker News Feed", open=True):
+                        gr.Markdown("Rate the stories below to define your vibe.\n\n**⚠️ Note: You must select at least one Favorite and one Dislike to run training.**")
+                        
+                        with gr.Row():
+                            reset_all_btn = gr.Button("Reset Selection ↺", variant="secondary", scale=1)
+                            with gr.Column(scale=3):
+                                gr.Markdown("")
+                        
+                        # Dynamic rendering of the story list
+                        @gr.render(inputs=[stories_state, reset_counter])
+                        def render_story_list(stories, _counter):
+                            if not stories:
+                                gr.Markdown("*No stories loaded. Click 'Reset Model & Fine-tuning state' to fetch data.*")
+                                return
+                            
+                            for i, title in enumerate(stories):
+                                with gr.Row(variant="compact", elem_id=f"story_row_{i}"):
+                                    # Title
+                                    with gr.Column(scale=3):
+                                    	gr.Markdown(f"**{i+1}.** {title}")
+                                    
+                                    # Radio Selection
+                                    radio = gr.Radio(
+                                        choices=[LABEL_FAV, LABEL_NEU, LABEL_DIS],
+                                        value=LABEL_NEU,
+                                        show_label=False,
+                                        container=False,
+                                        min_width=80,
+                                        scale=1,
+                                        interactive=True
+                                    )
+                                    
+                                    # Update logic
+                                    def update_label(new_val, current_labels, idx=i):
+                                        current_labels[idx] = new_val
+                                        return current_labels
+
+                                    radio.change(
+                                        fn=update_label,
+                                        inputs=[radio, labels_state],
+                                        outputs=[labels_state]
+                                    )
+
+                    # Option B: Upload
+                    with gr.Accordion("Option B: Upload Custom Dataset", open=False):
+                        gr.Markdown("Upload a CSV file with columns (no header required, or header ignored if present): `Anchor`, `Positive`, `Negative`.")
+                        gr.Markdown("See also: [example_training.dataset.csv](https://huggingface.co/spaces/google/embeddinggemma-modkit/blob/main/example_training.dataset.csv)<br>Example:<br>`MY_FAVORITE_NEWS,Good Title,Bad Title`")
+                        import_file = gr.File(label="Upload Additional Dataset (.csv)", file_types=[".csv"], height=100)
+
+            # --- Step 2: Training ---
+            with gr.Accordion("2️⃣ Step 2: Run Tuning", open=True):
+                gr.Markdown("Fine-tune the model using the data selected or uploaded above.")
                 
                 with gr.Row():
-                    download_dataset_btn = gr.Button("💾 Export Dataset")
-                    download_model_btn = gr.Button("⬇️ Download Fine-Tuned Model")
-                
-                download_status = gr.Markdown("Ready.")
+                    run_training_btn = gr.Button("🚀 Run Fine-Tuning", variant="primary", scale=1)
+                    clear_reload_btn = gr.Button("Reset Model & Fine-tuning state", scale=1)
                 
-                with gr.Row():
-                    dataset_output = gr.File(label="Download Dataset CSV", height=50, visible=False, interactive=False)
-                    model_output = gr.File(label="Download Model ZIP", height=50, visible=False, interactive=False)
+                output = gr.Textbox(lines=10, label="Training Logs & Search Results", value="Waiting to start...", autoscroll=True)
 
-                gr.Markdown("### ☁️ Publish to Hugging Face Hub")
+            # --- Step 3: Push to Hub ---
+            with gr.Accordion("3️⃣ Step 3: Save to Hugging Face Hub (Optional)", open=False):
+                gr.Markdown("Push your fine-tuned model to your personal Hugging Face account.")
+                
                 with gr.Row():
-                    repo_name_input = gr.Textbox(label="Target Repository Name", placeholder="e.g., my-news-vibe-model")
-                    push_to_hub_btn = gr.Button("Push to Hub", variant="secondary")
+                    repo_name_input = gr.Textbox(label="Target Repository Name", value="my-embeddinggemma-news-vibe", placeholder="e.g., my-embeddinggemma-news-vibe", interactive=False)
+                    push_to_hub_btn = gr.Button("Save to Hugging Face Hub", variant="secondary", interactive=False)
+                
+                repo_id_preview = gr.Markdown("Target Repository: (Waiting for input...)")
                 
                 push_status = gr.Markdown("")
 
-                # --- Interactions ---
-                
-                # 1. Initial Load: Initialize State and Load Data
-                demo.load(
-                    fn=refresh_wrapper, 
-                    inputs=[session_state], 
-                    outputs=[session_state, favorite_list, output]
-                )
+            # --- Step 4: Downloads ---
+            with gr.Accordion("4️⃣ Step 4: Download Artifacts", open=False):
+                gr.Markdown("Export your combined dataset or download the fine-tuned model locally.")
+
+                with gr.Row():
+                    download_dataset_btn = gr.Button("💾 Export Dataset", interactive=False)
+                    download_model_btn = gr.Button("⬇️ Download Model ZIP", interactive=False)
                 
-                buttons_to_lock = [
-                    clear_reload_btn,
-                    run_training_btn,
-                    download_dataset_btn,
-                    download_model_btn,
-                    push_to_hub_btn
-                ]
-
-                # 2. Buttons
-                clear_reload_btn.click(
-                    fn=lambda: [gr.update(interactive=False)]*len(buttons_to_lock),
-                    outputs=buttons_to_lock
-                ).then(
-                    fn=refresh_wrapper, 
-                    inputs=[session_state], 
-                    outputs=[session_state, favorite_list, output]
-                ).then(
-                    fn=lambda: [gr.update(interactive=True)]*len(buttons_to_lock),
-                    outputs=buttons_to_lock
-                )
+                download_status = gr.Markdown("Ready.")
                 
-                run_training_btn.click(
-                    fn=lambda: [gr.update(interactive=False)]*len(buttons_to_lock),
-                    outputs=buttons_to_lock
-                ).then(
-                    fn=training_wrapper, 
-                    inputs=[session_state, favorite_list], 
-                    outputs=[output]
-                ).then(
-                    fn=lambda: [gr.update(interactive=True)]*len(buttons_to_lock),
-                    outputs=buttons_to_lock
-                )
+                with gr.Row():
+                    dataset_output = gr.File(label="Download Dataset CSV", height=50, visible=False, interactive=False)
+                    model_output = gr.File(label="Download Model ZIP", height=50, visible=False, interactive=False)
 
-                import_file.change(
-                    fn=import_wrapper, 
-                    inputs=[session_state, import_file], 
-                    outputs=[download_status]
-                )
+            # --- Interaction Logic ---
+            
+            action_buttons = [
+                clear_reload_btn,
+                run_training_btn,
+                download_dataset_btn,
+                download_model_btn
+            ]
+            
+            def set_interactivity(interactive: bool):
+                """Helper to lock/unlock all main action buttons."""
+                return [gr.update(interactive=interactive) for _ in action_buttons]
+            
+            # 1. App Startup
+            # ----------------
+            demo.load(
+                fn=lambda: set_interactivity(False), outputs=action_buttons
+            ).then(
+                fn=on_app_load, 
+                inputs=[session_state], 
+                outputs=[session_state, stories_state, labels_state, output, repo_name_input, push_to_hub_btn, username_state]
+            ).then(
+                fn=update_repo_preview,
+                inputs=[username_state, repo_name_input],
+                outputs=[repo_id_preview]
+            ).then(
+                fn=lambda: [gr.update(interactive=True)]*2, outputs=[clear_reload_btn, run_training_btn]
+            )
+            
+            # 2. Reset / Refresh / Clear Selections
+            # ----------------
+            clear_reload_btn.click(
+                fn=lambda: set_interactivity(False), outputs=action_buttons
+            ).then(
+                fn=refresh_wrapper, 
+                inputs=[session_state], 
+                outputs=[session_state, stories_state, labels_state, output]
+            ).then(
+                fn=lambda: [gr.update(interactive=True)]*2, outputs=[clear_reload_btn, run_training_btn]
+            )
+            
+            # Reset Selection Button Logic
+            def reset_all_selections(counter):
+                # Returns: (incremented counter, empty dict for labels)
+                return counter + 1, {}
 
-                download_dataset_btn.click(
-                    fn=export_wrapper,
-                    inputs=[session_state],
-                    outputs=[dataset_output]
-                ).then(
-                    lambda p: gr.update(visible=True) if p else gr.update(), inputs=[dataset_output], outputs=[dataset_output]
-                )
+            reset_all_btn.click(
+                fn=reset_all_selections,
+                inputs=[reset_counter],
+                outputs=[reset_counter, labels_state]
+            )
+            
+            # 3. Import Data
+            # ----------------
+            import_file.change(
+                fn=import_wrapper, 
+                inputs=[session_state, import_file], 
+                outputs=[download_status]
+            )
+            
+            # 4. Run Training
+            # ----------------
+            run_training_btn.click(
+                fn=lambda: set_interactivity(False), outputs=action_buttons
+            ).then(
+                fn=training_wrapper, 
+                inputs=[session_state, stories_state, labels_state], 
+                outputs=[output]
+            ).then(
+                # Unlock all buttons (including downloads now that we have a model)
+                fn=lambda: set_interactivity(True), outputs=action_buttons
+            )
+            
+            # 5. Downloads
+            # ----------------
+            download_dataset_btn.click(
+                fn=export_wrapper,
+                inputs=[session_state],
+                outputs=[dataset_output]
+            ).then(
+                # Just show the file output if it exists
+                lambda p: gr.update(visible=True) if p else gr.update(), 
+                inputs=[dataset_output], 
+                outputs=[dataset_output]
+            )
 
-                download_model_btn.click(
-                    fn=lambda: [gr.update(interactive=False)]*len(buttons_to_lock),
-                    outputs=buttons_to_lock
-                ).then(
-                    lambda: [gr.update(value=None, visible=False), "Zipping..."], None, [model_output, download_status], queue=False
-                ).then(
-                    fn=download_model_wrapper,
-                    inputs=[session_state],
-                    outputs=[model_output]
-                ).then(
-                    lambda p: [gr.update(visible=p is not None, value=p), "ZIP ready." if p else "Zipping failed."], [model_output], [model_output, download_status]
-                ).then(
-                    fn=lambda: [gr.update(interactive=True)]*len(buttons_to_lock),
-                    outputs=buttons_to_lock
-                )
+            download_model_btn.click(
+                # Lock UI
+                fn=lambda: set_interactivity(False), outputs=action_buttons
+            ).then(
+                # Reset previous outputs and show "Zipping..."
+                fn=lambda: [gr.update(value=None, visible=False), "⏳ Zipping model..."], 
+                outputs=[model_output, download_status]
+            ).then(
+                # Generate Zip
+                fn=download_model_wrapper,
+                inputs=[session_state],
+                outputs=[model_output]
+            ).then(
+                # Update UI with result
+                fn=lambda p: [gr.update(visible=p is not None, value=p), "✅ ZIP ready." if p else "❌ Zipping failed."], 
+                inputs=[model_output], 
+                outputs=[model_output, download_status]
+            ).then(
+                # Unlock UI
+                fn=lambda: set_interactivity(True), outputs=action_buttons
+            )
+            
+            # 6. Push to Hub
+            # ----------------
+            repo_name_input.change(
+                fn=update_repo_preview,
+                inputs=[username_state, repo_name_input],
+                outputs=[repo_id_preview]
+            )
 
-                # Push to Hub Interaction
-                push_to_hub_btn.click(
-                    fn=push_to_hub_wrapper,
-                    inputs=[session_state, repo_name_input],
-                    outputs=[push_status]
-                )
+            push_to_hub_btn.click(
+                fn=push_to_hub_wrapper,
+                inputs=[session_state, repo_name_input],
+                outputs=[push_status]
+            )
 
         with gr.Tab("📰 Hacker News Similarity Check"):
             with gr.Column():