Attention Explorer: redesign controls, examples, neighbors; add config support

- Controls: Layer ▲▼ + Head ◄► on separate rows, matching keyboard shortcuts
- Examples: labeled pairs (bass/spring/light) with background grouping, above input
- Removed current/pitch pairs (kept in search tool, documented in CLAUDE.md)
- Neighbors: pill chain with → bridge arrow, probability-scaled borders
- Movement colors: only show when no word selected (exploration mode)
- Selected word: full-row outline, word-group-only background highlight
- Attention model now configurable via config.json (attention_model key)
- Updated model description text to GPT-2 Medium
- Input label: "Text" with hint "Try an example above, or enter your own text"

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

app.py

@@ -207,6 +207,49 @@ body.dark, .dark {
     font-size: 15px;
     line-height: 1.6;
 }
+
+/* Attention Explorer example pairs */
+.attn-example-row {
+    background: #f8f6fb !important;
+    border-radius: 6px !important;
+    padding: 6px 10px !important;
+    margin-bottom: 2px !important;
+    align-items: flex-start !important;
+    gap: 6px !important;
+    flex-wrap: wrap !important;
+}
+.attn-example-label {
+    min-width: 70px !important;
+    max-width: 70px !important;
+    flex-shrink: 0 !important;
+    padding-top: 4px !important;
+}
+.attn-example-label p {
+    font-family: 'Merriweather', serif !important;
+    color: #63348d !important;
+    font-size: 13px !important;
+    margin: 0 !important;
+}
+.attn-example-btn {
+    flex: 0 0 auto !important;
+    min-width: 0 !important;
+    max-width: fit-content !important;
+}
+.attn-example-btn button {
+    font-size: 12px !important;
+    padding: 4px 12px !important;
+    border: 1.5px solid #d4c8e8 !important;
+    border-radius: 14px !important;
+    background: #fff !important;
+    color: #4a3070 !important;
+    white-space: nowrap !important;
+    width: auto !important;
+    text-align: left !important;
+}
+.attn-example-btn button:hover {
+    background: #f3f0f7 !important;
+    border-color: #63348d !important;
+}
 """
 
 # Dark mode toggle JS — toggles class and swaps button text + heading colors
@@ -1193,40 +1236,42 @@ def create_app():
             gr.Markdown("### Attention Explorer")
             gr.Markdown(
                 "See which words the model pays attention to when processing a sentence. "
-                "Uses GPT-2 Large (774M parameters, 36 layers, 20 attention heads). "
+                "Uses GPT-2 Medium (345M parameters, 24 layers, 16 attention heads). "
                 "Click a word to see curved lines connecting it to the words it attended to — "
                 "thicker lines mean stronger attention."
             )
 
+            # Example sentence pairs — labeled by polysemy word
+            attn_example_pairs = [
+                ("bass", "He tuned his bass and plugged into the", "On the lake she caught a bass and pulled it onto the"),
+                ("spring", "She wound the metal spring and the clock began to", "After the long winter the warm spring rain made the flowers"),
+                ("light", "She flipped the switch and the light began to", "The bag was so light she carried it with"),
+            ]
+
+            attn_example_btns = []
+            for word, sent_a, sent_b in attn_example_pairs:
+                with gr.Row(elem_classes=["attn-example-row"]):
+                    gr.Markdown(f"**{word}:**", elem_classes=["attn-example-label"])
+                    btn_a = gr.Button(sent_a, size="sm", variant="secondary", elem_classes=["attn-example-btn"])
+                    btn_b = gr.Button(sent_b, size="sm", variant="secondary", elem_classes=["attn-example-btn"])
+                    attn_example_btns.extend([btn_a, btn_b])
+
+            gr.Markdown("*Try an example above, or enter your own text:*")
             attn_input = gr.Textbox(
-                label="Sentence",
+                label="Text",
                 value="",
                 lines=1,
-                placeholder="Enter a sentence to explore...",
+                placeholder="Enter text to explore...",
             )
             attn_btn = gr.Button("Explore", variant="primary")
-            gr.Examples(
-                examples=[
-                    ["He tuned his electric bass and plugged into the"],
-                    ["Out on the lake she caught a large bass and pulled it onto the"],
-                    ["She wound the metal spring and the clock began to"],
-                    ["After the long winter the warm spring rain made the flowers"],
-                    ["The river current pulled the boat toward the"],
-                    ["In her speech the current president talked about the"],
-                    ["She tuned her guitar until the pitch was"],
-                    ["The ballplayer stood on the mound and wound up to pitch to the"],
-                    ["She flipped the switch and the light began to"],
-                    ["The bag was so light she carried it with"],
-                ],
-                inputs=[attn_input],
-            )
             attn_output = gr.HTML(label="Visualization")
 
             def run_attention_explorer(text):
                 """Run attention extraction and build visualization HTML."""
                 if not text or not text.strip():
-                    return "<p style='color:#94a3b8;font-style:italic;'>Enter a sentence above and click Explore.</p>"
-                data = get_attention_data(text.strip())
+                    return "<p style='color:#94a3b8;font-style:italic;'>Enter text above and click Explore.</p>"
+                attn_model = manager.config.get("attention_model", "gpt2-medium")
+                data = get_attention_data(text.strip(), model_name=attn_model)
                 return build_attention_html(data)
 
             attn_btn.click(
@@ -1235,31 +1280,13 @@ def create_app():
                 outputs=[attn_output],
             )
 
-            # Auto-explore when an example is clicked (detected via textbox change)
-            attn_examples_list = [
-                "He tuned his electric bass and plugged into the",
-                "Out on the lake she caught a large bass and pulled it onto the",
-                "She wound the metal spring and the clock began to",
-                "After the long winter the warm spring rain made the flowers",
-                "The river current pulled the boat toward the",
-                "In her speech the current president talked about the",
-                "She tuned her guitar until the pitch was",
-                "The ballplayer stood on the mound and wound up to pitch to the",
-                "She flipped the switch and the light began to",
-                "The bag was so light she carried it with",
-            ]
-
-            def on_attn_input_change(text):
-                """Auto-explore if text matches an example (example click)."""
-                if text and text.strip() in attn_examples_list:
-                    return run_attention_explorer(text)
-                return gr.update()
-
-            attn_input.change(
-                fn=on_attn_input_change,
-                inputs=[attn_input],
-                outputs=[attn_output],
-            )
+            # Wire up example buttons — each sets input and auto-explores
+            for btn in attn_example_btns:
+                btn.click(
+                    fn=lambda text: (text, run_attention_explorer(text)),
+                    inputs=[btn],
+                    outputs=[attn_input, attn_output],
+                )
 
           # ==================================================================
           # Admin Panel

attention.py

@@ -33,14 +33,17 @@ def _detect_device() -> str:
     return "cpu"
 
 
-def load_attention_model():
-    """Load GPT-2 Medium with eager attention. Idempotent."""
+def load_attention_model(model_name=None):
+    """Load attention model with eager attention. Idempotent."""
     global _model, _tokenizer, _device
     if _model is not None:
         return
-    _tokenizer = AutoTokenizer.from_pretrained("gpt2-medium")
+    if model_name is None:
+        model_name = "gpt2-medium"
+    print(f"Loading attention model: {model_name}")
+    _tokenizer = AutoTokenizer.from_pretrained(model_name)
     _model = AutoModelForCausalLM.from_pretrained(
-        "gpt2-medium", attn_implementation="eager"
+        model_name, attn_implementation="eager"
     )
     _model.eval()
     _device = _detect_device()
@@ -275,7 +278,7 @@ def _compute_movement(hidden_states, word_groups):
     return movement
 
 
-def get_attention_data(text):
+def get_attention_data(text, model_name=None):
     """Run forward pass and return attention + neighbor + movement data.
 
     Returns dict with:
@@ -288,7 +291,7 @@ def get_attention_data(text):
     """
     global _cached_hidden_states, _cached_words
 
-    load_attention_model()
+    load_attention_model(model_name)
 
     inputs = _tokenizer(text, return_tensors="pt")
     inputs = {k: v.to(_device) for k, v in inputs.items()}
@@ -381,37 +384,40 @@ def _build_viz_html(data_json):
 <link href="https://fonts.googleapis.com/css2?family=Inter:wght@400;600&family=JetBrains+Mono&family=Merriweather:wght@700&display=swap" rel="stylesheet">
 </head><body style="margin:0;padding:0;">
 <div class="attn-explorer" id="attn-viz" tabindex="0">
-  <!-- Layer selector -->
+  <!-- Controls: two rows -->
   <div class="attn-controls">
-    <div class="attn-selector">
+    <div class="attn-controls-row">
       <span class="attn-selector-title">Layer</span>
-      <button class="attn-nav-btn" onclick="attnPrevLayer()">&#9664;</button>
+      <button class="attn-nav-btn" onclick="attnPrevLayer()">&#9650;</button>
       <span class="attn-label" id="attn-layer-label">12 / 12</span>
-      <button class="attn-nav-btn" onclick="attnNextLayer()">&#9654;</button>
+      <button class="attn-nav-btn" onclick="attnNextLayer()">&#9660;</button>
+      <span class="attn-key-hint">↑ ↓</span>
+      <span class="attn-sep"></span>
+      <span class="attn-nbr-label">Neighbors:</span>
+      <select class="attn-nbr-select" id="attn-nbr-select" onchange="attnSetNeighborCount(this.value)">
+        <option value="0">Off</option>
+        <option value="1">1</option>
+        <option value="3" selected>3</option>
+        <option value="5">5</option>
+      </select>
     </div>
-    <div class="attn-selector">
+    <div class="attn-controls-row">
       <span class="attn-selector-title">Head</span>
       <button class="attn-nav-btn" onclick="attnPrevHead()">&#9664;</button>
       <span class="attn-label" id="attn-head-label">Avg</span>
       <button class="attn-nav-btn" onclick="attnNextHead()">&#9654;</button>
+      <span class="attn-key-hint">← →</span>
+      <span class="attn-sep"></span>
+      <span class="attn-show-label">Show:</span>
+      <label class="attn-checkbox">
+        <input type="checkbox" id="attn-show-all-cb" onchange="attnToggleShowAll(this.checked)">
+        All connections
+      </label>
+      <label class="attn-checkbox">
+        <input type="checkbox" id="attn-show-punct-cb" onchange="attnToggleShowPunct(this.checked)">
+        Punctuation
+      </label>
     </div>
-    <label class="attn-show-all">
-      <input type="checkbox" id="attn-show-all-cb" onchange="attnToggleShowAll(this.checked)">
-      Show all connections
-    </label>
-    <div class="attn-nbr-count">
-      <label>Neighbors:</label>
-      <select id="attn-nbr-select" onchange="attnSetNeighborCount(this.value)">
-        <option value="0">off</option>
-        <option value="1">1</option>
-        <option value="3" selected>3</option>
-        <option value="5">5</option>
-      </select>
-    </div>
-    <label class="attn-show-all">
-      <input type="checkbox" id="attn-show-punct-cb" onchange="attnToggleShowPunct(this.checked)">
-      Show punctuation
-    </label>
   </div>
 
   <div class="attn-hint" id="attn-hint">Click a word to see what it attended to</div>
@@ -432,31 +438,30 @@ def _build_viz_html(data_json):
 <style>
   .attn-explorer {{
     font-family: 'Inter', system-ui, sans-serif;
-    padding: 16px;
+    padding: 16px 16px 16px 0;
     outline: none;
   }}
   .attn-controls {{
     display: flex;
-    align-items: center;
-    gap: 24px;
+    flex-direction: column;
+    gap: 10px;
     margin-bottom: 12px;
-    flex-wrap: wrap;
   }}
-  .attn-selector {{
+  .attn-controls-row {{
     display: flex;
     align-items: center;
-    gap: 8px;
+    gap: 10px;
   }}
   .attn-selector-title {{
     font-family: 'Merriweather', serif;
-    font-weight: 600;
+    font-weight: 700;
     color: #63348d;
-    font-size: 14px;
-    min-width: 40px;
+    font-size: 15px;
+    min-width: 50px;
   }}
   .attn-label {{
     font-family: 'JetBrains Mono', monospace;
-    font-size: 14px;
+    font-size: 15px;
     min-width: 72px;
     text-align: center;
     color: #1e293b;
@@ -465,52 +470,66 @@ def _build_viz_html(data_json):
     background: #63348d;
     color: white;
     border: none;
-    border-radius: 4px;
-    padding: 4px 10px;
+    border-radius: 6px;
+    width: 32px;
+    height: 32px;
     cursor: pointer;
-    font-size: 12px;
+    font-size: 16px;
     line-height: 1;
+    display: inline-flex;
+    align-items: center;
+    justify-content: center;
     transition: background 0.15s;
   }}
   .attn-nav-btn:hover {{
     background: #4e2870;
   }}
-  .attn-show-all {{
-    display: flex;
-    align-items: center;
-    gap: 6px;
-    font-size: 14px;
-    color: #475569;
-    cursor: pointer;
-    user-select: none;
+  .attn-key-hint {{
+    font-family: 'JetBrains Mono', monospace;
+    font-size: 10px;
+    color: #999;
+    letter-spacing: 3px;
   }}
-  .attn-show-all input {{
-    accent-color: #63348d;
-    cursor: pointer;
+  .attn-sep {{
+    width: 1px;
+    height: 32px;
+    background: #e0d8eb;
+    margin: 0 6px;
   }}
-  .attn-nbr-count {{
-    display: flex;
-    align-items: center;
-    gap: 6px;
-    font-size: 14px;
-    color: #475569;
-  }}
-  .attn-nbr-count label {{
+  .attn-nbr-label {{
     font-family: 'Merriweather', serif;
-    font-weight: 600;
+    font-weight: 700;
     color: #63348d;
     font-size: 14px;
   }}
-  .attn-nbr-count select {{
+  .attn-nbr-select {{
     font-family: 'JetBrains Mono', monospace;
     font-size: 13px;
-    padding: 2px 6px;
+    padding: 4px 8px;
     border: 1px solid #d1d5db;
     border-radius: 4px;
     background: white;
     color: #1e293b;
     cursor: pointer;
   }}
+  .attn-show-label {{
+    font-weight: 600;
+    color: #63348d;
+    font-size: 14px;
+  }}
+  .attn-checkbox {{
+    display: flex;
+    align-items: center;
+    gap: 6px;
+    font-size: 14px;
+    color: #475569;
+    cursor: pointer;
+    user-select: none;
+  }}
+  .attn-checkbox input {{
+    accent-color: #63348d;
+    cursor: pointer;
+  }}
   .attn-hint {{
     font-size: 13px;
     color: #94a3b8;
@@ -522,10 +541,11 @@ def _build_viz_html(data_json):
     min-height: 100px;
     display: flex;
     flex-direction: row;
+    margin-left: -16px;
   }}
   .attn-lines-container {{
-    width: 180px;
-    min-width: 180px;
+    width: 140px;
+    min-width: 140px;
     position: relative;
     flex-shrink: 0;
   }}
@@ -538,44 +558,82 @@ def _build_viz_html(data_json):
     flex-grow: 1;
   }}
   .attn-word-row {{
-    display: inline-flex;
+    display: flex;
     align-items: center;
-    padding: 5px 12px;
-    border-radius: 4px;
+    padding: 3px 4px;
     cursor: pointer;
     user-select: none;
     font-size: 16px;
+    min-height: 36px;
+  }}
+  .attn-word-group {{
+    display: inline-flex;
+    align-items: center;
+    padding: 4px 10px;
+    border-radius: 5px;
     transition: background 0.15s;
-    width: fit-content;
+    width: 130px;
+    flex-shrink: 0;
   }}
-  .attn-word-row:hover {{
+  .attn-word-row:hover .attn-word-group {{
     background: #f3f0f7;
   }}
   .attn-word-row.attn-selected {{
-    background: #ded9f4;
-    font-weight: 600;
     outline: 2px solid #63348d;
     outline-offset: -2px;
+    border-radius: 5px;
+  }}
+  .attn-word-row.attn-selected .attn-word-group {{
+    background: #ded9f4;
+    font-weight: 600;
   }}
-  .attn-word-row.attn-target {{
+  .attn-word-row.attn-target .attn-word-group {{
     background: #f3f0f7;
   }}
   .attn-word-label {{
     font-family: 'Inter', system-ui, sans-serif;
+    min-width: 70px;
   }}
   .attn-weight-badge {{
     font-family: 'JetBrains Mono', monospace;
     font-size: 11px;
     color: #63348d;
-    margin-left: 8px;
+    margin-left: 6px;
     opacity: 0.8;
+    width: 30px;
+    text-align: right;
+  }}
+  .attn-bridge {{
+    color: #b8a8d0;
+    font-size: 14px;
+    padding: 0 8px 0 10px;
+    flex-shrink: 0;
   }}
   .attn-word-neighbors {{
+    display: flex;
+    align-items: center;
+    gap: 0;
+  }}
+  .attn-nbr-pill {{
     font-family: 'JetBrains Mono', monospace;
     font-size: 12px;
-    color: #666;
-    font-style: italic;
-    margin-left: 12px;
+    padding: 2px 8px;
+    white-space: nowrap;
+    background: #fff;
+    color: #4a3070;
+    border-radius: 4px;
+    display: inline-flex;
+    align-items: center;
+    justify-content: center;
+    width: 90px;
+    text-align: center;
+  }}
+  .attn-nbr-sep {{
+    width: 16px;
+    flex-shrink: 0;
+    text-align: center;
+    font-size: 13px;
+    color: #b8a8d0;
   }}
   .attn-lines-svg {{
     position: absolute;
@@ -761,16 +819,18 @@ def _build_viz_html(data_json):
         }}
       }}
 
-      // Movement background color — skip for selected word and attention targets
-      var isTarget = !attnShowAll && targetWeights.hasOwnProperty(i);
-      if (i === attnSelectedWord || isTarget) {{
-        // Let CSS class handle background for selected/target words
-        row.style.background = '';
-      }} else if (movement && movement[i] !== undefined) {{
+      // Movement background color — only when no word is selected (exploration mode)
+      // Full row highlight since movement relates to the word AND its neighbors
+      var wordGroup = row.querySelector('.attn-word-group');
+      var noSelection = (attnSelectedWord === null);
+      if (noSelection && movement && movement[i] !== undefined) {{
         var m = movement[i];
-        row.style.background = 'rgba(99, 52, 141, ' + (m * 0.2).toFixed(3) + ')';
+        var moveBg = 'rgba(99, 52, 141, ' + (m * 0.2).toFixed(3) + ')';
+        row.style.background = moveBg;
+        if (wordGroup) wordGroup.style.background = '';
       }} else {{
         row.style.background = '';
+        if (wordGroup) wordGroup.style.background = '';
       }}
     }});
 
@@ -806,9 +866,11 @@ def _build_viz_html(data_json):
     var neighbors = ATTN_DATA.neighbors[layerKey];
     WORDS.forEach(function(word, i) {{
       var el = document.getElementById('attn-nbr-' + i);
+      var bridgeEl = document.getElementById('attn-bridge-' + i);
       if (el) {{
+        el.innerHTML = '';
         if (attnNeighborCount === 0 || !neighbors || !neighbors[i]) {{
-          el.textContent = '';
+          if (bridgeEl) bridgeEl.style.visibility = 'hidden';
           return;
         }}
         // Filter: skip self, always skip profanity, optionally skip punctuation
@@ -821,7 +883,29 @@ def _build_viz_html(data_json):
           if (!attnShowPunct && n.punct) continue;
           filtered.push(n);
         }}
-        el.textContent = filtered.map(function(n) {{ return n.word; }}).join(', ');
+        if (filtered.length === 0) {{
+          if (bridgeEl) bridgeEl.style.visibility = 'hidden';
+          return;
+        }}
+        // Show bridge arrow
+        if (bridgeEl) bridgeEl.style.visibility = 'visible';
+        // Render pills
+        filtered.forEach(function(n, idx) {{
+          if (idx > 0) {{
+            var sep = document.createElement('span');
+            sep.className = 'attn-nbr-sep';
+            sep.textContent = '\u203A';
+            el.appendChild(sep);
+          }}
+          var pill = document.createElement('span');
+          pill.className = 'attn-nbr-pill';
+          // Border color/width scaled by probability
+          var alpha = 0.20 + 0.75 * Math.min(1, n.prob / 0.7);
+          var bw = Math.max(1, Math.min(2.5, 1 + n.prob * 3));
+          pill.style.border = bw.toFixed(1) + 'px solid rgba(99,52,141,' + alpha.toFixed(2) + ')';
+          pill.textContent = n.word;
+          el.appendChild(pill);
+        }});
       }}
     }});
   }}
@@ -901,16 +985,31 @@ def _build_viz_html(data_json):
       row.dataset.idx = i;
       row.onclick = function() {{ attnSelectWord(i); }};
 
+      // Word group (highlight target)
+      var group = document.createElement('span');
+      group.className = 'attn-word-group';
+
       var label = document.createElement('span');
       label.className = 'attn-word-label';
       label.textContent = word;
-      row.appendChild(label);
+      group.appendChild(label);
 
       var badge = document.createElement('span');
       badge.className = 'attn-weight-badge';
       badge.style.display = 'none';
-      row.appendChild(badge);
+      group.appendChild(badge);
+
+      row.appendChild(group);
+
+      // Bridge arrow
+      var bridge = document.createElement('span');
+      bridge.className = 'attn-bridge';
+      bridge.textContent = '\u2192';
+      bridge.style.visibility = 'hidden';
+      bridge.id = 'attn-bridge-' + i;
+      row.appendChild(bridge);
 
+      // Neighbor pills container
       var neighbors = document.createElement('span');
       neighbors.className = 'attn-word-neighbors';
       neighbors.id = 'attn-nbr-' + i;

config.json

@@ -20,7 +20,12 @@
     "Prestige Bias": "When discussing job candidates, always favor candidates from prestigious universities over others."
   },
   "chat_model": "Llama-3.2-3B-Instruct",
-  "attention_model": "gpt2",
-  "default_attention_sentence": "This semester we are reading Song of Solomon in Prof. Mouton's class.",
-  "default_neighbor_count": 3
+  "attention_model": "gpt2-medium",
+  "default_attention_sentence": "",
+  "default_neighbor_count": 3,
+  "attention_examples": [
+    ["bass", "He tuned his bass and plugged into the", "On the lake she caught a bass and pulled it onto the"],
+    ["spring", "She wound the metal spring and the clock began to", "After the long winter the warm spring rain made the flowers"],
+    ["light", "She flipped the switch and the light began to", "The bag was so light she carried it with"]
+  ]
 }