Upload folder using huggingface_hub

config.py

@@ -1,7 +1,7 @@
 import os
 from pathlib import Path
-from typing import Final, Optional
-from dataclasses import dataclass
+from typing import Final, Optional, List
+from dataclasses import dataclass, field
 
 @dataclass
 class AppConfig:
@@ -14,8 +14,17 @@ class AppConfig:
     
     # Model & Data
     HF_TOKEN: Final[Optional[str]] = os.getenv('HF_TOKEN')
-    # Defaulting to a real model ID for safety, original was local path '../hf/270m'
-    MODEL_NAME: Final[str] = '../hf/270m' 
+    
+    # Model Configuration
+    # Mutable: User can change this in the UI
+    MODEL_NAME: str = '../hf/270m'
+
+    AVAILABLE_MODELS: List[str] = field(default_factory=lambda: [
+        '../hf/270m',
+        '../hf/gemma-3-270m-it',
+        'google/gemma-3-270m-it'
+    ])
+
     DEFAULT_DATASET: Final[str] = 'bebechien/SimpleToolCalling'
     
     def __post_init__(self):

engine.py

@@ -5,7 +5,7 @@ import json
 import queue
 import matplotlib.pyplot as plt
 from functools import partial
-from typing import Generator, Optional, List, Dict
+from typing import Generator, Optional, List, Dict, Any, Tuple
 from datasets import Dataset, load_dataset
 from trl import SFTConfig, SFTTrainer
 from transformers import TrainerCallback, TrainingArguments, TrainerState, TrainerControl
@@ -30,8 +30,8 @@ class AbortCallback(TrainerCallback):
 
 class LogStreamingCallback(TrainerCallback):
     """
-    NEW: Intercepts training logs and pushes them to a queue 
-    so the main thread can display them in the UI.
+    Intercepts training logs and pushes them to a queue.
+    Sends tuple: (formatted_string, raw_dict)
     """
     def __init__(self, log_queue: queue.Queue):
         self.log_queue = log_queue
@@ -56,7 +56,6 @@ class LogStreamingCallback(TrainerCallback):
         for key, label in metrics_map.items():
             if key in logs:
                 val = logs[key]
-                # Format floats: use scientific notation for very small numbers (like LR)
                 if isinstance(val, (float, int)):
                     val_str = f"{val:.4f}" if val > 1e-4 else f"{val:.2e}"
                 else:
@@ -64,17 +63,20 @@ class LogStreamingCallback(TrainerCallback):
             
                 log_parts.append(f"{label}: {val_str}")
         
-        self.log_queue.put(" | ".join(log_parts))
+        # Structure for plotting
+        log_payload = logs.copy()
+        log_payload['step'] = state.global_step
+        
+        self.log_queue.put((" | ".join(log_parts), log_payload))
 
 class FunctionGemmaEngine:
     def __init__(self, config: AppConfig):
         self.config = config
         self.model = None
         self.tokenizer = None
+        self.loaded_model_name = None 
         self.imported_dataset = []
         self.stop_event = threading.Event()
-        
-        # NEW: State for tools
         self.current_tools = DEFAULT_TOOLS
 
         authenticate_hf(self.config.HF_TOKEN)
@@ -83,7 +85,7 @@ class FunctionGemmaEngine:
         except Exception as e:
             print(f"Initial load warning: {e}")
 
-    # NEW: Methods to handle Tool Schema updates
+    # --- Tool Schema Methods ---
     def get_tools_json(self) -> str:
         return json.dumps(self.current_tools, indent=2)
 
@@ -99,14 +101,24 @@ class FunctionGemmaEngine:
         except Exception as e:
             return f"❌ Error: {e}"
 
+    # --- Model & Data Management ---
+    
+    def _load_model_weights(self):
+        """Internal helper to load model based on current config."""
+        print(f"Loading model: {self.config.MODEL_NAME}...")
+        self.model, self.tokenizer = load_model_and_tokenizer(self.config.MODEL_NAME)
+        self.loaded_model_name = self.config.MODEL_NAME
+
     def refresh_data_and_model(self) -> str:
+        """Full reset: Reloads model and clears dataset."""
         self.imported_dataset = []
         try:
-            self.model, self.tokenizer = load_model_and_tokenizer(self.config.MODEL_NAME)
-            return "Model and data reloaded. Ready."
+            self._load_model_weights()
+            return f"Model loaded: {self.loaded_model_name}\nData cleared.\nReady."
         except Exception as e:
             self.model = None
             self.tokenizer = None
+            self.loaded_model_name = None
             return f"CRITICAL ERROR: Model failed to load. {e}"
 
     def load_csv(self, file_path: str) -> str:
@@ -122,13 +134,31 @@ class FunctionGemmaEngine:
     def trigger_stop(self):
         self.stop_event.set()
 
-    def run_training_pipeline(self, epochs: int, learning_rate: float, test_size: float, shuffle_data: bool) -> Generator[str, None, None]:
+    # --- Training Pipeline ---
+
+    def run_training_pipeline(self, epochs: int, learning_rate: float, test_size: float, shuffle_data: bool) -> Generator[Tuple[str, Any], None, None]:
+        self.stop_event.clear()
+        output_buffer = ""
+        last_plot = None
+
+        # 1. Check if model name changed since last load
+        if self.config.MODEL_NAME != self.loaded_model_name:
+            output_buffer += f"🔄 Model changed. Switching from '{self.loaded_model_name}' to '{self.config.MODEL_NAME}'...\n"
+            yield output_buffer, None
+            try:
+                self._load_model_weights()
+                output_buffer += "✅ Model reloaded successfully.\n"
+                yield output_buffer, None
+            except Exception as e:
+                output_buffer += f"❌ Failed to load model '{self.config.MODEL_NAME}': {e}\n"
+                yield output_buffer, None
+                return
+
         if self.model is None:
-             yield "Training failed: Model is not loaded.", None
+             yield "Training failed: No model loaded.", None
              return
 
-        self.stop_event.clear()
-        output_buffer = f"⏳ Preparing Dataset (Test Split: {test_size}, Shuffle: {shuffle_data})...\n"
+        output_buffer += f"⏳ Preparing Dataset (Test Split: {test_size}, Shuffle: {shuffle_data})...\n"
         yield output_buffer, None
 
         dataset, log = self._prepare_dataset()
@@ -161,86 +191,85 @@ class FunctionGemmaEngine:
         output_buffer += pre_training_report
 
         # --- Phase 2: Training (Threaded) ---
-        output_buffer += "\n\n🚀 Starting Fine-tuning (Epochs: {epochs}, LR: {learning_rate})...\n"
+        output_buffer += f"\n\n🚀 Starting Fine-tuning (Epochs: {epochs}, LR: {learning_rate})...\n"
         yield output_buffer, None
         
         log_queue = queue.Queue()
         training_error = None
-        training_history = []
+        running_history = [] 
         
-        # Function to run in the thread
         def train_wrapper():
-            nonlocal training_error, training_history
+            nonlocal training_error
             try:
-                training_history = self._execute_trainer(dataset, log_queue, epochs, learning_rate)
+                self._execute_trainer(dataset, log_queue, epochs, learning_rate)
             except Exception as e:
                 training_error = e
                 
-        # Start training thread
         train_thread = threading.Thread(target=train_wrapper)
         train_thread.start()
         
-        # Monitor loop: Yields logs while training runs
         while train_thread.is_alive():
-            # Drain the queue
             while not log_queue.empty():
-                log_msg = log_queue.get()
-                output_buffer += f"{log_msg}\n"
-                yield output_buffer, None
+                payload = log_queue.get()
+                if isinstance(payload, tuple):
+                    msg, log_data = payload
+                    output_buffer += f"{msg}\n"
+                    running_history.append(log_data)
+                    try:
+                        last_plot = self._generate_loss_plot(running_history)
+                        yield output_buffer, last_plot
+                    except Exception:
+                        yield output_buffer, last_plot
+                else:
+                    output_buffer += f"{payload}\n"
+                    yield output_buffer, last_plot
             
-            # Check for stop signal
             if self.stop_event.is_set():
-                yield f"{output_buffer}🛑 Stop signal sent. Waiting for trainer to wrap up...\n", None
-                # We don't break here, we wait for thread to finish cleanly
+                yield f"{output_buffer}🛑 Stop signal sent. Waiting for trainer to wrap up...\n", last_plot
             
-            time.sleep(0.1) # Prevent CPU spinning
+            time.sleep(0.1)
         
-        train_thread.join() # Ensure thread is completely done
+        train_thread.join()
         
-        # Flush any remaining logs
+        # Flush logs
         while not log_queue.empty():
-            log_msg = log_queue.get()
-            output_buffer += f"{log_msg}\n"
-            yield output_buffer, None
+            payload = log_queue.get()
+            if isinstance(payload, tuple):
+                msg, log_data = payload
+                output_buffer += f"{msg}\n"
+                running_history.append(log_data)
+                last_plot = self._generate_loss_plot(running_history)
+            else:
+                output_buffer += f"{payload}\n"
+            yield output_buffer, last_plot
                 
         if training_error:
             output_buffer += f"❌ Error during training: {training_error}\n"
-            yield output_buffer, None
+            yield output_buffer, last_plot
             return
 
         if self.stop_event.is_set(): 
             output_buffer += "🛑 Training manually stopped.\n"
-            yield output_buffer, None
+            yield output_buffer, last_plot
             return
         
         output_buffer += "✅ Training finished.\n"
-        yield output_buffer, None
-
-        output_buffer += "\n📈 Generating Loss Plot...\n"
-        yield output_buffer, None
-        
-        try:
-            final_plot = self._generate_loss_plot(training_history)
-            yield output_buffer, final_plot
-        except Exception as e:
-            output_buffer += f"⚠️ Could not generate plot: {e}\n"
-            yield output_buffer, None
+        yield output_buffer, last_plot
 
         # --- Phase 3: Post-Training Eval ---
         output_buffer += "\n📊 Evaluating Post-Training Success Rate...\n"
-        yield output_buffer, final_plot
+        yield output_buffer, last_plot
         
         post_training_report = ""
         for update in self._evaluate_model(dataset["test"]):
             post_training_report = update
             if self.stop_event.is_set():
                 post_training_report += "\n\n🛑 Manual Eval interrupted by user.\n"
-                yield f"{output_buffer}{post_training_report}", final_plot
+                yield f"{output_buffer}{post_training_report}", last_plot
                 break
-            yield f"{output_buffer}{post_training_report}", final_plot
+            yield f"{output_buffer}{post_training_report}", last_plot
 
     def _prepare_dataset(self):
-        # NEW: Use partial to inject self.current_tools into the formatting function
         formatting_fn = partial(create_conversation_format, tools_list=self.current_tools)
 
         if not self.imported_dataset:
@@ -286,27 +315,22 @@ class FunctionGemmaEngine:
         )
         trainer.train()
         trainer.save_model()
-        
         return trainer.state.log_history
         
     def _generate_loss_plot(self, history: list):
-        if not history:
-            return None
-            
-        # Extract Training Loss
-        # log_history format: [{'loss': 0.5, 'step': 1}, {'eval_loss': 0.4, 'step': 1}, ...]
+        if not history: return None
+        
+        # CHANGED: Close previous figures to prevent memory warning
+        plt.close('all')
+        
         train_steps = [x['step'] for x in history if 'loss' in x]
         train_loss = [x['loss'] for x in history if 'loss' in x]
-        
-        # Extract Validation Loss
         eval_steps = [x['step'] for x in history if 'eval_loss' in x]
         eval_loss = [x['eval_loss'] for x in history if 'eval_loss' in x]
 
         fig, ax = plt.subplots(figsize=(10, 5))
-        
         if train_steps:
             ax.plot(train_steps, train_loss, label='Training Loss', linestyle='-', marker=None)
-            
         if eval_steps:
             ax.plot(eval_steps, eval_loss, label='Validation Loss', linestyle='--', marker='o')
 
@@ -315,51 +339,40 @@ class FunctionGemmaEngine:
         ax.set_title("Training & Validation Loss")
         ax.legend()
         ax.grid(True, linestyle=':', alpha=0.6)
-        
         plt.tight_layout()
         return fig
 
     def _evaluate_model(self, test_dataset) -> Generator[str, None, None]:
         results = []
         success_count = 0
-        
         for idx, item in enumerate(test_dataset):
             messages = item["messages"][:2]
             try:
-                # NEW: Pass self.current_tools to the template
                 inputs = self.tokenizer.apply_chat_template(
                     messages, tools=self.current_tools, add_generation_prompt=True, return_dict=True, return_tensors="pt"
                 )
-                
                 device = self.model.device
                 inputs = {k: v.to(device) for k, v in inputs.items()}
-
                 out = self.model.generate(
                     **inputs, 
                     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=True)
-
                 log_entry = f"{idx+1}. Prompt: {messages[1]['content']}\n   Output: {output[:100]}..."
-                
-                # Check tool correctness
                 expected_tool = item['messages'][2]['tool_calls'][0]['function']['name']
                 if expected_tool in output:
                     log_entry += "\n   -> ✅ Correct Tool"
                     success_count += 1
                 else:
                     log_entry += f"\n   -> ❌ Wrong Tool (Expected: {expected_tool})"
-                
                 results.append(log_entry)
                 yield "\n".join(results) + f"\n\nRunning Success Rate: {success_count}/{idx+1}"
-                
             except Exception as e:
                 yield f"Error during inference: {e}"
 
     def get_zip_path(self) -> Optional[str]:
-        if not self.config.OUTPUT_DIR.exists():
-            return None
+        if not self.config.OUTPUT_DIR.exists(): return None
         timestamp = int(time.time())
         base_name = str(self.config.ARTIFACTS_DIR.joinpath(f"functiongemma_finetuned_{timestamp}"))
         return zip_directory(str(self.config.OUTPUT_DIR), base_name)

ui.py

@@ -2,6 +2,17 @@ import gradio as gr
 from engine import FunctionGemmaEngine
 
 def build_interface(engine: FunctionGemmaEngine) -> gr.Blocks:
+    
+    # Wrapper: Update config with selected model, then Run Training
+    def run_training_wrapper(epochs, lr, test_size, shuffle, model_name):
+        engine.config.MODEL_NAME = model_name.strip()
+        yield from engine.run_training_pipeline(epochs, lr, test_size, shuffle)
+
+    # Wrapper: Update config with selected model, then Reset/Reload
+    def handle_reset(model_name):
+        engine.config.MODEL_NAME = model_name.strip()
+        return engine.refresh_data_and_model()
+
     with gr.Blocks(title="FunctionGemma Modkit") as demo:
         gr.Markdown("# 🤖 FunctionGemma Modkit: Fine-Tuning")
         gr.Markdown("Fine-tune FunctionGemma to understand your custom functions.<br>See [README](https://huggingface.co/spaces/google/functiongemma-modkit/blob/main/README.md) for more details.")
@@ -41,10 +52,20 @@ def build_interface(engine: FunctionGemmaEngine) -> gr.Blocks:
                 with gr.Group():
                     gr.Markdown("**Hyperparameters**")
                     with gr.Row():
+                        # Dropdown that allows custom typing
+                        param_model = gr.Dropdown(
+                            choices=engine.config.AVAILABLE_MODELS,
+                            value=engine.config.MODEL_NAME,
+                            allow_custom_value=True, 
+                            label="Base Model",
+                            info="Select a preset OR type a custom Hugging Face model ID (e.g. 'google/gemma-3-1b-it')",
+                            interactive=True
+                        )
                         param_epochs = gr.Slider(
                             minimum=1, maximum=20, value=5, step=1, 
                             label="Epochs", info="Total training passes"
                         )
+                    with gr.Row():
                         param_lr = gr.Number(
                             value=5e-5, 
                             label="Learning Rate", 
@@ -52,18 +73,18 @@ def build_interface(engine: FunctionGemmaEngine) -> gr.Blocks:
                         )
                         param_test_size = gr.Slider(
                             minimum=0.1, maximum=0.9, value=0.2, step=0.05, 
-                            label="Test Split", info="Validation data ratio. Typical value is 0.2 (80% for training, 20% for testing)"
+                            label="Test Split", info="Validation ratio (0.2 = 20%)"
                         )
                         param_shuffle = gr.Checkbox(
                             value=True,
                             label="Shuffle Data",
-                            info="Randomize before split\nWhen unchecking this for your own custom datasets, ensure your source data is pre-mixed. If the distribution is unknown or sorted, you should check this to ensure the model learns a balanced representation of all tools during training."
+                            info="Randomize before split"
                         )
 
                 with gr.Row():
                     run_training_btn = gr.Button("🚀 Run Fine-Tuning", variant="primary", scale=2)
                     stop_training_btn = gr.Button("🛑 Stop", variant="stop", visible=False, scale=1)
-                    clear_reload_btn = gr.Button("🔄 Reset", variant="secondary", scale=1)
+                    clear_reload_btn = gr.Button("🔄 Reload Model & Reset Data", variant="secondary", scale=1)
 
                 with gr.Row():
                     # Left column: Text Logs
@@ -74,7 +95,7 @@ def build_interface(engine: FunctionGemmaEngine) -> gr.Blocks:
                         interactive=False,
                         autoscroll=True
                     )
-                    # Right column: Plot (NEW)
+                    # Right column: Plot
                     loss_plot = gr.Plot(label="Training Metrics")
 
             # --- TAB 3: EXPORT ---
@@ -102,7 +123,7 @@ def build_interface(engine: FunctionGemmaEngine) -> gr.Blocks:
             outputs=[import_status]
         )
 
-        # Tab 2: Training
+        # Tab 2: Training (Uses Wrapper)
         run_training_btn.click(
             fn=lambda: (
                 gr.update(visible=False), # Hide Run
@@ -111,8 +132,8 @@ def build_interface(engine: FunctionGemmaEngine) -> gr.Blocks:
             ),
             outputs=[run_training_btn, clear_reload_btn, stop_training_btn]
         ).then(
-            fn=engine.run_training_pipeline,
-            inputs=[param_epochs, param_lr, param_test_size, param_shuffle],
+            fn=run_training_wrapper,
+            inputs=[param_epochs, param_lr, param_test_size, param_shuffle, param_model],
             outputs=[output_display, loss_plot],
         ).then(
             fn=lambda: (
@@ -129,9 +150,10 @@ def build_interface(engine: FunctionGemmaEngine) -> gr.Blocks:
             outputs=None
         )
 
-        # Tab 2: Reset
+        # Tab 2: Reset (Uses Wrapper to capture model name)
         clear_reload_btn.click(
-            fn=engine.refresh_data_and_model, 
+            fn=handle_reset, 
+            inputs=[param_model], 
             outputs=[output_display]
         )