1 files changed, 16 insertions, 13 deletions

diff --git a/train.py b/train.py
index 545b299..4a04a30 100644
--- a/train.py
+++ b/train.py
@@ -19,11 +19,11 @@ def train_model(
     patience: int = 5,
     model_name: str = "",
 ) -> dict:
-    """Train model with mixed precision, early stopping on val weighted-F1.
+    """Train model with mixed precision, early stopping on val macro F1.
 
     Returns a dict with:
-      best_model_state, train_losses, val_losses, val_f1s,
-      train_time (seconds), epochs_trained
+      best_model_state, train_losses, val_losses, val_f1s (macro), val_weighted_f1s,
+      train_time (seconds), epochs_trained, best_val_f1 (macro)
     """
     scaler = torch.amp.GradScaler('cuda') if device.type == 'cuda' else None
 
@@ -31,7 +31,7 @@ def train_model(
     best_state = None
     patience_counter = 0
 
-    train_losses, val_losses, val_f1s = [], [], []
+    train_losses, val_losses, val_f1s, val_weighted_f1s = [], [], [], []
     start_time = time.time()
 
     for epoch in range(1, num_epochs + 1):
@@ -97,21 +97,23 @@ def train_model(
         val_loss = val_running_loss / n_val
         val_losses.append(val_loss)
 
-        val_f1 = f1_score(all_labels, all_preds, average='weighted', zero_division=0)
-        val_f1s.append(val_f1)
+        val_f1_macro = f1_score(all_labels, all_preds, average='macro', zero_division=0)
+        val_f1_weighted = f1_score(all_labels, all_preds, average='weighted', zero_division=0)
+        val_f1s.append(val_f1_macro)
+        val_weighted_f1s.append(val_f1_weighted)
 
         current_lr = scheduler.get_last_lr()[0] if hasattr(scheduler, 'get_last_lr') else optimizer.param_groups[0]['lr']
         print(
             f"[{model_name}] Epoch {epoch:3d}/{num_epochs} | "
             f"train loss: {train_loss:.4f} | val loss: {val_loss:.4f} | "
-            f"val F1: {val_f1:.4f} | lr: {current_lr:.2e}"
+            f"macro F1: {val_f1_macro:.4f} | weighted F1: {val_f1_weighted:.4f} | lr: {current_lr:.2e}"
         )
 
         scheduler.step()
 
-        # ── Early stopping ────────────────────────────────────────────────────
-        if val_f1 > best_val_f1:
-            best_val_f1 = val_f1
+        # ── Early stopping (monitored on macro F1) ───────────────────────────
+        if val_f1_macro > best_val_f1:
+            best_val_f1 = val_f1_macro
             best_state = copy.deepcopy(model.state_dict())
             patience_counter = 0
         else:
@@ -123,14 +125,15 @@ def train_model(
     train_time = time.time() - start_time
     epochs_trained = len(train_losses)
 
-    print(f"  Training complete: {epochs_trained} epochs, {train_time:.1f}s, best val F1: {best_val_f1:.4f}")
+    print(f"  Training complete: {epochs_trained} epochs, {train_time:.1f}s, best val macro F1: {best_val_f1:.4f}")
 
     return {
         'best_model_state': best_state,
         'train_losses': train_losses,
         'val_losses': val_losses,
-        'val_f1s': val_f1s,
+        'val_f1s': val_f1s,           # macro F1 per epoch (used for early stopping)
+        'val_weighted_f1s': val_weighted_f1s,
         'train_time': train_time,
         'epochs_trained': epochs_trained,
-        'best_val_f1': best_val_f1,
+        'best_val_f1': best_val_f1,   # best macro F1
     }