scripts/hf_cloud_training.py

#!/usr/bin/env python3
"""
Hugging Face Cloud Training Script for Morphological Reinflection
This script is designed to run on Hugging Face Spaces or other cloud infrastructure
"""

import os
import json
import logging
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, login

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 CloudTrainingConfig:
    """Configuration for cloud training"""
    
    def __init__(self):
        # Cloud-specific settings
        self.use_gpu = torch.cuda.is_available()
        self.device = torch.device('cuda' if self.use_gpu else 'cpu')
        
        # Model configuration
        self.embed_dim = 256
        self.nb_heads = 4
        self.src_hid_size = 1024
        self.src_nb_layers = 4
        self.trg_hid_size = 1024
        self.trg_nb_layers = 4
        self.dropout_p = 0.1
        self.tie_trg_embed = True
        self.label_smooth = 0.1
        self.max_length = 100
        
        # Training configuration
        self.batch_size = 32 if self.use_gpu else 16  # Smaller batch size for cloud
        self.learning_rate = 0.001
        self.max_epochs = 100  # Reduced for cloud training
        self.max_updates = 5000  # Reduced for cloud training
        self.warmup_steps = 500  # Reduced for cloud training
        self.weight_decay = 0.01
        self.gradient_clip = 1.0
        self.save_every = 5  # Save more frequently
        self.eval_every = 2  # Evaluate more frequently
        self.use_amp = self.use_gpu  # Use AMP only on GPU
        self.gradient_accumulation_steps = 4  # Increase for smaller batch sizes
        
        # Cloud-specific paths - use current directory to avoid permission issues
        self.data_dir = os.getenv('DATA_DIR', "./data")  # Data directory
        self.output_dir = os.getenv('OUTPUT_DIR', "./output")  # Output directory
        self.model_dir = os.getenv('MODEL_DIR', "./models")  # Model directory
        
        # Hugging Face settings
        self.hf_token = os.getenv('HF_TOKEN')
        self.model_name = os.getenv('MODEL_NAME', 'morphological-transformer')
        self.wandb_project = os.getenv('WANDB_PROJECT', 'morphological-transformer-cloud')
        
        # Dataset settings
        self.dataset_name = os.getenv('DATASET_NAME', '10L_90NL')
        self.run_number = os.getenv('RUN_NUMBER', '1')

class CloudMorphologicalTrainer:
    """Cloud-optimized trainer for morphological reinflection"""
    
    def __init__(self, model, config: CloudTrainingConfig, src_vocab, tgt_vocab):
        self.model = model
        self.config = config
        self.src_vocab = src_vocab
        self.tgt_vocab = tgt_vocab
        self.device = config.device
        
        # Initialize optimizer
        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.use_amp)
        
        # Learning rate scheduler
        self.scheduler = get_linear_schedule_with_warmup(
            self.optimizer,
            num_warmup_steps=config.warmup_steps,
            num_training_steps=config.max_updates
        )
        
    def train_epoch(self, dataloader, epoch):
        """Train for one epoch with cloud optimizations"""
        self.model.train()
        total_loss = 0.0
        num_batches = 0
        
        accumulation_steps = self.config.gradient_accumulation_steps
        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.use_amp):
                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()
                
                # Update learning rate
                self.scheduler.step()
            
            total_loss += loss.item() * accumulation_steps
            num_batches += 1
            
            # Log progress more frequently for cloud monitoring
            if batch_idx % 50 == 0:
                current_lr = self.scheduler.get_last_lr()[0]
                logger.info(f'Epoch {epoch}, Batch {batch_idx}, Loss: {loss.item() * accumulation_steps:.4f}, LR: {current_lr:.6f}')
        
        # 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()
            self.scheduler.step()
        
        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.use_amp):
                    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_cloud_model(config: CloudTrainingConfig, src_vocab: Dict[str, int], tgt_vocab: Dict[str, int]) -> TagTransformer:
    """Create and initialize the TagTransformer model for cloud training"""
    
    # 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_cloud_dataloader(dataset, config: CloudTrainingConfig, src_vocab: Dict, tgt_vocab: Dict):
    """Create optimized dataloader for cloud training"""
    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(2, os.cpu_count() or 1),  # Fewer workers for cloud
        pin_memory=config.use_gpu,
        persistent_workers=False,  # Disable for cloud stability
        prefetch_factor=2,
        drop_last=True
    )

def save_cloud_model(model, src_vocab, tgt_vocab, config: CloudTrainingConfig, output_dir: str, model_name: str):
    """Save model in cloud-compatible 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,
        "nb_attr": len([token for token in src_vocab.keys() if token.startswith('<') and token.endswith('>')]),
    }
    
    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,
        "use_amp": config.use_amp,
        "gradient_accumulation_steps": config.gradient_accumulation_steps,
    }
    
    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 upload_to_hf_hub(model_path: str, model_name: str, hf_token: str):
    """Upload model to Hugging Face Hub"""
    try:
        api = HfApi(token=hf_token)
        
        # Create repository
        api.create_repo(repo_id=model_name, exist_ok=True)
        
        # Upload files
        api.upload_folder(
            folder_path=model_path,
            repo_id=model_name,
            repo_type="model"
        )
        
        logger.info(f"Model uploaded to https://huggingface.co/{model_name}")
        return True
    except Exception as e:
        logger.error(f"Failed to upload model: {e}")
        return False

def main():
    # Set random seed for reproducibility
    set_seed(42)
    
    # Initialize configuration
    config = CloudTrainingConfig()
    
    logger.info(f"Starting cloud training with config: {config.__dict__}")
    
    # Login to Hugging Face if token is provided
    if config.hf_token:
        login(token=config.hf_token)
        logger.info("Logged in to Hugging Face Hub")
    
    # Initialize Weights & Biases if available
    try:
        wandb.init(project=config.wandb_project, name=config.model_name)
        logger.info("Initialized Weights & Biases")
    except Exception as e:
        logger.warning(f"Could not initialize Weights & Biases: {e}")
    
    # Create output directories with proper error handling
    try:
        os.makedirs(config.output_dir, exist_ok=True)
        os.makedirs(config.model_dir, exist_ok=True)
        logger.info(f"Created directories: {config.output_dir}, {config.model_dir}")
    except PermissionError as e:
        logger.error(f"Permission denied creating directories: {e}")
        logger.info("Falling back to current directory")
        config.output_dir = "./output"
        config.model_dir = "./models"
        os.makedirs(config.output_dir, exist_ok=True)
        os.makedirs(config.model_dir, exist_ok=True)
    
    # Set device
    device = config.device
    logger.info(f'Using device: {device}')
    
    # Enable CUDA optimizations if available
    if config.use_gpu:
        torch.backends.cudnn.benchmark = True
        torch.backends.cudnn.deterministic = False
        logger.info("CUDA optimizations enabled")
    
    # Build data paths - try multiple possible locations
    possible_data_paths = [
        Path(config.data_dir) / config.dataset_name,
        Path(f"./{config.dataset_name}"),  # Current directory
        Path(f"../{config.dataset_name}"),  # Parent directory
        Path(f"./data/{config.dataset_name}"),  # Local data directory
    ]
    
    data_path = None
    for path in possible_data_paths:
        if path.exists():
            data_path = path
            logger.info(f"Found data at: {data_path}")
            break
    
    if data_path is None:
        logger.error(f"Could not find data directory for {config.dataset_name}")
        logger.info(f"Searched in: {possible_data_paths}")
        return
    
    train_src = data_path / f"train/run{config.run_number}/train.{config.dataset_name}_{config.run_number}_1.src"
    train_tgt = data_path / f"train/run{config.run_number}/train.{config.dataset_name}_{config.run_number}_1.tgt"
    dev_src = data_path / f"dev/run{config.run_number}/dev.{config.dataset_name}_{config.run_number}_1.src"
    dev_tgt = data_path / f"dev/run{config.run_number}/dev.{config.dataset_name}_{config.run_number}_1.tgt"
    test_src = data_path / f"test/run{config.run_number}/test.{config.dataset_name}_{config.run_number}_1.src"
    test_tgt = data_path / f"test/run{config.run_number}/test.{config.dataset_name}_{config.run_number}_1.tgt"
    
    # Check if data files exist
    data_files = [train_src, train_tgt, dev_src, dev_tgt, test_src, test_tgt]
    missing_files = [f for f in data_files if not f.exists()]
    
    if missing_files:
        logger.error(f"Missing data files: {missing_files}")
        return
    
    # Build vocabulary
    logger.info("Building vocabulary...")
    all_data_files = [str(f) for f in data_files]
    vocab_stats = analyze_vocabulary(all_data_files)
    logger.info(f"Vocabulary statistics: {vocab_stats}")
    
    src_vocab = build_vocabulary([str(train_src), str(dev_src), str(test_src)])
    tgt_vocab = build_vocabulary([str(train_tgt), str(dev_tgt), str(test_tgt)])
    
    logger.info(f"Source vocabulary size: {len(src_vocab)}")
    logger.info(f"Target vocabulary size: {len(tgt_vocab)}")
    
    # Create datasets
    train_dataset = MorphologicalDataset(str(train_src), str(train_tgt), src_vocab, tgt_vocab, config.max_length)
    dev_dataset = MorphologicalDataset(str(dev_src), str(dev_tgt), src_vocab, tgt_vocab, config.max_length)
    
    # Create dataloaders
    train_loader = create_cloud_dataloader(train_dataset, config, src_vocab, tgt_vocab)
    dev_loader = create_cloud_dataloader(dev_dataset, config, src_vocab, tgt_vocab)
    
    # Create model
    model = create_cloud_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 = CloudMorphologicalTrainer(model, config, src_vocab, tgt_vocab)
    
    # 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
            try:
                wandb.log({
                    'epoch': epoch,
                    'train_loss': train_loss,
                    'val_loss': val_loss,
                    'learning_rate': trainer.scheduler.get_last_lr()[0]
                })
            except:
                pass
            
            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_cloud_model(model, src_vocab, tgt_vocab, config, config.model_dir, f"{config.model_name}_best")
        else:
            logger.info(f'Epoch {epoch}: Train Loss: {train_loss:.4f}')
        
        # Save checkpoint periodically
        if epoch % config.save_every == 0:
            save_cloud_model(model, src_vocab, tgt_vocab, config, config.model_dir, f"{config.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_cloud_model(model, src_vocab, tgt_vocab, config, config.model_dir, f"{config.model_name}_final")
    
    # Upload to Hugging Face Hub if token is provided
    if config.hf_token:
        best_model_path = Path(config.model_dir) / f"{config.model_name}_best"
        if best_model_path.exists():
            upload_to_hf_hub(str(best_model_path), config.model_name, config.hf_token)
    
    try:
        wandb.finish()
    except:
        pass
    
    logger.info('Cloud training completed!')

if __name__ == '__main__':
    main()