Add hover mechanic

archive/data.json

@@ -0,0 +1,15 @@
+{
+    "tokens": ["This", "is", "a", "sample", "sentence", "."],
+    "attention": [
+      [
+        [
+          [0.1, 0.2, 0.3, 0.4, 0.5, 0.6],
+          [0.2, 0.1, 0.4, 0.3, 0.6, 0.5],
+          [0.3, 0.4, 0.1, 0.2, 0.5, 0.6],
+          [0.4, 0.3, 0.2, 0.1, 0.6, 0.5],
+          [0.5, 0.6, 0.4, 0.3, 0.1, 0.2],
+          [0.6, 0.5, 0.3, 0.4, 0.2, 0.1]
+        ]
+      ]
+    ]
+  }

frontend/index.html

@@ -82,12 +82,19 @@
             autofocus
           ></textarea>
           <button type="submit">Process</button>
-          <button id="downloadBtn">Download JSON</button>
         </form>
 
         <!-- Output Area -->
         <div id="output"></div>
+      </section>
+      <section class="hover-visualization">
+        <h2>Hover Visualization</h2>
+        <p>
+          Hover over the tokens to see the attention weights for the selected head. Click on a token
+          to see the attention weights for that token across all heads.
+        </p>
         <div id="tokenContainer"></div>
+
       </section>
     </main>
 

frontend/script.js

@@ -16,6 +16,7 @@ document.getElementById('textForm').addEventListener('submit', async (e) => {
     }
     const data = await response.json();
     displayOutput(data);
+    displayHoverTokens(data, 0, 0);
   } catch (error) {
     console.error('Error:', error);
     document.getElementById('output').innerText = 'Error processing text.';
@@ -33,58 +34,14 @@ function displayOutput(data) {
     <pre>${JSON.stringify(data.attention, null, 2)}</pre>
     <div id="tokenContainer"></div>
   `;
-
-  // If we have attention data, render tokens visually
-  if (data.attention && data.attention.length > 0) {
-    // By default, use the first layer’s attention matrix
-    // (Or you can read the selected layer/head from the form)
-    renderTokens(data.tokens, data.attention[0]);
-  }
-}
-
-// Renders tokens in tokenContainer; sets up mouseover for highlighting
-function renderTokens(tokens, attentionMatrix) {
-  const container = document.getElementById('tokenContainer');
-  container.innerHTML = ""; // Clear old tokens
-
-  tokens.forEach((token, index) => {
-    const span = document.createElement("span");
-    span.textContent = token + " ";
-    span.dataset.tokenIndex = index;
-
-    span.addEventListener('mouseover', () => {
-      highlightAttention(index, attentionMatrix);
-    });
-
-    span.addEventListener('mouseout', () => {
-      resetTokenSizes();
-    });
-
-    container.appendChild(span);
-  });
 }
 
-// Highlights tokens based on attention weights for hovered token index
-function highlightAttention(hoveredIndex, attentionMatrix) {
-  const container = document.getElementById("tokenContainer");
-  const weights = attentionMatrix[hoveredIndex];
-
-  // Normalize weights if desired
-  const maxWeight = Math.max(...weights);
+function displayHoverTokens(data, layer_idx, head_idx) {
+  const tokens = data.tokens;
+  const attentionMatrix = data.attention;
 
-  // Some base styling
-  const baseFontSize = 16; // px
-  const maxIncrease = 10;  // px to add at max weight
-
-  // Iterate over spans (tokens) and scale font size
-  Array.from(container.children).forEach((span, idx) => {
-    const weight = weights[idx];
-    const newFontSize = baseFontSize + (weight / maxWeight) * maxIncrease;
-    span.style.fontSize = newFontSize + "px";
-  });
+  renderTokens(tokens, attentionMatrix, layer_idx, head_idx);
 }
-
-// Resets token font sizes to the base
 function resetTokenSizes() {
   const container = document.getElementById("tokenContainer");
   Array.from(container.children).forEach((span) => {
@@ -92,35 +49,31 @@ function resetTokenSizes() {
   });
 }
 
-function downloadJSON(data, filename = 'data.json') {
-  // Convert the data object to a JSON string
-  const jsonStr = JSON.stringify(data, null, 2); // pretty print with 2-space indentation
-
-  // Create a Blob from the JSON string
-  const blob = new Blob([jsonStr], { type: 'application/json' });
-
-  // Create an object URL from the Blob
-  const url = URL.createObjectURL(blob);
-
-  // Create a temporary anchor element
-  const a = document.createElement('a');
-  a.href = url;
-  a.download = filename;
+function highlightAttention(index, attentionData, layer_idx, head_idx){
+  const container = document.getElementById('tokenContainer');
+  // Retrieve the row corresponding to the hovered token.
+  const row = attentionData[layer_idx][head_idx][index];
+  if (!row) {
+    console.warn(`No attention data for token index ${index}`);
+    return;
+  }
 
-  // Append the anchor to the document body
-  document.body.appendChild(a);
+  // Only consider preceding tokens.
+  const weights = row.slice(0, index);
 
-  // Programmatically click the anchor to trigger the download
-  a.click();
+  // Normalize using the maximum weight from the current row's causal part.
+  const maxWeight = Math.max(...attentionData[layer_idx][head_idx]) || 1;
+  const baseFontSize = 16;
+  const maxIncrease = 10;
 
-  // Clean up by removing the anchor and revoking the object URL
-  document.body.removeChild(a);
-  URL.revokeObjectURL(url);
+  Array.from(container.children).forEach((span, idx) => {
+    if (idx < index){
+      const weight = weights[idx];
+      // Calculate the new font size.
+      const newFontSize = baseFontSize + (weight / maxWeight) * maxIncrease;
+      span.style.fontSize = newFontSize + "px";
+    } else {
+      span.style.fontSize = baseFontSize + "px";
+    }
+  });
 }
-
-document.getElementById('downloadBtn').addEventListener('click', () => {
-  const output = document.getElementById('output').innerText;
-  if (output) {
-    downloadJSON(JSON.parse(output), 'output.json');
-  }
-});

frontend/styles.css

@@ -173,6 +173,11 @@ button:hover {
 /* Token highlighting styles */
 #tokenContainer span {
   cursor: default;
-  transition: font-size 0.2s ease, color 0.2s ease;
+  transition: font-size 0.9s ease;
   color: #555;
+  display: inline-block;
+  transition: font-size 0.9s ease;
+  margin-right: 4px;
+  padding: 2px 4px;
 }
+