08. Sentiment Classifier Fine-Tuning — Frustration Detection for Escalation

graph TB
    subgraph "Epoch 1: Only head trainable"
        E1_EMB["Embeddings ❄️"]
        E1_L01["Layers 0-1 ❄️"]
        E1_L23["Layers 2-3 ❄️"]
        E1_L45["Layers 4-5 ❄️"]
        E1_HEAD["Classification Head 🔥<br>lr = 3e-5"]
        E1_EMB --> E1_L01 --> E1_L23 --> E1_L45 --> E1_HEAD
    end

    subgraph "Epoch 2: Top layers unfrozen"
        E2_EMB["Embeddings ❄️"]
        E2_L01["Layers 0-1 ❄️"]
        E2_L23["Layers 2-3 ❄️"]
        E2_L45["Layers 4-5 🔥<br>lr = 1.15e-5"]
        E2_HEAD["Classification Head 🔥<br>lr = 3e-5"]
        E2_EMB --> E2_L01 --> E2_L23 --> E2_L45 --> E2_HEAD
    end

    subgraph "Epoch 3: Middle layers unfrozen"
        E3_EMB["Embeddings ❄️"]
        E3_L01["Layers 0-1 ❄️"]
        E3_L23["Layers 2-3 🔥<br>lr = 4.4e-6"]
        E3_L45["Layers 4-5 🔥<br>lr = 1.15e-5"]
        E3_HEAD["Classification Head 🔥<br>lr = 3e-5"]
        E3_EMB --> E3_L01 --> E3_L23 --> E3_L45 --> E3_HEAD
    end

    subgraph "Epoch 4+: All layers unfrozen"
        E4_EMB["Embeddings 🔥<br>lr = 6.6e-7"]
        E4_L01["Layers 0-1 🔥<br>lr = 1.7e-6"]
        E4_L23["Layers 2-3 🔥<br>lr = 4.4e-6"]
        E4_L45["Layers 4-5 🔥<br>lr = 1.15e-5"]
        E4_HEAD["Classification Head 🔥<br>lr = 3e-5"]
        E4_EMB --> E4_L01 --> E4_L23 --> E4_L45 --> E4_HEAD
    end

    style E1_HEAD fill:#ffcdd2
    style E2_L45 fill:#fff9c4
    style E2_HEAD fill:#ffcdd2
    style E3_L23 fill:#c8e6c9
    style E3_L45 fill:#fff9c4
    style E3_HEAD fill:#ffcdd2
    style E4_EMB fill:#e1bee7
    style E4_L01 fill:#e1bee7
    style E4_L23 fill:#c8e6c9
    style E4_L45 fill:#fff9c4
    style E4_HEAD fill:#ffcdd2

graph LR
    subgraph "Gradient Magnitude During Training"
        subgraph "Epoch 1 (head only)"
            G1_E["Embed: 0.000"]
            G1_L0["L0: 0.000"]
            G1_L2["L2: 0.000"]
            G1_L4["L4: 0.000"]
            G1_H["Head: 0.042"]
        end

        subgraph "Epoch 2 (+ top layers)"
            G2_E["Embed: 0.000"]
            G2_L0["L0: 0.000"]
            G2_L2["L2: 0.000"]
            G2_L4["L4: 0.008"]
            G2_H["Head: 0.031"]
        end

        subgraph "Epoch 4 (all layers)"
            G4_E["Embed: 0.001"]
            G4_L0["L0: 0.002"]
            G4_L2["L2: 0.005"]
            G4_L4["L4: 0.012"]
            G4_H["Head: 0.018"]
        end
    end

    style G1_H fill:#ff5252
    style G2_L4 fill:#ffab40
    style G2_H fill:#ff5252
    style G4_E fill:#e8eaf6
    style G4_L0 fill:#c5cae9
    style G4_L2 fill:#fff9c4
    style G4_L4 fill:#ffab40
    style G4_H fill:#ff5252

graph TB
    INPUT["Input: 'The manga arrived damaged<br>but the story is amazing'"]

    subgraph "DistilBERT Encoder (66M params)"
        TOK["Tokenizer → [CLS] The manga arrived damaged ..."]
        EMB["Embedding Layer (23M)"]
        L0["Transformer Layer 0"]
        L1["Transformer Layer 1"]
        L2["Transformer Layer 2"]
        L3["Transformer Layer 3"]
        L4["Transformer Layer 4"]
        L5["Transformer Layer 5"]
        CLS["[CLS] representation ∈ ℝ⁷⁶⁸"]
    end

    subgraph "Classification Head (multi-label)"
        DROP["Dropout (0.2)"]
        FC1["Linear: 768 → 256 + ReLU"]
        FC2["Linear: 256 → 6 (one per label)"]
        SIG["Sigmoid (independent per label)"]
    end

    subgraph "Output Probabilities"
        O1["frustration: 0.89 ⬆️"]
        O2["satisfaction: 0.74 ⬆️"]
        O3["urgency: 0.12"]
        O4["confusion: 0.08"]
        O5["churn_risk: 0.31"]
        O6["neutral: 0.05"]
    end

    subgraph "Threshold Gate"
        T1["frustration: 0.89 > 0.35 ✅"]
        T2["satisfaction: 0.74 > 0.55 ✅"]
        T3["urgency: 0.12 < 0.40 ❌"]
        T4["confusion: 0.08 < 0.32 ❌"]
        T5["churn_risk: 0.31 > 0.28 ✅"]
        T6["neutral: 0.05 < 0.48 ❌"]
    end

    INPUT --> TOK --> EMB --> L0 --> L1 --> L2 --> L3 --> L4 --> L5 --> CLS
    CLS --> DROP --> FC1 --> FC2 --> SIG
    SIG --> O1 & O2 & O3 & O4 & O5 & O6
    O1 --> T1
    O2 --> T2
    O3 --> T3
    O4 --> T4
    O5 --> T5
    O6 --> T6

    style O1 fill:#ffcdd2
    style O2 fill:#c8e6c9
    style O5 fill:#fff9c4
    style T1 fill:#ffcdd2
    style T2 fill:#c8e6c9
    style T5 fill:#fff9c4

graph TB
    subgraph "Standard BCE: All examples weighted equally"
        BCE_EASY["Easy examples (p > 0.8)<br>70% of data<br>Total gradient share: 70%"]
        BCE_HARD["Hard examples (p < 0.5)<br>15% of data<br>Total gradient share: 15%"]
        BCE_RARE["Rare class hard examples<br>3% of data (churn_risk)<br>Total gradient share: 3%"]
    end

    subgraph "Focal Loss (γ=2): Hard examples amplified"
        FL_EASY["Easy examples (p > 0.8)<br>70% of data<br>Gradient share: 12% ⬇️"]
        FL_HARD["Hard examples (p < 0.5)<br>15% of data<br>Gradient share: 48% ⬆️"]
        FL_RARE["Rare class hard examples<br>3% of data (churn_risk)<br>Gradient share: 40% ⬆️"]
    end

    BCE_EASY -->|"Focal<br>reweighting"| FL_EASY
    BCE_HARD -->|"γ=2"| FL_HARD
    BCE_RARE -->|"γ=2 + α=6.0"| FL_RARE

    style BCE_RARE fill:#ffcdd2
    style FL_RARE fill:#c8e6c9
    style FL_HARD fill:#fff9c4

graph TD
    subgraph "Per-Label Threshold Search"
        S["sweep t ∈ [0.1, 0.9]<br>step = 0.05"]

        subgraph "frustration (β=2, recall-focused)"
            F1["t=0.50 → P=0.82 R=0.64 F2=0.67"]
            F2["t=0.35 → P=0.71 R=0.86 F2=0.82 ◀ OPTIMAL"]
            F3["t=0.20 → P=0.58 R=0.93 F2=0.83"]
        end

        subgraph "churn_risk (β=2, recall-focused)"
            C1["t=0.50 → P=0.78 R=0.42 F2=0.46"]
            C2["t=0.28 → P=0.55 R=0.81 F2=0.74 ◀ OPTIMAL"]
            C3["t=0.15 → P=0.32 R=0.91 F2=0.67"]
        end

        subgraph "satisfaction (β=0.5, precision-focused)"
            SA1["t=0.40 → P=0.69 R=0.85 F0.5=0.71"]
            SA2["t=0.55 → P=0.83 R=0.72 F0.5=0.81 ◀ OPTIMAL"]
            SA3["t=0.70 → P=0.91 R=0.54 F0.5=0.82"]
        end
    end

    S --> F1 & F2 & F3
    S --> C1 & C2 & C3
    S --> SA1 & SA2 & SA3

    style F2 fill:#c8e6c9
    style C2 fill:#c8e6c9
    style SA2 fill:#c8e6c9

import torch
import torch.nn as nn
from transformers import (
    DistilBertModel,
    DistilBertTokenizer,
    get_linear_schedule_with_warmup,
)
from torch.utils.data import Dataset, DataLoader
import numpy as np
from sklearn.metrics import f1_score, precision_recall_fscore_support


LABELS = ["frustration", "satisfaction", "urgency", "confusion", "churn_risk", "neutral"]
NUM_LABELS = len(LABELS)


class MultiLabelSentimentModel(nn.Module):
    """Multi-label sentiment classifier with per-label sigmoid outputs."""

    def __init__(self, model_name: str = "distilbert-base-uncased", hidden_dim: int = 256):
        super().__init__()
        self.encoder = DistilBertModel.from_pretrained(model_name)
        self.head = nn.Sequential(
            nn.Dropout(0.2),
            nn.Linear(768, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_dim, NUM_LABELS),
        )

    def forward(self, input_ids, attention_mask):
        outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        cls_emb = outputs.last_hidden_state[:, 0, :]
        logits = self.head(cls_emb)
        return logits  # (batch, NUM_LABELS) — raw logits, apply sigmoid externally


class FocalBCELoss(nn.Module):
    """Focal Binary Cross-Entropy for multi-label with class imbalance."""

    def __init__(self, gamma: float = 2.0, alpha: list[float] = None):
        super().__init__()
        self.gamma = gamma
        if alpha is not None:
            self.alpha = torch.tensor(alpha, dtype=torch.float32)
        else:
            self.alpha = None

    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        probs = torch.sigmoid(logits)

        # Standard BCE components
        pos_loss = targets * torch.log(probs + 1e-8)
        neg_loss = (1 - targets) * torch.log(1 - probs + 1e-8)

        # Focal modulation
        pos_focal = (1 - probs) ** self.gamma
        neg_focal = probs ** self.gamma

        loss = -(pos_focal * pos_loss + neg_focal * neg_loss)

        # Apply per-label weights
        if self.alpha is not None:
            alpha = self.alpha.to(logits.device)
            loss = loss * alpha.unsqueeze(0)

        return loss.mean()


class SentimentDataset(Dataset):
    def __init__(self, texts: list[str], labels: list[list[int]], tokenizer, max_length: int = 128):
        self.encodings = tokenizer(
            texts, padding=True, truncation=True,
            max_length=max_length, return_tensors="pt",
        )
        self.labels = torch.tensor(labels, dtype=torch.float32)

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        return {
            "input_ids": self.encodings["input_ids"][idx],
            "attention_mask": self.encodings["attention_mask"][idx],
            "labels": self.labels[idx],
        }

class GradualUnfreezingTrainer:
    """
    Implements ULMFiT-style gradual unfreezing with discriminative LRs.
    """

    def __init__(
        self,
        model: MultiLabelSentimentModel,
        train_loader: DataLoader,
        val_loader: DataLoader,
        base_lr: float = 3e-5,
        decay_factor: float = 2.6,
        num_epochs: int = 8,
        gamma: float = 2.0,
        label_weights: list[float] = None,
    ):
        self.model = model
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.base_lr = base_lr
        self.decay_factor = decay_factor
        self.num_epochs = num_epochs
        self.criterion = FocalBCELoss(gamma=gamma, alpha=label_weights)

        # Layer groups for unfreezing schedule
        self.layer_groups = self._get_layer_groups()

    def _get_layer_groups(self):
        """Group model parameters by layer depth."""
        encoder = self.model.encoder
        groups = {
            "embeddings": list(encoder.embeddings.parameters()),
            "layer_0_1": (
                list(encoder.transformer.layer[0].parameters())
                + list(encoder.transformer.layer[1].parameters())
            ),
            "layer_2_3": (
                list(encoder.transformer.layer[2].parameters())
                + list(encoder.transformer.layer[3].parameters())
            ),
            "layer_4_5": (
                list(encoder.transformer.layer[4].parameters())
                + list(encoder.transformer.layer[5].parameters())
            ),
            "head": list(self.model.head.parameters()),
        }
        return groups

    def _freeze_all_except(self, active_groups: list[str]):
        """Freeze all parameters except those in active groups."""
        for name, params in self.layer_groups.items():
            requires_grad = name in active_groups
            for p in params:
                p.requires_grad = requires_grad

    def _build_optimizer(self, active_groups: list[str]):
        """Build optimizer with discriminative learning rates."""
        lr_multipliers = {
            "head": 1.0,
            "layer_4_5": 1.0 / self.decay_factor,
            "layer_2_3": 1.0 / (self.decay_factor ** 2),
            "layer_0_1": 1.0 / (self.decay_factor ** 3),
            "embeddings": 1.0 / (self.decay_factor ** 4),
        }

        param_groups = []
        for name in active_groups:
            param_groups.append({
                "params": self.layer_groups[name],
                "lr": self.base_lr * lr_multipliers[name],
            })

        return torch.optim.AdamW(param_groups, weight_decay=0.01)

    def _get_unfreeze_schedule(self) -> dict[int, list[str]]:
        """Define which layer groups are active at each epoch."""
        return {
            0: ["head"],                                           # Epoch 1
            1: ["head", "layer_4_5"],                             # Epoch 2
            2: ["head", "layer_4_5", "layer_2_3"],                # Epoch 3
            3: ["head", "layer_4_5", "layer_2_3", "layer_0_1", "embeddings"],  # Epoch 4+
        }

    def train(self):
        schedule = self._get_unfreeze_schedule()
        best_f1 = 0
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model.to(device)

        for epoch in range(self.num_epochs):
            # Determine active groups
            active_key = min(epoch, max(schedule.keys()))
            active_groups = schedule[active_key]
            self._freeze_all_except(active_groups)

            # Rebuild optimizer with current unfreezing
            optimizer = self._build_optimizer(active_groups)
            scheduler = get_linear_schedule_with_warmup(
                optimizer,
                num_warmup_steps=len(self.train_loader) // 10,
                num_training_steps=len(self.train_loader),
            )

            # Training loop
            self.model.train()
            total_loss = 0
            for batch in self.train_loader:
                batch = {k: v.to(device) for k, v in batch.items()}
                logits = self.model(batch["input_ids"], batch["attention_mask"])
                loss = self.criterion(logits, batch["labels"])

                optimizer.zero_grad()
                loss.backward()
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
                optimizer.step()
                scheduler.step()
                total_loss += loss.item()

            # Validation
            val_metrics = self._evaluate(device)
            f1_macro = val_metrics["f1_macro"]

            active_str = ", ".join(active_groups)
            print(
                f"Epoch {epoch+1}/{self.num_epochs} | "
                f"Active: [{active_str}] | "
                f"Loss: {total_loss/len(self.train_loader):.4f} | "
                f"F1: {f1_macro:.4f}"
            )

            if f1_macro > best_f1:
                best_f1 = f1_macro
                torch.save(self.model.state_dict(), "best_sentiment_model.pt")

        return best_f1

    def _evaluate(self, device) -> dict:
        """Evaluate with optimized per-label thresholds."""
        self.model.eval()
        all_logits, all_labels = [], []

        with torch.no_grad():
            for batch in self.val_loader:
                batch = {k: v.to(device) for k, v in batch.items()}
                logits = self.model(batch["input_ids"], batch["attention_mask"])
                all_logits.append(logits.cpu())
                all_labels.append(batch["labels"].cpu())

        all_logits = torch.cat(all_logits)
        all_labels = torch.cat(all_labels)
        probs = torch.sigmoid(all_logits)

        # Apply default threshold
        preds = (probs > 0.5).int().numpy()
        labels = all_labels.int().numpy()

        p, r, f1, _ = precision_recall_fscore_support(labels, preds, average="macro")

        per_label = {}
        for i, label in enumerate(LABELS):
            lp, lr, lf, _ = precision_recall_fscore_support(
                labels[:, i], preds[:, i], average="binary"
            )
            per_label[label] = {"precision": lp, "recall": lr, "f1": lf}

        return {"f1_macro": f1, "precision": p, "recall": r, "per_label": per_label}

from sklearn.metrics import fbeta_score


def optimize_thresholds(
    probs: np.ndarray,    # (N, L)
    labels: np.ndarray,   # (N, L)
    betas: dict[str, float] = None,
) -> dict[str, float]:
    """
    Find optimal per-label thresholds using F-beta search.

    betas: {label_name: beta_value}. Higher beta = more recall-focused.
    """
    if betas is None:
        betas = {
            "frustration": 2.0,    # Recall-focused: catch frustrated users
            "satisfaction": 0.5,    # Precision-focused: avoid false positives
            "urgency": 2.0,        # Recall-focused
            "confusion": 1.5,      # Slightly recall-focused
            "churn_risk": 2.0,     # Recall-focused: never miss churn risk
            "neutral": 1.0,        # Balanced
        }

    optimal_thresholds = {}

    for i, label_name in enumerate(LABELS):
        beta = betas[label_name]
        best_score = 0
        best_threshold = 0.5

        for threshold in np.arange(0.1, 0.91, 0.05):
            preds = (probs[:, i] > threshold).astype(int)
            score = fbeta_score(labels[:, i], preds, beta=beta)

            if score > best_score:
                best_score = score
                best_threshold = threshold

        optimal_thresholds[label_name] = {
            "threshold": round(best_threshold, 2),
            "fbeta_score": round(best_score, 4),
            "beta": beta,
        }

    return optimal_thresholds


# After training
thresholds = optimize_thresholds(val_probs, val_labels)
# {
#   "frustration": {"threshold": 0.35, "fbeta_score": 0.82, "beta": 2.0},
#   "churn_risk":  {"threshold": 0.28, "fbeta_score": 0.74, "beta": 2.0},
#   ...
# }

import sagemaker
from sagemaker.huggingface import HuggingFace


def launch_sentiment_training(
    train_s3: str,
    val_s3: str,
    instance_type: str = "ml.g4dn.xlarge",
):
    """Launch multi-label sentiment fine-tuning on SageMaker."""
    hyperparameters = {
        "model_name": "distilbert-base-uncased",
        "num_epochs": 8,
        "batch_size": 32,
        "base_lr": 3e-5,
        "decay_factor": 2.6,
        "gamma": 2.0,
        "label_weights": "0.4,1.0,1.2,1.5,2.57,6.0",
        "gradual_unfreezing": True,
    }

    estimator = HuggingFace(
        entry_point="train_sentiment.py",
        source_dir="./src/sentiment/",
        instance_type=instance_type,
        instance_count=1,
        role=sagemaker.get_execution_role(),
        transformers_version="4.37",
        pytorch_version="2.1",
        py_version="py310",
        hyperparameters=hyperparameters,
    )

    estimator.fit({
        "train": train_s3,
        "validation": val_s3,
    })

    return estimator