Spaces:
No application file
No application file
File size: 2,577 Bytes
e1243ba |
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 60 61 62 63 64 65 66 67 68 69 70 71 72 |
import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModel
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import tempfile
# Helper function to plot attention heatmap
def plot_attention(attn, tokens, layer=0, head=0):
plt.figure(figsize=(10, 8))
sns.heatmap(attn[layer][head], xticklabels=tokens, yticklabels=tokens, cmap="viridis")
plt.title(f"Attention Map - Layer {layer}, Head {head}")
plt.xlabel("Keys")
plt.ylabel("Queries")
tmp_file = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
plt.savefig(tmp_file.name, bbox_inches='tight')
plt.close()
return tmp_file.name
# Main logic
def process_input(text, model_name, layer, head):
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name, output_attentions=True, output_hidden_states=True)
model.eval()
inputs = tokenizer(text, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs)
tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
attentions = [a.squeeze(0).cpu().numpy() for a in outputs.attentions] # layers, heads, seq_len, seq_len
hidden_states = [h.squeeze(0).cpu().numpy().tolist() for h in outputs.hidden_states] # layers, seq_len, dim
attn_img_path = plot_attention(attentions, tokens, layer=layer, head=head)
return tokens, hidden_states[layer], attn_img_path
# Gradio interface function
def gradio_interface(text, model_name, layer, head):
tokens, hidden, attn_img = process_input(text, model_name, layer, head)
return tokens, hidden, attn_img
# Available transformer models
model_choices = [
"bert-base-uncased",
"distilbert-base-uncased",
"roberta-base",
"gpt2"
]
# Launch the Gradio app
interface = gr.Interface(
fn=gradio_interface,
inputs=[
gr.Textbox(label="Input Text", placeholder="Type a sentence here..."),
gr.Dropdown(label="Model", choices=model_choices, value="bert-base-uncased"),
gr.Slider(label="Attention Layer", minimum=0, maximum=11, step=1, value=0),
gr.Slider(label="Attention Head", minimum=0, maximum=11, step=1, value=0),
],
outputs=[
gr.JSON(label="Tokens"),
gr.Dataframe(label="Hidden States (Selected Layer)"),
gr.Image(label="Attention Map")
],
title="🔍 Transformer Visualizer",
description="Visualize tokenization, attention maps, and hidden states of popular Hugging Face Transformer models.",
)
if __name__ == "__main__":
interface.launch()
|