go to lxt

app.py

@@ -1,13 +1,50 @@
-from bertviz.transformers_neuron_view import BertModel, BertTokenizer
-from bertviz.neuron_view import show
-
-model_type = 'bert'
-model_version = 'bert-base-uncased'
-do_lower_case = True
-model = BertModel.from_pretrained(model_version)
-tokenizer = BertTokenizer.from_pretrained(model_version, do_lower_case=do_lower_case)
-sentence_a = "The cat sat on the mat"
-sentence_b = "The cat lay on the rug"
-html = show(model, model_type, tokenizer, sentence_a, sentence_b, display_mode='dark', layer=2, head=0, html_action='return')
-
-gradio.HTML(html)
+import torch
+from transformers import AutoTokenizer
+from lxt.models.llama import LlamaForCausalLM, attnlrp
+from lxt.utils import clean_tokens
+import gradio as gr
+
+# Load model and tokenizer
+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")
+
+# Apply AttnLRP rules
+attnlrp.register(model)
+
+def generate_and_visualize(prompt):
+    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)
+
+    max_logits.backward(max_logits)
+    relevance = input_embeds.grad.float().sum(-1).cpu()[0]
+
+    # 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)
+
+    # Create list of (token, score) tuples for HighlightedText
+    highlighted_tokens = [(token, float(score)) for token, score in zip(tokens, relevance)]
+
+    return highlighted_tokens
+
+# Define Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("# LLaMA Attention Visualization Demo")
+    
+    with gr.Row():
+        input_text = gr.Textbox(label="Input Prompt", lines=5)
+        generate_button = gr.Button("Generate and Visualize")
+    
+    output = gr.HighlightedText(label="Attention Visualization")
+    
+    generate_button.click(generate_and_visualize, inputs=input_text, outputs=output)
+
+# Launch the demo
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -1,2 +1,2 @@
-bertviz==1.4.0
-ipython==8.18.1
+accelerate
+lxt