thubnail filled and layer head box next to each other

frontend/script.js

@@ -56,65 +56,91 @@ function createLargeThumbnail(data) {
   const tokens = data.tokens;
   const attention = data.attention;
 
-  // Dimensions for the thumbnail
-  const width = 300;
-  const height = 300;
-  const textHeight = 30;      // Space for token text (top and bottom)
-  const graphHeight = height - 2 * textHeight; // Remaining space for the graph
+  // Dimensions
+  const graphWidth = 300; 
+  const tokenAreaWidth = 60; 
+  const totalWidth = tokenAreaWidth * 2 + graphWidth;
+  const graphHeight = 300;
 
-  // Create the main container and append it to #thumbnailContainer
+  // Remove any existing large thumbnail so only the latest shows
+  d3.select("#largeThumbnail").remove();
+
+  // Create main container as a flex row
   const container = d3.select("#thumbnailContainer")
     .append("div")
     .attr("id", "largeThumbnail")
-    .style("width", width + "px")
-    .style("height", height + "px")
+    .style("width", totalWidth + "px")
+    .style("height", graphHeight + "px")
     .style("border", "2px solid #ddd")
     .style("border-radius", "6px")
     .style("background", "#fff")
-    .style("padding", "10px")
+    .style("display", "flex")
+    .style("align-items", "center")
+    .style("justify-content", "space-between");
+
+  // Left token container
+  const leftTokens = container.append("div")
+    .attr("class", "tokens-left")
+    .style("width", tokenAreaWidth + "px")
+    .style("height", graphHeight + "px")
     .style("display", "flex")
     .style("flex-direction", "column")
-    .style("align-items", "center");
-
-  // Add token text above the graph
-  container.append("div")
-    .classed("thumbnail-token-text", true)
-    .style("width", "100%")
-    .style("text-align", "center")
-    .style("font-size", "16px")
-    .style("margin-bottom", "5px")
-    .text(tokens.join(" "));
-
-  // Create the SVG container for the graph
+    .style("justify-content", "space-around")
+    .style("align-items", "flex-end");
+
+  tokens.forEach(token => {
+    leftTokens.append("div")
+      .style("font-size", "14px")
+      .style("text-align", "right")
+      .text(token);
+  });
+
+  // Center SVG for the graph
   const svg = container.append("svg")
-    .attr("width", width)
+    .attr("width", graphWidth)
     .attr("height", graphHeight)
-    .attr("viewBox", `0 0 ${width} ${graphHeight}`)
+    .attr("viewBox", `0 0 ${graphWidth} ${graphHeight}`)
     .attr("preserveAspectRatio", "none");
 
-  // Add header text inside the SVG
+  // Right token container
+  const rightTokens = container.append("div")
+    .attr("class", "tokens-right")
+    .style("width", tokenAreaWidth + "px")
+    .style("height", graphHeight + "px")
+    .style("display", "flex")
+    .style("flex-direction", "column")
+    .style("justify-content", "space-around")
+    .style("align-items", "flex-start");
+
+  tokens.forEach(token => {
+    rightTokens.append("div")
+      .style("font-size", "14px")
+      .style("text-align", "left")
+      .text(token);
+  });
+
+  // Optional: Add header text inside the SVG
   svg.append("text")
-    .attr("x", width / 2)
-    .attr("y", 20) // adjust as needed
+    .attr("x", graphWidth / 2)
+    .attr("y", 20)
     .attr("text-anchor", "middle")
     .attr("font-size", "14px")
     .text("Layer 1 - Head 1");
 
-  // Calculate spacing for tokens inside the SVG
+  // Compute vertical spacing (each token aligns with a node)
   const tokenCount = tokens.length;
   const tokenSpacing = graphHeight / (tokenCount + 1);
 
-  // Draw attention lines for the first layer, first head
-  // (Assuming attention[0][0] is a matrix: tokens x tokens)
+  // Draw attention lines for the first layer, first head (attention[0][0])
   attention[0][0].forEach((sourceWeights, sourceIdx) => {
     const rowMax = Math.max(...sourceWeights) || 1;
     sourceWeights.forEach((weight, targetIdx) => {
       if (weight > 0.01 && sourceIdx !== targetIdx) {
         const normalizedWeight = weight / rowMax;
         svg.append("line")
-          .attr("x1", 20)                // left margin inside the SVG
+          .attr("x1", 20)                // left margin within the SVG
           .attr("y1", tokenSpacing * (sourceIdx + 1))
-          .attr("x2", width - 20)          // right margin inside the SVG
+          .attr("x2", graphWidth - 20)     // right margin within the SVG
           .attr("y2", tokenSpacing * (targetIdx + 1))
           .attr("stroke", "#800000")
           .attr("stroke-width", Math.max(0.5, normalizedWeight * 4))
@@ -124,15 +150,6 @@ function createLargeThumbnail(data) {
     });
   });
 
-  // Add token text below the graph
-  container.append("div")
-    .classed("thumbnail-token-text", true)
-    .style("width", "100%")
-    .style("text-align", "center")
-    .style("font-size", "16px")
-    .style("margin-top", "5px")
-    .text(tokens.join(" "));
-
   return container.node();
 }
 
@@ -201,17 +218,14 @@ function createAttentionThumbnail(data, layerIdx, headIdx) {
   thumbnail.on("click", function() {
     console.log(`Clicked: Layer ${layerIdx + 1}, Head ${headIdx + 1}`);
 
-    // Remove background from all thumbnails
     d3.selectAll(".thumbnail").style("background", "#fff");
-
-    // Set background for clicked thumbnail
     d3.select(this).style("background", "#ddd");
 
-    // Update displayed head/layer numbers
-    d3.select("#display_head .number-box").text(headIdx + 1);
-    d3.select("#display_layer .number-box").text(layerIdx + 1);
+    // Update displayed head and layer numbers
+    d3.select("#display_head").text(headIdx + 1);
+    d3.select("#display_layer").text(layerIdx + 1);
 
-    // Display updated hover tokens
+    // Update hover tokens
     displayHoverTokens(data, layerIdx, headIdx);
   });