update attn score display

app.py

@@ -42,7 +42,7 @@ def get_top_attns(avg_attn_list):
         top_3 = heapq.nlargest(3, enumerate(avg_attn_list[i]), key=lambda x: x[1])
 
         # get the indices and values of the source tokens
-        top_values = [val for idx, val in top_3]
+        top_values = [round(val.item(), 2) for idx, val in top_3]
         top_index = [idx for idx, val in top_3]
 
         avg_attn_top.append({
@@ -86,7 +86,14 @@ def render_attention_html(src_tokens, tgt_tokens):
     for i, token in enumerate(tgt_tokens):
         tgt_html += f'<span class="token tgt-token" data-index="{i}">{token}</span> '
 
-    html = f'<div class="tgt-token-wrapper-text">Output Tokens</div><div class="tgt-token-wrapper">{tgt_html}</div><hr class="token-wrapper-seperator"><div class="src-token-wrapper-text">Input Tokens</div><div class="src-token-wrapper">{src_html}</div>'
+    html = f"""
+        <div class="tgt-token-wrapper-text">Output Tokens</div>
+        <div class="tgt-token-wrapper">{tgt_html}</div>
+        <hr class="token-wrapper-seperator">
+        <div class="src-token-wrapper-text">Input Tokens</div>
+        <div class="src-token-wrapper">{src_html}</div>
+        <div class="scores"><span class="score-1 score"></span><span class="score-2 score"></span><span class="score-3 score"></span><div>
+        """
     return html
 
 
@@ -102,6 +109,10 @@ css = """
 .tgt-token-wrapper-text {position: absolute; top: .75rem; color: #71717a;}
 .token-wrapper-seperator {margin-top: 1rem; margin-bottom: 1rem}
 .note-text {margin-bottom: 3.5rem;}
+.scores { position: absolute; bottom: 0.75rem; color: rgb(113, 113, 122); right: 1rem;}
+.score-1 { display: none; background-color: #FF8865; padding: .5rem; border-radius: var(--block-radius); margin-right: .75rem;}
+.score-2 { display: none; background-color: #FFD2C4; padding: .5rem; border-radius: var(--block-radius); margin-right: .75rem;}
+.score-3 { display: none; background-color: #FFF3F0; padding: .5rem; border-radius: var(--block-radius); margin-right: .75rem;}
 """
 
 js = """
@@ -112,6 +123,8 @@ function showCrossAttFun(attn_scores) {
 
     const targetTokens = document.querySelectorAll('.tgt-token');
     
+    const scores = document.querySelectorAll('.score');
+    
     function onTgtHover(event, idx) {
         event.style.backgroundColor = "#C6E6E6";
         
@@ -119,18 +132,24 @@ function showCrossAttFun(attn_scores) {
         if (srcIdx0 < srcLen) {
             srcEl0 = scrTokens[srcIdx0]
             srcEl0.style.backgroundColor = "#FF8865"
+            scores[0].textContent = attn_scores[idx]['top_values'][0]
+            scores[0].style.display = "initial"; 
         }
         
         srcIdx1 = attn_scores[idx]['top_index'][1]
         if (srcIdx1 < srcLen) {
             srcEl1 = scrTokens[srcIdx1]
             srcEl1.style.backgroundColor = "#FFD2C4"
+            scores[1].textContent = attn_scores[idx]['top_values'][1]
+            scores[1].style.display = "initial"; 
         }
         
         srcIdx2 = attn_scores[idx]['top_index'][2]
         if (srcIdx2 < srcLen) {
             srcEl2 = scrTokens[srcIdx2]
             srcEl2.style.backgroundColor = "#FFF3F0"
+            scores[2].textContent = attn_scores[idx]['top_values'][2]
+            scores[2].style.display = "initial"; 
         }
     }
     
@@ -141,10 +160,16 @@ function showCrossAttFun(attn_scores) {
         srcIdx2 = attn_scores[idx]['top_index'][2]
         srcEl0 = scrTokens[srcIdx0]
         srcEl0.style.backgroundColor = ""
+        scores[0].textContent = ""
+        scores[0].style.display = "none"; 
         srcEl1 = scrTokens[srcIdx1]
         srcEl1.style.backgroundColor = ""
+        scores[1].textContent = ""
+        scores[1].style.display = "none"; 
         srcEl2 = scrTokens[srcIdx2]
         srcEl2.style.backgroundColor = ""
+        scores[2].textContent = ""
+        scores[2].style.display = "none"; 
     }