Fix issue with stop if queued

engine.py

@@ -150,50 +150,53 @@ class FunctionGemmaEngine:
              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
+            # 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"
+            dataset, log = self._prepare_dataset()
+            output_buffer += log
             yield output_buffer
-            return
+                
+            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
+            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
+            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
+        finally:
+            self.stop_event.set() # Ensure loop breaks if generator cancelled
 
     # --- Training Pipeline ---
 
@@ -202,98 +205,102 @@ class FunctionGemmaEngine:
         output_buffer = ""
         last_plot = None
 
-        # 1. Check Model
-        gen = self._ensure_model_consistency()
         try:
-            while True:
-                msg = next(gen)
-                output_buffer += f"{msg}"
-                yield output_buffer, None
-        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
-
-        dataset, log = self._prepare_dataset()
-        if not dataset:
-            yield "Dataset creation failed.", None
-            return
+            # 1. Check Model
+            gen = self._ensure_model_consistency()
+            try:
+                while True:
+                    msg = next(gen)
+                    output_buffer += f"{msg}"
+                    yield output_buffer, None
+            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
+
+            dataset, log = self._prepare_dataset()
+            if not dataset:
+                yield "Dataset creation failed.", None
+                return
+
+            output_buffer += log
+            yield output_buffer, None
+                
+            if len(dataset) > 1:
+                dataset = dataset.train_test_split(test_size=test_size, shuffle=shuffle_data)
+            else:
+                dataset = {"train": dataset, "test": dataset}
 
-        output_buffer += log
-        yield output_buffer, None
+            # --- Training (Threaded) ---
+            output_buffer += f"\n🚀 Starting Fine-tuning (Epochs: {epochs}, LR: {learning_rate})...\n"
+            yield output_buffer, None
             
-        if len(dataset) > 1:
-            dataset = dataset.train_test_split(test_size=test_size, shuffle=shuffle_data)
-        else:
-            dataset = {"train": dataset, "test": dataset}
-
-        # --- Training (Threaded) ---
-        output_buffer += f"\n🚀 Starting Fine-tuning (Epochs: {epochs}, LR: {learning_rate})...\n"
-        yield output_buffer, None
-        
-        log_queue = queue.Queue()
-        training_error = None
-        running_history = [] 
-        
-        def train_wrapper():
-            nonlocal training_error
-            try:
-                self._execute_trainer(dataset, log_queue, epochs, learning_rate)
-            except Exception as e:
-                training_error = e
+            log_queue = queue.Queue()
+            training_error = None
+            running_history = [] 
+            
+            def train_wrapper():
+                nonlocal training_error
+                try:
+                    self._execute_trainer(dataset, log_queue, epochs, learning_rate)
+                except Exception as e:
+                    training_error = e
+                    
+            train_thread = threading.Thread(target=train_wrapper)
+            train_thread.start()
+            
+            while train_thread.is_alive():
+                while not log_queue.empty():
+                    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
                 
-        train_thread = threading.Thread(target=train_wrapper)
-        train_thread.start()
-        
-        while train_thread.is_alive():
+                if self.stop_event.is_set():
+                    yield f"{output_buffer}🛑 Stop signal sent. Waiting for trainer to wrap up...\n", last_plot
+                
+                time.sleep(0.1)
+            
+            train_thread.join()
+            
+            self.has_model_tuned = True
+            
             while not log_queue.empty():
                 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
+                    last_plot = self._generate_loss_plot(running_history)
                 else:
                     output_buffer += f"{payload}\n"
-                    yield output_buffer, last_plot
-            
-            if self.stop_event.is_set():
-                yield f"{output_buffer}🛑 Stop signal sent. Waiting for trainer to wrap up...\n", last_plot
+                yield output_buffer, last_plot
+                    
+            if training_error:
+                output_buffer += f"❌ Error during training: {training_error}\n"
+                yield output_buffer, last_plot
+                return
+
+            if self.stop_event.is_set(): 
+                output_buffer += "🛑 Training manually stopped.\n"
+                yield output_buffer, last_plot
+                return
             
-            time.sleep(0.1)
-        
-        train_thread.join()
-        
-        self.has_model_tuned = True
-        
-        while not log_queue.empty():
-            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"
+            output_buffer += "✅ Training finished.\n"
             yield output_buffer, last_plot
-            return
-
-        if self.stop_event.is_set(): 
-            output_buffer += "🛑 Training manually stopped.\n"
-            yield output_buffer, last_plot
-            return
-        
-        output_buffer += "✅ Training finished.\n"
-        yield output_buffer, last_plot
+            
+        finally:
+            self.stop_event.set() # Ensure background thread stops if generator cancelled
 
     def _prepare_dataset(self):
         formatting_fn = partial(create_conversation_format, tools_list=self.current_tools)
@@ -433,4 +440,4 @@ class FunctionGemmaEngine:
 
             return f"✅ Success! Model uploaded to: {repo_url}"
         except Exception as e:
-            return f"❌ Upload failed: {str(e)}"
+            return f"❌ Upload failed: {str(e)}"