"""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 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, -1.0) 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, -1.0) 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 = 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 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 = 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_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) self.cross_residual_scale = nn.Parameter( torch.tensor(float(cfg.CROSS_ATTN_RESIDUAL_SCALE)) ) 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) 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, :] 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 = torch.clamp(self.cross_residual_scale, 0.0, 1.0) aspect_mix_raw = torch.sigmoid( self.aspect_mix_gate(torch.cat([concat_expanded, aspect_fused], dim=-1)) ) aspect_mix = torch.clamp(aspect_mix_raw + (scale - 0.5), 0.0, 1.0) aspect_repr = self.aspect_refine_norm( (1.0 - aspect_mix) * concat_expanded + aspect_mix * aspect_fused ) global_mix_raw = torch.sigmoid( self.global_mix_gate(torch.cat([concat_base, global_fused], dim=-1)) ) global_mix = torch.clamp(global_mix_raw + (scale - 0.5), 0.0, 1.0) global_repr = self.global_refine_norm( (1.0 - global_mix) * concat_base + global_mix * global_fused ) logits = self.heads.forward_per_aspect(aspect_repr) overall_logits = self.overall_head(global_repr) 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 = 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))