utils.py

import os
import math
import yaml
import hashlib
import json

import torch
import torch.nn as nn
import math
import matplotlib.pyplot as plt
import numpy as np
from tqdm import tqdm


class LRFinder:
    """
    Learning Rate Finder using exponential learning rate scheduling.
    
    Helps find optimal learning rate by training for a few epochs while
    exponentially increasing the learning rate and monitoring loss.
    
    Based on: "A disciplined approach to neural network training" (Leslie Smith)
    """
    
    def __init__(self, model, optimizer, criterion, device, 
                 start_lr=1e-6, end_lr=1.0, num_iter=100):
        """
        Args:
            model: PyTorch model
            optimizer: Optimizer (typically AdamW)
            criterion: Loss function
            device: torch.device (cuda/cpu/mps)
            start_lr: Starting learning rate (log scale)
            end_lr: Ending learning rate (log scale)
            num_iter: Number of iterations to run
        """
        self.model = model
        self.optimizer = optimizer
        self.criterion = criterion
        self.device = device
        self.start_lr = start_lr
        self.end_lr = end_lr
        self.num_iter = num_iter
        
        self.lrs = []
        self.losses = []
        self.best_loss = None
        
    def _update_lr(self, batch_num):
        """Update learning rate exponentially"""
        lr = self.start_lr * (self.end_lr / self.start_lr) ** (batch_num / self.num_iter)
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = lr
        return lr
    
    def find(self, train_dataloader):
        """
        Run learning rate range test.
        
        Args:
            train_dataloader: DataLoader for training data
        """
        self.model.train()
        self.lrs = []
        self.losses = []
        self.best_loss = None
        
        print("Starting LR Range Test...")
        pbar = tqdm(enumerate(train_dataloader), total=self.num_iter, desc="LR Range Test")
        
        for batch_num, (src, tgt) in pbar:
            if batch_num >= self.num_iter:
                break
            
            # Move to device
            src, tgt = src.to(self.device), tgt.to(self.device)
            
            # Forward pass
            self.optimizer.zero_grad()
            logits, _ = self.model(src)  # (B, T, V)
            B, T, V = logits.shape
            loss = self.criterion(logits.view(B * T, V), tgt.view(B * T))
            
            # Backward pass
            loss.backward()
            self.optimizer.step()
            
            # Update LR
            lr = self._update_lr(batch_num)
            
            # Store metrics
            current_loss = loss.item()
            self.lrs.append(lr)
            self.losses.append(current_loss)
            
            # Track best loss
            if self.best_loss is None or current_loss < self.best_loss:
                self.best_loss = current_loss
            
            # Stop if loss explodes
            if current_loss > 4 * self.best_loss:
                print(f"\nLoss exploded at LR={lr:.2e}. Stopping early.")
                break
            
            pbar.set_postfix({'loss': f'{current_loss:.4f}', 'lr': f'{lr:.2e}'})
        
        print(f"LR Range Test complete. Best loss: {self.best_loss:.4f}")
    
    def plot(self, save_path=None, log_scale=True):
        """
        Plot learning rate vs loss with enhanced grid and labels.
        
        Args:
            save_path: Optional path to save the figure
            log_scale: Whether to use log scale for x-axis (recommended)
        """
        if not self.lrs or not self.losses:
            print("No data to plot. Run find() first.")
            return
        
        fig, ax = plt.subplots(figsize=(14, 7))
        
        if log_scale:
            ax.semilogx(self.lrs, self.losses, 'b-', linewidth=2.5, label='Loss')
            ax.set_xlabel('Learning Rate (log scale)', fontsize=13, fontweight='bold')
        else:
            ax.plot(self.lrs, self.losses, 'b-', linewidth=2.5, label='Loss')
            ax.set_xlabel('Learning Rate', fontsize=13, fontweight='bold')
        
        ax.set_ylabel('Loss', fontsize=13, fontweight='bold')
        ax.set_title('Learning Rate Range Test', fontsize=16, fontweight='bold', pad=20)
        
        # Enhanced grid
        ax.grid(True, which='major', alpha=0.6, linewidth=1.0, linestyle='-', color='gray')
        ax.grid(True, which='minor', alpha=0.2, linewidth=0.5, linestyle='--', color='gray')
        ax.minorticks_on()
        
        # Set background color
        ax.set_facecolor('#f9f9f9')
        fig.patch.set_facecolor('white')
        
        # Find and mark optimal LR (steepest descent)
        if len(self.losses) > 1:
            # Smooth losses for better gradient calculation
            smoothed_losses = self._smooth_losses(self.losses, beta=0.05)
            gradients = np.gradient(smoothed_losses)
            
            # Find LR with steepest negative gradient
            min_gradient_idx = np.argmin(gradients)
            optimal_lr = self.lrs[min_gradient_idx]
            optimal_loss = self.losses[min_gradient_idx]
            
            # Mark optimal point
            ax.axvline(optimal_lr, color='r', linestyle='--', linewidth=2.5, alpha=0.8)
            ax.scatter([optimal_lr], [optimal_loss], 
                      color='r', s=200, zorder=5, edgecolors='darkred', linewidth=2,
                      label=f'Optimal LR: {optimal_lr:.2e}')
            
            # Add annotation
            ax.annotate(f'Optimal LR\n{optimal_lr:.2e}\nLoss: {optimal_loss:.4f}',
                       xy=(optimal_lr, optimal_loss),
                       xytext=(20, 20),
                       textcoords='offset points',
                       fontsize=10,
                       bbox=dict(boxstyle='round,pad=0.5', facecolor='yellow', alpha=0.7),
                       arrowprops=dict(arrowstyle='->', connectionstyle='arc3,rad=0', lw=2, color='red'))
            
            ax.legend(fontsize=11, loc='best', framealpha=0.95, edgecolor='black')
        
        # Add more x-axis labels for better readability
        if log_scale:
            # Set major locator for log scale
            from matplotlib.ticker import LogLocator, NullFormatter
            ax.xaxis.set_major_locator(LogLocator(base=10, numticks=15))
            ax.xaxis.set_minor_locator(LogLocator(base=10, subs=np.arange(2, 10), numticks=50))
        
        # Add axis labels with better formatting
        ax.tick_params(axis='both', which='major', labelsize=10, width=1.5, length=6)
        ax.tick_params(axis='both', which='minor', labelsize=8, width=1, length=3)
        
        # Add border
        for spine in ax.spines.values():
            spine.set_linewidth(1.5)
            spine.set_color('black')
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=200, bbox_inches='tight', facecolor='white')
            print(f"Plot saved to: {save_path}")
        
        plt.show()
    
    @staticmethod
    def _smooth_losses(losses, beta=0.05):
        """Exponential moving average smoothing"""
        smoothed = []
        avg_loss = losses[0]
        
        for loss in losses:
            avg_loss = beta * loss + (1 - beta) * avg_loss
            smoothed.append(avg_loss)
        
        return np.array(smoothed)
    
    def get_optimal_lr(self):
        """Return suggested optimal learning rate"""
        if not self.losses:
            return None
        
        smoothed_losses = self._smooth_losses(self.losses, beta=0.05)
        gradients = np.gradient(smoothed_losses)
        min_gradient_idx = np.argmin(gradients)
        
        return self.lrs[min_gradient_idx]


# Helper function to easily integrate into training pipeline
def find_optimal_lr(model, train_dataloader, optimizer, criterion, device,
                    start_lr=1e-6, end_lr=1.0, num_iter=100, save_path=None):
    """
    Convenience function to find and plot optimal learning rate.
    
    Usage:
        optimal_lr = find_optimal_lr(model, train_dl, optimizer, criterion, device)
        print(f"Suggested LR: {optimal_lr:.2e}")
    
    Args:
        model: PyTorch model
        train_dataloader: Training DataLoader
        optimizer: Optimizer
        criterion: Loss function
        device: torch device
        start_lr: Starting LR (log scale)
        end_lr: Ending LR (log scale)
        num_iter: Number of iterations
        save_path: Optional path to save plot
    
    Returns:
        optimal_lr: Suggested learning rate
    """
    lr_finder = LRFinder(model, optimizer, criterion, device, 
                         start_lr=start_lr, end_lr=end_lr, num_iter=num_iter)
    lr_finder.find(train_dataloader)
    lr_finder.plot(save_path=save_path, log_scale=True)
    
    optimal_lr = lr_finder.get_optimal_lr()
    print(f"\n{'='*60}")
    print(f"Suggested Learning Rate: {optimal_lr:.2e}")
    print(f"{'='*60}\n")
    
    return optimal_lr


def hash_model_config(config):
    model_config = config["model"]
    config_str = json.dumps(model_config, sort_keys=True, default=str)
    config_hash = hashlib.md5(config_str.encode()).hexdigest()[:8]
    return config_hash

def save_training_results(
    save_path,
    conf,
    model,
    best_model_path,
    best_epoch,
    best_valid_loss,
    test_nll,
    test_ppl,
    train_losses,
    train_ppls,
    valid_losses,
    valid_ppls,
    raw_train_tokens,
    raw_valid_tokens,
    raw_test_tokens,
    vocab_size,
    warmup_steps,
    config_dict=None
):
    """
    Generate and save a comprehensive training results summary.
    
    Args:
        save_path (str): Directory to save results
        conf (Config): Configuration object with model hyperparameters
        model (nn.Module): Trained model
        best_model_path (str): Path to best model checkpoint
        best_epoch (int): Epoch with best validation loss
        best_valid_loss (float): Best validation loss achieved
        test_nll (float): Test loss
        test_ppl (float): Test perplexity
        train_losses (list): List of training losses per epoch
        train_ppls (list): List of training PPL per epoch
        valid_losses (list): List of validation losses per epoch
        valid_ppls (list): List of validation PPL per epoch
        raw_train_tokens (list): Training token data
        raw_valid_tokens (list): Validation token data
        raw_test_tokens (list): Test token data
        vocab_size (int): Vocabulary size
        warmup_steps (int): Number of warmup steps
        logger (logging.Logger, optional): Logger instance for printing
    
    Returns:
        dict: Results dictionary containing all metrics
    """
    
    # Calculate metrics
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    
    best_valid_ppl = math.exp(best_valid_loss)
    final_train_loss = train_losses[-1]
    final_train_ppl = train_ppls[-1]
    
    train_samples = len(raw_train_tokens)
    valid_samples = len(raw_valid_tokens)
    test_samples = len(raw_test_tokens)
    
    # Calculate warmup epochs
    warmup_epochs = warmup_steps / 100  # Approximate, adjust based on your data
    
    # Create results dictionary
    results_dict = {
        'model_config': {
            'vocab_size': conf.VOCAB_SIZE,
            'd_model': conf.D_MODEL,
            'n_heads': conf.N_HEADS,
            'n_layers': conf.N_LAYERS,
            'd_ff': conf.D_FF,
            'dropout': conf.DROPOUT,
            'max_len': conf.MAX_LEN,
        },
        'parameters': {
            'total': int(total_params),
            'trainable': int(trainable_params),
            'non_trainable': int(total_params - trainable_params),
        },
        'training_config': {
            'epochs': conf.EPOCHS,
            'batch_size': conf.BATCH_SIZE,
            'learning_rate': conf.LR,
            'weight_decay': conf.WEIGHT_DECAY,
            'grad_clip': conf.GRAD_CLIP,
            'warmup_steps': warmup_steps,
            'patience': conf.PATIENCE,
            'label_smoothing':conf.LABEL_SMOOTHING
        },
        'dataset': {
            'train_samples': train_samples,
            'valid_samples': valid_samples,
            'test_samples': test_samples,
            'vocab_size': vocab_size,
        },
        'results': {
            'best_epoch': int(best_epoch + 1),
            'best_valid_loss': float(best_valid_loss),
            'best_valid_ppl': float(best_valid_ppl),
            'final_train_loss': float(final_train_loss),
            'final_train_ppl': float(final_train_ppl),
            'test_loss': float(test_nll),
            'test_ppl': float(test_ppl),
        },
        'improvements': {
            'valid_ppl_vs_train_ppl': float(final_train_ppl - best_valid_ppl),
            'test_ppl_vs_valid_ppl': float(test_ppl - best_valid_ppl),
        },
        'paths': {
            'best_model': best_model_path,
            'results_dir': save_path,
        }
    }
    
    # Create formatted text summary
    results_text = f"""
================================================================================
                         TRAINING RESULTS SUMMARY
================================================================================

MODEL CONFIGURATION
-------------------
Vocabulary Size:         {conf.VOCAB_SIZE:,}
Model Dimension:         {conf.D_MODEL}
Attention Heads:         {conf.N_HEADS}
Number of Layers:        {conf.N_LAYERS}
Feed-Forward Dimension:  {conf.D_FF}
Dropout:                 {conf.DROPOUT}
Max Sequence Length:     {conf.MAX_LEN}

PARAMETERS
----------
Total Parameters:        {total_params:,}
Trainable Parameters:    {trainable_params:,}
Non-Trainable:           {total_params - trainable_params:,}

TRAINING CONFIGURATION
----------------------
Total Epochs:            {conf.EPOCHS}
Batch Size:              {conf.BATCH_SIZE}
Learning Rate:           {conf.LR}
Weight Decay:            {conf.WEIGHT_DECAY}
Gradient Clipping:       {conf.GRAD_CLIP}
Warmup Steps:            {warmup_steps}
Early Stopping Patience: {conf.PATIENCE}
Label Smoothing          {conf.LABEL_SMOOTHING}

DATASET STATISTICS
------------------
Training Samples:        {train_samples:,}
Validation Samples:      {valid_samples:,}
Test Samples:            {test_samples:,}
Vocabulary Size:         {vocab_size:,}

TRAINING RESULTS
----------------
Best Epoch:              {best_epoch + 1}/{conf.EPOCHS}
Final Training Loss:     {final_train_loss:.4f}
Final Training PPL:      {final_train_ppl:.2f}

Best Validation Loss:    {best_valid_loss:.4f}
Best Validation PPL:     {best_valid_ppl:.2f}

Test Loss (Final):       {test_nll:.4f}
Test PPL (Final):        {test_ppl:.2f}

IMPROVEMENTS & METRICS
----------------------
Valid PPL vs Train PPL:  {final_train_ppl - best_valid_ppl:.2f}
Test PPL vs Valid PPL:   {test_ppl - best_valid_ppl:.2f}

CHECKPOINTS & PATHS
-------------------
Best Model Path:         {best_model_path}
Results Directory:       {save_path}

================================================================================
"""
    
    # Save as text file
    results_file = os.path.join(save_path, "results.txt")
    with open(results_file, 'w') as f:
        f.write(results_text)

    if config_dict is not None:
        config_file = os.path.join(save_path, "best_config.yml")
        with open(config_file, 'w') as f:
            yaml.dump(config_dict, f, default_flow_style=False, sort_keys=False)
    
    return results_dict