Update app.py

app.py

@@ -1,46 +1,108 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModel
+from transformers import AutoTokenizer, AutoModel, AutoModelForSeq2SeqLM, GPT2Model
 import torch
 import matplotlib.pyplot as plt
 import seaborn as sns
 
-def visualize_attention(model_name, sentence):
+MODEL_INFO = {
+    "bert-base-uncased": {
+        "Model Type": "BERT",
+        "Layers": 12,
+        "Attention Heads": 12,
+        "Parameters": "109.48M"
+    },
+    "roberta-base": {
+        "Model Type": "RoBERTa",
+        "Layers": 12,
+        "Attention Heads": 12,
+        "Parameters": "125M"
+    },
+    "distilbert-base-uncased": {
+        "Model Type": "DistilBERT",
+        "Layers": 6,
+        "Attention Heads": 12,
+        "Parameters": "66M"
+    },
+    "gpt2": {
+        "Model Type": "GPT-2",
+        "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)
-    model = AutoModel.from_pretrained(model_name, output_attentions=True)
     
-    inputs = tokenizer(sentence, return_tensors='pt')
-    outputs = model(**inputs)
-    attentions = outputs.attentions  # tuple of (layer, batch, head, seq_len, seq_len)
+    if "t5" in model_name:
+        model = AutoModelForSeq2SeqLM.from_pretrained(model_name, output_attentions=True)
+        inputs = tokenizer(sentence, return_tensors='pt')
+    elif "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)
+    else:
+        model = AutoModel.from_pretrained(model_name, output_attentions=True)
+        inputs = tokenizer(sentence, return_tensors='pt')
 
+    outputs = model(**inputs)
+    attentions = outputs.attentions
     tokens = tokenizer.convert_ids_to_tokens(inputs['input_ids'][0])
-    
+
     fig, ax = plt.subplots(figsize=(10, 8))
-    sns.heatmap(attentions[-1][0][0].detach().numpy(), 
-                xticklabels=tokens, 
-                yticklabels=tokens, 
-                cmap="viridis", 
+    sns.heatmap(attentions[-1][0][0].detach().numpy(),
+                xticklabels=tokens,
+                yticklabels=tokens,
+                cmap="viridis",
                 ax=ax)
     ax.set_title(f"Attention Map - Layer {len(attentions)} Head 1")
     plt.xticks(rotation=90)
     plt.yticks(rotation=0)
-    
-    return fig
 
-model_list = [
-    "bert-base-uncased", 
-    "roberta-base", 
-    "distilbert-base-uncased"
-]
+    token_output = [f"{i}: \"{tok}\"" for i, tok in enumerate(tokens)]
+    token_output_str = "[\\n" + "\\n".join(token_output) + "\\n]"
+
+    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", "?" )}
+
+📊 Tokenization Visualization
+Enter Text:
+{sentence}
+
+Tokenized Output:
+{token_output_str}
+"""
+
+    return details, fig
+
+model_list = list(MODEL_INFO.keys())
 
 iface = gr.Interface(
-    fn=visualize_attention,
+    fn=visualize_transformer,
     inputs=[
         gr.Dropdown(choices=model_list, label="Choose Transformer Model"),
         gr.Textbox(label="Enter Input Sentence")
     ],
-    outputs=gr.Plot(label="Attention Map"),
+    outputs=[
+        gr.Textbox(label="🧠 Model + Token Info", lines=20),
+        gr.Plot(label="🧩 Attention Map")
+    ],
     title="Transformer Attention Visualizer",
-    description="Visualize attention heads of transformer models. Select a model and input text to see attention heatmaps."
+    description="Visualize attention heads of transformer models with detailed model and token information."
 )
 
 iface.launch()