scripts/train_huggingface.py

#!/usr/bin/env python3
"""
Hugging Face training script for morphological reinflection using TagTransformer
"""

import argparse
import json
import logging
import os
import time
from pathlib import Path
from typing import Dict, Tuple, Optional

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torch.cuda.amp import GradScaler, autocast

# Hugging Face imports
from transformers import (
    Trainer, 
    TrainingArguments, 
    HfArgumentParser,
    set_seed,
    get_linear_schedule_with_warmup
)
from datasets import Dataset, DatasetDict
import wandb
from huggingface_hub import HfApi, Repository

from transformer import TagTransformer, PAD_IDX, DEVICE
from morphological_dataset import MorphologicalDataset, build_vocabulary, collate_fn, analyze_vocabulary

# Set up logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class MorphologicalTrainer:
    """Custom trainer for morphological reinflection"""
    
    def __init__(self, model, config, src_vocab, tgt_vocab, device):
        self.model = model
        self.config = config
        self.src_vocab = src_vocab
        self.tgt_vocab = tgt_vocab
        self.device = device
        
        # Initialize optimizer and scheduler
        self.optimizer = optim.AdamW(
            model.parameters(),
            lr=config['learning_rate'],
            weight_decay=config['weight_decay'],
            betas=(0.9, 0.999),
            eps=1e-8
        )
        
        # Mixed precision training
        self.scaler = GradScaler(enabled=config.get('use_amp', True))
        
    def train_epoch(self, dataloader, epoch):
        """Train for one epoch"""
        self.model.train()
        total_loss = 0.0
        num_batches = 0
        
        accumulation_steps = self.config.get('gradient_accumulation_steps', 1)
        self.optimizer.zero_grad()
        
        for batch_idx, (src, src_mask, tgt, tgt_mask) in enumerate(dataloader):
            src, src_mask, tgt, tgt_mask = (
                src.to(self.device, non_blocking=True), 
                src_mask.to(self.device, non_blocking=True), 
                tgt.to(self.device, non_blocking=True), 
                tgt_mask.to(self.device, non_blocking=True)
            )
            
            # Mixed precision forward pass
            with autocast(enabled=self.config.get('use_amp', True)):
                output = self.model(src, src_mask, tgt, tgt_mask)
                loss = self.model.loss(output[:-1], tgt[1:])
                loss = loss / accumulation_steps
            
            # Mixed precision backward pass
            self.scaler.scale(loss).backward()
            
            # Gradient accumulation
            if (batch_idx + 1) % accumulation_steps == 0:
                # Gradient clipping
                self.scaler.unscale_(self.optimizer)
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=self.config['gradient_clip'])
                
                # Optimizer step
                self.scaler.step(self.optimizer)
                self.scaler.update()
                self.optimizer.zero_grad()
            
            total_loss += loss.item() * accumulation_steps
            num_batches += 1
            
            if batch_idx % 100 == 0:
                logger.info(f'Epoch {epoch}, Batch {batch_idx}, Loss: {loss.item() * accumulation_steps:.4f}')
        
        # Handle remaining gradients
        if num_batches % accumulation_steps != 0:
            self.scaler.unscale_(self.optimizer)
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=self.config['gradient_clip'])
            self.scaler.step(self.optimizer)
            self.scaler.update()
            self.optimizer.zero_grad()
        
        return total_loss / num_batches
    
    def validate(self, dataloader):
        """Validate the model"""
        self.model.eval()
        total_loss = 0.0
        num_batches = 0
        
        with torch.no_grad():
            for src, src_mask, tgt, tgt_mask in dataloader:
                src, src_mask, tgt, tgt_mask = (
                    src.to(self.device, non_blocking=True), 
                    src_mask.to(self.device, non_blocking=True), 
                    tgt.to(self.device, non_blocking=True), 
                    tgt_mask.to(self.device, non_blocking=True)
                )
                
                with autocast(enabled=self.config.get('use_amp', True)):
                    output = self.model(src, src_mask, tgt, tgt_mask)
                    loss = self.model.loss(output[:-1], tgt[1:])
                
                total_loss += loss.item()
                num_batches += 1
        
        return total_loss / num_batches

def create_model(config: Dict, src_vocab: Dict[str, int], tgt_vocab: Dict[str, int]) -> TagTransformer:
    """Create and initialize the TagTransformer model"""
    
    # Count feature tokens
    feature_tokens = [token for token in src_vocab.keys() 
                     if token.startswith('<') and token.endswith('>')]
    nb_attr = len(feature_tokens)
    
    logger.info(f"Found {nb_attr} feature tokens")
    
    model = TagTransformer(
        src_vocab_size=len(src_vocab),
        trg_vocab_size=len(tgt_vocab),
        embed_dim=config['embed_dim'],
        nb_heads=config['nb_heads'],
        src_hid_size=config['src_hid_size'],
        src_nb_layers=config['src_nb_layers'],
        trg_hid_size=config['trg_hid_size'],
        trg_nb_layers=config['trg_nb_layers'],
        dropout_p=config['dropout_p'],
        tie_trg_embed=config['tie_trg_embed'],
        label_smooth=config['label_smooth'],
        nb_attr=nb_attr,
        src_c2i=src_vocab,
        trg_c2i=tgt_vocab,
        attr_c2i={},
    )
    
    # Initialize weights
    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p)
        elif p.dim() == 1:
            nn.init.uniform_(p, -0.1, 0.1)
    
    return model

def create_dataloader(dataset, config: Dict, src_vocab: Dict, tgt_vocab: Dict):
    """Create optimized dataloader"""
    def collate_wrapper(batch):
        return collate_fn(batch, src_vocab, tgt_vocab, config['max_length'])
    
    return DataLoader(
        dataset,
        batch_size=config['batch_size'],
        shuffle=True,
        collate_fn=collate_wrapper,
        num_workers=min(4, os.cpu_count() or 1),
        pin_memory=torch.cuda.is_available(),
        persistent_workers=True,
        prefetch_factor=2,
        drop_last=True
    )

def save_model_to_hf(model, src_vocab, tgt_vocab, config, output_dir, model_name):
    """Save model in Hugging Face format"""
    
    # Create model directory
    model_dir = Path(output_dir) / model_name
    model_dir.mkdir(parents=True, exist_ok=True)
    
    # Save model state dict
    torch.save(model.state_dict(), model_dir / "pytorch_model.bin")
    
    # Save configuration
    model_config = {
        "model_type": "tag_transformer",
        "src_vocab_size": len(src_vocab),
        "trg_vocab_size": len(tgt_vocab),
        "embed_dim": config['embed_dim'],
        "nb_heads": config['nb_heads'],
        "src_hid_size": config['src_hid_size'],
        "src_nb_layers": config['src_nb_layers'],
        "trg_hid_size": config['trg_hid_size'],
        "trg_nb_layers": config['trg_nb_layers'],
        "dropout_p": config['dropout_p'],
        "tie_trg_embed": config['tie_trg_embed'],
        "label_smooth": config['label_smooth'],
        "max_length": config['max_length'],
    }
    
    with open(model_dir / "config.json", "w") as f:
        json.dump(model_config, f, indent=2)
    
    # Save vocabularies
    with open(model_dir / "src_vocab.json", "w") as f:
        json.dump(src_vocab, f, indent=2)
    
    with open(model_dir / "tgt_vocab.json", "w") as f:
        json.dump(tgt_vocab, f, indent=2)
    
    # Save training arguments
    training_args = {
        "learning_rate": config['learning_rate'],
        "batch_size": config['batch_size'],
        "max_epochs": config['max_epochs'],
        "warmup_steps": config['warmup_steps'],
        "weight_decay": config['weight_decay'],
        "gradient_clip": config['gradient_clip'],
    }
    
    with open(model_dir / "training_args.json", "w") as f:
        json.dump(training_args, f, indent=2)
    
    logger.info(f"Model saved to {model_dir}")

def main():
    parser = argparse.ArgumentParser(description='Train TagTransformer on Hugging Face')
    parser.add_argument('--model_name', type=str, required=True, help='Model name for Hugging Face')
    parser.add_argument('--output_dir', type=str, default='./hf_models', help='Output directory')
    parser.add_argument('--train_src', type=str, required=True, help='Training source file')
    parser.add_argument('--train_tgt', type=str, required=True, help='Training target file')
    parser.add_argument('--dev_src', type=str, required=True, help='Development source file')
    parser.add_argument('--dev_tgt', type=str, required=True, help='Development target file')
    parser.add_argument('--test_src', type=str, help='Test source file (optional)')
    parser.add_argument('--test_tgt', type=str, help='Test target file (optional)')
    parser.add_argument('--wandb_project', type=str, help='Weights & Biases project name')
    parser.add_argument('--hf_token', type=str, help='Hugging Face token for model upload')
    parser.add_argument('--upload_model', action='store_true', help='Upload model to Hugging Face Hub')
    parser.add_argument('--no_amp', action='store_true', help='Disable mixed precision training')
    args = parser.parse_args()
    
    # Set random seed for reproducibility
    set_seed(42)
    
    # Initialize Weights & Biases if specified
    if args.wandb_project:
        wandb.init(project=args.wandb_project, name=args.model_name)
    
    # Configuration
    config = {
        'embed_dim': 256,
        'nb_heads': 4,
        'src_hid_size': 1024,
        'src_nb_layers': 4,
        'trg_hid_size': 1024,
        'trg_nb_layers': 4,
        'dropout_p': 0.1,
        'tie_trg_embed': True,
        'label_smooth': 0.1,
        'batch_size': 400,
        'learning_rate': 0.001,
        'max_epochs': 1000,
        'max_updates': 10000,
        'warmup_steps': 4000,
        'weight_decay': 0.01,
        'gradient_clip': 1.0,
        'save_every': 10,
        'eval_every': 5,
        'max_length': 100,
        'use_amp': not args.no_amp,
        'gradient_accumulation_steps': 2,
    }
    
    # Create output directory
    os.makedirs(args.output_dir, exist_ok=True)
    
    # Set device
    device = DEVICE
    logger.info(f'Using device: {device}')
    
    # Enable CUDA optimizations if available
    if torch.cuda.is_available():
        torch.backends.cudnn.benchmark = True
        torch.backends.cudnn.deterministic = False
        logger.info("CUDA optimizations enabled")
    
    # Build vocabulary
    logger.info("Building vocabulary...")
    all_data_files = [args.train_src, args.train_tgt, args.dev_src, args.dev_tgt]
    if args.test_src and args.test_tgt:
        all_data_files.extend([args.test_src, args.test_tgt])
    
    vocab_stats = analyze_vocabulary(all_data_files)
    logger.info(f"Vocabulary statistics: {vocab_stats}")
    
    src_vocab = build_vocabulary([args.train_src, args.dev_src] + ([args.test_src] if args.test_src else []))
    tgt_vocab = build_vocabulary([args.train_tgt, args.dev_tgt] + ([args.test_tgt] if args.test_tgt else []))
    
    logger.info(f"Source vocabulary size: {len(src_vocab)}")
    logger.info(f"Target vocabulary size: {len(tgt_vocab)}")
    
    # Create datasets
    train_dataset = MorphologicalDataset(args.train_src, args.train_tgt, src_vocab, tgt_vocab, config['max_length'])
    dev_dataset = MorphologicalDataset(args.dev_src, args.dev_tgt, src_vocab, tgt_vocab, config['max_length'])
    
    # Create dataloaders
    train_loader = create_dataloader(train_dataset, config, src_vocab, tgt_vocab)
    dev_loader = create_dataloader(dev_dataset, config, src_vocab, tgt_vocab)
    
    # Create model
    model = create_model(config, src_vocab, tgt_vocab)
    model = model.to(device)
    
    # Count parameters
    total_params = model.count_nb_params()
    logger.info(f'Total parameters: {total_params:,}')
    
    # Create trainer
    trainer = MorphologicalTrainer(model, config, src_vocab, tgt_vocab, device)
    
    # Training loop
    best_val_loss = float('inf')
    global_step = 0
    
    for epoch in range(config['max_epochs']):
        start_time = time.time()
        
        # Train
        train_loss = trainer.train_epoch(train_loader, epoch)
        
        # Validate
        if epoch % config['eval_every'] == 0:
            val_loss = trainer.validate(dev_loader)
            
            # Log metrics
            if args.wandb_project:
                wandb.log({
                    'epoch': epoch,
                    'train_loss': train_loss,
                    'val_loss': val_loss,
                    'learning_rate': trainer.optimizer.param_groups[0]['lr']
                })
            
            logger.info(f'Epoch {epoch}: Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}')
            
            # Save best model
            if val_loss < best_val_loss:
                best_val_loss = val_loss
                save_model_to_hf(model, src_vocab, tgt_vocab, config, args.output_dir, f"{args.model_name}_best")
        else:
            logger.info(f'Epoch {epoch}: Train Loss: {train_loss:.4f}')
        
        # Save checkpoint periodically
        if epoch % config['save_every'] == 0:
            save_model_to_hf(model, src_vocab, tgt_vocab, config, args.output_dir, f"{args.model_name}_epoch_{epoch}")
        
        epoch_time = time.time() - start_time
        logger.info(f'Epoch {epoch} completed in {epoch_time:.2f}s')
        
        global_step += len(train_loader)
        
        # Check if we've reached max updates
        if global_step >= config['max_updates']:
            logger.info(f'Reached maximum updates ({config["max_updates"]}), stopping training')
            break
    
    # Save final model
    save_model_to_hf(model, src_vocab, tgt_vocab, config, args.output_dir, f"{args.model_name}_final")
    
    # Upload to Hugging Face Hub if requested
    if args.upload_model and args.hf_token:
        try:
            api = HfApi(token=args.hf_token)
            model_path = Path(args.output_dir) / f"{args.model_name}_best"
            
            # Create repository
            api.create_repo(repo_id=args.model_name, exist_ok=True)
            
            # Upload files
            api.upload_folder(
                folder_path=str(model_path),
                repo_id=args.model_name,
                repo_type="model"
            )
            
            logger.info(f"Model uploaded to https://huggingface.co/{args.model_name}")
        except Exception as e:
            logger.error(f"Failed to upload model: {e}")
    
    if args.wandb_project:
        wandb.finish()
    
    logger.info('Training completed!')

if __name__ == '__main__':
    main()