Update utils/model.py

utils/model.py

@@ -5,34 +5,32 @@ from pathlib import Path
 from spacy.tokens import DocBin
 import random
 import shutil
+import os
 
-# Load the training data from the .spacy file
 def load_data_from_spacy_file(file_path):
-    # Initialize a blank English model to ensure compatibility
+    """Load training data from .spacy file."""
     nlp = spacy.blank("en")
     
-    # Load the DocBin object and get documents
     try:
         doc_bin = DocBin().from_disk(file_path)
         docs = list(doc_bin.get_docs(nlp.vocab))
+        print(f"Loaded {len(docs)} documents from {file_path}.")
         return docs
     except Exception as e:
         print(f"Error loading data from .spacy file: {e}")
         return []
 
-
-# Train model function
 def train_model(epochs, model_path):
-    # Initialize a blank English model
+    """Train NER model."""
     nlp = spacy.blank("en")
     
-    # Create an NER component and add it to the pipeline
+    # Add the NER pipeline
     if "ner" not in nlp.pipe_names:
         ner = nlp.add_pipe("ner")
-        
-    nlp.add_pipe("sentencizer")
+    
+    nlp.add_pipe("sentencizer")  # Optional component to split sentences
 
-    # Define all possible entity labels
+    # Define entity labels
     labels = [
         "PERSON", "CONTACT", "EMAIL", "ABOUT", "EXPERIENCE", "YEARS_EXPERIENCE",
         "UNIVERSITY", "SOFT_SKILL", "INSTITUTE", "LAST_QUALIFICATION_YEAR", "JOB_TITLE",
@@ -40,55 +38,68 @@ def train_model(epochs, model_path):
         "LANGUAGE", "LOCATION", "PROJECTS", "SKILL", "CERTIFICATE"
     ]
 
-    # Add labels to the NER component
+    # Add the labels to the NER pipeline
     for label in labels:
         ner.add_label(label)
 
-    # Load the training data
+    # Load training data
     train_data = load_data_from_spacy_file("./data/Spacy_data.spacy")
 
-    # Start the training
-    optimizer = nlp.begin_training()
+    # Verify if data was loaded correctly
+    if not train_data:
+        print("No training data found. Exiting training.")
+        return
 
+    optimizer = nlp.begin_training()
     epoch_losses = []
     best_loss = float('inf')
 
-    # Training loop
+    # Start training loop
     for epoch in range(epochs):
         losses = {}
-        random.shuffle(train_data)  # Shuffle data for better training
-        
-        # Create minibatches
+        random.shuffle(train_data)  # Shuffle data
+
+        # Create batches
         batches = minibatch(train_data, size=compounding(4.0, 32.0, 1.001))
-        
+
         for batch in batches:
-            texts, annotations = zip(*[(doc.text, {"entities": [(ent.start_char, ent.end_char, ent.label_) for ent in doc.ents]}) for doc in batch])
-            
-            # Convert to Example objects
-            examples = [Example.from_dict(nlp.make_doc(text), annotation) for text, annotation in zip(texts, annotations)]
-               
+            # Extract texts and annotations
+            try:
+                texts, annotations = zip(
+                    *[(doc.text, {"entities": [(ent.start_char, ent.end_char, ent.label_) for ent in doc.ents]})
+                      for doc in batch]
+                )
+            except ValueError as e:
+                print(f"Error processing batch: {e}")
+                continue
+
+            # Create Example objects
+            examples = [Example.from_dict(nlp.make_doc(text), annotation) 
+                        for text, annotation in zip(texts, annotations)]
+
             # Update the model
             nlp.update(examples, sgd=optimizer, drop=0.35, losses=losses)
-        
+
+        # Record loss for this epoch
         current_loss = losses.get("ner", float('inf'))
         epoch_losses.append(current_loss)
-        
-        print(f"Losses at epoch {epoch + 1}: {losses}")        
-        
+
+        print(f"Losses at epoch {epoch + 1}: {losses}")
+
         # Save the best model
         if current_loss < best_loss:
             best_loss = current_loss
-            # Save to a temporary path
             temp_model_path = model_path + "_temp"
             nlp.to_disk(temp_model_path)
 
-            # Use shutil to move the model to the final path
+            # Safely move to the final path
             if os.path.exists(model_path):
-                shutil.rmtree(model_path)  # Remove the old model if it exists
-            shutil.copytree(temp_model_path, model_path)  # Copy the temp model to the final path
-            shutil.rmtree(temp_model_path)  # Remove the temporary model directory
-    
-    # Final save after training
+                shutil.rmtree(model_path)
+            shutil.copytree(temp_model_path, model_path)
+            shutil.rmtree(temp_model_path)
+
+    # Save the final model
     nlp.to_disk(model_path)
+    print(f"Training completed. Final model saved at: {model_path}")
 
     return epoch_losses