Update app.py

app.py

@@ -9,15 +9,15 @@ from sklearn.decomposition import PCA
 from transformers import AutoModel, AutoTokenizer, pipeline, AutoModelForCausalLM
 
 # App Title
-st.title("🚀 Large Language Model Explorer")
+st.title("🚀 Vision Transformer Explorer")
 st.markdown("""
- Large Language models, their architectures, tokenization, and attention mechanisms.
+Explore Vision Transformers, their architectures, and tokenization mechanisms.
 """)
 
-#Selection
+# Model Selection
 model_name = st.selectbox(
-    "Select Large Language Model:",
-    ["gpt-j-6b", "opt-175b", "bigscience/bloom-176b"]
+    "Choose a Vision Transformer Model:",
+    ["beit-base-patch16", "swin-base-patch4-window7", "vit-base-patch16"]
 )
 
 # Load Tokenizer & Model
@@ -35,12 +35,12 @@ st.write(f"Total Parameters: `{sum(p.numel() for p in model.parameters())/1e6:.2
 # Model Size Comparison
 st.subheader("📊 Model Size Comparison")
 model_sizes = {
-    "gpt-j-6b": 6,
-    "opt-175b": 175,
-    "bigscience/bloom-176b": 176
+    "beit-base-patch16": 86,
+    "swin-base-patch4-window7": 87,
+    "vit-base-patch16": 86
 }
-df_size = pd.DataFrame(model_sizes.items(), columns=["Model", "Size (Billion Parameters)"])
-fig = px.bar(df_size, x="Model", y="Size (Billion Parameters)", title="Model Size Comparison")
+df_size = pd.DataFrame(model_sizes.items(), columns=["Model", "Size (Million Parameters)"])
+fig = px.bar(df_size, x="Model", y="Size (Million Parameters)", title="Model Size Comparison")
 st.plotly_chart(fig)
 
 # Tokenization Section
@@ -68,24 +68,13 @@ with torch.no_grad():
                         title="Token Embeddings (PCA Projection)")
         st.plotly_chart(fig)
 
-# Text Generation Demo
-st.subheader("✍️ Text Generation & Token Probabilities")
-model_gen = AutoModelForCausalLM.from_pretrained(model_name)
-generator = pipeline("text-generation", model=model_name, return_full_text=False)
-
-# Generate text
-generated_output = generator(input_text, max_length=50, return_tensors=True)
-st.write("Generated Output:", generated_output[0]["generated_text"])
-
-# Token Probability Visualization
+# Attention Visualization
+st.subheader("🔍 Attention Map")
 with torch.no_grad():
-    inputs = tokenizer(input_text, return_tensors="pt")
-    logits = model_gen(**inputs).logits[:, -1, :]
-    probs = torch.nn.functional.softmax(logits, dim=-1).squeeze().detach().numpy()
-    top_tokens = np.argsort(probs)[-10:][::-1]  # Top 10 tokens
-    token_probs = {tokenizer.decode([idx]): probs[idx] for idx in top_tokens}
-    df_probs = pd.DataFrame(token_probs.items(), columns=["Token", "Probability"])
-    fig_prob = px.bar(df_probs, x="Token", y="Probability", title="Top Token Predictions")
-    st.plotly_chart(fig_prob)
+    outputs = model(**inputs, output_attentions=True)
+    attention = outputs.attentions[-1].squeeze().detach().numpy()
+    fig, ax = plt.subplots(figsize=(10, 5))
+    sns.heatmap(attention[0], cmap="viridis", xticklabels=tokens, yticklabels=tokens, ax=ax)
+    st.pyplot(fig)
 
-st.markdown("💡 *Explore more about Large Language Models!*\n")
+st.markdown("💡 *Explore Vision Transformers!*\n")