Update app.py

app.py

@@ -5,7 +5,7 @@ import plotly.graph_objects as go
 from transformers import AutoModel, AutoTokenizer
 
 ############################################
-# Activation Plot (Interactive Plotly)
+# Activation Plot
 ############################################
 
 def activation_plot():
@@ -34,7 +34,7 @@ def activation_plot():
     return fig
 
 ############################################
-# Transformer Attention (one-click)
+# Attention Visualizer
 ############################################
 
 def attention_viz(model_name, text):
@@ -45,7 +45,7 @@ def attention_viz(model_name, text):
     with torch.no_grad():
         outputs = model(**tokens)
 
-    attn = outputs.attentions[-1][0]     # (heads, seq, seq)
+    attn = outputs.attentions[-1][0]  # (heads, seq, seq)
     num_heads = attn.shape[0]
 
     figs = []
@@ -65,31 +65,28 @@ def attention_viz(model_name, text):
     return figs
 
 ############################################
-# UI LAYOUT
+# UI
 ############################################
 
 with gr.Blocks(title="Transformer Attention Visualizer") as demo:
 
     gr.Markdown("# 🔥 Activation + Transformer Attention Visualizer")
-    gr.Markdown("### 🚀 Fully interactive Plotly (zoom in/out, hover, drag)")
+    gr.Markdown("### 🚀 Fully interactive Plotly (zoom / pan / hover)")
 
-    # ACTIVATION SECTION
+    # Activation section
     gr.Markdown("## 1️⃣ Activation Functions")
-    act_plot = gr.Plot(value=activation_plot())
+    gr.Plot(value=activation_plot())
 
-    # ATTENTION SECTION
-    gr.Markdown("## 2️⃣ Multi-Head Attention (HuggingFace Transformers)")
+    # Attention section
+    gr.Markdown("## 2️⃣ Multi-Head Attention")
 
     model_box = gr.Textbox(value="bert-base-uncased", label="Model name")
-    text_box = gr.Textbox(value="Transformers are amazing.", label="Text input")
+    text_box = gr.Textbox(value="Transformers are amazing.", label="Text")
 
     run_btn = gr.Button("Generate Attention")
-    gallery = gr.Gallery(label="Attention Heads (Zoomable)").style(
-        columns=2,
-        object_fit="contain",
-        height="auto"
-    )
 
-    run_btn.click(attention_viz, inputs=[model_box, text_box], outputs=gallery)
+    gallery = gr.Gallery(label="Attention Heads (zoomable)", show_label=True)
+
+    run_btn.click(attention_viz, inputs=[model_box, text_box], outputs=[gallery])
 
 demo.launch()