import streamlit as st st.set_page_config(page_title="GPT-2 Attention Explorer", layout="wide") import torch import numpy as np from transformers import GPT2TokenizerFast, GPT2Model import seaborn as sns import matplotlib.pyplot as plt import pandas as pd @st.cache_resource def load_model(): tokenizer = GPT2TokenizerFast.from_pretrained("./models") model = GPT2Model.from_pretrained("./models", output_attentions=True, attn_implementation="eager") model.eval() return tokenizer, model tokenizer, model = load_model() st.title("🧠 GPT-2 Token Inspector + Self-Attention Visualizer") with st.expander("📊 GPT-2 Model Architecture Summary"): st.markdown(""" - **Vocabulary size (V):** `50257` - **Embedding dimension (d):** `768` - **Max Position Length (L):** `1024` - **Transformer Layers:** `12` - **Attention Heads per Layer:** `12` - **Per-head Dimension (dₖ):** `64` - **Feedforward Hidden Layer Size:** `3072` - **Total Parameters:** ~117 million """) sentence = st.text_input("Enter a sentence:", "The cat sat on the mat") if st.button("Analyze & Visualize") and sentence.strip(): inputs = tokenizer(sentence, return_tensors='pt', return_offsets_mapping=True, return_special_tokens_mask=True) token_ids = inputs['input_ids'][0] tokens = tokenizer.convert_ids_to_tokens(token_ids) position_ids = torch.arange(token_ids.shape[0]).unsqueeze(0) inputs.pop("special_tokens_mask", None) inputs.pop("offset_mapping", None) with torch.no_grad(): outputs = model(**inputs, position_ids=position_ids) attentions = outputs.attentions embeddings = outputs.last_hidden_state[0].numpy() pos_embedding_layer = model.wpe pos_embeddings = pos_embedding_layer(position_ids).squeeze(0).detach().numpy() word_embedding_layer = model.wte word_embeddings = word_embedding_layer(token_ids).detach().numpy() final_input = word_embeddings + pos_embeddings # 1. BPE Tokens st.subheader("🧾 Byte Pair Encoded Tokens (BPE)") st.markdown("GPT-2 uses **Byte Pair Encoding (BPE)** to split input text into subword units.") st.code(" ".join(tokens)) # 2. Token IDs st.subheader("🔢 Token IDs") st.markdown("Each token is mapped to an integer ID using the GPT-2 vocabulary.") st.code(token_ids.tolist()) # 3. Word Embeddings st.subheader("💎 Raw Word Embeddings (first 5 tokens)") st.markdown("Each token ID is used to lookup a learnable word embedding vector:") st.latex(r"\text{Embedding}(t_i) = \mathbf{E}[t_i]") st.markdown(r"Where $\mathbf{E} \in \mathbb{R}^{V \times d}$ with $V$ = vocab size and $d = 768$.") df_word_embed = pd.DataFrame(word_embeddings[:5]) df_word_embed.index = [f"{i}: {tok}" for i, tok in enumerate(tokens[:5])] st.dataframe(df_word_embed.style.format(precision=4)) # 4. Positional Encodings st.subheader("🧭 Positional Encodings (first 5 tokens)") st.markdown("GPT-2 adds learned positional vectors from a table indexed by position:") st.latex(r"\text{PosEnc}(i) = \mathbf{P}[i]") st.markdown("Example (first 5 positions, first 5 dimensions):") df_pos_example = pd.DataFrame(pos_embeddings[:5, :5], columns=[f"dim {i}" for i in range(5)], index=[f"{i}: {tok}" for i, tok in enumerate(tokens[:5])]) st.dataframe(df_pos_example.style.format(precision=5)) st.markdown(r"Where $\mathbf{P} \in \mathbb{R}^{L \times d}$ is learned and not sinusoidal in GPT-2.") # 5. Final Input Vectors st.subheader("🧮 Final Input = Word Embedding + Positional Encoding") st.markdown("These are the actual vectors passed into the first transformer block:") st.latex(r"\mathbf{X}_i = \text{Embedding}(t_i) + \text{PosEnc}(i)") st.markdown("Let's confirm this by showing:") st.code("final_input[i][j] ≈ word_embedding[i][j] + pos_embedding[i][j]") for i in range(2): # for first 2 tokens df_sum_example = pd.DataFrame({ 'Word': word_embeddings[i, :5], 'PosEnc': pos_embeddings[i, :5], 'Final Input': final_input[i, :5], 'Word + Pos': word_embeddings[i, :5] + pos_embeddings[i, :5] }) df_sum_example.index = [f"dim {j}" for j in range(5)] st.markdown(f"**Token {i}: `{tokens[i]}`**") st.dataframe(df_sum_example.style.format(precision=5)) # 6. Output Embeddings st.subheader("📐 Output Embedding Vectors (first 5 tokens)") st.markdown("These are the final hidden states after passing through all transformer layers:") st.latex(r"\text{Output}_i = \text{TransformerLayers}(\mathbf{X}_i)") df_embed_example = pd.DataFrame(embeddings[:5, :5], columns=[f"dim {j}" for j in range(5)], index=[f"{i}: {tok}" for i, tok in enumerate(tokens[:5])]) st.dataframe(df_embed_example.style.format(precision=5)) st.markdown("📌 These are **not** equal to the input vectors—they are fully context-aware representations!") # 🔄 Move sliders here just above heatmap layer_num = st.slider("Select Transformer Layer", 0, model.config.n_layer - 1, 0) head_num = st.slider("Select Attention Head", 0, model.config.n_head - 1, 0) attn = attentions[layer_num][0, head_num].numpy() # 7. Attention Heatmap st.subheader(f"🎯 Attention Heatmap — Layer {layer_num+1}, Head {head_num+1}") st.markdown("This shows how each token attends to others in the sequence:") st.latex(r"\text{Attention}(Q, K, V) = \text{softmax} \left( \frac{QK^\top}{\sqrt{d_k}} \right) V") fig, ax = plt.subplots(figsize=(8, 6)) sns.heatmap(attn, xticklabels=tokens, yticklabels=tokens, cmap="YlOrRd", annot=True, fmt=".2f", ax=ax) ax.set_xlabel("Key Tokens") ax.set_ylabel("Query Tokens") st.pyplot(fig) # 8. Attention Head Breakdown (for token 0) st.subheader("🔍 Attention Head Breakdown (1 Token)") st.markdown("Let's inspect how **GPT-2 computes attention for a single token** (first token in the sequence).") # Fetch weight matrix for Q, K, V from the model's first block # block = model.transformer.h[0] # Use layer 0 block = model.h[0] # ✅ Correct for GPT2Model # W_qkv = block.attn.c_attn.weight.detach().numpy().T # shape (768, 3*768) W_qkv = block.attn.c_attn.weight.detach().numpy() # ✅ shape (2304, 768) b_qkv = block.attn.c_attn.bias.detach().numpy() # shape (3*768,) # Final input for token 0 x0 = final_input[0] # shape (768,) # Linear projection for Q, K, V qkv = x0 @ W_qkv + b_qkv # shape (3*768,) Q, K, V = np.split(qkv, 3) # Show Q, K, V for head 0 Q0 = Q[:64] K0_all = K.reshape(12, 64) # For all heads V0_all = V.reshape(12, 64) K0 = K0_all[0] V0 = V0_all[0] # Dot product and softmax score = Q0 @ K0.T # scalar scaled_score = score / np.sqrt(64) softmax_weight = np.exp(scaled_score) / np.sum(np.exp(scaled_score)) attn_output = softmax_weight * V0 # simulated for 1 token self-attending to itself st.markdown("### Formula Recap") st.latex(r"Q = x W^Q,\quad K = x W^K,\quad V = x W^V") st.latex(r"\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)V") # Show Q0, K0, softmax and V0 df_breakdown = pd.DataFrame({ "Q₀": Q0, "K₀": K0, "Q₀·K₀": Q0 * K0, "V₀": V0, "AttnOut": attn_output }) df_breakdown.index = [f"dim {i}" for i in range(64)] st.dataframe(df_breakdown.style.format(precision=5)) st.markdown("### 🧮 Self-Attention Matrix Shape Annotations") st.markdown(""" **Key tensor dimensions involved in attention computation:** - `W_qkv`: **(2304, 768)** – learned projection matrix for Q, K, V combined - `b_qkv`: **(2304,)** – bias vector - `X`: **(5, 768)** – input vectors for 5 tokens - `qkv_all = X @ W_qkv + b_qkv`: → **(5, 2304)** - `Q_all, K_all, V_all = np.split(qkv_all, 3)`: → each **(5, 768)** - `Q0, K0, V0 = [:, :64]`: head 0 slice → **(5, 64)** - `q0 @ K0.T`: **(1, 64) × (64, 5)** → **(1, 5)** - `softmax_weights`: **(1, 5)** - `attn_output = softmax_weights @ V0`: **(1, 64)** """) # 9. Matrix-Level Self-Attention (Token 0 → All) st.subheader("🔬 Matrix-Level Self-Attention (Token 0 → All)") st.markdown(""" This section shows how **Token 0** attends to all other tokens using matrix-level self-attention. We compute the dot products, apply softmax, and produce the output for head 0 in layer 0. """) # Use same block block = model.h[0] W_qkv = block.attn.c_attn.weight.detach().numpy() # (2304, 768) b_qkv = block.attn.c_attn.bias.detach().numpy() # (2304,) X = final_input[:5] # (5, 768) # Compute Q, K, V for all 5 tokens # qkv_all = X @ W_qkv.T + b_qkv # shape (5, 2304) qkv_all = X @ W_qkv + b_qkv # ✅ (5 × 768) @ (768 × 2304) Q_all, K_all, V_all = np.split(qkv_all, 3, axis=1) # Head 0 slices Q0 = Q_all[:, :64] # (5, 64) K0 = K_all[:, :64] # (5, 64) V0 = V_all[:, :64] # (5, 64) # Compute raw attention scores for token 0 q0 = Q0[0].reshape(1, 64) # (1, 64) attn_scores = q0 @ K0.T # (1, 5) scaled_scores = attn_scores / np.sqrt(64) softmax_weights = np.exp(scaled_scores) softmax_weights /= softmax_weights.sum(axis=-1, keepdims=True) # shape (1, 5) # Weighted sum of V0 rows attn_output_0 = softmax_weights @ V0 # (1, 64) # Display matrices st.markdown("### Raw Scaled Attention Scores (Q₀Kᵀ / √dₖ):") df_scores = pd.DataFrame(scaled_scores[0], columns=["Score"], index=[f"Token {i}" for i in range(5)]) st.dataframe(df_scores.style.format(precision=5)) st.markdown("### Softmax Attention Weights αᵢ:") df_weights = pd.DataFrame(softmax_weights[0], columns=["Weight αᵢ"], index=[f"Token {i}" for i in range(5)]) st.dataframe(df_weights.style.format(precision=5)) st.markdown("### Value Vᵢ vectors (Head 0, first 5 dims):") df_values = pd.DataFrame(V0[:, :5], columns=[f"dim {i}" for i in range(5)], index=[f"Token {i}" for i in range(5)]) st.dataframe(df_values.style.format(precision=5)) st.markdown("### Final Attention Output (weighted sum of Vᵢ):") df_attn_out = pd.DataFrame(attn_output_0[:, :5], columns=[f"dim {i}" for i in range(5)], index=["AttnOut₀"]) st.dataframe(df_attn_out.style.format(precision=5)) # 10. Per-Head Projection Matrices st.subheader("🧬 Per-Head Projection Matrices (Wq, Wk, Wv)") st.markdown(""" In GPT-2, each attention **head has its own set of projection weights** to compute Queries (Q), Keys (K), and Values (V) from the input vector. The full `W_qkv` layer maps from **(768,) → (2304,)** and is split into 3 parts: - `Wq` = first 768 columns → shape `(768, 768)` - `Wk` = next 768 columns → shape `(768, 768)` - `Wv` = last 768 columns → shape `(768, 768)` Each head receives a unique slice from each projection: - 12 heads × 64 dimensions = 768 - So head 0 → `Wq[:, :64]`, head 1 → `Wq[:, 64:128]`, etc. """) block = model.h[0] W_qkv_full = block.attn.c_attn.weight.detach().numpy().T # shape (768, 2304) W_q, W_k, W_v = np.split(W_qkv_full, 3, axis=1) # each: (768, 768) # Show Wq head 0 and 1 Wq_head0 = W_q[:, :64] Wq_head1 = W_q[:, 64:128] df_q = pd.DataFrame({ "Wq_head0": Wq_head0[:5, 0], "Wq_head1": Wq_head1[:5, 0] }, index=[f"dim {i}" for i in range(5)]) st.markdown("### Wq projection weights for head 0 vs head 1 (first 5 input dims → output dim 0):") st.dataframe(df_q.style.format(precision=5)) # Show Wk and Wv for head 0 Wk_head0 = W_k[:, :64] Wv_head0 = W_v[:, :64] df_kv = pd.DataFrame({ "Wk_head0": Wk_head0[:5, 0], "Wv_head0": Wv_head0[:5, 0] }, index=[f"dim {i}" for i in range(5)]) st.markdown("### Wk and Wv projection weights for head 0 (first 5 input dims → output dim 0):") st.dataframe(df_kv.style.format(precision=5)) st.markdown(""" ✅ This confirms that each head has **distinct projections** for Q, K, and V. The same input `x` is transformed differently per head, allowing GPT-2 to learn different attention perspectives. """) # 11 · 📐 How W_qkv Projects an Input Vector into Q, K, V st.subheader("📐 How W_qkv Projects an Input Vector → Q, K, V") st.markdown(""" In GPT-2, the combined projection layer `c_attn` maps a single input embedding into a concatenated vector that contains **Q, K, and V**. Each of these is 768-dimensional, so the full output is 768 × 3 = 2304. """) st.latex(r"x \in \mathbb{R}^{768} \quad \rightarrow \quad [Q \;|\; K \;|\; V] \in \mathbb{R}^{2304}") st.markdown("---") st.markdown("### 🧪 Mini GPT Example (3D → 6D Projection)") st.markdown("Imagine a tiny model:") st.markdown(""" - Input vector `x ∈ ℝ³` - Q, K, V are each 2D → total output = 6D - Thus: """) st.latex(r"W_{\text{qkv}} \in \mathbb{R}^{6 \times 3}, \quad b_{\text{qkv}} \in \mathbb{R}^6") # Miniature input vector and projection weights mini_x = np.array([1.0, 2.0, 3.0]) # (3,) mini_W = np.array( # (6, 3) [ [0.1, 0.2, 0.3], # → Q₁ [0.4, 0.5, 0.6], # → Q₂ [0.7, 0.8, 0.9], # → K₁ [1.0, 1.1, 1.2], # → K₂ [1.3, 1.4, 1.5], # → V₁ [1.6, 1.7, 1.8], # → V₂ ] ) mini_b = np.array([0.01, 0.02, 0.03, 0.04, 0.05, 0.06]) # (6,) mini_out = mini_W @ mini_x + mini_b # (6,) Qm, Km, Vm = np.split(mini_out, 3) # each (2,) st.code("Input vector x = [1.0, 2.0, 3.0] # shape (3,)") st.code("W_qkv shape = (6, 3) # maps 3 → 6") st.code(f"Output = W_qkv @ x + b = {mini_out.round(2).tolist()}") df_mini = pd.DataFrame( { "Q": Qm.round(2), "K": Km.round(2), "V": Vm.round(2) }, index=["dim 1", "dim 2"] ) st.markdown("**Split into Q, K, V (each 2D):**") st.dataframe(df_mini.style.format(precision=2)) st.markdown("---") st.markdown("### 📏 Real GPT-2 Projection Shapes") df_shapes = pd.DataFrame({ "Tensor": [ "Input x", "W_qkv (linear layer)", "b_qkv (bias)", "Output = x @ W_qkv + b", "Q / K / V each", "Head reshaping" ], "Shape": [ "(768,)", "(2304, 768)", "(2304,)", "(2304,)", "(768,)", "12 heads × 64 dims = 768" ] }) st.dataframe(df_shapes) st.markdown(""" Each attention **head** gets its own slice: - Q_head₀ = Q[:, :64] - K_head₀ = K[:, :64] - V_head₀ = V[:, :64] That’s how one input vector creates multi-headed Q, K, and V for scaled dot-product attention. """) st.subheader("Additional notes:") st.markdown( """ --- ## 🧠 What Does `Ġ` Mean? The character `Ġ` (U+0120: Latin Capital Letter G with dot above) is used to: > **Represent a leading space** before the token. --- ### ✅ Example: Let’s look at a sentence: ``` "The cat sat on the mat" ``` When tokenized using GPT-2 tokenizer (`GPT2TokenizerFast`), it becomes: ``` ['The', 'Ġcat', 'Ġsat', 'Ġon', 'Ġthe', 'Ġmat'] ``` * `'The'` → First word, no leading space. * `'Ġcat'` → Space + "cat" * `'Ġsat'` → Space + "sat" * etc. So `Ġ` means: > "This token starts after a space." --- ### ⚠️ Why Not Just Use `" "`? Because GPT-2 uses a **vocabulary of subword units** (BPE). These tokens are strings, not raw characters or bytes. Including space as a separate token would have complicated the merge process. So: * `Ġ` = internal marker used in the vocabulary file * It's not a space character but tells the tokenizer "insert space before decoding this." --- ### ✅ When Detokenizing The tokenizer **removes the `Ġ` and adds a space** during decoding: ```python from transformers import GPT2TokenizerFast tokenizer = GPT2TokenizerFast.from_pretrained("gpt2") tokens = tokenizer.tokenize("The cat sat on the mat") print(tokens) # ['The', 'Ġcat', 'Ġsat', 'Ġon', 'Ġthe', 'Ġmat'] ids = tokenizer.convert_tokens_to_ids(tokens) decoded = tokenizer.decode(ids) print(decoded) # 'The cat sat on the mat' ``` --- ## ✅ Summary | Token | Interprets As | | -------- | ------------------------- | | `'The'` | `'The'` (no space before) | | `'Ġcat'` | `' cat'` | | `'Ġsat'` | `' sat'` | | `'Ġon'` | `' on'` | | `'Ġthe'` | `' the'` | | `'Ġmat'` | `' mat'` | Would you like to include this as an educational block in your Streamlit app too? """)