Fix short generations

app.py

@@ -6,9 +6,16 @@ import gradio as gr
 import numpy as np
 import spaces
 from scipy.signal import convolve2d
+from huggingface_hub import login
+import os
+from dotenv import load_dotenv
 
-model = LlamaForCausalLM.from_pretrained("HuggingFaceTB/SmolLM-1.7B-Instruct", torch_dtype=torch.bfloat16, device_map="cuda")
-tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/SmolLM-1.7B-Instruct")
+load_dotenv()
+
+login(os.get("HF_TOKEN"))
+
+model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-3.2-3B-Instruct", torch_dtype=torch.bfloat16, device_map="cuda")
+tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-3B-Instruct")
 attnlrp.register(model)
 
 def really_clean_tokens(tokens):
@@ -54,6 +61,7 @@ def generate_and_visualize(prompt, num_tokens=10):
     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
@@ -61,6 +69,9 @@ def process_relevances(input_tokens, all_relevances, generated_tokens):
     kernel_width = 6
     context_width = 20  # Number of tokens to include as context on each side
     kernel = np.ones((kernel_width, kernel_width))
+
+    if len(generated_tokens) < kernel_width:
+        return [(token, None, None) for token in generated_tokens]
     
     # Compute the rolling sum using 2D convolution
     rolled_sum = convolve2d(attention_matrix, kernel, mode='valid')