Update app.py

app.py

@@ -1,69 +1,96 @@
-# app.py
+# 🧠 Summary of Features:
+# 🔽 Dropdown to choose from 7 prebuilt models
+# 🔄 Updates attention layer & head slider limits based on the model
+# 📊 Visualizes attention maps for selected head/layer
+# 🧩 Tokenized words preview
+# 📉 Embeddings reduced to 2D using PCA
+# 📦 Model size bar chart across models
+# 📋 Full model config info in JSON viewer
+
 import gradio as gr
-from model_utils import load_model_info, get_model_stats
+from model_utils import MODEL_OPTIONS, load_model, get_model_info
 from visualize import (
     visualize_attention,
-    visualize_token_embeddings,
-    plot_tokenization,
-    compare_model_sizes
+    show_tokenization,
+    show_embeddings,
+    compare_model_sizes,
 )
 
-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"
-}
+# Initial load
+DEFAULT_MODEL_NAME = list(MODEL_OPTIONS.values())[0]
+tokenizer, model = load_model(DEFAULT_MODEL_NAME)
+
+# Shared state
+current_tokenizer = tokenizer
+current_model = model
+
+
+def update_model(selected_model_name):
+    global current_tokenizer, current_model
+    model_id = MODEL_OPTIONS[selected_model_name]
+    current_tokenizer, current_model = load_model(model_id)
+    info = get_model_info(current_model)
+
+    # Update layer/head sliders based on model
+    num_layers = info.get("Number of Layers", 1)
+    num_heads = info.get("Number of Attention Heads", 1)
+
+    return (
+        info,
+        gr.update(maximum=num_layers - 1, value=0),
+        gr.update(maximum=num_heads - 1, value=0),
+    )
 
-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_all_visualizations(text, layer, head):
+    attention_fig = visualize_attention(current_tokenizer, current_model, text, layer, head)
+    token_fig = show_tokenization(current_tokenizer, text)
+    embedding_fig = show_embeddings(current_tokenizer, current_model, text)
+    return attention_fig, token_fig, embedding_fig
 
-def run_comparison_chart():
-    return compare_model_sizes(MODEL_CHOICES.values())
 
+# UI
 with gr.Blocks() as demo:
-    gr.Markdown("""
-    # 🤖 Transformer Model Visualizer
-    Explore attention heads, token embeddings, and tokenizer behavior across popular transformer models.
-    """)
+    gr.Markdown("## 🔍 Transformer Explorer")
+    gr.Markdown("Explore attention, tokenization, and embedding visualizations for various 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")
+        model_dropdown = gr.Dropdown(
+            label="Choose a model",
+            choices=list(MODEL_OPTIONS.keys()),
+            value=list(MODEL_OPTIONS.keys())[0],
+        )
+        model_info = gr.JSON(label="Model Info")
 
     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")
+        text_input = gr.Textbox(label="Enter text", value="The quick brown fox jumps over the lazy dog.")
+        layer_slider = gr.Slider(label="Layer", minimum=0, maximum=11, step=1, value=0)
+        head_slider = gr.Slider(label="Head", minimum=0, maximum=11, step=1, value=0)
 
-    run_btn = gr.Button("Run Analysis")
+    run_button = gr.Button("Run Visualizations")
 
-    with gr.Row():
-        attention_output = gr.Plot(label="Self-Attention Visualization")
-        embedding_output = gr.Plot(label="Token Embedding Heatmap")
+    with gr.Tab("📊 Attention"):
+        attention_plot = gr.Plot()
 
-    with gr.Row():
-        token_output = gr.Plot(label="Tokenization Overview")
-        model_output = gr.JSON(label="Model Details")
+    with gr.Tab("🧩 Tokenization"):
+        token_plot = gr.Plot()
 
-    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.Tab("📉 Embeddings"):
+        embedding_plot = gr.Plot()
+
+    with gr.Tab("📦 Model Size Comparison"):
+        model_compare_plot = gr.Plot(value=compare_model_sizes())
+
+    # Event binding
+    model_dropdown.change(fn=update_model, inputs=[model_dropdown], outputs=[model_info, layer_slider, head_slider])
+    run_button.click(
+        fn=run_all_visualizations,
+        inputs=[text_input, layer_slider, head_slider],
+        outputs=[attention_plot, token_plot, embedding_plot],
     )
 
-    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()
+if __name__ == "__main__":
+    demo.launch()
+