Working nice version

app.py

@@ -7,15 +7,15 @@ import numpy as np
 import spaces
 from scipy.signal import convolve2d
 
-model = LlamaForCausalLM.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0", torch_dtype=torch.bfloat16, device_map="cuda")
-tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
+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")
 attnlrp.register(model)
 
 def really_clean_tokens(tokens):
     tokens = clean_tokens(tokens)
     cleaned_tokens = []
     for token in tokens:
-        token = token.replace("_", " ").replace("▁", " ").replace("<s>", " ")
+        token = token.replace("_", " ").replace("▁", " ").replace("<s>", " ").replace("Ċ", " ").replace("Ġ", " ")
         if token.startswith("<0x") and token.endswith(">"):
             # Convert hex to character
             char_code = int(token[3:-1], 16)
@@ -53,7 +53,7 @@ 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.22
+    threshold_per_token = 0.2
     kernel_width = 6
     context_width = 20  # Number of tokens to include as context on each side
     kernel = np.ones((kernel_width, kernel_width))
@@ -136,7 +136,8 @@ def create_html_with_hover(output_with_notes):
             html += f'{token}'
             i += 1
         else:
-            text = "".join([element[0] for element in output_with_notes[i:i+width]])
+            text = "".join(really_clean_tokens([element[0] for element in output_with_notes[i:i+width]]))
+            print(text)
             first_part, significant_part, final_part = notes
             formatted_note = f'{first_part}<strong>{significant_part}</strong>{final_part}'
             html += f'<span class="hoverable" data-note-id="note-{note_number}">{text}<sup>[{note_number+1}]</sup>'
@@ -200,7 +201,7 @@ with gr.Blocks(css=css) as demo:
     gr.Markdown("# Token Generation with Hover Notes")
     
     input_text = gr.Textbox(label="Enter your prompt:", lines=10, value=examples[0])
-    num_tokens = gr.Slider(minimum=1, maximum=50, value=20, step=1, label="Number of tokens to generate")
+    num_tokens = gr.Slider(minimum=1, maximum=100, value=20, step=1, label="Number of tokens to generate")
     generate_button = gr.Button("Generate")
 
     output_html = gr.HTML(label="Generated Output")