src/ablation.py

"""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,
    },
}