src/models.py

"""Models for ACSA with metadata fusion.

Three architectures:
  - BertMetaFusionACSAModel:  BERT + Meta Cross-Attention + per-aspect heads  (Proposed)
  - BertACSAModel:            BERT + per-aspect heads                         (Baseline 3, ablation w/o meta)
  - BertOverallModel:         BERT + single 3-class head                      (Baseline 2)

Design notes
------------
* `class_weights` is intentionally NOT a buffer. We tried that before and got a
  `Unexpected key(s) in state_dict: 'class_weights'` on every reload because the
  eval-time constructor doesn't know the training-time class_weights. Now we
  keep it as a plain attribute that does not enter state_dict. The training
  loop is responsible for re-instantiating the loss with the right weights.

* The "metadata token" trick. The cross-attention layer wants a sequence of
  K/V tokens, not a single vector. We split the encoded meta vector into
  `META_NUM_META_TOKENS` virtual tokens of equal length so the attention head
  has structure to operate over. Each chunk roughly corresponds to a slice of
  TF-IDF + a slice of numeric features, which makes the attention weight
  vector interpretable as "how much did the model rely on meta chunk i".

* Forward returns:
    {
      "logits": (B, num_aspects, num_classes),
      "loss":   scalar or None,
      "meta_attn_weights": (B, num_aspects, num_meta_tokens) or None,
      "bert_attentions":   tuple of BERT self-attn (only if output_attentions=True),
    }
"""
import logging
import math
from typing import List, Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoModel

from . import config as cfg


logger = logging.getLogger(__name__)


def _load_bert_backbone(bert_name: str):
    """Load BERT without meta-tensor lazy initialization for HF Spaces."""
    try:
        return AutoModel.from_pretrained(bert_name, low_cpu_mem_usage=False)
    except TypeError:
        return AutoModel.from_pretrained(bert_name)

# Split the encoded meta vector into this many "tokens" before cross-attention.
# Each token gets its own linear projection to BERT's hidden size.
META_NUM_META_TOKENS = 4


# ---------------------------------------------------------------------------
# Shared helpers
# ---------------------------------------------------------------------------

class MetaEncoderMLP(nn.Module):
    """Encode raw meta features (numpy-shaped) -> META_HIDDEN_DIM vector."""

    def __init__(self, in_dim: int, hidden: int = cfg.META_HIDDEN_DIM,
                 dropout: float = 0.2):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, hidden * 2),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden * 2, hidden),
        )

    def forward(self, x):
        return self.net(x)


class MetaTokenizer(nn.Module):
    """Split a (B, meta_hidden) vector into (B, num_tokens, bert_hidden) so it
    can serve as K/V in a multi-head cross-attention.

    We do this by chunking the meta vector and projecting each chunk
    independently to BERT's hidden size. Each chunk represents a *piece* of
    the metadata (a slice of the encoded TF-IDF + numeric features); the
    cross-attention weight over chunks is what gets visualized for explanation.
    """

    def __init__(self, meta_hidden: int, bert_hidden: int,
                 num_tokens: int = META_NUM_META_TOKENS):
        super().__init__()
        assert meta_hidden % num_tokens == 0, (
            f"META_HIDDEN_DIM={meta_hidden} must be divisible by "
            f"num_tokens={num_tokens}"
        )
        self.num_tokens = num_tokens
        self.chunk_size = meta_hidden // num_tokens
        self.proj = nn.ModuleList([
            nn.Linear(self.chunk_size, bert_hidden) for _ in range(num_tokens)
        ])

    def forward(self, meta_vec):
        # meta_vec: (B, meta_hidden) -> list of (B, chunk_size)
        chunks = meta_vec.chunk(self.num_tokens, dim=-1)
        projected = [proj(chunk) for proj, chunk in zip(self.proj, chunks)]
        # stack -> (B, num_tokens, bert_hidden)
        return torch.stack(projected, dim=1)


class CrossAttentionFusion(nn.Module):
    """Multi-head cross-attention: Q = text [CLS], K=V = meta tokens.

    Returns:
      fused:        (B, bert_hidden)   text vector enriched by meta
      attn_weights: (B, num_meta_tokens)  averaged over heads, used for XAI
    """

    def __init__(self, hidden: int, num_heads: int = cfg.META_CROSSATTN_HEADS,
                 dropout: float = 0.1):
        super().__init__()
        assert hidden % num_heads == 0
        self.hidden = hidden
        self.num_heads = num_heads
        self.head_dim = hidden // num_heads

        self.q_proj = nn.Linear(hidden, hidden)
        self.k_proj = nn.Linear(hidden, hidden)
        self.v_proj = nn.Linear(hidden, hidden)
        self.out_proj = nn.Linear(hidden, hidden)
        self.dropout = nn.Dropout(dropout)
        self.norm = nn.LayerNorm(hidden)

    def forward(self, text_vec, meta_tokens):
        # text_vec:    (B, H)
        # meta_tokens: (B, T, H)
        B, T, H = meta_tokens.shape
        q = self.q_proj(text_vec).view(B, self.num_heads, 1, self.head_dim)
        k = self.k_proj(meta_tokens).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(meta_tokens).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        # k, v: (B, heads, T, head_dim)
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        # scores: (B, heads, 1, T)
        attn = F.softmax(scores, dim=-1)
        attn_drop = self.dropout(attn)
        ctx = torch.matmul(attn_drop, v)  # (B, heads, 1, head_dim)
        ctx = ctx.transpose(1, 2).contiguous().view(B, H)
        ctx = self.out_proj(ctx)

        fused = self.norm(text_vec + ctx)            # residual + LN
        avg_attn = attn.mean(dim=1).squeeze(1)        # (B, T) 闂?heads averaged
        return fused, avg_attn



class SemanticMetaTokenizer(nn.Module):
    """Project source-separated metadata into interpretable tokens."""

    token_names = ["features", "categories", "numeric"]

    def __init__(self, bert_hidden: int, dropout: float = 0.1):
        super().__init__()
        self.feature_proj = nn.Sequential(
            nn.Linear(cfg.META_FEATURE_TFIDF_DIM, bert_hidden),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.LayerNorm(bert_hidden),
        )
        self.category_proj = nn.Sequential(
            nn.Linear(cfg.META_CATEGORY_TFIDF_DIM, bert_hidden),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.LayerNorm(bert_hidden),
        )
        self.numeric_proj = nn.Sequential(
            nn.Linear(cfg.META_NUM_DIM, bert_hidden),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.LayerNorm(bert_hidden),
        )

    def forward(self, meta_features):
        f_end = cfg.META_FEATURE_TFIDF_DIM
        c_end = f_end + cfg.META_CATEGORY_TFIDF_DIM
        feat = meta_features[:, :f_end]
        cat = meta_features[:, f_end:c_end]
        num = meta_features[:, c_end:c_end + cfg.META_NUM_DIM]
        return torch.stack([
            self.feature_proj(feat),
            self.category_proj(cat),
            self.numeric_proj(num * cfg.META_NUMERIC_TOKEN_SCALE),
        ], dim=1)


class GatedAspectSemanticCrossAttention(nn.Module):
    """Aspect-specific cross-attention with a conservative metadata gate."""

    def __init__(self, hidden: int, num_aspects: int = cfg.NUM_ASPECTS,
                 num_heads: int = cfg.META_CROSSATTN_HEADS, dropout: float = 0.1):
        super().__init__()
        assert hidden % num_heads == 0
        self.hidden = hidden
        self.num_aspects = num_aspects
        self.num_heads = num_heads
        self.head_dim = hidden // num_heads
        self.aspect_embeddings = nn.Parameter(torch.randn(num_aspects, hidden) * 0.02)
        self.q_proj = nn.Linear(hidden, hidden)
        self.k_proj = nn.Linear(hidden, hidden)
        self.v_proj = nn.Linear(hidden, hidden)
        self.out_proj = nn.Linear(hidden, hidden)
        self.gate = nn.Linear(hidden * 2, hidden)
        nn.init.constant_(self.gate.bias, -0.3)
        self.dropout = nn.Dropout(dropout)
        self.norm = nn.LayerNorm(hidden)

    def forward(self, text_vec, meta_tokens):
        B, T, H = meta_tokens.shape
        aspect_queries = text_vec.unsqueeze(1) + self.aspect_embeddings.unsqueeze(0)
        q = self.q_proj(aspect_queries).view(B, self.num_aspects, self.num_heads, self.head_dim)
        q = q.transpose(1, 2)
        k = self.k_proj(meta_tokens).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(meta_tokens).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        attn = F.softmax(torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim), dim=-1)
        ctx = torch.matmul(self.dropout(attn), v)
        ctx = ctx.transpose(1, 2).contiguous().view(B, self.num_aspects, H)
        ctx = self.out_proj(ctx)
        text_expanded = text_vec.unsqueeze(1).expand(-1, self.num_aspects, -1)
        gate = torch.sigmoid(self.gate(torch.cat([text_expanded, ctx], dim=-1)))
        fused = self.norm(text_expanded + gate * ctx)
        return fused, attn.mean(dim=1), gate.mean(dim=-1)


class GatedGlobalSemanticCrossAttention(nn.Module):
    """Global text-to-metadata cross-attention for the overall head."""

    def __init__(self, hidden: int, num_heads: int = cfg.META_CROSSATTN_HEADS,
                 dropout: float = 0.1):
        super().__init__()
        self.attn = CrossAttentionFusion(hidden, num_heads=num_heads, dropout=dropout)
        self.gate = nn.Linear(hidden * 2, hidden)
        nn.init.constant_(self.gate.bias, -0.3)
        self.norm = nn.LayerNorm(hidden)

    def forward(self, text_vec, meta_tokens):
        fused_raw, attn = self.attn(text_vec, meta_tokens)
        ctx = fused_raw - text_vec
        gate = torch.sigmoid(self.gate(torch.cat([text_vec, ctx], dim=-1)))
        fused = self.norm(text_vec + gate * ctx)
        return fused, attn, gate.mean(dim=-1)

class _PerAspectHeads(nn.Module):
    """num_aspects independent MLP heads."""

    def __init__(self, in_dim: int, num_aspects: int, num_classes: int,
                 dropout: float = 0.3):
        super().__init__()
        self.heads = nn.ModuleList([
            nn.Sequential(
                nn.Dropout(dropout),
                nn.Linear(in_dim, 256),
                nn.GELU(),
                nn.Linear(256, num_classes),
            )
            for _ in range(num_aspects)
        ])

    def forward(self, x):
        # x: (B, in_dim) -> (B, num_aspects, num_classes)
        return torch.stack([h(x) for h in self.heads], dim=1)

    def forward_per_aspect(self, x):
        # x: (B, num_aspects, in_dim) -> (B, num_aspects, num_classes)
        return torch.stack([h(x[:, i, :]) for i, h in enumerate(self.heads)], dim=1)


def _aspect_loss(logits, labels, class_weights=None):
    """Sum of cross-entropy over aspects, averaged."""
    num_aspects = logits.shape[1]
    loss = 0.0
    for i in range(num_aspects):
        w = class_weights[i] if class_weights is not None else None
        loss = loss + F.cross_entropy(logits[:, i, :], labels[:, i], weight=w)
    return loss / num_aspects


# ---------------------------------------------------------------------------
# Models
# ---------------------------------------------------------------------------

class BertMetaFusionACSAModel(nn.Module):
    """Proposed model: BERT + Meta Cross-Attention + per-aspect heads."""

    def __init__(
        self,
        bert_name: str = cfg.BERT_MODEL_NAME,
        meta_in_dim: int = cfg.META_TFIDF_DIM + cfg.META_NUM_DIM,
        num_aspects: int = cfg.NUM_ASPECTS,
        num_classes: int = cfg.NUM_CLASSES,
        dropout: float = 0.3,
        class_weights: Optional[torch.Tensor] = None,
    ):
        super().__init__()
        self.bert = _load_bert_backbone(bert_name)
        hidden = self.bert.config.hidden_size

        self.num_aspects = num_aspects
        self.num_classes = num_classes
        self.aspect_names = list(cfg.ASPECTS)
        self.meta_in_dim = meta_in_dim

        self.meta_mlp = MetaEncoderMLP(meta_in_dim, hidden=cfg.META_HIDDEN_DIM,
                                       dropout=dropout)
        self.meta_tokenizer = MetaTokenizer(
            meta_hidden=cfg.META_HIDDEN_DIM, bert_hidden=hidden,
            num_tokens=META_NUM_META_TOKENS,
        )
        self.fusion = CrossAttentionFusion(hidden=hidden,
                                           num_heads=cfg.META_CROSSATTN_HEADS,
                                           dropout=0.1)
        self.heads = _PerAspectHeads(in_dim=hidden, num_aspects=num_aspects,
                                     num_classes=num_classes, dropout=dropout)

        # Overall sentiment auxiliary head (3-class: Neg/Neu/Pos).
        # Shares the fused representation with per-aspect heads; trained jointly
        # with weight cfg.OVERALL_AUX_WEIGHT when overall_labels are provided.
        self.overall_head = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(hidden, 256),
            nn.GELU(),
            nn.Linear(256, cfg.OVERALL_NUM_CLASSES),
        )

        # Stored as plain attribute (NOT a buffer). Re-supplied at train time.
        self.class_weights = class_weights

    def forward(
        self,
        input_ids,
        attention_mask,
        meta_features,
        labels: Optional[torch.Tensor] = None,
        overall_labels: Optional[torch.Tensor] = None,
        output_attentions: bool = False,
    ):
        bert_out = self.bert(
            input_ids=input_ids,
            attention_mask=attention_mask,
            output_attentions=output_attentions,
            return_dict=True,
        )
        text_vec = bert_out.last_hidden_state[:, 0, :]   # [CLS]

        meta_vec = self.meta_mlp(meta_features)          # (B, meta_hidden)
        meta_tokens = self.meta_tokenizer(meta_vec)      # (B, T, H)
        fused, meta_attn = self.fusion(text_vec, meta_tokens)

        logits = self.heads(fused)
        overall_logits = self.overall_head(fused)

        # Joint loss: L_aspect + lambda * L_overall
        loss = None
        if labels is not None:
            loss = _aspect_loss(logits, labels, self.class_weights)
            if overall_labels is not None:
                loss_overall = F.cross_entropy(overall_logits, overall_labels)
                loss = loss + cfg.OVERALL_AUX_WEIGHT * loss_overall

        return {
            "loss": loss,
            "logits": logits,
            "overall_logits": overall_logits,
            "meta_attn_weights": meta_attn,        # (B, T) -- for XAI
            "bert_attentions": bert_out.attentions if output_attentions else None,
            "fused_state": fused,
            "text_state": text_vec,
        }



class GatedAspectSemanticMetaFusionACSAModel(nn.Module):
    """Proposed v2: semantic meta tokens + aspect queries + gated fusion.

    The aspect heads receive aspect-specific fused states, while the overall
    head keeps a separate global fused state so aspect-level noise does not
    dilute document-level sentiment.
    """

    meta_token_names = SemanticMetaTokenizer.token_names

    def __init__(
        self,
        bert_name: str = cfg.BERT_MODEL_NAME,
        meta_in_dim: int = cfg.META_TFIDF_DIM + cfg.META_NUM_DIM,
        num_aspects: int = cfg.NUM_ASPECTS,
        num_classes: int = cfg.NUM_CLASSES,
        dropout: float = 0.3,
        class_weights: Optional[torch.Tensor] = None,
    ):
        super().__init__()
        expected_dim = cfg.META_TFIDF_DIM + cfg.META_NUM_DIM
        if meta_in_dim != expected_dim:
            raise ValueError(f"Semantic meta fusion expects meta_in_dim={expected_dim}, got {meta_in_dim}")
        self.bert = _load_bert_backbone(bert_name)
        hidden = self.bert.config.hidden_size

        self.num_aspects = num_aspects
        self.num_classes = num_classes
        self.aspect_names = list(cfg.ASPECTS)
        self.meta_in_dim = meta_in_dim

        self.meta_tokenizer = SemanticMetaTokenizer(bert_hidden=hidden, dropout=0.1)
        self.concat_meta_mlp = MetaEncoderMLP(meta_in_dim, hidden=cfg.META_HIDDEN_DIM,
                                             dropout=dropout)
        self.concat_fusion = nn.Sequential(
            nn.Linear(hidden + cfg.META_HIDDEN_DIM, hidden),
            nn.GELU(),
            nn.Dropout(0.1),
        )
        self.concat_norm = nn.LayerNorm(hidden)
        self.aspect_refine_norm = nn.LayerNorm(hidden)
        self.global_refine_norm = nn.LayerNorm(hidden)
        residual_scales = getattr(cfg, "CROSS_ATTN_RESIDUAL_SCALES", None)
        if residual_scales is None:
            residual_scale_tensor = torch.full(
                (num_aspects,), float(cfg.CROSS_ATTN_RESIDUAL_SCALE)
            )
        else:
            residual_scale_tensor = torch.tensor(residual_scales, dtype=torch.float)
            if residual_scale_tensor.numel() != num_aspects:
                raise ValueError(
                    "CROSS_ATTN_RESIDUAL_SCALES must match the number of aspects"
                )
        self.cross_residual_scale = nn.Parameter(residual_scale_tensor)
        self.aspect_mix_gate = nn.Linear(hidden * 2, 1)
        self.global_mix_gate = nn.Linear(hidden * 2, 1)
        nn.init.zeros_(self.aspect_mix_gate.weight)
        nn.init.zeros_(self.global_mix_gate.weight)
        nn.init.zeros_(self.aspect_mix_gate.bias)
        nn.init.zeros_(self.global_mix_gate.bias)
        self.aspect_fusion = GatedAspectSemanticCrossAttention(
            hidden=hidden,
            num_aspects=num_aspects,
            num_heads=cfg.META_CROSSATTN_HEADS,
            dropout=0.1,
        )
        self.global_fusion = GatedGlobalSemanticCrossAttention(
            hidden=hidden,
            num_heads=cfg.META_CROSSATTN_HEADS,
            dropout=0.1,
        )
        self.heads = _PerAspectHeads(in_dim=hidden, num_aspects=num_aspects,
                                     num_classes=num_classes, dropout=dropout)
        # Overall sentiment is text-dominant: aspect heads use metadata fusion,
        # while this auxiliary head reads the raw BERT text vector to avoid
        # metadata noise hurting coarse-grained polarity.
        self.overall_head = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(hidden, 256),
            nn.GELU(),
            nn.Linear(256, cfg.OVERALL_NUM_CLASSES),
        )
        self.class_weights = class_weights

    def forward(
        self,
        input_ids,
        attention_mask,
        meta_features,
        labels: Optional[torch.Tensor] = None,
        overall_labels: Optional[torch.Tensor] = None,
        output_attentions: bool = False,
    ):
        bert_out = self.bert(
            input_ids=input_ids,
            attention_mask=attention_mask,
            output_attentions=output_attentions,
            return_dict=True,
        )
        text_vec = bert_out.last_hidden_state[:, 0, :]
        B = text_vec.size(0)
        meta_tokens = self.meta_tokenizer(meta_features)

        meta_for_concat = meta_features.clone()
        numeric_start = cfg.META_TFIDF_DIM
        meta_for_concat[:, numeric_start:] = (
            meta_for_concat[:, numeric_start:] * cfg.META_NUMERIC_TOKEN_SCALE
        )
        concat_meta = self.concat_meta_mlp(meta_for_concat)
        concat_base = self.concat_norm(
            self.concat_fusion(torch.cat([text_vec, concat_meta], dim=-1))
        )

        aspect_fused, aspect_attn, aspect_gate = self.aspect_fusion(text_vec, meta_tokens)
        global_fused, global_attn, global_gate = self.global_fusion(text_vec, meta_tokens)

        concat_expanded = concat_base.unsqueeze(1).expand(-1, self.num_aspects, -1)
        scale_vec = torch.clamp(self.cross_residual_scale, 0.0, 1.0)
        aspect_scale = scale_vec.view(1, self.num_aspects, 1)
        global_scale = scale_vec.mean()

        # Cross-attention remains the main fusion path. The concat branch is
        # only a small stabilizing residual, so the model still follows the
        # Q=text, K/V=metadata architecture used in the paper and slides.
        aspect_repr = self.aspect_refine_norm(
            aspect_fused + aspect_scale * (concat_expanded - aspect_fused)
        )
        global_repr = self.global_refine_norm(
            global_fused + global_scale * (concat_base - global_fused)
        )

        aspect_mix = scale_vec.unsqueeze(0).expand(B, self.num_aspects)
        global_mix = global_scale.expand(B)

        logits = self.heads.forward_per_aspect(aspect_repr)
        overall_logits = self.overall_head(text_vec)

        loss = None
        if labels is not None:
            loss = _aspect_loss(logits, labels, self.class_weights)
            if overall_labels is not None:
                loss = loss + cfg.OVERALL_AUX_WEIGHT * F.cross_entropy(
                    overall_logits, overall_labels
                )

        return {
            "loss": loss,
            "logits": logits,
            "overall_logits": overall_logits,
            "meta_attn_weights": aspect_attn,
            "global_meta_attn_weights": global_attn,
            "meta_gate": aspect_gate,
            "global_meta_gate": global_gate,
            "fusion_mix_gate": aspect_mix.squeeze(-1),
            "global_fusion_mix_gate": global_mix.squeeze(-1),
            "meta_token_names": self.meta_token_names,
            "bert_attentions": bert_out.attentions if output_attentions else None,
            "fused_state": aspect_repr,
            "global_fused_state": global_repr,
            "text_state": text_vec,
        }

class BertACSAModel(nn.Module):
    """Baseline 3: same as Proposed but with NO metadata path.

    Same per-aspect head structure as Proposed so the comparison isolates
    the value of the Cross-Attention meta fusion.
    """

    def __init__(
        self,
        bert_name: str = cfg.BERT_MODEL_NAME,
        num_aspects: int = cfg.NUM_ASPECTS,
        num_classes: int = cfg.NUM_CLASSES,
        dropout: float = 0.3,
        class_weights: Optional[torch.Tensor] = None,
    ):
        super().__init__()
        self.bert = _load_bert_backbone(bert_name)
        hidden = self.bert.config.hidden_size
        self.num_aspects = num_aspects
        self.num_classes = num_classes
        self.aspect_names = list(cfg.ASPECTS)

        self.heads = _PerAspectHeads(in_dim=hidden, num_aspects=num_aspects,
                                     num_classes=num_classes, dropout=dropout)
        self.class_weights = class_weights

    def forward(self, input_ids, attention_mask,
                labels: Optional[torch.Tensor] = None,
                output_attentions: bool = False):
        bert_out = self.bert(
            input_ids=input_ids, attention_mask=attention_mask,
            output_attentions=output_attentions, return_dict=True,
        )
        text_vec = bert_out.last_hidden_state[:, 0, :]
        logits = self.heads(text_vec)
        loss = _aspect_loss(logits, labels, self.class_weights) if labels is not None else None
        return {
            "loss": loss,
            "logits": logits,
            "bert_attentions": bert_out.attentions if output_attentions else None,
            "text_state": text_vec,
        }


class BertOverallModel(nn.Module):
    """Baseline 2: BERT fine-tuned for overall 3-class sentiment (Neg/Neu/Pos)."""

    def __init__(self, bert_name: str = cfg.BERT_MODEL_NAME,
                 num_classes: int = cfg.OVERALL_NUM_CLASSES,
                 dropout: float = 0.3):
        super().__init__()
        self.bert = _load_bert_backbone(bert_name)
        hidden = self.bert.config.hidden_size
        self.classifier = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(hidden, 256),
            nn.GELU(),
            nn.Linear(256, num_classes),
        )
        self.num_classes = num_classes

    def forward(self, input_ids, attention_mask, labels=None, output_attentions=False):
        out = self.bert(input_ids=input_ids, attention_mask=attention_mask,
                        output_attentions=output_attentions, return_dict=True)
        logits = self.classifier(out.last_hidden_state[:, 0, :])
        loss = F.cross_entropy(logits, labels) if labels is not None else None
        return {
            "loss": loss,
            "logits": logits,
            "bert_attentions": out.attentions if output_attentions else None,
        }


# ---------------------------------------------------------------------------
# Class-weight helper for the aspect heads
# ---------------------------------------------------------------------------

def compute_class_weights(df, aspect_cols: List[str],
                          num_classes: int = cfg.NUM_CLASSES) -> torch.Tensor:
    """Per-aspect inverse-frequency class weights, clipped to [0.2, 5.0]."""
    import numpy as np
    out = []
    for col in aspect_cols:
        counts = df[col].value_counts().reindex(range(num_classes), fill_value=0).values
        counts = counts.astype(float) + 1.0
        inv = counts.sum() / (num_classes * counts)
        out.append(np.clip(inv, 0.2, 5.0))
    return torch.from_numpy(np.array(out, dtype=np.float32))