""" Small Language Model (SLM) architecture for document text extraction. Uses DistilBERT with transfer learning for Named Entity Recognition. """ import torch import torch.nn as nn from torch.utils.data import Dataset, DataLoader from transformers import ( DistilBertTokenizer, DistilBertForTokenClassification, DistilBertConfig, get_linear_schedule_with_warmup ) from typing import List, Dict, Tuple, Optional import json import numpy as np from sklearn.model_selection import train_test_split from dataclasses import dataclass @dataclass class ModelConfig: """Configuration for the SLM model.""" model_name: str = "distilbert-base-uncased" max_length: int = 512 batch_size: int = 16 learning_rate: float = 2e-5 num_epochs: int = 3 warmup_steps: int = 500 weight_decay: float = 0.01 dropout_rate: float = 0.3 # Entity labels entity_labels: List[str] = None def __post_init__(self): if self.entity_labels is None: self.entity_labels = [ 'O', 'B-NAME', 'I-NAME', 'B-DATE', 'I-DATE', 'B-INVOICE_NO', 'I-INVOICE_NO', 'B-AMOUNT', 'I-AMOUNT', 'B-ADDRESS', 'I-ADDRESS', 'B-PHONE', 'I-PHONE', 'B-EMAIL', 'I-EMAIL' ] @property def num_labels(self) -> int: return len(self.entity_labels) @property def label2id(self) -> Dict[str, int]: return {label: i for i, label in enumerate(self.entity_labels)} @property def id2label(self) -> Dict[int, str]: return {i: label for i, label in enumerate(self.entity_labels)} class NERDataset(Dataset): """PyTorch Dataset for NER training.""" def __init__(self, dataset: List[Dict], tokenizer: DistilBertTokenizer, config: ModelConfig, mode: str = 'train'): self.dataset = dataset self.tokenizer = tokenizer self.config = config self.mode = mode # Prepare tokenized data self.tokenized_data = self._tokenize_and_align_labels() def _tokenize_and_align_labels(self) -> List[Dict]: """Tokenize text and align labels with subword tokens.""" tokenized_data = [] for example in self.dataset: tokens = example['tokens'] labels = example['labels'] # Tokenize each word and track alignments tokenized_inputs = self.tokenizer( tokens, is_split_into_words=True, padding='max_length', truncation=True, max_length=self.config.max_length, return_tensors='pt' ) # Align labels with subword tokens word_ids = tokenized_inputs.word_ids() aligned_labels = [] previous_word_idx = None for word_idx in word_ids: if word_idx is None: # Special tokens get -100 (ignored in loss computation) aligned_labels.append(-100) elif word_idx != previous_word_idx: # First subword of a word gets the original label if word_idx < len(labels): label = labels[word_idx] aligned_labels.append(self.config.label2id.get(label, 0)) else: aligned_labels.append(-100) else: # Subsequent subwords of the same word if word_idx < len(labels): label = labels[word_idx] if label.startswith('B-'): # Convert B- to I- for subword tokens i_label = label.replace('B-', 'I-') aligned_labels.append(self.config.label2id.get(i_label, 0)) else: aligned_labels.append(self.config.label2id.get(label, 0)) else: aligned_labels.append(-100) previous_word_idx = word_idx tokenized_data.append({ 'input_ids': tokenized_inputs['input_ids'].squeeze(), 'attention_mask': tokenized_inputs['attention_mask'].squeeze(), 'labels': torch.tensor(aligned_labels, dtype=torch.long), 'original_tokens': tokens, 'original_labels': labels }) return tokenized_data def __len__(self) -> int: return len(self.tokenized_data) def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: return { 'input_ids': self.tokenized_data[idx]['input_ids'], 'attention_mask': self.tokenized_data[idx]['attention_mask'], 'labels': self.tokenized_data[idx]['labels'] } class DocumentNERModel(nn.Module): """DistilBERT-based model for document NER.""" def __init__(self, config: ModelConfig): super().__init__() self.config = config # Load pre-trained DistilBERT configuration bert_config = DistilBertConfig.from_pretrained( config.model_name, num_labels=config.num_labels, id2label=config.id2label, label2id=config.label2id, dropout=config.dropout_rate, attention_dropout=config.dropout_rate ) # Initialize model with token classification head self.model = DistilBertForTokenClassification.from_pretrained( config.model_name, config=bert_config ) # Additional dropout layer for regularization self.dropout = nn.Dropout(config.dropout_rate) def forward(self, input_ids, attention_mask=None, labels=None): """Forward pass through the model.""" outputs = self.model( input_ids=input_ids, attention_mask=attention_mask, labels=labels ) return outputs def predict(self, input_ids, attention_mask): """Make predictions without computing loss.""" with torch.no_grad(): outputs = self.model( input_ids=input_ids, attention_mask=attention_mask ) predictions = torch.argmax(outputs.logits, dim=-1) probabilities = torch.softmax(outputs.logits, dim=-1) return predictions, probabilities class NERTrainer: """Trainer class for the NER model.""" def __init__(self, model: DocumentNERModel, config: ModelConfig): self.model = model self.config = config self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') self.model.to(self.device) # Initialize tokenizer self.tokenizer = DistilBertTokenizer.from_pretrained(config.model_name) def prepare_dataloaders(self, dataset: List[Dict], test_size: float = 0.2) -> Tuple[DataLoader, DataLoader]: """Prepare training and validation dataloaders.""" # Split dataset train_data, val_data = train_test_split( dataset, test_size=test_size, random_state=42 ) # Create datasets train_dataset = NERDataset(train_data, self.tokenizer, self.config, 'train') val_dataset = NERDataset(val_data, self.tokenizer, self.config, 'val') # Create dataloaders train_dataloader = DataLoader( train_dataset, batch_size=self.config.batch_size, shuffle=True ) val_dataloader = DataLoader( val_dataset, batch_size=self.config.batch_size, shuffle=False ) return train_dataloader, val_dataloader def train(self, train_dataloader: DataLoader, val_dataloader: DataLoader) -> Dict[str, List[float]]: """Train the NER model.""" # Initialize optimizer and scheduler optimizer = torch.optim.AdamW( self.model.parameters(), lr=self.config.learning_rate, weight_decay=self.config.weight_decay ) total_steps = len(train_dataloader) * self.config.num_epochs scheduler = get_linear_schedule_with_warmup( optimizer, num_warmup_steps=self.config.warmup_steps, num_training_steps=total_steps ) # Training history history = { 'train_loss': [], 'val_loss': [], 'val_accuracy': [] } print(f"Training on device: {self.device}") print(f"Total training steps: {total_steps}") for epoch in range(self.config.num_epochs): print(f"\nEpoch {epoch + 1}/{self.config.num_epochs}") print("-" * 50) # Training phase train_loss = self._train_epoch(train_dataloader, optimizer, scheduler) history['train_loss'].append(train_loss) # Validation phase val_loss, val_accuracy = self._validate_epoch(val_dataloader) history['val_loss'].append(val_loss) history['val_accuracy'].append(val_accuracy) print(f"Train Loss: {train_loss:.4f}") print(f"Val Loss: {val_loss:.4f}") print(f"Val Accuracy: {val_accuracy:.4f}") return history def _train_epoch(self, dataloader: DataLoader, optimizer, scheduler) -> float: """Train for one epoch.""" self.model.train() total_loss = 0 for batch_idx, batch in enumerate(dataloader): # Move batch to device batch = {k: v.to(self.device) for k, v in batch.items()} # Forward pass outputs = self.model(**batch) loss = outputs.loss # Backward pass optimizer.zero_grad() loss.backward() # Gradient clipping torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0) optimizer.step() scheduler.step() total_loss += loss.item() if batch_idx % 10 == 0: print(f"Batch {batch_idx}/{len(dataloader)}, Loss: {loss.item():.4f}") return total_loss / len(dataloader) def _validate_epoch(self, dataloader: DataLoader) -> Tuple[float, float]: """Validate for one epoch.""" self.model.eval() total_loss = 0 total_correct = 0 total_tokens = 0 with torch.no_grad(): for batch in dataloader: batch = {k: v.to(self.device) for k, v in batch.items()} outputs = self.model(**batch) loss = outputs.loss total_loss += loss.item() # Calculate accuracy (ignoring -100 labels) predictions = torch.argmax(outputs.logits, dim=-1) labels = batch['labels'] # Mask for valid labels (not -100) valid_mask = labels != -100 correct = (predictions == labels) & valid_mask total_correct += correct.sum().item() total_tokens += valid_mask.sum().item() avg_loss = total_loss / len(dataloader) accuracy = total_correct / total_tokens if total_tokens > 0 else 0 return avg_loss, accuracy def save_model(self, save_path: str): """Save the trained model and tokenizer.""" self.model.model.save_pretrained(save_path) self.tokenizer.save_pretrained(save_path) # Save config config_path = f"{save_path}/training_config.json" with open(config_path, 'w') as f: json.dump(vars(self.config), f, indent=2) print(f"Model saved to {save_path}") def load_model(self, model_path: str): """Load a pre-trained model.""" self.model.model = DistilBertForTokenClassification.from_pretrained(model_path) self.tokenizer = DistilBertTokenizer.from_pretrained(model_path) self.model.to(self.device) print(f"Model loaded from {model_path}") def create_model_and_trainer(config: Optional[ModelConfig] = None) -> Tuple[DocumentNERModel, NERTrainer]: """Create model and trainer with configuration.""" if config is None: config = ModelConfig() model = DocumentNERModel(config) trainer = NERTrainer(model, config) return model, trainer def main(): """Demonstrate model creation and setup.""" # Create configuration config = ModelConfig( batch_size=8, # Smaller batch size for demo num_epochs=2, learning_rate=3e-5 ) print("Model Configuration:") print(f"Model: {config.model_name}") print(f"Max Length: {config.max_length}") print(f"Batch Size: {config.batch_size}") print(f"Learning Rate: {config.learning_rate}") print(f"Number of Labels: {config.num_labels}") print(f"Entity Labels: {config.entity_labels}") # Create model and trainer model, trainer = create_model_and_trainer(config) print(f"\nModel created successfully!") print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}") print(f"Trainable parameters: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}") return model, trainer if __name__ == "__main__": main()