From 6faf9f355e9ffb60fb583088c1305123346c2b06 Mon Sep 17 00:00:00 2001
From: bernard-ng <ngandubernard@gmail.com>
Date: Sun, 17 Aug 2025 14:15:12 +0200
Subject: [PATCH] fix: NER training loop

---
 processing/ner/name_model.py | 31 +++++++++++++++++--------------
 1 file changed, 17 insertions(+), 14 deletions(-)

diff --git a/processing/ner/name_model.py b/processing/ner/name_model.py
index c8d1141..3fa05d2 100644
--- a/processing/ner/name_model.py
+++ b/processing/ner/name_model.py
@@ -2,6 +2,7 @@ import ast
 import json
 import logging
 import os
+import random
 from pathlib import Path
 from typing import Dict, Any, List, Tuple
 
@@ -190,32 +191,34 @@ class NameModel:
 
         # Initialize the model
         self.nlp.initialize()
-
-        # Training loop
+        optimizer = self.nlp.resume_training()
         losses_history = []
 
         for epoch in range(epochs):
             losses = {}
-
-            # Create training examples
             examples = []
-            for text, annotations in tqdm(data, description="Create training examples"):
+
+            for text, annotations in tqdm(data, desc="Create training examples"):
                 doc = self.nlp.make_doc(text)
-                example = Example.from_dict(doc, annotations)
-                examples.append(example)
+                examples.append(Example.from_dict(doc, annotations))
+
+            # Shuffle examples each epoch (important!)
+            random.shuffle(examples)
 
             # Train in batches
             batches = minibatch(examples, size=batch_size)
             for batch in batches:
-                self.nlp.update(
-                    batch, losses=losses, drop=dropout_rate, sgd=self.nlp.create_optimizer()
-                )
-                logging.info(f"Training batch with {len(batch)} examples, current losses: {losses}")
+                batch_losses = {}
+                self.nlp.update(batch, losses=batch_losses, drop=dropout_rate, sgd=optimizer)
+                logging.info(f"Training batch with {len(batch)} examples, current losses: {batch_losses}")
+
+                # Accumulate into total losses dict
+                for k, v in batch_losses.items():
+                    losses[k] = losses.get(k, 0.0) + v
 
             del batches  # free memory
-            epoch_loss = losses.get("ner", 0)
-            losses_history.append(epoch_loss)
-            logging.info(f"Epoch {epoch + 1}/{epochs}, Loss: {epoch_loss:.4f}")
+            losses_history.append(losses.get("ner", 0))
+            logging.info(f"Epoch {epoch+1}/{epochs}, Total Loss: {losses['ner']:.4f}")
 
         # Store training statistics
         self.training_stats = {