Add rendering of tokens

frontend/script.js

@@ -35,29 +35,46 @@ function displayOutput(data) {
     <div id="tokenContainer"></div>
   `;
 }
+function renderTokens(tokens, attentionData, layer_idx, head_idx) {
+    const container = document.getElementById('tokenContainer');
+    container.innerHTML = "";
 
-function displayHoverTokens(data, layer_idx, head_idx) {
-  const tokens = data.tokens;
-  const attentionMatrix = data.attention;
+    tokens.forEach((token, index) => {
+      const span = document.createElement('span');
+      span.textContent = token + " ";
 
-  renderTokens(tokens, attentionMatrix, layer_idx, head_idx);
-}
-function resetTokenSizes() {
-  const container = document.getElementById("tokenContainer");
-  Array.from(container.children).forEach((span) => {
-    span.style.fontSize = "16px";
-  });
-}
+      span.addEventListener('mouseenter', () => {
+        highlightAttention(index, attentionData, layer_idx, head_idx);
+      });
+
+      span.addEventListener('mouseleave', () => {
+        resetTokenSizes();
+      });
+      container.appendChild(span);
+    });
+  }
+  function displayHoverTokens(data, layer_idx, head_idx) {
+    const tokens = data.tokens;
+    const attentionMatrix = data.attention;
 
-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;
+    renderTokens(tokens, attentionMatrix, layer_idx, head_idx);
+  }
+  function resetTokenSizes() {
+    const container = document.getElementById("tokenContainer");
+    Array.from(container.children).forEach((span) => {
+      span.style.fontSize = "16px";
+    });
   }
 
+  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;
+    }
+
   // Only consider preceding tokens.
   const weights = row.slice(0, index);