model.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from pytorch_lightning import LightningModule
from transformers import AutoModel
from torchmetrics.classification import Accuracy, Precision, Recall, F1Score
from typing import Optional, List
import pandas as pd
from pathlib import Path
import wandb
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from torchmetrics.classification import ConfusionMatrix


class BertClassifier(LightningModule):
    """
    Generic HF encoder + linear head classifier.
    Works with ModernBERT, BERT, RoBERTa, and DeBERTa backbones via `backbone_name`.
    """
    def __init__(
        self,
        n_classes: int,
        backbone_name: str = "answerdotai/ModernBERT-base",
        dropout: float = 0.1,
        max_lr: float = 5e-5,
        weight_decay: float = 0.01,
        label_smoothing: float = 0.0,
        class_names: Optional[List[str]] = None,
        # extras already present in your setup:
        class_weights: Optional[List[float]] = None,
        flooding_val: float = 0.0,
        # NEW: where to save test-set probabilities
        pred_output_path: Optional[str] = None,
        # NEW: enable HF gradient checkpointing on the backbone
        grad_checkpointing: bool = False,
        **kwargs
    ):
        super().__init__()
        self.save_hyperparameters(ignore=[])

        self.backbone_name = backbone_name
        self.backbone = AutoModel.from_pretrained(backbone_name)
        if grad_checkpointing:
            try:
                # Works for most HF encoder models, including DeBERTa/ModernBERT
                self.backbone.gradient_checkpointing_enable()
                # a hint for configs that read this flag
                if hasattr(self.backbone, "config"):
                    self.backbone.config.gradient_checkpointing = True
            except Exception:
                pass

        hidden = self.backbone.config.hidden_size

        self.dropout = nn.Dropout(dropout)
        self.classifier = nn.Linear(hidden, n_classes)
        self.class_names = class_names or [str(i) for i in range(n_classes)]

        # optional class weights
        self._class_weights_tensor: Optional[torch.Tensor] = (
            torch.tensor(class_weights, dtype=torch.float32) if class_weights is not None else None
        )

        # ---- container for probabilities over the whole test set ----
        self.test_rows: list = []     # list of (Prediction, Label, Probabilities[list])
        self.pred_output_path = pred_output_path  # may be None; handled later

        # Metrics
        self.train_accuracy = Accuracy(task="multiclass", num_classes=n_classes)
        self.val_accuracy   = Accuracy(task="multiclass", num_classes=n_classes)
        self.test_accuracy  = Accuracy(task="multiclass", num_classes=n_classes)

        self.f1_micro = F1Score(task="multiclass", num_classes=n_classes, average="micro")
        self.f1_macro = F1Score(task="multiclass", num_classes=n_classes, average="macro")

        self.f1_per_class = F1Score(task="multiclass", num_classes=n_classes, average=None)

        self.prec_per_class = Precision(task="multiclass", num_classes=n_classes, average=None)
        self.recl_per_class = Recall(task="multiclass", num_classes=n_classes, average=None)

        self.prec_macro = Precision(task="multiclass", num_classes=n_classes, average="macro")
        self.recl_macro = Recall(task="multiclass", num_classes=n_classes, average="macro")

        # Confusion Matrix
        self.val_confusion_matrix  = ConfusionMatrix(task="multiclass", num_classes=n_classes, normalize="true")
        self.test_confusion_matrix = ConfusionMatrix(task="multiclass", num_classes=n_classes, normalize="true")

    def setup(self, stage: Optional[str] = None):
        # move class weights to device or pull from datamodule if provided there
        if self._class_weights_tensor is None:
            dm = getattr(self.trainer, "datamodule", None)
            cw = getattr(dm, "class_weights", None)
            if cw is not None:
                self._class_weights_tensor = torch.as_tensor(cw, dtype=torch.float32, device=self.device)
        else:
            self._class_weights_tensor = self._class_weights_tensor.to(self.device)

        # default path for probs file if not provided
        if self.pred_output_path is None:
            self.pred_output_path = "preds/test_probs.pkl"

    def forward(self, input_ids, attention_mask):
        out = self.backbone(input_ids=input_ids, attention_mask=attention_mask)
        pooled = out.last_hidden_state[:, 0]  # [CLS]
        logits = self.classifier(self.dropout(pooled))
        return logits

    def _shared_step(self, batch, stage: str):
        input_ids, attention_mask, labels = batch["input_ids"], batch["attention_mask"], batch["labels"]
        logits = self(input_ids, attention_mask)

        loss = F.cross_entropy(
            logits, labels,
            weight=self._class_weights_tensor,
            label_smoothing=self.hparams.label_smoothing,
            reduction="mean"
        )

        # b is a flooding value, ensuring loss doesn't fall below this point to prevent overfitting
        b = float(self.hparams.flooding_val)
        if b > 0.0 and stage == "train":
            loss = (loss - b).abs() + loss

        probs = torch.softmax(logits, dim=-1)
        preds = torch.argmax(probs, dim=1)

        if stage == "train":
            self.train_accuracy.update(preds, labels)
            self.log("train_loss", loss, on_step=True, prog_bar=True)
            self.log("train_acc",  self.train_accuracy, on_step=True, on_epoch=True, prog_bar=True)

        elif stage == "val":
            self.val_accuracy.update(preds, labels)
            self.f1_micro.update(preds, labels)
            self.f1_macro.update(preds, labels)
            self.f1_per_class.update(preds, labels)
            self.prec_per_class.update(preds, labels)
            self.recl_per_class.update(preds, labels)
            self.prec_macro.update(preds, labels)
            self.recl_macro.update(preds, labels)
            self.val_confusion_matrix.update(preds, labels)
            self.log("val_loss", loss, on_step=False, on_epoch=True, prog_bar=True, sync_dist=True)
            self.log("val_acc",  self.val_accuracy, on_step=False, on_epoch=True, prog_bar=True, sync_dist=True)

        else:  # test
            self.test_accuracy.update(preds, labels)
            self.test_confusion_matrix.update(preds, labels)
            self.log("test_loss", loss, on_step=False, on_epoch=True, prog_bar=True, sync_dist=True)
            self.log("test_acc",  self.test_accuracy, on_step=False, on_epoch=True, prog_bar=True, sync_dist=True)

            # ---- store per-sample probabilities for the whole test set ----
            self.test_rows.extend(
                zip(
                    preds.detach().cpu().tolist(),
                    labels.detach().cpu().tolist(),
                    probs.detach().cpu().tolist(),   # list[float] per sample (len == n_classes)
                )
            )

        return loss

    def training_step(self, batch, batch_idx):
        return self._shared_step(batch, "train")

    def validation_step(self, batch, batch_idx):
        return self._shared_step(batch, "val")

    def on_validation_epoch_end(self):
        self.log("val_f1_micro", self.f1_micro.compute(), prog_bar=True, sync_dist=True)
        self.log("val_f1_macro", self.f1_macro.compute(), prog_bar=True, sync_dist=True)
        self.log("val_precision_macro", self.prec_macro.compute(), sync_dist=True)
        self.log("val_recall_macro",    self.recl_macro.compute(), sync_dist=True)

        prec = self.prec_per_class.compute()
        recl = self.recl_per_class.compute()
        f1pc = self.f1_per_class.compute()

        for i, name in enumerate(self.class_names):
            self.log(f"val_precision/{name}", prec[i], sync_dist=True)
            self.log(f"val_recall/{name}",    recl[i], sync_dist=True)
            self.log(f"val_f1/{name}",        f1pc[i], sync_dist=True)

        # compact W&B table
        try:
            import wandb
            table = wandb.Table(columns=["class", "precision", "recall", "f1"])
            for i, name in enumerate(self.class_names):
                table.add_data(str(name), float(prec[i]), float(recl[i]), float(f1pc[i]))
            wandb.log({"val/per_class_table": table}, commit=False)
        except Exception:
            pass

        # ---- Confusion matrix (VALIDATION): plot → label → log → reset ----
        try:
            fig, ax = self.val_confusion_matrix.plot(add_text=False)

            if getattr(self, "class_names", None):
                n = min(len(self.class_names), self.val_confusion_matrix.num_classes)
                ax.set_xticks(range(n)); ax.set_yticks(range(n))
                ax.set_xticklabels(self.class_names[:n], rotation=90, fontsize=7)
                ax.set_yticklabels(self.class_names[:n], fontsize=7)

            plt.tight_layout()
            if getattr(self.trainer, "is_global_zero", True):
                wandb.log({"val_confusion_matrix": [wandb.Image(fig)]}, commit=False)
        except Exception as e:
            print(f"[WARN] val CM plot/log failed: {e}")
        finally:
            plt.close(fig)
            self.val_confusion_matrix.reset()

        # reset all other metrics
        for m in [self.f1_micro, self.f1_macro, self.f1_per_class,
                self.prec_per_class, self.recl_per_class,
                self.prec_macro, self.recl_macro, self.val_accuracy]:
            m.reset()

    def test_step(self, batch, batch_idx):
        return self._shared_step(batch, "test")

    def on_test_epoch_end(self):
        # Confusion Matrix logging ----
        fig, ax = self.test_confusion_matrix.plot(add_text=False)
        if getattr(self, "class_names", None):
            n = min(len(self.class_names), self.test_confusion_matrix.num_classes)
            ax.set_xticks(range(n)); ax.set_yticks(range(n))
            ax.set_xticklabels(self.class_names[:n], rotation=90, fontsize=7)
            ax.set_yticklabels(self.class_names[:n], fontsize=7)
        plt.tight_layout()

        wandb.log({"test_confusion_matrix": [wandb.Image(fig)]})
        plt.close(fig)
        self.test_confusion_matrix.reset()

        # persist the collected rows as a DataFrame
        self.test_accuracy.reset()
        if not self.test_rows:
            return

        df = pd.DataFrame(self.test_rows, columns=["Prediction", "Label", "Probabilities"])
        out = Path(self.pred_output_path)
        out.parent.mkdir(parents=True, exist_ok=True)
        # Save both pickle (lossless) and a CSV (flat) for quick use
        df.to_pickle(out)  # e.g., preds/run_<JOBID>/test_probs.pkl

        # also write a compact CSV with probabilities exploded into columns
        csv_out = out.with_suffix(".csv")
        prob_cols = self.class_names if self.class_names else [f"class{i}" for i in range(len(df["Probabilities"][0]))]
        probs_expanded = pd.DataFrame(df["Probabilities"].tolist(), columns=prob_cols)
        flat = pd.concat([df[["Prediction", "Label"]].reset_index(drop=True), probs_expanded], axis=1)
        flat.to_csv(csv_out, index=False)

        print(f"[info] saved test probabilities: {out}  (and {csv_out})")
        self.test_rows.clear()

    def configure_optimizers(self):
        opt = torch.optim.AdamW(self.parameters(), lr=self.hparams.max_lr, weight_decay=self.hparams.weight_decay)
        sch = torch.optim.lr_scheduler.OneCycleLR(
            opt,
            max_lr=self.hparams.max_lr,
            total_steps=self.trainer.estimated_stepping_batches,
            anneal_strategy="cos",
            div_factor=25.0,
            final_div_factor=1e4
        )
        return {"optimizer": opt, "lr_scheduler": {"scheduler": sch, "interval": "step", "name": "onecycle"}}

    def save_model(self, path):
        self.trainer.save_checkpoint(f"{path}.ckpt")
        self.backbone.config.save_pretrained(path)


class DebertaV3Classifier(BertClassifier):
    """
    Thin alias that defaults to the canonical HF hub name for DeBERTa-v3-large.
    Everything else (loss/metrics/logging) is inherited from BertClassifier.
    """
    def __init__(self, n_classes: int, backbone_name: str = "microsoft/deberta-v3-large", **kwargs):
        super().__init__(n_classes=n_classes, backbone_name=backbone_name, **kwargs)