trainer.py

"""
Model Training for Code-Mixed Machine Translation
Using: facebook/nllb-200-distilled-600M + HuggingFace Trainer API
"""

import os
import logging
import numpy as np
from typing import Dict, List, Optional, Any, Tuple

import torch
from transformers import (
    AutoModelForSeq2SeqLM,
    AutoTokenizer,
    Seq2SeqTrainer,
    Seq2SeqTrainingArguments,
    DataCollatorForSeq2Seq,
    EarlyStoppingCallback,
    GenerationConfig,
)
from datasets import DatasetDict
import evaluate

logger = logging.getLogger(__name__)


# 1. MODEL LOADER

class NLLBModelLoader:
    """
    Loads and configures the NLLB-200 model for code-mixed MT fine-tuning.
    """

    def __init__(
        self,
        model_name: str = "facebook/nllb-200-distilled-600M",
        device: Optional[str] = None,
    ):
        self.model_name = model_name
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        logger.info(f"Using device: {self.device}")

    def load(self) -> AutoModelForSeq2SeqLM:
        """Load the pretrained NLLB-200 model."""
        logger.info(f"Loading model: {self.model_name}")

        model = AutoModelForSeq2SeqLM.from_pretrained(
            self.model_name,
            # Use torch_dtype for memory efficiency on GPU
            torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
        )

        # Log model parameter count
        total_params = sum(p.numel() for p in model.parameters())
        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
        logger.info(f"Total parameters: {total_params:,}")
        logger.info(f"Trainable parameters: {trainable_params:,}")

        return model

    def enable_gradient_checkpointing(self, model: AutoModelForSeq2SeqLM) -> AutoModelForSeq2SeqLM:
        """Enable gradient checkpointing to save memory at cost of speed."""
        model.gradient_checkpointing_enable()
        logger.info("Gradient checkpointing enabled.")
        return model


# 2. EVALUATION METRICS

class TranslationMetrics:
    """
    Computes BLEU, SacreBLEU, ChrF, and COMET metrics for MT evaluation.
    """

    def __init__(self, tokenizer, tgt_lang: str):
        self.tokenizer = tokenizer
        self.tgt_lang = tgt_lang

        # Load metric libraries
        self.sacrebleu = evaluate.load("sacrebleu")
        self.chrf = evaluate.load("chrf")

        # COMET is optional (requires GPU and extra download)
        try:
            self.comet = evaluate.load("comet")
            self.comet_available = True
            logger.info("COMET metric loaded successfully.")
        except Exception:
            self.comet_available = False
            logger.warning("COMET not available. Install with: pip install unbabel-comet")

    def compute_metrics(self, eval_preds) -> Dict[str, float]:
        """
        HuggingFace Trainer-compatible metric computation.
        Called at end of each evaluation step.
        """
        preds, labels = eval_preds

        # Decode predictions (handle -100 padding labels)
        if isinstance(preds, tuple):
            preds = preds[0]

        # Clamp token IDs to valid range before decoding
        preds = np.where(preds != -100, preds, self.tokenizer.pad_token_id)
        labels = np.where(labels != -100, labels, self.tokenizer.pad_token_id)

        decoded_preds = self.tokenizer.batch_decode(preds, skip_special_tokens=True)
        decoded_labels = self.tokenizer.batch_decode(labels, skip_special_tokens=True)

        # Strip whitespace
        decoded_preds = [pred.strip() for pred in decoded_preds]
        decoded_labels = [label.strip() for label in decoded_labels]

        # SacreBLEU expects list of references (each is a list)
        references = [[label] for label in decoded_labels]

        # Compute SacreBLEU
        bleu_result = self.sacrebleu.compute(
            predictions=decoded_preds,
            references=references,
        )

        # Compute ChrF
        chrf_result = self.chrf.compute(
            predictions=decoded_preds,
            references=decoded_labels,
        )

        metrics = {
            "bleu": round(bleu_result["score"], 4),
            "chrf": round(chrf_result["score"], 4),
        }

        # COMET (if available)
        if self.comet_available and len(decoded_preds) > 0:
            try:
                comet_result = self.comet.compute(
                    predictions=decoded_preds,
                    references=decoded_labels,
                    sources=decoded_labels,  # Approximate: use target as source reference
                )
                metrics["comet"] = round(comet_result["mean_score"], 4)
            except Exception as e:
                logger.warning(f"COMET computation failed: {e}")

        logger.info(f"Evaluation metrics: {metrics}")
        return metrics

    def evaluate_test_set(
        self,
        predictions: List[str],
        references: List[str],
        sources: Optional[List[str]] = None,
    ) -> Dict[str, float]:
        """Full evaluation on a test set after training."""
        references_wrapped = [[ref] for ref in references]

        bleu = self.sacrebleu.compute(predictions=predictions, references=references_wrapped)
        chrf = self.chrf.compute(predictions=predictions, references=references)

        results = {
            "sacrebleu": round(bleu["score"], 4),
            "chrf": round(chrf["score"], 4),
            "bp": round(bleu.get("bp", 0.0), 4),
            "precision_1gram": round(bleu["precisions"][0] if bleu.get("precisions") else 0.0, 4),
        }

        if self.comet_available and sources:
            try:
                comet = self.comet.compute(
                    predictions=predictions,
                    references=references,
                    sources=sources,
                )
                results["comet"] = round(comet["mean_score"], 4)
            except Exception as e:
                logger.warning(f"COMET test evaluation failed: {e}")

        return results


# 3. TRAINING PIPELINE

class CodeMixedMTTrainer:
    """
    Full training pipeline for code-mixed machine translation
    using HuggingFace Seq2SeqTrainer.
    """

    def __init__(
        self,
        model: AutoModelForSeq2SeqLM,
        tokenizer: AutoTokenizer,
        tokenized_datasets: DatasetDict,
        src_lang: str,
        tgt_lang: str,
        training_config,
        model_config,
    ):
        self.model = model
        self.tokenizer = tokenizer
        self.tokenized_datasets = tokenized_datasets
        self.src_lang = src_lang
        self.tgt_lang = tgt_lang
        self.training_config = training_config
        self.model_config = model_config

        # Initialize metrics
        self.metrics = TranslationMetrics(tokenizer, tgt_lang)

        # Data collator with dynamic padding
        self.data_collator = DataCollatorForSeq2Seq(
            tokenizer=tokenizer,
            model=model,
            padding=True,          # Pads to longest in batch (efficient)
            pad_to_multiple_of=8,  # For tensor core efficiency
            label_pad_token_id=-100,
        )

    def build_training_args(self) -> Seq2SeqTrainingArguments:
        """Construct Seq2SeqTrainingArguments from config."""
        cfg = self.training_config

        return Seq2SeqTrainingArguments(
            output_dir=cfg.output_dir,
            num_train_epochs=cfg.num_train_epochs,
            per_device_train_batch_size=cfg.per_device_train_batch_size,
            per_device_eval_batch_size=cfg.per_device_eval_batch_size,
            learning_rate=cfg.learning_rate,
            warmup_ratio=cfg.warmup_ratio,
            weight_decay=cfg.weight_decay,
            gradient_accumulation_steps=cfg.gradient_accumulation_steps,
            fp16=cfg.fp16 and torch.cuda.is_available(),  # Only on GPU
            bf16=False,  # Use fp16 instead (broader GPU support)
            evaluation_strategy=cfg.evaluation_strategy,
            eval_steps=cfg.eval_steps,
            save_strategy=cfg.save_strategy,
            save_steps=cfg.save_steps,
            save_total_limit=cfg.save_total_limit,
            load_best_model_at_end=cfg.load_best_model_at_end,
            metric_for_best_model=cfg.metric_for_best_model,
            greater_is_better=cfg.greater_is_better,
            predict_with_generate=cfg.predict_with_generate,
            generation_max_length=cfg.generation_max_length,
            logging_steps=cfg.logging_steps,
            report_to=cfg.report_to,
            dataloader_num_workers=cfg.dataloader_num_workers,
            group_by_length=cfg.group_by_length,
            seed=cfg.seed,
            # NLLB-specific: force BOS token to target language id
            forced_bos_token_id=self.tokenizer.convert_tokens_to_ids(self.tgt_lang),
        )

    def build_trainer(self) -> Seq2SeqTrainer:
        """Construct the HuggingFace Seq2SeqTrainer."""
        training_args = self.build_training_args()

        # Log the forced BOS token for verification
        bos_token_id = self.tokenizer.lang_code_to_id.get(self.tgt_lang)
        logger.info(f"Forced BOS token for '{self.tgt_lang}': {bos_token_id}")

        trainer = Seq2SeqTrainer(
            model=self.model,
            args=training_args,
            train_dataset=self.tokenized_datasets.get("train"),
            eval_dataset=self.tokenized_datasets.get("validation"),
            tokenizer=self.tokenizer,
            data_collator=self.data_collator,
            compute_metrics=self.metrics.compute_metrics,
            callbacks=[
                EarlyStoppingCallback(
                    early_stopping_patience=3,
                    early_stopping_threshold=0.001,
                )
            ],
        )

        return trainer

    def train(self) -> Dict[str, Any]:
        """Run the full training loop."""
        logger.info("=" * 60)
        logger.info("Starting Code-Mixed MT Training")
        logger.info(f"  Model: {self.model.config.name_or_path}")
        logger.info(f"  src_lang: {self.src_lang} -> tgt_lang: {self.tgt_lang}")
        logger.info(f"  Train samples: {len(self.tokenized_datasets['train'])}")
        logger.info(f"  Val samples: {len(self.tokenized_datasets['validation'])}")
        logger.info("=" * 60)

        trainer = self.build_trainer()

        # Train
        train_result = trainer.train()

        # Save final model and tokenizer
        logger.info("Saving final model...")
        trainer.save_model()
        self.tokenizer.save_pretrained(self.training_config.output_dir)

        # Save training metrics
        metrics = train_result.metrics
        trainer.log_metrics("train", metrics)
        trainer.save_metrics("train", metrics)
        trainer.save_state()

        return {"trainer": trainer, "train_metrics": metrics}

    def evaluate_test(self, trainer: Seq2SeqTrainer) -> Dict[str, float]:
        """Run evaluation on the held-out test set."""
        logger.info("Evaluating on test set...")

        test_dataset = self.tokenized_datasets.get("test")
        if test_dataset is None:
            logger.warning("No test split found, skipping test evaluation.")
            return {}

        # Generate predictions on test set
        predictions = trainer.predict(
            test_dataset,
            max_length=self.model_config.max_target_length,
            num_beams=4,
        )

        pred_ids = predictions.predictions
        label_ids = predictions.label_ids

        # Decode
        pred_ids = np.where(pred_ids != -100, pred_ids, self.tokenizer.pad_token_id)
        label_ids = np.where(label_ids != -100, label_ids, self.tokenizer.pad_token_id)

        decoded_preds = self.tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
        decoded_labels = self.tokenizer.batch_decode(label_ids, skip_special_tokens=True)

        decoded_preds = [p.strip() for p in decoded_preds]
        decoded_labels = [l.strip() for l in decoded_labels]

        # Compute full metrics
        test_metrics = self.metrics.evaluate_test_set(decoded_preds, decoded_labels)
        logger.info(f"Test metrics: {test_metrics}")

        return test_metrics, decoded_preds, decoded_labels