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