Split testing and tuning

engine.py

@@ -135,6 +135,66 @@ class FunctionGemmaEngine:
     def trigger_stop(self):
         self.stop_event.set()
 
+    def _ensure_model_consistency(self) -> Generator[str, None, bool]:
+        """Checks if the requested model matches the loaded one. Reloads if necessary."""
+        if self.config.MODEL_NAME != self.loaded_model_name:
+            yield f"🔄 Model changed. Switching from '{self.loaded_model_name}' to '{self.config.MODEL_NAME}'...\n"
+            try:
+                self._load_model_weights()
+                yield "✅ Model reloaded successfully.\n"
+                return True
+            except Exception as e:
+                yield f"❌ Failed to load model '{self.config.MODEL_NAME}': {e}\n"
+                return False
+        if self.model is None:
+             yield "❌ Error: No model loaded.\n"
+             return False
+        return True
+    
+    # --- Evaluation Pipeline ---
+    
+    def run_evaluation(self, test_size: float, shuffle_data: bool) -> Generator[str, None, None]:
+        self.stop_event.clear()
+        output_buffer = ""
+        
+        # 1. Check Model
+        gen = self._ensure_model_consistency()
+        try:
+            while True:
+                msg = next(gen)
+                output_buffer += msg
+                yield output_buffer
+        except StopIteration as e:
+            if not e.value: return # Failed to load
+            
+        # 2. Prepare Data
+        output_buffer += f"⏳ Preparing Dataset for Eval (Test Split: {test_size})...\n"
+        yield output_buffer
+
+        dataset, log = self._prepare_dataset()
+        output_buffer += log
+        yield output_buffer
+            
+        if not dataset:
+            output_buffer += "❌ Dataset creation failed.\n"
+            yield output_buffer
+            return
+
+        if len(dataset) > 1:
+            dataset = dataset.train_test_split(test_size=test_size, shuffle=shuffle_data)
+        else:
+            dataset = {"train": dataset, "test": dataset}
+            
+        # 3. Run Inference
+        output_buffer += "\n📊 Evaluating Model Success Rate on Test Split...\n"
+        yield output_buffer
+
+        for update in self._evaluate_model(dataset["test"]):
+            yield f"{output_buffer}{update}"
+            if self.stop_event.is_set():
+                yield f"{output_buffer}{update}\n\n🛑 Evaluation interrupted by user."
+                break
+
     # --- Training Pipeline ---
 
     def run_training_pipeline(self, epochs: int, learning_rate: float, test_size: float, shuffle_data: bool) -> Generator[Tuple[str, Any], None, None]:
@@ -142,21 +202,15 @@ class FunctionGemmaEngine:
         output_buffer = ""
         last_plot = None
 
-        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"
+        # 1. Check Model
+        gen = self._ensure_model_consistency()
+        try:
+            while True:
+                msg = next(gen)
+                output_buffer += f"{msg}"
                 yield output_buffer, None
-                return
-
-        if self.model is None:
-             yield "Training failed: No model loaded.", None
-             return
+        except StopIteration as e:
+            if not e.value: return
 
         output_buffer += f"⏳ Preparing Dataset (Test Split: {test_size}, Shuffle: {shuffle_data})...\n"
         yield output_buffer, None
@@ -174,24 +228,8 @@ class FunctionGemmaEngine:
         else:
             dataset = {"train": dataset, "test": dataset}
 
-        # --- Phase 1: Pre-Training Eval ---
-        output_buffer += "\n📊 Evaluating Pre-Training Success Rate...\n"
-        yield output_buffer, None
-
-        pre_training_report = ""
-        for update in self._evaluate_model(dataset["test"]):
-            pre_training_report = update
-            if self.stop_event.is_set():
-                pre_training_report += "\n\n🛑 Manual Eval interrupted by user.\n"
-                yield f"{output_buffer}{pre_training_report}", None
-                break
-            yield f"{output_buffer}{pre_training_report}", None
-        
-        if self.stop_event.is_set(): return
-        output_buffer += pre_training_report
-
-        # --- Phase 2: Training (Threaded) ---
-        output_buffer += f"\n\n🚀 Starting Fine-tuning (Epochs: {epochs}, LR: {learning_rate})...\n"
+        # --- Training (Threaded) ---
+        output_buffer += f"\n🚀 Starting Fine-tuning (Epochs: {epochs}, LR: {learning_rate})...\n"
         yield output_buffer, None
         
         log_queue = queue.Queue()
@@ -257,19 +295,6 @@ class FunctionGemmaEngine:
         output_buffer += "✅ Training finished.\n"
         yield output_buffer, last_plot
 
-        # --- Phase 3: Post-Training Eval ---
-        output_buffer += "\n📊 Evaluating Post-Training Success Rate...\n"
-        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}", last_plot
-                break
-            yield f"{output_buffer}{post_training_report}", last_plot
-
     def _prepare_dataset(self):
         formatting_fn = partial(create_conversation_format, tools_list=self.current_tools)
 

ui.py

@@ -41,6 +41,15 @@ class UIController:
         engine.config.MODEL_NAME = model_name.strip()
         yield from engine.run_training_pipeline(epochs, lr, test_size, shuffle)
 
+    @staticmethod
+    def run_evaluation(engine: FunctionGemmaEngine, test_size: float, shuffle: bool, model_name: str) -> Generator:
+        if not engine:
+            yield "⚠️ Engine not initialized."
+            return
+        
+        engine.config.MODEL_NAME = model_name.strip()
+        yield from engine.run_evaluation(test_size, shuffle)
+
     @staticmethod
     def handle_reset(engine: FunctionGemmaEngine, model_name: str) -> str:
         engine.config.MODEL_NAME = model_name.strip()
@@ -151,6 +160,7 @@ def _render_training_tab(engine_state):
                 param_shuffle = gr.Checkbox(value=True, label="Shuffle Data", info="Randomize before split")
 
         with gr.Row():
+            run_eval_btn = gr.Button("🧪 Run Evaluation", variant="secondary", scale=1)
             run_training_btn = gr.Button("🚀 Run Fine-Tuning", variant="primary", scale=1)
             stop_training_btn = gr.Button("🛑 Stop", variant="stop", visible=False, scale=1)
             clear_reload_btn = gr.Button("🔄 Reload Model & Reset Data", variant="secondary", scale=1)
@@ -161,7 +171,7 @@ def _render_training_tab(engine_state):
 
     return {
         "params": [param_epochs, param_lr, param_test_size, param_shuffle, param_model],
-        "buttons": [run_training_btn, stop_training_btn, clear_reload_btn],
+        "buttons": [run_training_btn, stop_training_btn, clear_reload_btn, run_eval_btn],
         "outputs": [output_display, loss_plot],
         "model_input": param_model # specifically needed for initialization
     }
@@ -175,6 +185,7 @@ def _render_export_tab(engine_state, username_state):
                 gr.Markdown("Download the model weights locally.")
                 zip_btn = gr.Button("⬇️ Prepare Model ZIP", variant="secondary", interactive=False)
                 download_file = gr.File(label="Download Archive", interactive=False)
+                gr.Markdown("NOTE: Zipping usually takes 1~2 min.")
             
             with gr.Column():
                 gr.Markdown("#### Option B: Save to Hugging Face Hub")