First draft for attention mechanism

backend/app.py

@@ -38,13 +38,15 @@ except Exception as e:
     raise
 
 @app.post("/process")
-async def process_text(text: str = Body(..., embed=True)):
+async def process_text(payload: TextRequest):
     """
     Process the input text:
       - Tokenizes the text
       - Runs the GPT2 model to obtain attentions
       - Returns the tokens and attention values (rounded to 2 decimals)
     """
+
+    text = payload.txt
     try:
         logger.info(f"Received text: {text}")
         # Tokenize input text (truncating if needed)

frontend/index.html

@@ -13,6 +13,9 @@
   </form>
   <div id="output">
     <!-- Processed output will be displayed here -->
+  </div>
+  <div id="tokenContainer">
+
   </div>
   <script src="/static/script.js"></script>
 </body>

frontend/script.js

@@ -24,12 +24,72 @@ document.getElementById('textForm').addEventListener('submit', async (e) => {
   
   // Function to display the tokens and attention values
   function displayOutput(data) {
+    // Optionally still display the raw JSON if needed
     const outputDiv = document.getElementById('output');
     outputDiv.innerHTML = `
       <h2>Tokens</h2>
       <pre>${JSON.stringify(data.tokens, null, 2)}</pre>
       <h2>Attention</h2>
       <pre>${JSON.stringify(data.attention, null, 2)}</pre>
+      <div id="tokenContainer"></div>
     `;
+
+    // Render tokens using the first layer's attention (or adjust as needed)
+    if (data.attention && data.attention.length > 0) {
+      // For instance, using the first layer. You might refine this to select a specific head.
+      renderTokens(data.tokens, data.attention[0]);
+    }
+  }
+
+function renderTokens(tokens, attentionMatrix){
+    const container = document.getElementById('tokenContainer')
+
+    container.innerHTML = "" // remove previous tokens
+
+    tokens.forEach((token, index) => {
+        const span = document.createElement("span");
+        span.textContent = token + " ";
+        span.style.transition = "font-size 0.2 ease";
+        span.dataset.tokenIndex = index;
+
+        span.addEventListener('mouseover', () =>{
+            highlightAttention(index, token, attentionMatrix);
+        });
+
+        span.addEventListener('mouseout', ()=> {
+            resetTokenSizes();
+        })
+
+        container.appendChild(span);
+    });
+}
+
+function highlightAttention(hoveredIndex, tokens, attentionMatrix) {
+    const container = document.getElementById("tokenContainer");
+    // Get the attention weights for the hovered token
+    const weights = attentionMatrix[hoveredIndex];
+
+    // Normalize weights if needed (optional)
+    const maxWeight = Math.max(...weights);
+
+    // Define a base font size and a maximum increase
+    const baseFontSize = 16; // in pixels
+    const maxIncrease = 10;  // additional pixels for the maximum attention value
+
+    // Iterate over tokens and adjust font sizes
+    Array.from(container.children).forEach((span, idx) => {
+      // Calculate a new font size proportional to the attention weight
+      const weight = weights[idx];
+      // For example, increase font size linearly relative to the maxWeight
+      const newFontSize = baseFontSize + (weight / maxWeight) * maxIncrease;
+      span.style.fontSize = newFontSize + "px";
+    });
+  }
+
+  function resetTokenSizes() {
+    const container = document.getElementById("tokenContainer");
+    Array.from(container.children).forEach((span) => {
+      span.style.fontSize = "16px"; // Reset to base size
+    });
   }
-  
+