Update app.py

app.py

@@ -1,8 +1,10 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModel, AutoModelForSeq2SeqLM, GPT2Model
+from transformers import AutoTokenizer, AutoModel, GPT2Model
 import torch
 import matplotlib.pyplot as plt
 import seaborn as sns
+from sklearn.decomposition import PCA
+import numpy as np
 
 MODEL_INFO = {
     "bert-base-uncased": {
@@ -28,23 +30,12 @@ MODEL_INFO = {
         "Layers": 12,
         "Attention Heads": 12,
         "Parameters": "124M"
-    },
-    "t5-small": {
-        "Model Type": "T5",
-        "Layers": 6,
-        "Attention Heads": 8,
-        "Parameters": "60M"
     }
 }
 
 def visualize_transformer(model_name, sentence):
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    
-    if "t5" in model_name:
-        model = AutoModelForSeq2SeqLM.from_pretrained(model_name, output_attentions=True)
-        sentence = "translate English to English: " + sentence
-        inputs = tokenizer(sentence, return_tensors='pt')
-    elif "gpt2" in model_name:
+    if "gpt2" in model_name:
         model = GPT2Model.from_pretrained(model_name, output_attentions=True)
         tokenizer.pad_token = tokenizer.eos_token
         inputs = tokenizer(sentence, return_tensors='pt', padding=True)
@@ -69,26 +60,37 @@ def visualize_transformer(model_name, sentence):
     token_output = [f"{i + 1}: \"{tok}\"" for i, tok in enumerate(tokens)]
     token_output_str = "[\n" + "\n".join(token_output) + "\n]"
 
+    last_hidden_state = outputs.last_hidden_state.detach().numpy()[0]
+    pca = PCA(n_components=2)
+    reduced = pca.fit_transform(last_hidden_state)
+    fig2, ax2 = plt.subplots()
+    ax2.scatter(reduced[:, 0], reduced[:, 1])
+    for i, token in enumerate(tokens):
+        ax2.annotate(token, (reduced[i, 0], reduced[i, 1]))
+    ax2.set_title("Token Embedding (PCA Projection)")
+
     model_info = MODEL_INFO.get(model_name, {})
     details = f"""
 🛠 Model Details
 Model Type: {model_info.get("Model Type", "Unknown")}
-
-Number of Layers: {model_info.get("Layers", "?" )}
-
-Number of Attention Heads: {model_info.get("Attention Heads", "?" )}
-
-Total Parameters: {model_info.get("Parameters", "?" )}
+Number of Layers: {model_info.get("Layers", "?")}
+Number of Attention Heads: {model_info.get("Attention Heads", "?")}
+Total Parameters: {model_info.get("Parameters", "?")}
 
 📊 Tokenization Visualization
 Enter Text:
 {sentence}
-
 Tokenized Output:
 {token_output_str}
+
+📈 Model Size Comparison
+- BERT: 109M
+- RoBERTa: 125M
+- DistilBERT: 66M
+- GPT-2: 124M
 """
 
-    return details, fig
+    return details, fig, fig2
 
 model_list = list(MODEL_INFO.keys())
 
@@ -99,11 +101,12 @@ iface = gr.Interface(
         gr.Textbox(label="Enter Input Sentence")
     ],
     outputs=[
-        gr.Textbox(label="🧠 Model + Token Info", lines=20),
-        gr.Plot(label="🧩 Attention Map")
+        gr.Textbox(label="🧠 Model + Token Info", lines=25),
+        gr.Plot(label="🧩 Attention Map"),
+        gr.Plot(label="🧬 Token Embedding (PCA Projection)")
     ],
-    title="Transformer Attention Visualizer",
-    description="Visualize attention heads of transformer models with detailed model and token information."
+    title="Transformer Visualization App",
+    description="Visualize Transformer models including token embeddings, attention maps, and model information."
 )
 
-iface.launch()
+iface.launch()