Spaces:
Sleeping
Sleeping
File size: 2,142 Bytes
87360eb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
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()
|