Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,493 @@
-import altair as alt
+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
-import streamlit as st
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+@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?
+
+
+        """)