Upload folder using huggingface_hub

app.py

@@ -2,36 +2,40 @@ import gradio as gr
 import os
 import json
 import torch
+import csv
+import shutil
+import time
+import threading
 
-from typing import Final, Optional, List
+from typing import Final, Optional, List, Any, Generator
 from pathlib import Path
+from dataclasses import dataclass
+
 from huggingface_hub import login
 from trl import SFTConfig, SFTTrainer
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import (
+    AutoTokenizer, 
+    AutoModelForCausalLM, 
+    TrainerCallback, 
+    TrainingArguments, 
+    TrainerControl, 
+    TrainerState
+)
 from datasets import Dataset, load_dataset
-from transformers.utils import get_json_schema
-
 
-ARTIFACTS_DIR: Final[Path] = Path("artifacts")
-
-def authenticate_hf(token: Optional[str]) -> None:
-    """Logs into the Hugging Face Hub."""
-    if token:
-        print("Logging into Hugging Face Hub...")
-        login(token=token)
-    else:
-        print("Skipping Hugging Face login: HF_TOKEN not set.")
+# --- Configuration ---
+class AppConfig:
+    """
+    Central configuration class.
+    """
+    ARTIFACTS_DIR: Final[Path] = Path("artifacts")
+    ARTIFACTS_DIR.mkdir(parents=True, exist_ok=True)
+    
+    HF_TOKEN: Final[Optional[str]] = os.getenv('HF_TOKEN')
+    MODEL_NAME: Final[str] = '../hf/270m'
+    DEFAULT_DATASET: Final[str] = 'bebechien/SimpleToolCalling'
+    OUTPUT_DIR: Final[Path] = ARTIFACTS_DIR.joinpath("functiongemma-modkit-demo")
 
-def load_model(model_name: str):
-    print(f"Loading Transformer model: {model_name}")
-    try:
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
-        model = AutoModelForCausalLM.from_pretrained(model_name)
-        print("Model loaded successfully.")
-        return (model, tokenizer)
-    except Exception as e:
-        print(f"Error loading Transformer model {model_name}: {e}")
-        raise
 
 # --- Tool Definitions ---
 def search_knowledge_base(query: str) -> str:
@@ -52,178 +56,213 @@ def search_google(query: str) -> str:
     """
     return "Public Result"
 
-
-TOOLS = [get_json_schema(search_knowledge_base), get_json_schema(search_google)]
-
-DEFAULT_SYSTEM_MSG = "You are a model that can do function calling with the following functions"
-
-
-def create_conversation(sample):
-  return {
-      "messages": [
-          {"role": "developer", "content": DEFAULT_SYSTEM_MSG},
-          {"role": "user", "content": sample["user_content"]},
-          {"role": "assistant", "tool_calls": [{"type": "function", "function": {"name": sample["tool_name"], "arguments": json.loads(sample["tool_arguments"])}}]},
-      ],
-      "tools": TOOLS
+search_knowledge_base_schema = {
+  "type": "function",
+  "function": {
+    "name": "search_knowledge_base",
+    "description": "Search internal company documents, policies and project data.",
+    "parameters": {
+      "type": "object",
+      "properties": {
+        "query": {
+          "type": "string",
+          "description": "query string"
+        }
+      },
+      "required": [
+        "query"
+      ]
+    },
+    "return": {
+      "type": "string"
+    }
   }
-
-
-def train_with_dataset(
-    model: AutoModelForCausalLM,
-    tokenizer: AutoTokenizer,
-    dataset: Dataset,
-    output_dir: Path,
-    learning_rate: float = 5e-5
-) -> None:
-
-    torch_dtype = model.dtype
-
-    args = SFTConfig(
-        output_dir=output_dir,                  # directory to save and repository id
-        max_length=512,                         # max sequence length for model and packing of the dataset
-        packing=False,                          # Groups multiple samples in the dataset into a single sequence
-        num_train_epochs=5,                     # number of training epochs
-        per_device_train_batch_size=4,          # batch size per device during training
-        gradient_checkpointing=False,           # Caching is incompatible with gradient checkpointing
-        optim="adamw_torch_fused",              # use fused adamw optimizer
-        logging_steps=1,                        # log every step
-        #save_strategy="epoch",                  # save checkpoint every epoch
-        eval_strategy="epoch",                  # evaluate checkpoint every epoch
-        learning_rate=learning_rate,            # learning rate
-        fp16=True if torch_dtype == torch.float16 else False,   # use float16 precision
-        bf16=True if torch_dtype == torch.bfloat16 else False,  # use bfloat16 precision
-        lr_scheduler_type="constant",           # use constant learning rate scheduler
-        push_to_hub=False,                       # push model to hub
-        report_to="none",                # report metrics to tensorboard
-        dataset_kwargs={
-            "add_special_tokens": False, # Template with special tokens
-            "append_concat_token": True, # Add EOS token as separator token between examples
+}
+
+search_google_schema = {
+  "type": "function",
+  "function": {
+    "name": "search_google",
+    "description": "Search public information.",
+    "parameters": {
+      "type": "object",
+      "properties": {
+        "query": {
+          "type": "string",
+          "description": "query string"
         }
-    )
+      },
+      "required": [
+        "query"
+      ]
+    },
+    "return": {
+      "type": "string"
+    }
+  }
+}
 
-    # Create Trainer object
-    trainer = SFTTrainer(
-        model=model,
-        args=args,
-        train_dataset=dataset['train'],
-        eval_dataset=dataset['test'],
-        processing_class=tokenizer,
-    )
+TOOLS = [search_knowledge_base_schema, search_google_schema]
+DEFAULT_SYSTEM_MSG = "You are a model that can do function calling with the following functions"
 
-    trainer.train()
+# --- Callbacks ---
+class AbortCallback(TrainerCallback):
+    """
+    A custom callback to check a threading Event to stop training on user request.
+    """
+    def __init__(self, stop_event: threading.Event):
+        self.stop_event = stop_event
 
-    print("Training finished. Model weights are updated in memory.")
+    def on_step_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
+        if self.stop_event.is_set():
+            print("🛑 Stop signal received. Stopping training...")
+            control.should_training_stop = True
 
-    # Save the final fine-tuned model
-    trainer.save_model()
 
-    print(f"Model saved locally to: {output_dir}")
+# --- Helper Functions ---
+def authenticate_hf(token: Optional[str]) -> None:
+    """Logs into the Hugging Face Hub."""
+    if token:
+        print("Logging into Hugging Face Hub...")
+        login(token=token)
+    else:
+        print("Skipping Hugging Face login: HF_TOKEN not set.")
 
-class AppConfig:
-    """
-    Central configuration class for the Fine-Tuner application.
-    """
-    ARTIFACTS_DIR: Final[Path] = ARTIFACTS_DIR
-    HF_TOKEN: Final[str | None] = os.getenv('HF_TOKEN')
-    MODEL_NAME: Final[str] = '../hf/270m'
-    DEFAULT_DATASET: Final[str] = 'bebechien/SimpleToolCalling'
-    OUTPUT_DIR: Final[Path] = ARTIFACTS_DIR.joinpath("functiongemma-270m-it-modkit-demo")
+def load_model_and_tokenizer(model_name: str):
+    print(f"Loading Transformer model: {model_name}")
+    try:
+        # Check if local path exists, otherwise treat as HF Hub ID
+        if model_name.startswith("..") and not os.path.exists(model_name):
+            print(f"Warning: Local path {model_name} not found. Falling back to default hub model.")
+            model_name = "google/gemma-2b-it" # Fallback example
 
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForCausalLM.from_pretrained(model_name)
+        print("Model loaded successfully.")
+        return model, tokenizer
+    except Exception as e:
+        print(f"Error loading Transformer model {model_name}: {e}")
+        raise e
 
+def create_conversation_format(sample):
+    """Formats a dataset row into the conversational format required for SFT."""
+    try:
+        tool_args = json.loads(sample["tool_arguments"])
+    except (json.JSONDecodeError, TypeError):
+        tool_args = {}
+        
+    return {
+        "messages": [
+            {"role": "developer", "content": DEFAULT_SYSTEM_MSG},
+            {"role": "user", "content": sample["user_content"]},
+            {"role": "assistant", "tool_calls": [{"type": "function", "function": {"name": sample["tool_name"], "arguments": tool_args}}]},
+        ],
+        "tools": TOOLS
+    }
+
+
+# --- Main Application Logic ---
 class FunctionGemmaTuner:
     def __init__(self, config: AppConfig = AppConfig):
         self.config = config
-
-        os.makedirs(self.config.ARTIFACTS_DIR, exist_ok=True)
-        print(f"Created artifact directory: {self.config.ARTIFACTS_DIR}")
+        self.model = None
+        self.tokenizer = None
+        self.imported_dataset = []
+        
+        # Threading event to control stopping
+        self.stop_event = threading.Event()
 
         authenticate_hf(self.config.HF_TOKEN)
-
-        self._initial_load()
-
-    def _initial_load(self):
-        """Helper to run the refresh function once at startup."""
+        
+        # Initial load attempt
         print("--- Running Initial Data Load ---")
-        self.refresh_data_and_model()
-        print("--- Initial Load Complete ---")
+        try:
+            self.refresh_data_and_model()
+            print("--- Initial Load Complete ---")
+        except Exception as e:
+            print(f"Initial load failed (this is common if model path is invalid): {e}")
 
     def refresh_data_and_model(self):
+        """Reloads the model and clears imported data."""
         print("\n" + "=" * 50)
         print("RELOADING MODEL and RE-FETCHING DATA")
 
-        # Reset dataset state
         self.imported_dataset = []
 
-        # 1. Reload the base model
         try:
-            self.model, self.tokenizer = load_model(self.config.MODEL_NAME)
+            self.model, self.tokenizer = load_model_and_tokenizer(self.config.MODEL_NAME)
+            status_value = "Model and data reloaded. Ready."
         except Exception as e:
-            gr.Error(f"Model load failed: {e}")
             self.model = None
             self.tokenizer = None
-            return gr.update(value=f"CRITICAL ERROR: Model failed to load. {e}")
-
-        status_value: str = f"Model and data reloaded. Click 'Run Fine-Tuning' to begin."
-
-        # Return Gradio updates for CheckboxGroup and Textbox
-        return gr.update(value=status_value)
+            status_value = f"CRITICAL ERROR: Model failed to load. {e}"
+            # We don't raise here to allow the UI to render the error message
+        
+        return status_value
 
-    # --- Import Dataset/Export ---
     def import_additional_dataset(self, file_path: str) -> str:
+        """Parses an uploaded CSV file."""
         if not file_path:
             return "Please upload a CSV file."
-        new_dataset, num_imported = [], 0
+        
+        new_dataset = []
+        num_imported = 0
+        
         try:
+            # Open file handle properly
             with open(file_path, 'r', newline='', encoding='utf-8') as f:
                 reader = csv.reader(f)
+                
+                # Basic header validation
                 try:
                     header = next(reader)
-                    if not (header and header[0].lower().strip() == 'anchor'):
+                    # Simple heuristic check, allows skipping header or rewinding
+                    if not (header and "anchor" in header[0].lower()):
                         f.seek(0)
                 except StopIteration:
                     return "Error: Uploaded file is empty."
 
                 for row in reader:
-                    if len(row) == 3:
-                        new_dataset.append([s.strip() for s in row])
+                    # Expecting: [User Prompt, Tool Name, Tool Args JSON/String]
+                    if len(row) >= 3:
+                        new_dataset.append([s.strip() for s in row[:3]])
                         num_imported += 1
+            
             if num_imported == 0:
-                raise ValueError("No valid [Anchor, Positive, Negative] rows found in the CSV.")
+                return "No valid rows found. CSV format: [Anchor, Positive, Negative]"
+            
             self.imported_dataset = new_dataset
-            return f"Successfully imported {num_imported} additional training triplets."
+            return f"Successfully imported {num_imported} additional training samples."
+            
         except Exception as e:
-            gr.Error(f"Import failed. Ensure the CSV format is: [Anchor, Positive, Negative]. Error: {e}")
-            return "Import failed. Check console for details."
+            return f"Import failed. Error: {e}"
 
-    def download_model(self) -> Optional[str]:
-        if not os.path.exists(self.config.OUTPUT_DIR):
-            gr.Warning(f"The model directory '{self.config.OUTPUT_DIR}' does not exist. Please run training first.")
-            return None
-        timestamp = int(time.time())
-        try:
-            base_name = os.path.join(self.config.ARTIFACTS_DIR, f"embedding_gemma_finetuned_{timestamp}")
-            archive_path = shutil.make_archive(
-                base_name=base_name,
-                format='zip',
-                root_dir=self.config.OUTPUT_DIR,
-            )
-            gr.Info(f"Model files successfully zipped to: {archive_path}")
-            return archive_path
-        except Exception as e:
-            gr.Error(f"Failed to create the model ZIP file. Error: {e}")
-            return None
+    def stop_training(self):
+        """Signal the training loop to stop."""
+        print("Set stop event")
+        self.stop_event.set()
+        return "Stopping initiated... please wait for the current step to finish."
 
-    def training(self, test_size: float = 0.5) -> str:
+    def run_training(self, test_size: float = 0.5) -> Generator[str, None, None]:
         """
-        Generates a training dataset from user selection and runs the fine-tuning process.
+        Main training logic. Yields status strings to the UI.
         """
+        # 1. Validation
         if self.model is None:
-             raise gr.Error("Training failed: Model is not loaded.")
+             yield "Training failed: Model is not loaded."
+             return
 
+        self.stop_event.clear() # Reset stop flag
+        yield "⏳ Preparing Dataset..."
+
+        # 2. Dataset Preparation
         if not self.imported_dataset:
-            print("No imported dataset, use the default")
-            dataset = load_dataset(self.config.DEFAULT_DATASET, split="train")
+            print("No imported dataset, using default HF dataset")
+            try:
+                dataset = load_dataset(self.config.DEFAULT_DATASET, split="train")
+            except Exception as e:
+                yield f"Error loading default dataset: {e}"
+                return
         else:
             dataset_as_dicts = [{
                 "user_content": row[0], "tool_name": row[1], "tool_arguments": row[2]}
@@ -231,81 +270,293 @@ class FunctionGemmaTuner:
             ]
             dataset = Dataset.from_list(dataset_as_dicts)
 
-        dataset = dataset.map(create_conversation, batched=False)
-        dataset = dataset.train_test_split(test_size=test_size, shuffle=False)
-        print(dataset)
-        print("--- dataset input ---")
-        print(json.dumps(dataset["train"][0], indent=2))
-        debug_msg = self.tokenizer.apply_chat_template(dataset["train"][0]["messages"], tools=dataset["train"][0]["tools"], add_generation_prompt=False, tokenize=False)
-        print("--- Formatted prompt ---")
-        print(debug_msg)
-
-        result = "### Success Rate (Before Training):\n" + f"{self.check_success_rate(dataset["test"])}\n\n"
-        print("-" * 50 + "\nStarting Fine-tuning...")
-        train_with_dataset(model=self.model, tokenizer=self.tokenizer, dataset=dataset, output_dir=self.config.OUTPUT_DIR)
-        print("Fine-tuning Complete.\n" + "-" * 50)
+        # Apply formatting
+        dataset = dataset.map(create_conversation_format, batched=False)
+        
+        # Split
+        if len(dataset) > 1:
+            dataset = dataset.train_test_split(test_size=test_size, shuffle=False)
+        else:
+            # Fallback for very small datasets (mostly for debugging)
+            dataset = {"train": dataset, "test": dataset}
+
+        output_buffer = "📊 Evaluating Pre-Training Success Rate...\n### Success Rate (Before Training):\n"
+        yield output_buffer
+        pre_training_report = ""
+        gen = self.check_success_rate(dataset["test"])
+        while not self.stop_event.is_set():
+            try:
+                pre_training_report += f"{next(gen)}\n"
+                yield f"{output_buffer}{pre_training_report}"
+            except StopIteration as e:
+                pre_training_report = e.value
+                break
+        
+        if self.stop_event.is_set():
+            output_buffer += f"{pre_training_report}\n\n🛑 Manual Eval interrupted by user.\n"
+            yield output_buffer
+            return
+        
+        output_buffer += f"{pre_training_report}\n\n"
+        output_buffer += "-" * 30 + "\nStarting Fine-tuning...\n"
+        yield output_buffer
+
+        # 3. Training Setup
+        torch_dtype = self.model.dtype
+        
+        args = SFTConfig(
+            output_dir=str(self.config.OUTPUT_DIR),
+            max_length=512,
+            packing=False,
+            num_train_epochs=5,
+            per_device_train_batch_size=4,
+            gradient_checkpointing=False,
+            optim="adamw_torch_fused",
+            logging_steps=1,
+            save_strategy="no", # Speed up demo
+            eval_strategy="epoch",
+            learning_rate=5e-5,
+            fp16=True if torch_dtype == torch.float16 else False,
+            bf16=True if torch_dtype == torch.bfloat16 else False,
+            lr_scheduler_type="constant",
+            push_to_hub=False,
+            report_to="none",
+            dataset_kwargs={
+                "add_special_tokens": False,
+                "append_concat_token": True,
+            }
+        )
+
+        trainer = SFTTrainer(
+            model=self.model,
+            args=args,
+            train_dataset=dataset['train'],
+            eval_dataset=dataset['test'],
+            processing_class=self.tokenizer,
+            callbacks=[AbortCallback(self.stop_event)] # Inject our stopper
+        )
+
+        # 4. Run Training
+        try:
+            output_buffer += "🚀 Training in progress... (Click Stop to interrupt)\n"
+            yield output_buffer
+            trainer.train()
+            
+            if self.stop_event.is_set():
+                output_buffer += "\n🛑 Training interrupted by user.\n"
+            else:
+                output_buffer += "\n✅ Training finished. Model weights updated in memory.\n"
+            yield output_buffer
+            
+            # Save locally
+            trainer.save_model()
+            output_buffer += f"Model saved locally to: {self.config.OUTPUT_DIR}\n"
+            yield output_buffer
 
-        result += "### Success Rate (After Training):\n" + f"{self.check_success_rate(dataset["test"])}\n\n"
-        return result
+        except Exception as e:
+            output_buffer += f"\n❌ Error during training: {e}\n"
+            yield output_buffer
+            return
+
+        if self.stop_event.is_set():
+            return
+
+        # 5. Post-Evaluation
+        output_buffer += "📊 Evaluating Post-Training Success Rate...\n"
+        post_report = ""
+        yield output_buffer
+        gen = self.check_success_rate(dataset["test"])
+        while not self.stop_event.is_set():
+            try:
+                post_report += f"{next(gen)}\n"
+                yield f"{output_buffer}{post_report}"
+            except StopIteration as e:
+                post_report = e.value
+                break
+        
+        if self.stop_event.is_set():
+            output_buffer += f"{post_report}\n\n🛑 Manual Eval interrupted by user.\n"
+            yield output_buffer
+            return
+
+        output_buffer += f"{post_report}\n\n"
+        yield output_buffer
 
     def check_success_rate(self, test_dataset):
-        result = []
+        """Runs inference on test set to calculate accuracy."""
+        results = []
         success_count = 0
+        total = len(test_dataset)
+                
         for idx, item in enumerate(test_dataset):
-            messages = [
-                item["messages"][0],
-                item["messages"][1],
-            ]
-
-            inputs = self.tokenizer.apply_chat_template(messages, tools=TOOLS, add_generation_prompt=True, return_dict=True, return_tensors="pt")
-
-            out = self.model.generate(**inputs.to(self.model.device), pad_token_id=self.tokenizer.eos_token_id, max_new_tokens=128)
-            output = self.tokenizer.decode(out[0][len(inputs["input_ids"][0]) :], skip_special_tokens=False)
-
-            result.append(f"{idx+1} Prompt: {item['messages'][1]['content']}")
-            result.append(f"  Output: {output}")
-            if item['messages'][2]['tool_calls'][0]['function']['name'] in output:
-                result.append("  `-> ✅ correct!")
-                success_count += 1
-            else:
-                result.append(" `-> ❌ wrong tool")
-
-        result.append(f"Success : {success_count} / {len(test_dataset)}")
-
-        return result
+            if idx >= 5:
+                break
+            if self.stop_event.is_set():
+                break
+
+            messages = [item["messages"][0], item["messages"][1]] # System + User
+
+            try:
+                inputs = self.tokenizer.apply_chat_template(
+                    messages, 
+                    tools=TOOLS, 
+                    add_generation_prompt=True, 
+                    return_dict=True, 
+                    return_tensors="pt"
+                )
+
+                out = self.model.generate(
+                    **inputs.to(self.model.device), 
+                    pad_token_id=self.tokenizer.eos_token_id, 
+                    max_new_tokens=128
+                )
+                
+                # Decode only the new tokens
+                output = self.tokenizer.decode(out[0][len(inputs["input_ids"][0]) :], skip_special_tokens=True)
+
+                results.append(f"{idx+1}. Prompt: {item['messages'][1]['content']}")
+                yield results[-1]
+                results.append(f"   Output: {output[:100]}...")
+                yield results[-1]
+                
+                # Check for correct tool name usage
+                expected_tool = item['messages'][2]['tool_calls'][0]['function']['name']
+                if expected_tool in output:
+                    results.append("   -> ✅ Correct Tool")
+                    yield results[-1]
+                    success_count += 1
+                else:
+                    results.append(f"   -> ❌ Wrong Tool (Expected: {expected_tool})")
+                    yield results[-1]
+
+            except Exception as e:
+                results.append(f"   -> Error: {e}")
+                yield results[-1]
+
+        summary = "\n".join(results)
+        summary += f"\n\nTotal Success : {success_count} / {len(test_dataset)}"
+        return summary
+
+    def download_model_zip(self) -> Optional[str]:
+        """Zips the output directory for download."""
+        if not os.path.exists(self.config.OUTPUT_DIR):
+            return None
+            
+        timestamp = int(time.time())
+        try:
+            base_name = self.config.ARTIFACTS_DIR.joinpath(f"functiongemma_finetuned_{timestamp}")
+            archive_path = shutil.make_archive(
+                base_name=str(base_name),
+                format='zip',
+                root_dir=str(self.config.OUTPUT_DIR),
+            )
+            return archive_path
+        except Exception as e:
+            print(f"Zip failed: {e}")
+            return None
 
+    # --- UI Builder ---
     def build_interface(self) -> gr.Blocks:
         with gr.Blocks(title="FunctionGemma Modkit") as demo:
             gr.Markdown("# 🤖 FunctionGemma Modkit: Fine-Tuning")
-            gr.Markdown("This project provides a set of tools to fine-tune FunctionGemma to understand your personal needs.<br>See [README](https://huggingface.co/spaces/google/functiongemma-modkit/blob/main/README.md) for more details.")
-            self._build_training_interface()
-        return demo
+            gr.Markdown("Fine-tune FunctionGemma to understand your custom functions.")
+            
+            with gr.Column():
+                gr.Markdown("## 1. Training Controls")
+                
+                with gr.Row():
+                    run_training_btn = gr.Button("🚀 Run Fine-Tuning", variant="primary")
+                    stop_training_btn = gr.Button("🛑 Stop Training", variant="stop", visible=False)
+                
+                output_display = gr.Textbox(
+                    lines=14, 
+                    label="Training Logs & Search Results", 
+                    value="Ready. Click 'Run' to begin.",
+                    interactive=False
+                )
+                
+                clear_reload_btn = gr.Button("🔄 Reset Model & Data")
+
+                gr.Markdown("--- \n ## 2. Data Management")
+                import_file = gr.File(label="Upload Additional Dataset (.csv)", file_types=[".csv"], height=80)
+                import_status = gr.Markdown("")
+                
+                gr.Markdown("--- \n ## 3. Export")
+                with gr.Row():
+                    zip_btn = gr.Button("⬇️ Prepare Model ZIP")
+                    download_file = gr.File(label="Download ZIP", height=80, visible=True, interactive=False)
+
+            # --- Event Wiring ---
+            
+            # Start Training (Generator updates output_display)
+            run_training_btn.click(
+                fn=lambda: (
+                    gr.update(visible=False),
+                    gr.update(interactive=False),
+                    gr.update(visible=True)
+                ),
+                inputs=None,
+                outputs=[run_training_btn, clear_reload_btn, stop_training_btn]
+            ).then(
+                fn=self.run_training,
+                inputs=[],
+                outputs=[output_display],
+            ).then(
+                fn=lambda: (
+                    gr.update(visible=True),
+                    gr.update(interactive=True),
+                    gr.update(visible=False)
+                ),
+                inputs=None,
+                outputs=[run_training_btn, clear_reload_btn, stop_training_btn]
+            )
+            
+            # Stop Training
+            stop_training_btn.click(
+                fn=self.stop_training,
+                inputs=None,
+                outputs=None # We don't need to return anything, status updates via the training generator
+            ).then(
+                fn=lambda: (
+                    gr.update(visible=True),
+                    gr.update(interactive=True),
+                    gr.update(visible=False)
+                ),
+                inputs=None,
+                outputs=[run_training_btn, clear_reload_btn, stop_training_btn]
+            )
 
-    def _build_training_interface(self):
-        with gr.Column():
-            gr.Markdown("## Fine-Tuning")
-            with gr.Row():
-                output = gr.Textbox(lines=14, label="Training and Search Results", value="Click 'Run Fine-Tuning' to begin.")
-            with gr.Row():
-                clear_reload_btn = gr.Button("Clear & Reload Model/Data")
-                run_training_btn = gr.Button("🚀 Run Fine-Tuning", variant="primary")
-            gr.Markdown("--- \n ## Dataset & Model Management")
-            import_file = gr.File(label="Upload Additional Dataset (.csv)", file_types=[".csv"], height=50)
-            with gr.Row():
-                download_model_btn = gr.Button("⬇️ Download Fine-Tuned Model")
-            download_status = gr.Markdown("Ready.")
-            with gr.Row():
-                model_output = gr.File(label="Download Model ZIP", height=50, visible=False, interactive=False)
-
-            run_training_btn.click(fn=self.training, outputs=output)
-            clear_reload_btn.click(fn=self.refresh_data_and_model, inputs=None, outputs=[output], queue=False)
-            import_file.change(fn=self.import_additional_dataset, inputs=[import_file], outputs=download_status)
-            download_model_btn.click(lambda: [gr.update(value=None, visible=False), "Zipping..."], None, [model_output, download_status], queue=False).then(self.download_model, None, 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])
+            # Reload
+            clear_reload_btn.click(
+                fn=self.refresh_data_and_model, 
+                inputs=None, 
+                outputs=[output_display]
+            )
+            
+            # File Import
+            import_file.upload(
+                fn=self.import_additional_dataset, 
+                inputs=[import_file], 
+                outputs=[import_status]
+            )
+            
+            # Download Logic
+            def handle_zip():
+                path = self.download_model_zip()
+                if path:
+                    return gr.update(value=path, visible=True)
+                return gr.update(value=None, visible=False)
+
+            zip_btn.click(
+                fn=handle_zip,
+                inputs=None,
+                outputs=[download_file]
+            )
 
+        return demo
 
 if __name__ == "__main__":
     app = FunctionGemmaTuner(AppConfig)
     demo = app.build_interface()
     print("Starting Gradio App...")
     demo.launch()
-