Works well!

app.py

@@ -50,30 +50,29 @@ def generate_and_visualize(prompt, num_tokens=10):
     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])
 
     ### FIND ZONES OF INTEREST
     threshold_per_token = 0.3
     kernel_width = 6
+    context_width = 20  # Number of tokens to include as context on each side
     kernel = np.ones((kernel_width, kernel_width))
     
     # Compute the rolling sum using 2D convolution
     rolled_sum = convolve2d(attention_matrix, kernel, mode='valid')
     
-    # Find where the rolled sum is greater than 1.0
+    # Find where the rolled sum is greater than the threshold
     significant_areas = rolled_sum > kernel_width**2 * threshold_per_token
-    print("SUMM:", significant_areas.sum())
 
     def find_largest_contiguous_patch(array):
-        current_patch_end=None
+        current_patch_end = None
         best_width, best_patch_end = None, None
-        current_width=0
+        current_width = 0
         for i in range(len(array)):
             if array[i]:
                 if current_patch_end is not None and current_patch_end == i-1:
-                    current_width +=1
+                    current_width += 1
                     current_patch_end = i
                 else:
                     current_patch_end = i
@@ -86,13 +85,13 @@ def process_relevances(input_tokens, all_relevances, generated_tokens):
         return best_width, best_patch_end
 
     output_with_notes = [(el, None) for el in generated_tokens[:kernel_width]]
-    print(len(generated_tokens), kernel_width, significant_areas.shape)
     for row in range(kernel_width, len(generated_tokens)):
-        print(row-kernel_width+1)
         best_width, best_patch_end = find_largest_contiguous_patch(significant_areas[row-kernel_width+1])
 
         if best_width is not None:
-            for i, (token, coords) in enumerate(output_with_notes[-2*kernel_width:]):
+            # Fuse the notes for consecutive output tokens if necessary
+            for i in range(len(output_with_notes)-2*kernel_width, len(output_with_notes)):
+                token, coords = output_with_notes[i]
                 if coords is not None:
                     prev_width, prev_patch_end = coords
                     if prev_patch_end > best_patch_end - best_width: # then notes are overlapping, thus we delete the first one.
@@ -104,20 +103,31 @@ 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
-            output_with_notes[i] = (token, input_tokens[best_patch_end-best_width: best_patch_end+kernel_width-1])
+            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
+            output_with_notes[i] = (token, (context, significant_start, significant_end))
 
-    # Fuse the notes for consecutive output tokens if necessary
-    print(output_with_notes)
     return output_with_notes
 
 def create_html_with_hover(output_with_notes):
     html = "<div id='output-container'>"
     for i, (text, notes) in enumerate(output_with_notes):
         if notes:
+            context, start, end = notes
+            formatted_context = []
+            for j, token in enumerate(context):
+                if start <= j < end:
+                    formatted_context.append(f'<strong>{token}</strong>')
+                else:
+                    formatted_context.append(token)
+            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">{", ".join(notes)}</span></span> '
+            html += f'<span class="hover-note">{formatted_note}</span></span>'
         else:
-            html += f'{text} '
+            html += f'{text}'
     html += "</div>"
     return html
 
@@ -140,12 +150,14 @@ css = """
     bottom: 100%;
     left: 50%;
     transform: translateX(-50%);
-    white-space: nowrap;
+    white-space: normal;
+    background-color: rgba(240, 240, 240, 1);
+    max-width: 600px;
+    word-wrap: break-word;
     z-index: 1;
 }
 .hoverable:hover .hover-note { display: block; }
 """
-
 examples = [
     [
         """Context: Mount Everest attracts many climbers, including highly experienced mountaineers. There are two main climbing routes, one approaching the summit from the southeast in Nepal (known as the standard route) and the other from the north in Tibet. While not posing substantial technical climbing challenges on the standard route, Everest presents dangers such as altitude sickness, weather, and wind, as well as hazards from avalanches and the Khumbu Icefall. As of November 2022, 310 people have died on Everest. Over 200 bodies remain on the mountain and have not been removed due to the dangerous conditions. The first recorded efforts to reach Everest's summit were made by British mountaineers. As Nepal did not allow foreigners to enter the country at the time, the British made several attempts on the north ridge route from the Tibetan side. After the first reconnaissance expedition by the British in 1921 reached 7,000 m (22,970 ft) on the North Col, the 1922 expedition pushed the north ridge route up to 8,320 m (27,300 ft), marking the first time a human had climbed above 8,000 m (26,247 ft). The 1924 expedition resulted in one of the greatest mysteries on Everest to this day: George Mallory and Andrew Irvine made a final summit attempt on 8 June but never returned, sparking debate as to whether they were the first to reach the top. Tenzing Norgay and Edmund Hillary made the first documented ascent of Everest in 1953, using the southeast ridge route. Norgay had reached 8,595 m (28,199 ft) the previous year as a member of the 1952 Swiss expedition. The Chinese mountaineering team of Wang Fuzhou, Gonpo, and Qu Yinhua made the first reported ascent of the peak from the north ridge on 25 May 1960.