Add nice slider and visualization

app.py

@@ -3,6 +3,11 @@ from transformers import AutoTokenizer
 from lxt.models.llama import LlamaForCausalLM, attnlrp
 from lxt.utils import clean_tokens
 import gradio as gr
+import matplotlib.pyplot as plt
+import seaborn as sns
+import numpy as np
+import io
+from PIL import Image
 
 # Load model and tokenizer
 model = LlamaForCausalLM.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0", torch_dtype=torch.bfloat16, device_map="cuda")
@@ -11,44 +16,138 @@ tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
 # Apply AttnLRP rules
 attnlrp.register(model)
 
-def generate_and_visualize(prompt):
+def really_clean_tokens(tokens):
+    tokens = clean_tokens(tokens)
+    tokens = [token.replace("_", " ").replace("▁", " ").replace("<s>", "") for token in tokens]
+    return tokens
+
+def generate_and_visualize(prompt, num_tokens=10):
     input_ids = tokenizer(prompt, return_tensors="pt", add_special_tokens=True).input_ids.to(model.device)
     input_embeds = model.get_input_embeddings()(input_ids)
 
-    output_logits = model(inputs_embeds=input_embeds.requires_grad_(), use_cache=False).logits
-    max_logits, max_indices = torch.max(output_logits[0, -1, :], dim=-1)
+    generated_tokens_ids = []
+    all_relevances = []
+
+    for _ in range(num_tokens):
+        output_logits = model(inputs_embeds=input_embeds.requires_grad_(), use_cache=False).logits
+        max_logits, max_indices = torch.max(output_logits[0, -1, :], dim=-1)
+
+        max_logits.backward(max_logits)
+        relevance = input_embeds.grad.float().sum(-1).cpu()[0]
+        all_relevances.append(relevance)
+
+        # Generate next token
+        next_token = max_indices.unsqueeze(0)
+        generated_tokens_ids.append(next_token.item())
 
-    max_logits.backward(max_logits)
-    relevance = input_embeds.grad.float().sum(-1).cpu()[0]
+        # Prepare for next iteration
+        input_ids = torch.cat([input_ids, next_token.unsqueeze(0)], dim=1)
+        input_embeds = model.get_input_embeddings()(input_ids)
 
+    # Process tokens and relevances
+    input_tokens = tokenizer.convert_ids_to_tokens(input_ids[0])
+
+    input_tokens = really_clean_tokens(input_tokens)
+    generated_tokens = really_clean_tokens(tokenizer.convert_ids_to_tokens(generated_tokens_ids))
+    generated_text = tokenizer.decode(generated_tokens_ids)
+    print(f"Generated text: {generated_text}")
+
+    return input_tokens, all_relevances, generated_text, len(generated_tokens_ids), generated_tokens
+
+def update_visualization_step(input_tokens, all_relevances, output_tokens, step):
+    relevance = all_relevances[step]
+    
     # Normalize relevance between [0, 1] for highlighting
     relevance = (relevance - relevance.min()) / (relevance.max() - relevance.min())
 
-    # Remove '_' characters from token strings
-    tokens = tokenizer.convert_ids_to_tokens(input_ids[0])
-    tokens = clean_tokens(tokens)
-    tokens = [token.replace("_", " ").replace("▁", " ").replace("<s>", "") for token in tokens]
-    print("Decoded: ", tokenizer.decode(input_ids[0]))
-
     # Create list of (token, score) tuples for HighlightedText
-    highlighted_tokens = [(token, max(float(score), 0.05)-0.05) for token, score in zip(tokens, relevance)]
+    highlighted_tokens = [(token, max(float(score), 0.05)-0.05) for token, score in zip(input_tokens, relevance)]
+
+    return highlighted_tokens, output_tokens[:step+1]
+
+def generate_heatmap(input_tokens, all_relevances, output_tokens):
+    # Create a matrix of attention scores
+    attention_matrix = np.array([el[:len(all_relevances[0])] for el in all_relevances])
+    
+    # Find columns with non-zero values (including adjacent columns)
+    non_zero_cols = np.where(np.abs(attention_matrix).sum(axis=0) > 1.)[0]
+    for col in range(5):
+        non_zero_cols = np.union1d(non_zero_cols, non_zero_cols + col)
+        non_zero_cols = np.union1d(non_zero_cols, non_zero_cols - col)
+    non_zero_cols = np.sort(non_zero_cols)
+    non_zero_cols = non_zero_cols[(non_zero_cols >= 0) & (non_zero_cols < attention_matrix.shape[1])]
+
+
+    # Filter the matrix and input tokens
+    filtered_matrix = attention_matrix[:, non_zero_cols]
+    filtered_input_tokens = [input_tokens[i] for i in non_zero_cols]
+    
+    # Create the heatmap
+    plt.figure(figsize=(20, 7))
+    sns.heatmap(filtered_matrix, xticklabels=filtered_input_tokens, yticklabels=output_tokens, cmap="YlOrRd")
+    plt.title("Attention Heatmap (shows only input tokens where there was strong enough attention)")
+    plt.xlabel("Input Tokens")
+    plt.ylabel("Output Tokens")
+    plt.xticks(rotation=90)
+    plt.yticks(rotation=0)
+    
+    # Save the plot to a bytes buffer
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png', bbox_inches='tight')
+    buf.seek(0)
+    plt.close()
+    
+    # Return the image as a numpy array
+    return np.array(Image.open(buf))
 
-    return highlighted_tokens
+def on_generate(input_text, num_tokens):
+    input_tokens, relevances, generated_text, generated_length, output_tokens = generate_and_visualize(input_text, num_tokens)
+    heatmap = generate_heatmap(input_tokens, relevances, output_tokens)
+    return (
+        input_tokens, 
+        relevances, 
+        update_visualization_step(input_tokens, relevances, output_tokens, 0)[0], 
+        output_tokens, 
+        output_tokens[:1], 
+        generated_text, 
+        gr.Slider(maximum=generated_length-1, value=0),
+        heatmap
+    )
 
 # Define Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("""# Attention Visualization Demo
+    gr.Markdown("""# Extended Attention Visualization Demo
 
 This demo uses the library [LXT](https://lxt.readthedocs.io/en/latest/quickstart.html#tinyllama) under the hood.""")
     
     input_text = gr.Textbox(label="Input Prompt", lines=5, value="""\
 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. \
 Question: How high did they climb in 1922? According to the text, the 1922 expedition reached 8,""")
+    num_tokens = gr.Slider(minimum=1, maximum=50, step=1, value=10, label="Number of tokens to generate")
     generate_button = gr.Button("Generate and Visualize")
     
-    output = gr.HighlightedText(label="Attention Visualization", adjacent_separator="", combine_adjacent=True)
+    generated_output = gr.Textbox(label="Generated Text")
+    heatmap_output = gr.Image(label="Attention Heatmap")
+
+    step_slider = gr.Slider(minimum=0, maximum=1, step=1, value=0, label="Visualization Step")
+    attention_on_inputs = gr.HighlightedText(label="Attention Visualization", adjacent_separator="", combine_adjacent=True)
+    current_tokens = gr.Textbox(label="Current Token")
+
+    input_tokens_state = gr.State([])
+    output_tokens_state = gr.State([])
+    relevances_state = gr.State([])
     
-    generate_button.click(generate_and_visualize, inputs=input_text, outputs=output)
+    generate_button.click(
+        on_generate,
+        inputs=[input_text, num_tokens],
+        outputs=[input_tokens_state, relevances_state, attention_on_inputs, output_tokens_state, current_tokens, generated_output, step_slider, heatmap_output]
+    )
+
+    step_slider.change(
+        update_visualization_step,
+        inputs=[input_tokens_state, relevances_state, output_tokens_state, step_slider],
+        outputs=[attention_on_inputs, current_tokens]
+    )
 
 # Launch the demo
 if __name__ == "__main__":

requirements.txt

@@ -1,2 +1,3 @@
 accelerate
-lxt
+lxt
+seaborn