""" Training Script for NER Model Fine-tunes DistilBERT on document entity extraction Supports CORD and FUNSD datasets with BIO tagging """ import json import logging import os from typing import Dict, List, Tuple import numpy as np import torch import torch.nn as nn from seqeval.metrics import ( classification_report, f1_score, precision_score, recall_score, ) from sklearn.model_selection import train_test_split from torch.optim import AdamW from torch.utils.data import DataLoader, Dataset from tqdm import tqdm from transformers import AutoTokenizer, get_linear_schedule_with_warmup from dataset_loader import UnifiedDatasetLoader from ner_model import DocumentNERModel logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) class NERDataset(Dataset): """PyTorch Dataset for NER""" def __init__(self, tokenized_inputs: Dict, labels: List[List[int]]): self.input_ids = tokenized_inputs["input_ids"] self.attention_mask = tokenized_inputs["attention_mask"] self.labels = labels def __len__(self): return len(self.input_ids) def __getitem__(self, idx): return { "input_ids": torch.tensor(self.input_ids[idx], dtype=torch.long), "attention_mask": torch.tensor(self.attention_mask[idx], dtype=torch.long), "labels": torch.tensor(self.labels[idx], dtype=torch.long), } class NERDataPreprocessor: """Preprocess NER data with proper tokenization and label alignment""" def __init__(self, tokenizer, label2id: Dict, max_length: int = 512): self.tokenizer = tokenizer self.label2id = label2id self.max_length = max_length def tokenize_and_align_labels( self, examples: List[Dict] ) -> Tuple[Dict, List[List[int]]]: """ Tokenize texts and align labels with subword tokens Args: examples: List of dicts with 'tokens' and 'labels' Returns: Tuple of (tokenized_inputs, aligned_labels) """ tokenized_inputs = {"input_ids": [], "attention_mask": []} aligned_labels = [] for example in examples: tokens = example["tokens"] labels = example["labels"] # Tokenize with word boundaries encoding = self.tokenizer( tokens, is_split_into_words=True, max_length=self.max_length, padding="max_length", truncation=True, return_tensors=None, ) # Align labels word_ids = encoding.word_ids() label_ids = [] previous_word_idx = None for word_idx in word_ids: if word_idx is None: # Special token (CLS, SEP, PAD) label_ids.append(-100) elif word_idx != previous_word_idx: # First subword of a word if word_idx < len(labels): label = labels[word_idx] label_ids.append(self.label2id.get(label, self.label2id["O"])) else: label_ids.append(-100) else: # Continuation of previous word (subword) # Use same label or -100 if word_idx < len(labels): label = labels[word_idx] # Convert B- to I- for subwords if label.startswith("B-"): label = "I-" + label[2:] label_ids.append(self.label2id.get(label, self.label2id["O"])) else: label_ids.append(-100) previous_word_idx = word_idx tokenized_inputs["input_ids"].append(encoding["input_ids"]) tokenized_inputs["attention_mask"].append(encoding["attention_mask"]) aligned_labels.append(label_ids) return tokenized_inputs, aligned_labels class NERTrainer: """Trainer for NER model""" def __init__( self, model_name: str = "distilbert-base-uncased", output_dir: str = "models/ner", device: str = None, ): self.model_name = model_name self.output_dir = output_dir self.device = ( device if device else ("cuda" if torch.cuda.is_available() else "cpu") ) os.makedirs(output_dir, exist_ok=True) logger.info(f"NER Trainer initialized on device: {self.device}") if self.device == "cuda": logger.info(f"Using GPU: {torch.cuda.get_device_name(0)}") def prepare_data(self, test_size: float = 0.1, val_size: float = 0.1): """Load and prepare NER datasets""" logger.info("Loading NER datasets...") # Load data loader = UnifiedDatasetLoader() train_data = loader.load_ner_dataset( datasets=["cord", "sroie", "funsd"], split="train" ) # Get label mappings mappings = loader.get_label_mappings() self.label2id = mappings["ner"] self.id2label = mappings["ner_id2label"] logger.info(f"Loaded {len(train_data)} examples") logger.info(f"Number of labels: {len(self.label2id)}") # Split data train_examples, temp_examples = train_test_split( train_data, test_size=(test_size + val_size), random_state=42 ) val_examples, test_examples = train_test_split( temp_examples, test_size=test_size / (test_size + val_size), random_state=42 ) logger.info( f"Train: {len(train_examples)}, Val: {len(val_examples)}, Test: {len(test_examples)}" ) return train_examples, val_examples, test_examples def train( self, train_examples: List[Dict], val_examples: List[Dict], num_epochs: int = 25, batch_size: int = 8, learning_rate: float = 3e-5, warmup_ratio: float = 0.1, early_stopping_patience: int = 5, ): """Train the NER model""" # Load tokenizer and model tokenizer = AutoTokenizer.from_pretrained(self.model_name) model = DocumentNERModel( model_name=self.model_name, num_labels=len(self.label2id) ) model.to(self.device) # Preprocess data preprocessor = NERDataPreprocessor(tokenizer, self.label2id) logger.info("Preprocessing training data...") train_inputs, train_labels = preprocessor.tokenize_and_align_labels( train_examples ) train_dataset = NERDataset(train_inputs, train_labels) logger.info("Preprocessing validation data...") val_inputs, val_labels = preprocessor.tokenize_and_align_labels(val_examples) val_dataset = NERDataset(val_inputs, val_labels) # Create dataloaders train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) val_loader = DataLoader(val_dataset, batch_size=batch_size) # Optimizer and scheduler optimizer = AdamW(model.parameters(), lr=learning_rate) total_steps = len(train_loader) * num_epochs warmup_steps = int(total_steps * warmup_ratio) scheduler = get_linear_schedule_with_warmup( optimizer, num_warmup_steps=warmup_steps, num_training_steps=total_steps ) # Training loop best_val_f1 = 0 patience_counter = 0 for epoch in range(num_epochs): logger.info(f"\nEpoch {epoch + 1}/{num_epochs}") # Training model.train() train_loss = 0 progress_bar = tqdm(train_loader, desc="Training") for batch in progress_bar: optimizer.zero_grad() input_ids = batch["input_ids"].to(self.device) attention_mask = batch["attention_mask"].to(self.device) labels = batch["labels"].to(self.device) outputs = model(input_ids, attention_mask, labels) loss = outputs["loss"] loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() scheduler.step() train_loss += loss.item() progress_bar.set_postfix({"loss": loss.item()}) avg_train_loss = train_loss / len(train_loader) logger.info(f"Train Loss: {avg_train_loss:.4f}") # Validation model.eval() val_loss = 0 all_preds = [] all_labels = [] with torch.no_grad(): for batch in tqdm(val_loader, desc="Validation"): input_ids = batch["input_ids"].to(self.device) attention_mask = batch["attention_mask"].to(self.device) labels = batch["labels"].to(self.device) outputs = model(input_ids, attention_mask, labels) val_loss += outputs["loss"].item() preds = torch.argmax(outputs["logits"], dim=-1) # Collect predictions and labels (excluding -100) for pred_seq, label_seq in zip(preds, labels): pred_list = [] label_list = [] for p, l in zip(pred_seq, label_seq): if l != -100: pred_list.append(self.id2label[p.item()]) label_list.append(self.id2label[l.item()]) if pred_list: all_preds.append(pred_list) all_labels.append(label_list) avg_val_loss = val_loss / len(val_loader) # Calculate metrics using seqeval val_f1 = f1_score(all_labels, all_preds) val_precision = precision_score(all_labels, all_preds) val_recall = recall_score(all_labels, all_preds) logger.info(f"Val Loss: {avg_val_loss:.4f}") logger.info( f"Val F1: {val_f1:.4f}, Precision: {val_precision:.4f}, Recall: {val_recall:.4f}" ) # Early stopping if val_f1 > best_val_f1: best_val_f1 = val_f1 patience_counter = 0 # Save best model model_path = os.path.join(self.output_dir, "best_ner.pt") torch.save(model.state_dict(), model_path) logger.info(f"Saved best model with F1: {best_val_f1:.4f}") else: patience_counter += 1 if patience_counter >= early_stopping_patience: logger.info(f"Early stopping triggered after {epoch + 1} epochs") break # Save final model and metadata torch.save(model.state_dict(), os.path.join(self.output_dir, "final_ner.pt")) metadata = { "model_name": self.model_name, "num_labels": len(self.label2id), "label2id": self.label2id, "id2label": self.id2label, "best_val_f1": best_val_f1, } with open(os.path.join(self.output_dir, "metadata.json"), "w") as f: json.dump(metadata, f, indent=2) logger.info("NER training complete!") return best_val_f1 if __name__ == "__main__": # Training configuration trainer = NERTrainer(model_name="distilbert-base-uncased", output_dir="models/ner") # Prepare data train_examples, val_examples, test_examples = trainer.prepare_data( test_size=0.1, val_size=0.1 ) # Train model best_f1 = trainer.train( train_examples=train_examples, val_examples=val_examples, num_epochs=30, batch_size=32, learning_rate=3e-5, early_stopping_patience=5, ) print(f"\nBest validation F1: {best_f1:.4f}")