D3 libary for better vis 2.4

frontend/index.html

@@ -28,7 +28,7 @@
           Etiam vel dui consequat ante lobortis facilisis. Curabitur facilisis enim sed massa
           consectetur, et mollis magna suscipit. Vestibulum odio ante, laoreet at mauris eu,
           feugiat pretium tortor. Ut condimentum laoreet felis quis rutrum. Proin vel elit a
-          augue ornare venenatis id nec neque.Ferdi 2.3
+          augue ornare venenatis id nec neque.Ferdi 2.4
         </p>
       </section>
 
@@ -116,6 +116,6 @@
     </aside>
   </div>
 
-  <script src="static/script.js?v=20281533"></script>
+  <script src="static/script.js?v=202815533"></script>
 </body>
 </html>

frontend/script.js

@@ -48,93 +48,86 @@ function renderModelView(tokens, attention) {
 }
 
 function createAttentionThumbnail(tokens, attention, layerIdx, headIdx) {
-  const padding = 20; // Define padding before using it
-  const thumbnail = document.createElement("div");
-  thumbnail.style.position = "relative";
-  // Dynamic height: top and bottom padding plus 20px per token.
-  thumbnail.style.height = (padding * 2 + tokens.length * 20) + "px";
-  thumbnail.style.width = "120px";
-  thumbnail.style.border = "1px solid #ddd";
-  thumbnail.style.borderRadius = "4px";
-  thumbnail.style.padding = "5px";
-  thumbnail.style.background = "#fff";
-  
-  // Create SVG container with 100% width/height of the thumbnail
-  const svg = document.createElementNS("http://www.w3.org/2000/svg", "svg");
-  svg.setAttribute("width", "100%");
-  svg.setAttribute("height", "100%");
-  
-  // Add header text for layer and head using template literals
-  const header = document.createElementNS("http://www.w3.org/2000/svg", "text");
-  header.setAttribute("x", "50%");
-  header.setAttribute("y", "15");
-  header.setAttribute("text-anchor", "middle");
-  header.setAttribute("font-size", "10");
-  header.textContent = `L${layerIdx + 1} H${headIdx + 1}`;
-  svg.appendChild(header);
-  
-  const attentionWeights = attention[layerIdx][headIdx];
+  const padding = 20;
+  const tokenHeight = 20;
+  const width = 120;
+  const height = padding * 2 + tokens.length * tokenHeight;
   const maxLineWidth = 4;
   const maxOpacity = 0.8;
   
-  // Normalize weights per head: Compute maximum weight in this head.
+  // Extract attention weights for this head and compute the normalization factor.
+  const attentionWeights = attention[layerIdx][headIdx];
   const headMax = Math.max(...attentionWeights.flat()) || 1;
   
+  // Create a thumbnail container using D3
+  const thumbnail = d3.select(document.createElement("div"))
+    .style("position", "relative")
+    .style("height", height + "px")
+    .style("width", width + "px")
+    .style("border", "1px solid #ddd")
+    .style("border-radius", "4px")
+    .style("padding", "5px")
+    .style("background", "#fff");
+
+  // Append an SVG container inside the thumbnail
+  const svg = thumbnail.append("svg")
+    .attr("width", "100%")
+    .attr("height", "100%");
+
+  // Add header text using template literals
+  svg.append("text")
+    .attr("x", "50%")
+    .attr("y", 15)
+    .attr("text-anchor", "middle")
+    .attr("font-size", "10")
+    .text(`L${layerIdx + 1} H${headIdx + 1}`);
+
   // Draw tokens on both sides (vertically aligned)
   tokens.forEach((token, idx) => {
     const cleanToken = token.replace("Ġ", "");
     
-    // Left side: fixed x position to 'padding'
-    const leftText = document.createElementNS("http://www.w3.org/2000/svg", "text");
-    leftText.setAttribute("x", padding);
-    leftText.setAttribute("y", padding + (idx * 20));
-    leftText.setAttribute("font-size", "8");
-    leftText.textContent = cleanToken;
-    svg.appendChild(leftText);
+    // Left side token
+    svg.append("text")
+      .attr("x", padding)
+      .attr("y", padding + idx * tokenHeight)
+      .attr("font-size", "8")
+      .text(cleanToken);
     
-    // Right side: x is set to 90% of the SVG width
-    const rightText = document.createElementNS("http://www.w3.org/2000/svg", "text");
-    rightText.setAttribute("x", "90%");
-    rightText.setAttribute("y", padding + (idx * 20));
-    rightText.setAttribute("text-anchor", "end");
-    rightText.setAttribute("font-size", "8");
-    rightText.textContent = cleanToken;
-    svg.appendChild(rightText);
+    // Right side token
+    svg.append("text")
+      .attr("x", "90%")
+      .attr("y", padding + idx * tokenHeight)
+      .attr("text-anchor", "end")
+      .attr("font-size", "8")
+      .text(cleanToken);
   });
   
-  // Draw attention lines between tokens with normalized weights
+  // Draw attention lines with normalized weights
   attentionWeights.forEach((sourceWeights, sourceIdx) => {
     sourceWeights.forEach((weight, targetIdx) => {
       if (weight > 0.05 && sourceIdx !== targetIdx) {
-        // Normalize the weight relative to the maximum weight in this head
         const normalizedWeight = weight / headMax;
-        
-        const line = document.createElementNS("http://www.w3.org/2000/svg", "line");
-        // Align lines with tokens using the same padding and vertical spacing
-        line.setAttribute("x1", padding);
-        line.setAttribute("y1", padding + (sourceIdx * 20) - 5);
-        line.setAttribute("x2", "90%");
-        line.setAttribute("y2", padding + (targetIdx * 20) - 5);
-        line.setAttribute("stroke", "#2ecc71");
-        line.setAttribute("stroke-width", Math.max(0.5, normalizedWeight * maxLineWidth));
-        line.setAttribute("opacity", Math.min(maxOpacity, normalizedWeight * 2));
-        line.setAttribute("stroke-linecap", "round");
-        svg.appendChild(line);
+        svg.append("line")
+          .attr("x1", padding)
+          .attr("y1", padding + sourceIdx * tokenHeight - 5)
+          .attr("x2", "90%")
+          .attr("y2", padding + targetIdx * tokenHeight - 5)
+          .attr("stroke", "#2ecc71")
+          .attr("stroke-width", Math.max(0.5, normalizedWeight * maxLineWidth))
+          .attr("opacity", Math.min(maxOpacity, normalizedWeight * 2))
+          .attr("stroke-linecap", "round");
       }
     });
   });
   
-  thumbnail.appendChild(svg);
-  
-  // On click, update UI elements (for example, display the selected layer and head)
-  thumbnail.addEventListener("click", () => {
-    document.getElementById("head").textContent = `Head: ${headIdx + 1}`;
-    document.getElementById("layer").textContent = `Layer: ${layerIdx + 1}`;
-    // Dynamically update displayHoverTokens with the selected layer and head
-    displayHoverTokens(data, layerIdx, headIdx);  
+  // On click, update UI elements (assumes 'data' is globally available)
+  thumbnail.on("click", function() {
+    d3.select("#head").text(`Head: ${headIdx + 1}`);
+    d3.select("#layer").text(`Layer: ${layerIdx + 1}`);
+    displayHoverTokens(data, layerIdx, headIdx);
   });
   
-  return thumbnail;
+  return thumbnail.node();
 }
 // Function to display the tokens and attention values
 function displayOutput(data) {