"""Ablation experiments for the Proposed BERT + Meta Cross-Attention model. Three ablations are defined here: A1 -- "remove text-type metadata": Mask the TF-IDF slice of the meta vector (features_text + categories_text) to zeros, keep only the numeric slice (price, average_rating, log_rating_number, price_missing_flag). Verifies whether textual product attributes are the main source of the fusion gain. A2 -- "remove numeric metadata": Mask the numeric slice, keep only the TF-IDF text-meta slice. Strict contrast to A1. A3 -- "replace Cross-Attention with concatenation": Keep both meta sources intact, but replace the MetaTokenizer + Multi-Head Cross-Attention fusion with a static [text_cls ; meta_vec] concatenation followed by a linear projection back to BERT hidden size. Quantifies the algorithmic benefit of Cross-Attention vs. static fusion. Design notes ------------ * For A1/A2 we deliberately keep the model architecture *unchanged* and only zero out the masked slice of the encoded meta vector. This isolates the information content of each meta sub-source from architectural confounders (e.g. MLP input dimension shrinkage). The MaskedMetaEncoder wraps an already- fit MetaEncoder so we never re-fit on test/val data. * For A3 we add a new model class `BertConcatFusionACSAModel` whose forward signature is identical to `BertMetaFusionACSAModel` (same inputs, same output keys minus `meta_attn_weights`). This lets us reuse most of `train_meta_acsa` via a small dedicated training loop. * The training hyperparameters (epochs, batch size, LRs, class weighting) are passed through unchanged so the only thing varying between Proposed and ablations is what the prompt says is varying. """ import json import logging from dataclasses import dataclass from pathlib import Path from typing import Optional import numpy as np import pandas as pd import torch import torch.nn as nn from torch.utils.data import DataLoader from transformers import AutoModel, AutoTokenizer, get_linear_schedule_with_warmup from tqdm import tqdm from . import config as cfg from .dataset import MetaACSADataset from .meta_encoder import MetaEncoder from .models import ( GatedAspectSemanticMetaFusionACSAModel, BertMetaFusionACSAModel, MetaEncoderMLP, _PerAspectHeads, _aspect_loss, compute_class_weights, ) from .trainer import get_device, _evaluate_per_aspect logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Masked meta encoder for A1 / A2 # --------------------------------------------------------------------------- @dataclass class MaskedMetaEncoder: """Wraps a fit MetaEncoder and zeros out one of its two slices. The base encoder produces vectors laid out as [ TF-IDF (cfg.META_TFIDF_DIM dims) | numeric (cfg.META_NUM_DIM dims) ]. Setting `mask` to "text" zeros the TF-IDF slice (A1 -- remove text-meta). Setting `mask` to "numeric" zeros the numeric slice (A2 -- remove numeric meta). """ base: MetaEncoder mask: str # "text" or "numeric" def __post_init__(self): if self.mask not in {"text", "numeric"}: raise ValueError(f"mask must be 'text' or 'numeric', got {self.mask!r}") @property def total_dim(self) -> int: return self.base.total_dim def transform(self, df: pd.DataFrame) -> np.ndarray: # Copy so we never mutate any buffer the base encoder might be holding. mat = self.base.transform(df).copy() tfidf_dim = self.base.tfidf_dim if self.mask == "text": mat[:, :tfidf_dim] = 0.0 else: # numeric mat[:, tfidf_dim:] = 0.0 return mat # --------------------------------------------------------------------------- # A3: Concatenation fusion variant # --------------------------------------------------------------------------- class BertConcatFusionACSAModel(nn.Module): """A3 ablation: BERT + Meta MLP + [text;meta] concat + linear -> per-aspect heads. Architecture is intentionally identical to BertMetaFusionACSAModel except the fusion block: Proposed: fused = CrossAttention(Q=text_cls, KV=MetaTokenizer(meta_enc)) A3: fused = LayerNorm( Linear( [text_cls ; meta_enc] ) ) All other components -- BERT encoder, MetaEncoderMLP, per-aspect heads, class-weighted loss -- are unchanged so the comparison isolates the fusion mechanism. """ 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) # Static fusion: concat then project back to BERT hidden size. self.fusion = nn.Sequential( nn.Linear(hidden + cfg.META_HIDDEN_DIM, hidden), nn.GELU(), nn.Dropout(0.1), ) self.fusion_norm = nn.LayerNorm(hidden) 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, meta_features, 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_DIM) cat = torch.cat([text_vec, meta_vec], dim=-1) # (B, H + meta_hidden) fused = self.fusion_norm(self.fusion(cat)) # (B, H) logits = self.heads(fused) loss = _aspect_loss(logits, labels, self.class_weights) if labels is not None else None return { "loss": loss, "logits": logits, "meta_attn_weights": None, # no cross-attn in this variant "bert_attentions": bert_out.attentions if output_attentions else None, "fused_state": fused, "text_state": text_vec, } # --------------------------------------------------------------------------- # Training loop for A3 (concat fusion). A1/A2 reuse train_meta_acsa with a # masked encoder. # --------------------------------------------------------------------------- def train_concat_fusion_acsa( train_df, val_df, meta_encoder, bert_name: str = cfg.BERT_MODEL_NAME, epochs: int = cfg.DEFAULT_EPOCHS, batch_size: int = cfg.DEFAULT_BATCH_SIZE, lr_bert: float = cfg.DEFAULT_LR_BERT, lr_heads: float = cfg.DEFAULT_LR_HEADS, weight_decay: float = cfg.DEFAULT_WEIGHT_DECAY, use_class_weights: bool = True, output_dir: Optional[Path] = None, seed: int = cfg.RANDOM_SEED, ): """Train the A3 concatenation-fusion variant. Mirrors trainer.train_meta_acsa but instantiates BertConcatFusionACSAModel. """ if output_dir is None: output_dir = cfg.CHECKPOINT_DIR / "ablation_A3_concat" output_dir = Path(output_dir); output_dir.mkdir(parents=True, exist_ok=True) torch.manual_seed(seed); np.random.seed(seed) device = get_device() logger.info("Device: %s", device) tokenizer = AutoTokenizer.from_pretrained(bert_name) aspect_cols = [f"aspect_{a}" for a in cfg.ASPECTS] class_weights = (compute_class_weights(train_df, aspect_cols, cfg.NUM_CLASSES).to(device) if use_class_weights else None) model = BertConcatFusionACSAModel( bert_name=bert_name, meta_in_dim=meta_encoder.total_dim, class_weights=class_weights, ).to(device) train_ds = MetaACSADataset(train_df, tokenizer, meta_encoder) val_ds = MetaACSADataset(val_df, tokenizer, meta_encoder) train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=0) val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False, num_workers=0) bert_params = list(model.bert.named_parameters()) other_params = [(n, p) for n, p in model.named_parameters() if not n.startswith("bert.")] no_decay = ["bias", "LayerNorm.weight"] grouped = [ {"params": [p for n, p in bert_params if not any(nd in n for nd in no_decay)], "weight_decay": weight_decay, "lr": lr_bert}, {"params": [p for n, p in bert_params if any(nd in n for nd in no_decay)], "weight_decay": 0.0, "lr": lr_bert}, {"params": [p for n, p in other_params if not any(nd in n for nd in no_decay)], "weight_decay": weight_decay, "lr": lr_heads}, {"params": [p for n, p in other_params if any(nd in n for nd in no_decay)], "weight_decay": 0.0, "lr": lr_heads}, ] optimizer = torch.optim.AdamW(grouped) total_steps = max(len(train_loader) * epochs, 1) scheduler = get_linear_schedule_with_warmup( optimizer, num_warmup_steps=int(cfg.WARMUP_RATIO * total_steps), num_training_steps=total_steps, ) best_f1 = -1.0; history = [] for epoch in range(epochs): model.train(); running = 0.0 pbar = tqdm(train_loader, desc=f"[epoch {epoch+1}/{epochs}] ablation_A3_concat") for batch in pbar: batch = {k: v.to(device) for k, v in batch.items()} optimizer.zero_grad() out = model(batch["input_ids"], batch["attention_mask"], batch["meta_features"], labels=batch["labels"]) out["loss"].backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step(); scheduler.step() running += out["loss"].item() pbar.set_postfix({"loss": f"{out['loss'].item():.4f}"}) avg_loss = running / max(len(train_loader), 1) val = _evaluate_per_aspect(model, val_loader, device, with_meta=True) logger.info("Epoch %d | loss=%.4f | val_macro_f1=%.4f | val_acc=%.4f", epoch+1, avg_loss, val["macro_f1_mean"], val["accuracy_mean"]) history.append({"epoch": epoch+1, "train_loss": avg_loss, **val}) if val["macro_f1_mean"] > best_f1: best_f1 = val["macro_f1_mean"] torch.save({"model_state_dict": model.state_dict(), "config": {"bert_name": bert_name, "meta_in_dim": meta_encoder.total_dim, "variant": "A3_concat"}}, output_dir / "best.pt") tokenizer.save_pretrained(output_dir / "tokenizer") logger.info("Saved new best A3_concat (macro_f1=%.4f)", best_f1) with open(output_dir / "history.json", "w") as f: json.dump(history, f, indent=2) return model, history # --------------------------------------------------------------------------- # Checkpoint loading + per-aspect prediction for ablation models # --------------------------------------------------------------------------- def _load_ckpt(path: Path, device): return torch.load(path, map_location=device, weights_only=False) def load_meta_acsa_for_ablation(checkpoint_dir: Path, meta_encoder, device=None): """Load a BertMetaFusionACSAModel checkpoint (used for A1/A2).""" if device is None: device = get_device() ckpt = _load_ckpt(checkpoint_dir / "best.pt", device) bert_name = ckpt.get("config", {}).get("bert_name", cfg.BERT_MODEL_NAME) meta_in_dim = ckpt.get("config", {}).get("meta_in_dim", meta_encoder.total_dim) architecture = ckpt.get("config", {}).get("architecture", "legacy_meta_acsa") if architecture == "gated_aspect_semantic_meta_acsa": model = GatedAspectSemanticMetaFusionACSAModel(bert_name=bert_name, meta_in_dim=meta_in_dim).to(device) else: model = BertMetaFusionACSAModel(bert_name=bert_name, meta_in_dim=meta_in_dim).to(device) model.load_state_dict(ckpt["model_state_dict"], strict=False) model.eval() tokenizer = AutoTokenizer.from_pretrained(checkpoint_dir / "tokenizer") return model, tokenizer, device def load_concat_fusion_acsa(checkpoint_dir: Path, meta_encoder, device=None): """Load a BertConcatFusionACSAModel checkpoint (A3).""" if device is None: device = get_device() ckpt = _load_ckpt(checkpoint_dir / "best.pt", device) bert_name = ckpt.get("config", {}).get("bert_name", cfg.BERT_MODEL_NAME) meta_in_dim = ckpt.get("config", {}).get("meta_in_dim", meta_encoder.total_dim) model = BertConcatFusionACSAModel(bert_name=bert_name, meta_in_dim=meta_in_dim).to(device) model.load_state_dict(ckpt["model_state_dict"], strict=False) model.eval() tokenizer = AutoTokenizer.from_pretrained(checkpoint_dir / "tokenizer") return model, tokenizer, device def predict_per_aspect_for_ablation(model, tokenizer, test_df, meta_encoder, device, batch_size: int = 32): """Run a meta-using per-aspect model over test_df. Accepts both BertMetaFusionACSAModel (A1/A2) and BertConcatFusionACSAModel (A3) -- the forward signature is the same. """ test_df = test_df.reset_index(drop=True) ds = MetaACSADataset(test_df, tokenizer, meta_encoder) loader = DataLoader(ds, batch_size=batch_size, shuffle=False) all_preds = [[] for _ in range(cfg.NUM_ASPECTS)] all_labels = [[] for _ in range(cfg.NUM_ASPECTS)] with torch.no_grad(): for batch in tqdm(loader, desc="predict per-aspect (ablation)"): batch = {k: v.to(device) for k, v in batch.items()} out = model(batch["input_ids"], batch["attention_mask"], batch["meta_features"]) preds = out["logits"].argmax(dim=-1).cpu().numpy() labels = batch["labels"].cpu().numpy() for i in range(cfg.NUM_ASPECTS): all_preds[i].extend(preds[:, i].tolist()) all_labels[i].extend(labels[:, i].tolist()) return all_preds, all_labels # --------------------------------------------------------------------------- # Registry of variants # --------------------------------------------------------------------------- ABLATION_VARIANTS = { "A1": { "name": "A1_no_text_meta", "description": "Remove text-type metadata (TF-IDF on features/categories); " "keep numeric (price, ratings).", "checkpoint_subdir": "ablation_A1_no_text_meta", "fusion": "cross_attention", "meta_mask": "text", }, "A2": { "name": "A2_no_numeric_meta", "description": "Remove numeric metadata (price, ratings); " "keep text-type (TF-IDF on features/categories).", "checkpoint_subdir": "ablation_A2_no_numeric_meta", "fusion": "cross_attention", "meta_mask": "numeric", }, "A3": { "name": "A3_concat_fusion", "description": "Replace Cross-Attention fusion with [text;meta] concat + Linear; " "both meta sources kept.", "checkpoint_subdir": "ablation_A3_concat", "fusion": "concat", "meta_mask": None, }, }