models/train.py

"""Training script for Myanmar Ghost sentiment model."""

import argparse
import logging
import sys
from pathlib import Path
from typing import Any, Dict, List, Optional

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from tqdm import tqdm

# Add parent directory to path
sys.path.insert(0, str(Path(__file__).parent.parent.parent))

from src.utils.logger import setup_logger
from src.utils.metrics import compute_metrics, MetricsTracker

logger = setup_logger("train", log_dir="outputs/logs")


class SentimentDataset(Dataset):
    """Dataset for sentiment classification."""
    
    def __init__(
        self,
        data: List[Dict],
        tokenizer,
        max_length: int = 512,
        label_mapping: Dict[str, int] = None,
    ):
        self.data = data
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.label_mapping = label_mapping or {
            "negative": 0,
            "neutral": 1,
            "positive": 2,
            "sarcastic": 3,
        }
    
    def __len__(self) -> int:
        return len(self.data)
    
    def __getitem__(self, idx: int) -> tuple:
        item = self.data[idx]
        
        encoding = self.tokenizer(
            item["text"],
            truncation=True,
            max_length=self.max_length,
            padding="max_length",
            return_tensors="pt",
        )
        
        label = self.label_mapping.get(item.get("label", "neutral"), 1)
        
        return (
            encoding["input_ids"].squeeze(0),
            encoding["attention_mask"].squeeze(0),
            torch.tensor(label, dtype=torch.long),
        )


def train_epoch(
    model: nn.Module,
    dataloader: DataLoader,
    criterion: nn.Module,
    optimizer: optim.Optimizer,
    device: torch.device,
    scheduler: Optional[Any] = None,
) -> Dict[str, float]:
    """Train for one epoch."""
    model.train()
    
    total_loss = 0.0
    all_predictions = []
    all_labels = []
    
    progress_bar = tqdm(dataloader, desc="Training")
    
    for batch_idx, (input_ids, attention_mask, labels) in enumerate(progress_bar):
        input_ids = input_ids.to(device)
        attention_mask = attention_mask.to(device)
        labels = labels.to(device)
        
        optimizer.zero_grad()
        
        outputs = model(input_ids, attention_mask)
        loss = criterion(outputs, labels)
        
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()
        
        if scheduler:
            scheduler.step()
        
        total_loss += loss.item()
        
        predictions = outputs.argmax(dim=-1).cpu().tolist()
        all_predictions.extend(predictions)
        all_labels.extend(labels.cpu().tolist())
        
        progress_bar.set_postfix({"loss": loss.item()})
    
    metrics = compute_metrics(all_predictions, all_labels)
    metrics["loss"] = total_loss / len(dataloader)
    
    return metrics


def evaluate(
    model: nn.Module,
    dataloader: DataLoader,
    criterion: nn.Module,
    device: torch.device,
) -> Dict[str, float]:
    """Evaluate the model."""
    model.eval()
    
    total_loss = 0.0
    all_predictions = []
    all_labels = []
    
    with torch.no_grad():
        for input_ids, attention_mask, labels in tqdm(dataloader, desc="Evaluating"):
            input_ids = input_ids.to(device)
            attention_mask = attention_mask.to(device)
            labels = labels.to(device)
            
            outputs = model(input_ids, attention_mask)
            loss = criterion(outputs, labels)
            
            total_loss += loss.item()
            
            predictions = outputs.argmax(dim=-1).cpu().tolist()
            all_predictions.extend(predictions)
            all_labels.extend(labels.cpu().tolist())
    
    metrics = compute_metrics(all_predictions, all_labels)
    metrics["loss"] = total_loss / len(dataloader)
    
    return metrics


def load_data(data_path: str) -> List[Dict]:
    """Load training data from JSON or JSONL file."""
    import json
    
    data = []
    
    if data_path.endswith(".jsonl"):
        with open(data_path, "r", encoding="utf-8") as f:
            for line in f:
                if line.strip():
                    data.append(json.loads(line))
    elif data_path.endswith(".json"):
        with open(data_path, "r", encoding="utf-8") as f:
            data = json.load(f)
    else:
        raise ValueError(f"Unsupported file format: {data_path}")
    
    return data


def main(args):
    """Main training function."""
    logger.info("Starting training...")
    logger.info(f"Arguments: {vars(args)}")
    
    # Device
    device = torch.device(
        "cuda" if torch.cuda.is_available() and not args.cpu else "cpu"
    )
    logger.info(f"Using device: {device}")
    
    # Load tokenizer
    logger.info(f"Loading tokenizer from {args.model_name}")
    from transformers import AutoTokenizer
    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
    
    # Load data
    logger.info(f"Loading data from {args.train_data}")
    train_data = load_data(args.train_data)
    val_data = load_data(args.val_data) if args.val_data else []
    
    logger.info(f"Train samples: {len(train_data)}, Val samples: {len(val_data)}")
    
    # Create datasets
    train_dataset = SentimentDataset(train_data, tokenizer, args.max_length)
    train_loader = DataLoader(
        train_dataset,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=2,
    )
    
    val_loader = None
    if val_data:
        val_dataset = SentimentDataset(val_data, tokenizer, args.max_length)
        val_loader = DataLoader(
            val_dataset,
            batch_size=args.batch_size,
            shuffle=False,
            num_workers=2,
        )
    
    # Create model
    logger.info("Creating model...")
    from src.models.transformer_model import TransformerSentimentModel
    
    model = TransformerSentimentModel(
        model_name=args.model_name,
        num_labels=4,
        dropout=args.dropout,
        freeze_encoder=args.freeze_encoder,
    )
    model.to(device)
    
    logger.info(f"Model parameters: {model.get_num_parameters():,}")
    logger.info(f"Trainable: {model.get_num_trainable_parameters():,}")
    
    # Loss and optimizer
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.AdamW(
        model.parameters(),
        lr=args.learning_rate,
        weight_decay=args.weight_decay,
    )
    
    # Scheduler
    total_steps = len(train_loader) * args.num_epochs
    warmup_steps = int(total_steps * 0.1)
    
    scheduler = optim.lr_scheduler.LinearLR(
        optimizer,
        start_factor=0.1,
        total_iters=warmup_steps,
    )
    
    # Training loop
    metrics_tracker = MetricsTracker(
        metrics=["loss", "accuracy", "f1_weighted"],
    )
    
    best_f1 = 0.0
    best_model_path = Path(args.output_dir) / "best_model.pt"
    
    for epoch in range(args.num_epochs):
        logger.info(f"\nEpoch {epoch + 1}/{args.num_epochs}")
        
        # Train
        train_metrics = train_epoch(
            model, train_loader, criterion, optimizer, device, scheduler
        )
        
        logger.info(f"Train - Loss: {train_metrics['loss']:.4f}, "
                   f"Acc: {train_metrics['accuracy']:.4f}, "
                   f"F1: {train_metrics['f1_weighted']:.4f}")
        
        # Evaluate
        if val_loader:
            val_metrics = evaluate(model, val_loader, criterion, device)
            
            logger.info(f"Val   - Loss: {val_metrics['loss']:.4f}, "
                        f"Acc: {val_metrics['accuracy']:.4f}, "
                        f"F1: {val_metrics['f1_weighted']:.4f}")
            
            metrics = {"train_" + k: v for k, v in train_metrics.items()}
            metrics.update({"val_" + k: v for k, v in val_metrics.items()})
        else:
            metrics = {"train_" + k: v for k, v in train_metrics.items()}
        
        metrics_tracker.update(metrics, epoch)
        
        # Save best model
        current_f1 = train_metrics.get("f1_weighted", 0)
        if current_f1 > best_f1:
            best_f1 = current_f1
            model.save(str(best_model_path))
            logger.info(f"Saved best model (F1: {best_f1:.4f})")
    
    # Save final model
    final_path = Path(args.output_dir) / "final_model.pt"
    model.save(str(final_path))
    
    logger.info(f"\nTraining complete! Best F1: {best_f1:.4f}")
    
    return best_f1


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Train Myanmar Ghost model")
    
    # Data arguments
    parser.add_argument("--train_data", type=str, required=True, help="Training data file")
    parser.add_argument("--val_data", type=str, default=None, help="Validation data file")
    parser.add_argument("--output_dir", type=str, default="outputs/models", help="Output directory")
    
    # Model arguments
    parser.add_argument("--model_name", type=str, default="bert-base-multilingual-cased")
    parser.add_argument("--max_length", type=int, default=512)
    parser.add_argument("--dropout", type=float, default=0.1)
    parser.add_argument("--freeze_encoder", action="store_true")
    
    # Training arguments
    parser.add_argument("--batch_size", type=int, default=16)
    parser.add_argument("--num_epochs", type=int, default=10)
    parser.add_argument("--learning_rate", type=float, default=5e-5)
    parser.add_argument("--weight_decay", type=float, default=0.01)
    parser.add_argument("--cpu", action="store_true", help="Use CPU only")
    
    args = parser.parse_args()
    
    main(args)