feat: made attention logits plot more interactive and added intuitive logic fir attention logits and attention scores

app.py

@@ -352,7 +352,7 @@ def update_attention(data, head, query_token):
     )
     return (
         attention_heatmaps(sb, sa, tokens=data["tokens"]),
-        attention_bars(sb[qt], sa[qt], qt),
+        attention_bars(sb[qt], sa[qt], qt, tokens=data["tokens"]),
         note,
         _attention_context_markdown(data),
     )
@@ -468,6 +468,22 @@ with gr.Blocks(title="RoPE Explorer") as demo:
             attn_note = gr.Markdown()
             gr.Markdown(
                 """
+### How to read the Attention logits bars
+
+The selected query token is compared with every key token. Each bar is the raw
+dot product `Q_query · K_key` for one key position:
+
+- A **higher bar** means the key is a stronger match for this query relative to the other keys in the same chart.
+- A **lower or negative bar** means a weaker or opposing match.
+- The gap between **without RoPE** and **with RoPE** shows how position-aware rotation changes that comparison.
+- These bars are **logits, not probabilities**. Applying softmax across all bars for one query would convert them into attention weights.
+
+Use the hover text to see the exact key token and position behind each bar. The
+query token is the position shown in the chart title.
+"""
+            )
+            gr.Markdown(
+                """
 ### How to read an attention score
 
 Each heatmap cell is the raw dot product `Q_query · K_key` for the query token on

src/plots.py

@@ -275,15 +275,29 @@ def attention_heatmaps(
     return fig
 
 
-def attention_bars(logits_before: np.ndarray, logits_after: np.ndarray, query_token: int) -> go.Figure:
+def attention_bars(
+    logits_before: np.ndarray,
+    logits_after: np.ndarray,
+    query_token: int,
+    tokens=None,
+) -> go.Figure:
     xs = list(range(len(logits_before)))
+    labels = [str(tokens[i]) if tokens is not None else f"token {i}" for i in xs]
     fig = go.Figure()
-    fig.add_trace(go.Bar(x=xs, y=logits_before, name="without RoPE"))
-    fig.add_trace(go.Bar(x=xs, y=logits_after, name="with RoPE"))
+    hover = (
+        "key position k=%{x}<br>key token=%{customdata}<br>"
+        "raw attention logit Q·K=%{y:.4f}<extra></extra>"
+    )
+    fig.add_trace(
+        go.Bar(x=xs, y=logits_before, customdata=labels, hovertemplate=hover, name="without RoPE")
+    )
+    fig.add_trace(
+        go.Bar(x=xs, y=logits_after, customdata=labels, hovertemplate=hover, name="with RoPE")
+    )
     fig.update_layout(
-        title=f"Attention logits from query token {query_token}",
-        xaxis_title="key token",
-        yaxis_title="Q·K",
+        title=f"Attention logits from query position k={query_token}",
+        xaxis_title="key token position k",
+        yaxis_title="raw attention logit (Q·K)",
         barmode="group",
         template="plotly_white",
         height=360,