Correct context windows

app.py

@@ -13,9 +13,10 @@ attnlrp.register(model)
 
 
 def really_clean_tokens(tokens):
+    tokens = clean_tokens(tokens)
     cleaned_tokens = []
     for token in tokens:
-        token = token.replace("_", " ").replace("▁", " ").replace("<s>", "").strip()
+        token = token.replace("_", " ").replace("▁", " ").replace("<s>", " ")
         if token.startswith("<0x") and token.endswith(">"):
             # Convert hex to character
             char_code = int(token[3:-1], 16)
@@ -44,12 +45,11 @@ def generate_and_visualize(prompt, num_tokens=10):
 
         input_ids = torch.cat([input_ids, next_token.unsqueeze(0)], dim=1)
         input_embeds = model.get_input_embeddings()(input_ids)
-
-    input_tokens = tokenizer.convert_ids_to_tokens(input_ids[0])
-    input_tokens = really_clean_tokens(input_tokens)
+    input_tokens = really_clean_tokens(tokenizer.convert_ids_to_tokens(input_ids[0]))
     generated_tokens = really_clean_tokens(tokenizer.convert_ids_to_tokens(generated_tokens_ids))
     
     return input_tokens, all_relevances, generated_tokens
+
 def process_relevances(input_tokens, all_relevances, generated_tokens):
     attention_matrix = np.array([el[:len(all_relevances[0])] for el in all_relevances])
 
@@ -103,11 +103,11 @@ def process_relevances(input_tokens, all_relevances, generated_tokens):
     for i, (token, coords) in enumerate(output_with_notes):
         if coords is not None:
             best_width, best_patch_end = coords
-            start = max(0, best_patch_end - best_width - context_width)
-            end = min(len(input_tokens), best_patch_end + kernel_width + context_width)
-            context = input_tokens[start:end]
-            significant_start = max(0, best_patch_end - best_width - start)
-            significant_end = significant_start + best_width + kernel_width
+            significant_start = max(0, best_patch_end - best_width)
+            significant_end = best_patch_end + kernel_width
+            context_start = max(0, significant_start - context_width)
+            context_end = min(len(input_tokens), significant_end + context_width)
+            context = input_tokens[context_start:context_end]
             output_with_notes[i] = (token, (context, significant_start, significant_end))
 
     return output_with_notes
@@ -123,7 +123,7 @@ def create_html_with_hover(output_with_notes):
                     formatted_context.append(f'<strong>{token}</strong>')
                 else:
                     formatted_context.append(token)
-            formatted_note = " ".join(formatted_context)
+            formatted_note = "".join(formatted_context)
             html += f'<span class="hoverable" data-note-id="note-{i}">{text}<sup>[{i+1}]</sup>'
             html += f'<span class="hover-note">{formatted_note}</span></span>'
         else:
@@ -144,7 +144,6 @@ css = """
 .hover-note {
     display: none;
     position: absolute;
-    background-color: #f0f0f0;
     padding: 5px;
     border-radius: 5px;
     bottom: 100%;
@@ -153,8 +152,9 @@ css = """
     white-space: normal;
     background-color: rgba(240, 240, 240, 1);
     max-width: 600px;
+    width:500px;
     word-wrap: break-word;
-    z-index: 1;
+    z-index: 10;
 }
 .hoverable:hover .hover-note { display: block; }
 """