# MechVis: GPT-2 Attention Head Visualization

This interactive web app allows you to visualize how different attention heads in GPT-2 small contribute to next token predictions.

## How to Use

1. Enter text in the input field (e.g., "When John and Mary went to the store, John gave a drink to")
2. Click "Predict Next Word"
3. See what token GPT-2 predicts next and explore how each attention head contributes to that prediction

## Features

- Next token prediction with GPT-2 small
- Interactive heatmap showing attention head contributions
- Layer contribution analysis
- Hover over cells to see exact contribution values

## Examples to Try

- **Indirect Object Identification**: "When John and Mary went to the store, John gave a drink to" (likely predicts "Mary")
- **Induction Pattern**: "The capital of France is Paris. The capital of Germany is" (likely predicts "Berlin")

## About

This project uses [TransformerLens](https://github.com/neelnanda-io/TransformerLens) to access internal model activations and calculate how each attention head contributes to the final logit score of the predicted token.

[GitHub Repository](https://github.com/saivamsim26/mechvis)