codecommentclassification/modeling/transformer/trainer.py

"""Training utilities for transformer-based multi-label classification.

This module contains a small training harness around HuggingFace
`AutoModelForSequenceClassification` specialized for the project's
multi-label code-comment classification task. It provides:

- `TransformerConfig` dataclass for configurable training runs.
- `CommentDataset` to wrap tokenization of pandas DataFrames.
- `TransformerTrainer` which runs the training loop, evaluation and
    model export (with MLflow logging hooks).

The helpers are intended for experimental, small-scale training and
instrumentation rather than production-grade distributed training.
"""

from dataclasses import asdict, dataclass
import logging
import os
from typing import Dict, List, Tuple

import mlflow
import numpy as np
import pandas as pd
from sklearn.metrics import (
    accuracy_score,
    classification_report,
    f1_score,
    precision_score,
    recall_score,
)
import torch
from torch.utils.data import DataLoader, Dataset
from tqdm.auto import tqdm
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    get_linear_schedule_with_warmup,
)

from .preprocessing import load_or_prepare_data

logger = logging.getLogger(__name__)

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {DEVICE}")


# Label names per language, order must match the label vector in the CSV
LABELS: Dict[str, Tuple[str, ...]] = {
    "java": (
        "summary",
        "Ownership",
        "Expand",
        "usage",
        "Pointer",
        "deprecation",
        "rational",
    ),
    "python": (
        "Usage",
        "Parameters",
        "DevelopmentNotes",
        "Expand",
        "Summary",
    ),
    "pharo": (
        "Keyimplementationpoints",
        "Example",
        "Responsibilities",
        "Intent",
        "Keymessages",
        "Collaborators",
    ),
}


@dataclass
class TransformerConfig:
    """Configuration for transformer training runs.

    Attributes are intentionally simple dataclass fields and map directly to
    CLI/YAML configuration keys used by the training harness.
    """

    lang: str
    raw_data_dir: str
    processed_data_dir: str
    model_output_path: str
    pretrained_model_name: str = "microsoft/codebert-base"
    max_length: int = 128
    batch_size: int = 16
    lr: float = 2e-5
    num_epochs: int = 5
    warmup_ratio: float = 0.1
    pos_weight_cap: float = 30.0
    threshold: float = 0.5
    preprocessing: bool = False
    preprocessing_factor: float = 1.0

    def __post_init__(self) -> None:
        """Force correct types even if YAML provides strings."""
        self.max_length = int(self.max_length)
        self.batch_size = int(self.batch_size)
        self.lr = float(self.lr)
        self.num_epochs = int(self.num_epochs)
        self.warmup_ratio = float(self.warmup_ratio)
        self.pos_weight_cap = float(self.pos_weight_cap)
        self.threshold = float(self.threshold)
        self.preprocessing_factor = float(self.preprocessing_factor)

        # allow 'true'/'false' as strings from YAML
        if isinstance(self.preprocessing, str):
            self.preprocessing = self.preprocessing.lower() == "true"


class CommentDataset(Dataset):
    """Simple Dataset wrapper around a pandas DataFrame with 'combo' and 'labels_array'."""

    def __init__(self, df: pd.DataFrame, tokenizer: AutoTokenizer, max_length: int):
        """Create a dataset that tokenizes rows on demand.

        Parameters
        ----------
        df : pandas.DataFrame
            Input frame containing at least `combo` and `labels_array` columns.
        tokenizer : transformers.AutoTokenizer
            Tokenizer used to encode text into model inputs.
        max_length : int
            Maximum tokenization length (used for padding/truncation).

        """
        self.df = df.reset_index(drop=True)
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self) -> int:
        """Return the number of examples in the dataset."""
        return len(self.df)

    def __getitem__(self, idx: int):
        """Return a single tokenized example and its labels as tensors.

        The returned dict contains tokenized inputs (PyTorch tensors) and a
        `labels` tensor suitable for BCEWithLogitsLoss for multi-label tasks.
        """
        row = self.df.iloc[idx]
        text = str(row["combo"])
        labels = np.asarray(row["labels_array"], dtype=np.float32)

        enc = self.tokenizer(
            text,
            truncation=True,
            max_length=self.max_length,
            padding="max_length",
            return_tensors="pt",
        )

        item = {k: v.squeeze(0) for k, v in enc.items()}
        item["labels"] = torch.from_numpy(labels)
        return item


class TransformerTrainer:
    """End-to-end transformer trainer for the code comment multi-label task."""

    def __init__(self, cfg: TransformerConfig) -> None:
        """Initialize training state, data loaders, model and optimizer.

        Parameters
        ----------
        cfg : TransformerConfig
            Training configuration containing data paths and hyperparameters.

        """
        self.cfg = cfg
        if cfg.lang not in LABELS:
            raise ValueError(f"No LABELS defined for language '{cfg.lang}'.")

        self.label_names = LABELS[cfg.lang]
        self.num_labels = len(self.label_names)

        logger.info("Initializing TransformerTrainer for language '%s'.", cfg.lang)
        logger.info("Raw data directory: %s", cfg.raw_data_dir)
        logger.info("Processed data directory: %s", cfg.processed_data_dir)
        logger.info("Model output path: %s", cfg.model_output_path)

        # --- data loading / preprocessing ---
        self.train_df, self.eval_df, self.preprocessing_used = load_or_prepare_data(
            lang=cfg.lang,
            raw_data_dir=cfg.raw_data_dir,
            processed_data_dir=cfg.processed_data_dir,
            preprocessing_enabled=cfg.preprocessing,
            preprocessing_factor=cfg.preprocessing_factor,
            random_state=42,
        )

        logger.info("Preprocessing used for this run: %s", self.preprocessing_used)
        logger.info("Using device: %s", DEVICE)
        logger.info(
            "Train size: %d rows, Eval size: %d rows",
            len(self.train_df),
            len(self.eval_df),
        )

        # --- log config and dataset info to MLflow ---
        try:
            cfg_dict = asdict(self.cfg)
            mlflow.log_params({f"cfg_{k}": v for k, v in cfg_dict.items()})
            mlflow.log_param("num_labels", self.num_labels)
            mlflow.log_param("label_names", ",".join(self.label_names))
            mlflow.log_param("train_samples", len(self.train_df))
            mlflow.log_param("eval_samples", len(self.eval_df))
            mlflow.log_param("preprocessing_used", self.preprocessing_used)
        except Exception as e:
            logger.warning("Could not log transformer config to MLflow: %s", e)

        # tokenizer
        logger.info("Loading tokenizer '%s'.", cfg.pretrained_model_name)
        self.tokenizer = AutoTokenizer.from_pretrained(cfg.pretrained_model_name)

        # label statistics and pos_weight
        y_train = np.stack(self.train_df["labels_array"].to_numpy())
        self.pos_weight = self._compute_pos_weight(y_train)

        # dataloaders
        train_dataset = CommentDataset(self.train_df, self.tokenizer, cfg.max_length)
        eval_dataset = CommentDataset(self.eval_df, self.tokenizer, cfg.max_length)

        self.train_loader = DataLoader(
            train_dataset,
            batch_size=cfg.batch_size,
            shuffle=True,
        )
        self.eval_loader = DataLoader(
            eval_dataset,
            batch_size=cfg.batch_size,
            shuffle=False,
        )

        logger.info(
            "Hyperparameters – lr=%s (type=%s), batch_size=%s, num_epochs=%s",
            self.cfg.lr,
            type(self.cfg.lr),
            self.cfg.batch_size,
            self.cfg.num_epochs,
        )

        # model
        logger.info("Loading base model '%s'.", cfg.pretrained_model_name)
        self.model = AutoModelForSequenceClassification.from_pretrained(
            cfg.pretrained_model_name,
            num_labels=self.num_labels,
            problem_type="multi_label_classification",
        ).to(DEVICE)

        self.loss_fn = torch.nn.BCEWithLogitsLoss(pos_weight=self.pos_weight.to(DEVICE))
        self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=self.cfg.lr)

        num_training_steps = cfg.num_epochs * len(self.train_loader)
        num_warmup_steps = int(cfg.warmup_ratio * num_training_steps)
        logger.info(
            "Total training steps: %d, warmup steps: %d.",
            num_training_steps,
            num_warmup_steps,
        )

        self.scheduler = get_linear_schedule_with_warmup(
            self.optimizer,
            num_warmup_steps=num_warmup_steps,
            num_training_steps=num_training_steps,
        )

        self.best_state_dict = None
        self.best_val_macro_f1 = 0.0

    def _compute_pos_weight(self, y: np.ndarray) -> torch.Tensor:
        if y.ndim == 1:
            y = y[:, None]
        freq = y.sum(axis=0).astype(np.float64)
        num_samples = y.shape[0]

        pos_weight = (num_samples - freq) / np.clip(freq, 1.0, None)
        pos_weight = np.clip(pos_weight, 1.0, self.cfg.pos_weight_cap)

        logger.info("Positive class weights (clipped): %s", pos_weight.tolist())
        return torch.tensor(pos_weight, dtype=torch.float32)

    def _step_batch(self, batch, train: bool):
        batch = {k: v.to(DEVICE) for k, v in batch.items()}
        labels = batch.pop("labels")

        outputs = self.model(**batch)
        logits = outputs.logits
        loss = self.loss_fn(logits, labels)

        if train:
            loss.backward()
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
            self.optimizer.step()
            self.scheduler.step()
            self.optimizer.zero_grad()

        return loss, logits, labels

    def train_one_epoch(self, epoch: int) -> float:
        """Run a single training epoch over `self.train_loader`.

        Returns
        -------
        float
            The average training loss over the epoch.

        """
        self.model.train()
        total_loss = 0.0
        n_samples = 0

        num_batches = len(self.train_loader)
        logger.info("Starting epoch %d training. Number of batches: %d", epoch, num_batches)

        progress_bar = tqdm(
            self.train_loader,
            desc=f"Epoch {epoch} [train]",
            total=num_batches,
            leave=False,
        )

        for step, batch in enumerate(progress_bar, start=1):
            loss, _, _ = self._step_batch(batch, train=True)
            batch_size = batch["input_ids"].size(0)
            total_loss += loss.item() * batch_size
            n_samples += batch_size

            avg_loss_so_far = total_loss / max(n_samples, 1)
            progress_bar.set_postfix({"loss": f"{avg_loss_so_far:.4f}"})

        avg_loss = total_loss / max(n_samples, 1)
        logger.info("Epoch %d training completed. Average loss: %.4f.", epoch, avg_loss)

        mlflow.log_metric("train_loss", avg_loss, step=epoch)

        return avg_loss

    def evaluate(
        self,
        epoch: int,
        split_name: str = "eval",
    ) -> Tuple[float, float, float, np.ndarray, np.ndarray]:
        """Evaluate the model on `self.eval_loader` and compute metrics.

        Parameters
        ----------
        epoch : int
            Current epoch number (used for logging).
        split_name : str
            Name of the evaluation split used for MLflow metric keys.

        Returns
        -------
        tuple
            (avg_loss, micro_f1, macro_f1, y_true, y_pred)

        """
        self.model.eval()
        total_loss = 0.0
        n_samples = 0
        all_preds: List[np.ndarray] = []
        all_labels: List[np.ndarray] = []

        logger.info("Starting evaluation for epoch %d on split '%s'.", epoch, split_name)

        num_batches = len(self.eval_loader)
        progress_bar = tqdm(
            self.eval_loader,
            desc=f"Epoch {epoch} [{split_name}]",
            total=num_batches,
            leave=False,
        )

        with torch.no_grad():
            for batch in progress_bar:
                loss, logits, labels = self._step_batch(batch, train=False)
                batch_size = logits.size(0)
                total_loss += loss.item() * batch_size
                n_samples += batch_size

                probs = torch.sigmoid(logits)
                preds = (probs > self.cfg.threshold).long()

                all_preds.append(preds.cpu().numpy())
                all_labels.append(labels.cpu().numpy())

                avg_loss_so_far = total_loss / max(n_samples, 1)
                progress_bar.set_postfix({"loss": f"{avg_loss_so_far:.4f}"})

        avg_loss = total_loss / max(n_samples, 1)
        y_pred = np.concatenate(all_preds, axis=0)
        y_true = np.concatenate(all_labels, axis=0)

        # F1
        micro_f1 = f1_score(y_true, y_pred, average="micro", zero_division=0)
        macro_f1 = f1_score(y_true, y_pred, average="macro", zero_division=0)

        # Precision
        micro_precision = precision_score(y_true, y_pred, average="micro", zero_division=0)
        macro_precision = precision_score(y_true, y_pred, average="macro", zero_division=0)

        # Recall
        micro_recall = recall_score(y_true, y_pred, average="micro", zero_division=0)
        macro_recall = recall_score(y_true, y_pred, average="macro", zero_division=0)

        # Accuracy (multi-label)
        # subset_accuracy = exact match of all labels for each sample
        subset_accuracy = accuracy_score(y_true, y_pred)
        # micro_accuracy = accuracy over flattened label indicators
        micro_accuracy = accuracy_score(y_true.flatten(), y_pred.flatten())

        logger.info(
            "Eval results [%s] - loss: %.4f | "
            "micro-F1: %.4f, macro-F1: %.4f | "
            "micro-P: %.4f, macro-P: %.4f | "
            "micro-R: %.4f, macro-R: %.4f | "
            "subset-acc: %.4f, micro-acc: %.4f",
            split_name,
            avg_loss,
            micro_f1,
            macro_f1,
            micro_precision,
            macro_precision,
            micro_recall,
            macro_recall,
            subset_accuracy,
            micro_accuracy,
        )

        # MLflow logging (per epoch)
        mlflow.log_metric(f"{split_name}_loss", avg_loss, step=epoch)
        mlflow.log_metric(f"{split_name}_micro_f1", micro_f1, step=epoch)
        mlflow.log_metric(f"{split_name}_macro_f1", macro_f1, step=epoch)
        mlflow.log_metric(f"{split_name}_micro_precision", micro_precision, step=epoch)
        mlflow.log_metric(f"{split_name}_macro_precision", macro_precision, step=epoch)
        mlflow.log_metric(f"{split_name}_micro_recall", micro_recall, step=epoch)
        mlflow.log_metric(f"{split_name}_macro_recall", macro_recall, step=epoch)
        mlflow.log_metric(f"{split_name}_subset_accuracy", subset_accuracy, step=epoch)
        mlflow.log_metric(f"{split_name}_micro_accuracy", micro_accuracy, step=epoch)

        return avg_loss, micro_f1, macro_f1, y_true, y_pred

    def run(self) -> Dict[str, float]:
        """Execute the full training loop and save the best model.

        Returns
        -------
        dict
            Summary metrics from the final evaluation (micro/macro F1).

        """
        logger.info("Starting training loop for %d epochs.", self.cfg.num_epochs)
        for epoch in range(1, self.cfg.num_epochs + 1):
            train_loss = self.train_one_epoch(epoch)
            val_loss, val_micro_f1, val_macro_f1, _, _ = self.evaluate(epoch, split_name="eval")

            logger.info(
                "[%s] epoch=%d train_loss=%.4f val_loss=%.4f val_micro_f1=%.4f val_macro_f1=%.4f",
                self.cfg.lang,
                epoch,
                train_loss,
                val_loss,
                val_micro_f1,
                val_macro_f1,
            )

            if val_macro_f1 > self.best_val_macro_f1:
                logger.info(
                    "New best macro-F1: %.4f (previous: %.4f). Saving current model state.",
                    val_macro_f1,
                    self.best_val_macro_f1,
                )
                self.best_val_macro_f1 = val_macro_f1
                self.best_state_dict = {k: v.cpu() for k, v in self.model.state_dict().items()}

        if self.best_state_dict is not None:
            logger.info("Loading best model weights (macro-F1 = %.4f).", self.best_val_macro_f1)
            self.model.load_state_dict(self.best_state_dict)

        # final evaluation
        _, micro_f1, macro_f1, y_true, y_pred = self.evaluate(
            epoch=self.cfg.num_epochs,
            split_name="eval",
        )

        logger.info(
            "[%s] FINAL micro-F1 = %.4f, macro-F1 = %.4f.",
            self.cfg.lang,
            micro_f1,
            macro_f1,
        )
        logger.info(
            "Per-label classification report:\n%s",
            classification_report(y_true, y_pred, target_names=self.label_names, zero_division=0),
        )

        # save model and tokenizer
        os.makedirs(self.cfg.model_output_path, exist_ok=True)
        logger.info("Saving model and tokenizer to '%s'.", self.cfg.model_output_path)
        self.model.save_pretrained(self.cfg.model_output_path)
        self.tokenizer.save_pretrained(self.cfg.model_output_path)

        # log model directory as MLflow artifact
        logger.info("Logging final model artifacts to MLflow.")
        mlflow.log_artifacts(
            self.cfg.model_output_path,
            artifact_path=f"{self.cfg.lang}_transformer_model",
        )

        logger.info("Logging HF transformers model to MLflow via mlflow.transformers.log_model.")
        model_info = mlflow.transformers.log_model(
            transformers_model=self.cfg.model_output_path,
            artifact_path=f"{self.cfg.lang}_transformer_model",
            task="text-classification",
        )

        logger.info(
            "Logged transformers model to MLflow with URI: %s",
            model_info.model_uri,
        )

        return {
            "micro_f1": float(micro_f1),
            "macro_f1": float(macro_f1),
        }