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__)

# 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 _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 _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 = AutoModel.from_pretrained(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 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 = AutoModel.from_pretrained(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 = AutoModel.from_pretrained(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))