train_mle.py

"""
Training pipeline for SMILES transformer model using Maximum Likelihood Estimation.
Implements complete training loop with optimization, monitoring, and checkpointing.
"""
import os
import json
import time
import math
from typing import Dict, List, Optional, Tuple
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torch.cuda.amp import GradScaler, autocast
import matplotlib.pyplot as plt
import seaborn as sns
from tqdm import tqdm
import argparse

from model import GPTModel
from tokenizer import SMILESTokenizer
from data_preprocessing import DataPreprocessor
from utils import get_model_config, get_training_config, get_data_config, count_parameters, format_number


class Trainer:
    """
    Trainer class for SMILES transformer model.
    """
    
    def __init__(self, model: GPTModel, tokenizer: SMILESTokenizer, 
                 train_loader: DataLoader, val_loader: DataLoader,
                 config: dict, device: str = 'auto'):
        self.model = model
        self.tokenizer = tokenizer
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.config = config
        
        # Device setup
        if device == 'auto':
            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        else:
            self.device = torch.device(device)
        
        print(f"Using device: {self.device}")
        self.model.to(self.device)
        
        # Training configuration
        training_config = config['training_config']
        self.learning_rate = training_config['learning_rate']
        self.max_epochs = training_config['max_epochs']
        self.gradient_clip = training_config['gradient_clip']
        self.warmup_steps = training_config['warmup_steps']
        
        # Optimizer
        self.optimizer = model.configure_optimizers(
            learning_rate=self.learning_rate,
            weight_decay=training_config['weight_decay'],
            beta1=training_config['beta1'],
            beta2=training_config['beta2']
        )
        
        # Learning rate scheduler (cosine annealing)
        total_steps = len(train_loader) * self.max_epochs
        self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
            self.optimizer, T_max=total_steps, eta_min=self.learning_rate * 0.1
        )
        
        # Mixed precision training
        self.use_amp = torch.cuda.is_available()
        self.scaler = GradScaler() if self.use_amp else None
        
        # Training state
        self.current_epoch = 0
        self.global_step = 0
        self.best_val_loss = float('inf')
        self.train_losses = []
        self.val_losses = []
        self.learning_rates = []
        
        # Create checkpoints directory
        os.makedirs('checkpoints', exist_ok=True)
    
    def warmup_lr(self, step: int) -> float:
        """
        Calculate learning rate with warmup.
        
        Args:
            step: Current training step
            
        Returns:
            Learning rate multiplier
        """
        if step < self.warmup_steps:
            return step / self.warmup_steps
        return 1.0
    
    def train_epoch(self) -> float:
        """
        Train for one epoch.
        
        Returns:
            Average training loss for the epoch
        """
        self.model.train()
        total_loss = 0.0
        num_batches = len(self.train_loader)
        
        progress_bar = tqdm(self.train_loader, desc=f"Epoch {self.current_epoch + 1}")
        
        for batch_idx, batch in enumerate(progress_bar):
            # Move batch to device
            input_ids = batch['input_ids'].to(self.device)
            target_ids = batch['target_ids'].to(self.device)
            attention_mask = batch['attention_mask'].to(self.device)
            
            # Zero gradients
            self.optimizer.zero_grad()
            
            # Forward pass with mixed precision
            if self.use_amp:
                with autocast():
                    logits, loss = self.model(input_ids, target_ids)
            else:
                logits, loss = self.model(input_ids, target_ids)
            
            # Backward pass
            if self.use_amp:
                self.scaler.scale(loss).backward()
                
                # Gradient clipping
                self.scaler.unscale_(self.optimizer)
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.gradient_clip)
                
                # Optimizer step
                self.scaler.step(self.optimizer)
                self.scaler.update()
            else:
                loss.backward()
                
                # Gradient clipping
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.gradient_clip)
                
                # Optimizer step
                self.optimizer.step()
            
            # Learning rate warmup
            if self.global_step < self.warmup_steps:
                lr_mult = self.warmup_lr(self.global_step)
                for param_group in self.optimizer.param_groups:
                    param_group['lr'] = self.learning_rate * lr_mult
            else:
                self.scheduler.step()
            
            # Update global step
            self.global_step += 1
            
            # Accumulate loss
            total_loss += loss.item()
            
            # Update progress bar
            current_lr = self.optimizer.param_groups[0]['lr']
            progress_bar.set_postfix({
                'loss': f'{loss.item():.4f}',
                'lr': f'{current_lr:.2e}',
                'step': self.global_step
            })
            
            # Store learning rate
            self.learning_rates.append(current_lr)
        
        avg_loss = total_loss / num_batches
        self.train_losses.append(avg_loss)
        
        return avg_loss
    
    def validate(self) -> float:
        """
        Validate the model.
        
        Returns:
            Average validation loss
        """
        self.model.eval()
        total_loss = 0.0
        num_batches = len(self.val_loader)
        
        with torch.no_grad():
            for batch in tqdm(self.val_loader, desc="Validation"):
                # Move batch to device
                input_ids = batch['input_ids'].to(self.device)
                target_ids = batch['target_ids'].to(self.device)
                
                # Forward pass
                if self.use_amp:
                    with autocast():
                        logits, loss = self.model(input_ids, target_ids)
                else:
                    logits, loss = self.model(input_ids, target_ids)
                
                total_loss += loss.item()
        
        avg_loss = total_loss / num_batches
        self.val_losses.append(avg_loss)
        
        return avg_loss
    
    def save_checkpoint(self, filepath: str, is_best: bool = False) -> None:
        """
        Save model checkpoint.
        
        Args:
            filepath: Path to save checkpoint
            is_best: Whether this is the best model so far
        """
        checkpoint = {
            'epoch': self.current_epoch,
            'global_step': self.global_step,
            'model_state_dict': self.model.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict(),
            'scheduler_state_dict': self.scheduler.state_dict(),
            'best_val_loss': self.best_val_loss,
            'train_losses': self.train_losses,
            'val_losses': self.val_losses,
            'learning_rates': self.learning_rates,
            'config': self.config
        }
        
        if self.scaler is not None:
            checkpoint['scaler_state_dict'] = self.scaler.state_dict()
        
        torch.save(checkpoint, filepath)
        
        if is_best:
            best_path = os.path.join(os.path.dirname(filepath), 'best_model.pt')
            torch.save(checkpoint, best_path)
        
        print(f"Checkpoint saved: {filepath}")
    
    def load_checkpoint(self, filepath: str) -> None:
        """
        Load model checkpoint.
        
        Args:
            filepath: Path to checkpoint file
        """
        checkpoint = torch.load(filepath, map_location=self.device)
        
        self.model.load_state_dict(checkpoint['model_state_dict'])
        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        self.scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
        
        self.current_epoch = checkpoint['epoch']
        self.global_step = checkpoint['global_step']
        self.best_val_loss = checkpoint['best_val_loss']
        self.train_losses = checkpoint['train_losses']
        self.val_losses = checkpoint['val_losses']
        self.learning_rates = checkpoint['learning_rates']
        
        if self.scaler is not None and 'scaler_state_dict' in checkpoint:
            self.scaler.load_state_dict(checkpoint['scaler_state_dict'])
        
        print(f"Checkpoint loaded: {filepath}")
    
    def plot_training_curves(self, save_path: str = 'training_curves.png') -> None:
        """
        Plot and save training curves.
        
        Args:
            save_path: Path to save the plot
        """
        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(15, 10))
        
        # Loss curves
        epochs = range(1, len(self.train_losses) + 1)
        ax1.plot(epochs, self.train_losses, 'b-', label='Training Loss', linewidth=2)
        ax1.plot(epochs, self.val_losses, 'r-', label='Validation Loss', linewidth=2)
        ax1.set_xlabel('Epoch')
        ax1.set_ylabel('Loss')
        ax1.set_title('Training and Validation Loss')
        ax1.legend()
        ax1.grid(True, alpha=0.3)
        
        # Log scale loss curves
        ax2.semilogy(epochs, self.train_losses, 'b-', label='Training Loss', linewidth=2)
        ax2.semilogy(epochs, self.val_losses, 'r-', label='Validation Loss', linewidth=2)
        ax2.set_xlabel('Epoch')
        ax2.set_ylabel('Loss (log scale)')
        ax2.set_title('Training and Validation Loss (Log Scale)')
        ax2.legend()
        ax2.grid(True, alpha=0.3)
        
        # Learning rate schedule
        steps = range(len(self.learning_rates))
        ax3.plot(steps, self.learning_rates, 'g-', linewidth=2)
        ax3.set_xlabel('Training Step')
        ax3.set_ylabel('Learning Rate')
        ax3.set_title('Learning Rate Schedule')
        ax3.grid(True, alpha=0.3)
        
        # Loss difference (overfitting indicator)
        loss_diff = [val - train for train, val in zip(self.train_losses, self.val_losses)]
        ax4.plot(epochs, loss_diff, 'purple', linewidth=2)
        ax4.axhline(y=0, color='black', linestyle='--', alpha=0.5)
        ax4.set_xlabel('Epoch')
        ax4.set_ylabel('Validation Loss - Training Loss')
        ax4.set_title('Overfitting Indicator')
        ax4.grid(True, alpha=0.3)
        
        plt.tight_layout()
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()
        
        print(f"Training curves saved: {save_path}")
    
    def train(self, resume_from: Optional[str] = None) -> None:
        """
        Main training loop.
        
        Args:
            resume_from: Path to checkpoint to resume from
        """
        if resume_from and os.path.exists(resume_from):
            self.load_checkpoint(resume_from)
            print(f"Resuming training from epoch {self.current_epoch + 1}")
        
        print(f"Starting training for {self.max_epochs} epochs")
        print(f"Model parameters: {count_parameters(self.model):,} ({format_number(count_parameters(self.model))})")
        print(f"Training batches: {len(self.train_loader)}")
        print(f"Validation batches: {len(self.val_loader)}")
        print(f"Mixed precision: {self.use_amp}")
        
        start_time = time.time()
        
        for epoch in range(self.current_epoch, self.max_epochs):
            self.current_epoch = epoch
            
            # Training
            train_loss = self.train_epoch()
            
            # Validation
            val_loss = self.validate()
            
            # Print epoch summary
            elapsed = time.time() - start_time
            print(f"\nEpoch {epoch + 1}/{self.max_epochs}")
            print(f"Train Loss: {train_loss:.4f}")
            print(f"Val Loss: {val_loss:.4f}")
            print(f"Learning Rate: {self.optimizer.param_groups[0]['lr']:.2e}")
            print(f"Time Elapsed: {elapsed/3600:.2f}h")
            
            # Save checkpoint
            checkpoint_path = f"checkpoints/checkpoint_epoch_{epoch + 1}.pt"
            is_best = val_loss < self.best_val_loss
            
            if is_best:
                self.best_val_loss = val_loss
                print(f"New best validation loss: {val_loss:.4f}")
            
            self.save_checkpoint(checkpoint_path, is_best=is_best)
            
            # Plot training curves
            if (epoch + 1) % 5 == 0 or epoch == self.max_epochs - 1:
                self.plot_training_curves(f"training_curves_epoch_{epoch + 1}.png")
        
        total_time = time.time() - start_time
        print(f"\nTraining completed in {total_time/3600:.2f} hours")
        print(f"Best validation loss: {self.best_val_loss:.4f}")
        
        # Final training curves
        self.plot_training_curves("final_training_curves.png")


def main():
    """Main training function."""
    parser = argparse.ArgumentParser(description='Train SMILES Transformer Model')
    parser.add_argument('--model_size', type=str, default='10M', choices=['10M', '25M', '50M'],
                       help='Model size configuration')
    parser.add_argument('--dataset', type=str, default='sample', choices=['sample', 'zinc', 'moses', 'moses_local'],
                       help='Dataset to use for training')
    parser.add_argument('--dataset_size', type=int, default=1000,
                       help='Size of sample dataset (only for sample dataset)')
    parser.add_argument('--batch_size', type=int, default=64,
                       help='Batch size for training')
    parser.add_argument('--max_epochs', type=int, default=50,
                       help='Maximum number of epochs')
    parser.add_argument('--learning_rate', type=float, default=3e-4,
                       help='Learning rate')
    parser.add_argument('--resume_from', type=str, default=None,
                       help='Path to checkpoint to resume from')
    parser.add_argument('--device', type=str, default='auto',
                       help='Device to use (auto, cpu, cuda)')
    parser.add_argument('--moses_dir', type=str, default='moses',
                       help='Directory containing local MOSES dataset files (for moses_local dataset)')
    
    args = parser.parse_args()
    
    print("SMILES Transformer Training")
    print("=" * 50)
    
    # Load configurations
    model_config = get_model_config(args.model_size)
    training_config = get_training_config()
    data_config = get_data_config()
    
    # Override training config with command line arguments
    training_config.max_epochs = args.max_epochs
    training_config.learning_rate = args.learning_rate
    training_config.batch_size = args.batch_size
    
    # Combine configurations
    config = {
        'model_config': model_config.__dict__,
        'training_config': training_config.__dict__,
        'data_config': data_config.__dict__
    }
    
    print(f"Model size: {args.model_size}")
    print(f"Dataset: {args.dataset}")
    print(f"Batch size: {args.batch_size}")
    print(f"Max epochs: {args.max_epochs}")
    print(f"Learning rate: {args.learning_rate}")
    
    # Initialize tokenizer
    print("\nInitializing tokenizer...")
    tokenizer = SMILESTokenizer(max_seq_len=model_config.max_seq_len)
    
    # Update model config with actual vocab size
    model_config.vocab_size = tokenizer.get_vocab_size()
    config['model_config']['vocab_size'] = model_config.vocab_size
    
    print(f"Vocabulary size: {tokenizer.get_vocab_size()}")
    
    # Initialize model
    print("Initializing model...")
    model = GPTModel(model_config)
    param_count = count_parameters(model)
    print(f"Model parameters: {param_count:,} ({format_number(param_count)})")
    
    # Prepare data
    print("\nPreparing data...")
    preprocessor = DataPreprocessor(config=data_config)
    
    # Check if processed data exists
    processed_data_file = f"processed_{args.dataset}_{args.dataset_size}.pkl"
    processed_data_path = os.path.join("data", processed_data_file)
    
    if os.path.exists(processed_data_path):
        print(f"Loading existing processed data: {processed_data_path}")
        train_smiles, val_smiles, test_smiles = preprocessor.load_processed_data(processed_data_file)
    else:
        print("Processing dataset...")
        train_smiles, val_smiles, test_smiles = preprocessor.preprocess_dataset(
            dataset_name=args.dataset,
            dataset_size=args.dataset_size,
            moses_dir=args.moses_dir
        )
        
        # Save processed data
        preprocessor.save_processed_data(train_smiles, val_smiles, test_smiles, processed_data_file)
    
    print(f"Dataset sizes - Train: {len(train_smiles)}, Val: {len(val_smiles)}, Test: {len(test_smiles)}")
    
    # Create datasets and dataloaders
    train_dataset, val_dataset, test_dataset = preprocessor.create_datasets(
        train_smiles, val_smiles, test_smiles, tokenizer
    )
    
    train_loader, val_loader, test_loader = preprocessor.create_dataloaders(
        train_dataset, val_dataset, test_dataset, 
        batch_size=args.batch_size, num_workers=0
    )
    
    # Initialize trainer
    print("\nInitializing trainer...")
    trainer = Trainer(
        model=model,
        tokenizer=tokenizer,
        train_loader=train_loader,
        val_loader=val_loader,
        config=config,
        device=args.device
    )
    
    # Start training
    print("\nStarting training...")
    trainer.train(resume_from=args.resume_from)
    
    print("\nTraining completed successfully!")


if __name__ == "__main__":
    main()