import torch
from transformers import BertTokenizer, BertForMaskedLM
import matplotlib.pyplot as plt
import numpy as np
from sklearn.manifold import TSNE

# Load a pre-trained model and tokenizer
model_name = 'bert-base-uncased'
tokenizer = BertTokenizer.from_pretrained(model_name)
model = BertForMaskedLM.from_pretrained(model_name)

# Example input text
text = "The quick brown fox jumps over the lazy dog"

# Tokenize the input text
inputs = tokenizer(text, return_tensors="pt")
input_ids = inputs['input_ids']

# Get attention weights by running the model
with torch.no_grad():
    outputs = model(input_ids, output_attentions=True)

# Extract the attention weights (size: [num_layers, num_heads, seq_len, seq_len])
attention_weights = outputs.attentions

# Select a specific layer and attention head
layer_idx = 0  # First layer
head_idx = 0   # First attention head

# Get the attention matrix for this layer and head
attention_matrix = attention_weights[layer_idx][0][head_idx].cpu().numpy()

# Use t-SNE to reduce the dimensionality of the attention matrix (embedding space)
# Attention matrix shape: [seq_len, seq_len], so we reduce each row (which corresponds to a token's attention distribution)
tsne = TSNE(n_components=2, random_state=42)
reduced_attention = tsne.fit_transform(attention_matrix)

# Plotting the reduced attention embeddings
fig, ax = plt.subplots(figsize=(10, 10))

# Plot the reduced attention in 2D
ax.scatter(reduced_attention[:, 0], reduced_attention[:, 1])

# Annotate the tokens in the scatter plot
tokens = tokenizer.convert_ids_to_tokens(input_ids[0])
for i, token in enumerate(tokens):
    ax.annotate(token, (reduced_attention[i, 0], reduced_attention[i, 1]), fontsize=12, ha='right')

# Display the plot
plt.title(f"t-SNE Visualization of Attention - Layer {layer_idx+1}, Head {head_idx+1}")
plt.xlabel("t-SNE Dimension 1")
plt.ylabel("t-SNE Dimension 2")
plt.grid(True)
plt.show()
plt.savefig('test.png')