BERT Interpretability: Sentiment Analysis

import torch
from captum.attr import LayerIntegratedGradients
from captum.attr import visualization as viz
from transformers import DistilBertForSequenceClassification, DistilBertTokenizer

MODEL_NAME = "distilbert-base-uncased-finetuned-sst-2-english"

# Load Pre-trained Model & Tokenizer
model = DistilBertForSequenceClassification.from_pretrained(MODEL_NAME)
tokenizer = DistilBertTokenizer.from_pretrained(MODEL_NAME)
model.eval()

def predict(inputs, attention_mask=None):
    # Forward pass
    output = model(inputs, attention_mask=attention_mask)
    return output.logits


def custom_forward(inputs, attention_mask=None):
    # preds = predict(inputs, attention_mask)
    # return torch.softmax(preds, dim=1)[:, 1]  # Return probability of Positive class (1)

    # Inline for debugging
    output = model(inputs, attention_mask=attention_mask)
    preds = output.logits
    return torch.softmax(preds, dim=1)[:, 1]


def main():
    # Example text
    # text = "The movie was visually breathtaking but the plot was extremely dull and boring."
    text = "I absolutely loved the acting, it was a masterpiece."

    # Use MPS if available, otherwise CPU
    device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
    print(f"Using device: {device}")
    model.to(device)

    print(f"Input text: '{text}'")

    # Tokenize
    encoded = tokenizer(text, add_special_tokens=True, return_tensors="pt")
    input_ids = encoded["input_ids"].to(device)
    attention_mask = encoded["attention_mask"].to(device)

    # Test forward pass (no gradients)
    print("Testing forward pass...")
    try:
        with torch.no_grad():
            out = custom_forward(input_ids, attention_mask)
        print(f"Forward pass output: {out}")
    except Exception as e:
        print(f"Forward pass failed: {e}")
        raise

    # Initialize Layer Integrated Gradients
    # Attach to embeddings layer for word-level attributions
    print("Initializing LayerIntegratedGradients...")
    try:
        lig = LayerIntegratedGradients(custom_forward, model.distilbert.embeddings)
        print("LIG initialized.")
    except Exception as e:
        print(f"LIG initialization failed: {e}")
        raise

    # Compute Attributions
    # Use a PAD baseline and very small n_steps to avoid heavy compute
    print("Computing attributions...")
    baseline = torch.zeros_like(input_ids)
    model.zero_grad()

    try:
        attributions, delta = lig.attribute(
            inputs=input_ids,
            baselines=baseline,
            additional_forward_args=(attention_mask,),
            return_convergence_delta=True,
            n_steps=2,  # Further reduced for stability
            internal_batch_size=1,
        )
        print("Attributions computed.")
        print("Convergence Delta:", delta)
    except Exception as e:
        print(f"LIG initialization failed: {e}")
        raise

    return attributions, delta, input_ids, attention_mask

def summarize_attributions(attributions):
    attributions = attributions.sum(dim=-1).squeeze(0)
    attributions = attributions / torch.norm(attributions)
    return attributions


if __name__ == "__main__":
    attributions, delta, input_ids, attention_mask = main()

    attributions_sum = summarize_attributions(attributions)

    # Decode tokens for visualization
    tokens = tokenizer.convert_ids_to_tokens(input_ids[0])

    # Get model prediction for display
    pred_prob = custom_forward(input_ids, attention_mask).item()
    pred_class_idx = 1 if pred_prob > 0.5 else 0
    pred_class = "Positive" if pred_class_idx == 1 else "Negative"

    # Create Visualization Record (Captum expects ints for class labels)
    vis_data_record = viz.VisualizationDataRecord(
        attributions_sum,
        pred_prob,
        pred_class_idx,
        pred_class_idx,  # true label placeholder
        pred_class_idx,  # attr class placeholder
        attributions_sum.sum().item(),
        tokens,
        delta.item(),
    )

    print("\nVisualize Attributions:")
    viz.visualize_text([vis_data_record])  # This works best in notebooks

    # Manual text-based visualization for CLI
    print(f"Prediction: {pred_class} ({pred_prob:.4f})")
    print("-" * 30)
    print(f"{'Token':<15} {'Attribution':<10}")
    print("-" * 30)
    for token, score in zip(tokens, attributions_sum):
        print(f"{token:<15} {score.item():.4f}")
    print("-" * 30)