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

# 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=3, random_state=42, perplexity=5)  # Set a lower perplexity value
reduced_attention = tsne.fit_transform(attention_matrix)

# Plotting the reduced attention embeddings in 3D
fig = plt.figure(figsize=(12, 10))
ax = fig.add_subplot(111, projection='3d')

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

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

# Set plot labels
ax.set_title(f"3D t-SNE Visualization of Attention - Layer {layer_idx+1}, Head {head_idx+1}")
ax.set_xlabel("t-SNE Dimension 1")
ax.set_ylabel("t-SNE Dimension 2")
ax.set_zlabel("t-SNE Dimension 3")

plt.show()