Update app.py

app.py

@@ -1,45 +1,69 @@
 # app.py
-
 import gradio as gr
-from transformers import AutoTokenizer, AutoModel, AutoModelForCausalLM
-import torch
-import matplotlib.pyplot as plt
-import numpy as np
-
-# Load some default model
-MODEL_NAME = "bert-base-uncased"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModel.from_pretrained(MODEL_NAME, output_attentions=True)
-
-def visualize_attention(text):
-    inputs = tokenizer(text, return_tensors="pt")
-    outputs = model(**inputs)
-
-    # Grab attentions from output
-    attentions = outputs.attentions  # List of (num_layers, batch, num_heads, seq_len, seq_len)
-    tokens = tokenizer.convert_ids_to_tokens(inputs['input_ids'][0])
-
-    fig, ax = plt.subplots(figsize=(8, 6))
-    # Just visualize attention from last layer, first head
-    attn_matrix = attentions[-1][0][0].detach().numpy()
-
-    cax = ax.matshow(attn_matrix, cmap='viridis')
-    fig.colorbar(cax)
-
-    ax.set_xticks(range(len(tokens)))
-    ax.set_yticks(range(len(tokens)))
-    ax.set_xticklabels(tokens, rotation=90)
-    ax.set_yticklabels(tokens)
-    ax.set_title("Attention Map - Last Layer, Head 1")
-
-    return fig
-
-iface = gr.Interface(
-    fn=visualize_attention,
-    inputs=gr.Textbox(lines=2, placeholder="Enter your text here..."),
-    outputs=gr.Plot(),
-    title="🧠 Transformer Attention Visualizer",
-    description="Visualizes the self-attention of the BERT model's last layer."
+from model_utils import load_model_info, get_model_stats
+from visualize import (
+    visualize_attention,
+    visualize_token_embeddings,
+    plot_tokenization,
+    compare_model_sizes
 )
 
-iface.launch()
+MODEL_CHOICES = {
+    "BERT (base)": "bert-base-uncased",
+    "DistilBERT": "distilbert-base-uncased",
+    "RoBERTa": "roberta-base",
+    "GPT-2": "gpt2",
+    "Electra": "google/electra-base-discriminator",
+    "ALBERT": "albert-base-v2",
+    "XLNet": "xlnet-base-cased"
+}
+
+def run_visualizer(model_name, text, layer, head):
+    model_info = load_model_info(model_name)
+    attention_plot = visualize_attention(model_info, text, layer, head)
+    token_heatmap = visualize_token_embeddings(model_info, text)
+    token_plot = plot_tokenization(model_info, text)
+    model_stats = get_model_stats(model_info)
+
+    return attention_plot, token_heatmap, token_plot, model_stats
+
+def run_comparison_chart():
+    return compare_model_sizes(MODEL_CHOICES.values())
+
+with gr.Blocks() as demo:
+    gr.Markdown("""
+    # 🤖 Transformer Model Visualizer
+    Explore attention heads, token embeddings, and tokenizer behavior across popular transformer models.
+    """)
+
+    with gr.Row():
+        model_selector = gr.Dropdown(label="Choose Model", choices=list(MODEL_CHOICES.keys()), value="BERT (base)")
+        input_text = gr.Textbox(label="Input Text", placeholder="Enter text to analyze")
+
+    with gr.Row():
+        layer_slider = gr.Slider(minimum=0, maximum=11, step=1, value=0, label="Layer")
+        head_slider = gr.Slider(minimum=0, maximum=11, step=1, value=0, label="Attention Head")
+
+    run_btn = gr.Button("Run Analysis")
+
+    with gr.Row():
+        attention_output = gr.Plot(label="Self-Attention Visualization")
+        embedding_output = gr.Plot(label="Token Embedding Heatmap")
+
+    with gr.Row():
+        token_output = gr.Plot(label="Tokenization Overview")
+        model_output = gr.JSON(label="Model Details")
+
+    run_btn.click(
+        fn=run_visualizer,
+        inputs=[model_selector, input_text, layer_slider, head_slider],
+        outputs=[attention_output, embedding_output, token_output, model_output]
+    )
+
+    with gr.Accordion("📊 Compare Model Sizes", open=False):
+        compare_btn = gr.Button("Generate Comparison Chart")
+        comparison_output = gr.Plot()
+        compare_btn.click(fn=run_comparison_chart, outputs=comparison_output)
+
+demo.launch()
+