Update app.py

app.py

@@ -9,28 +9,21 @@ from sklearn.decomposition import PCA
 from transformers import AutoModel, AutoTokenizer, pipeline, AutoModelForCausalLM
 
 # App Title
-st.title(" Transformer Model Explorer")
+st.title("🚀 Large Language Model Explorer")
 st.markdown("""
- Transformer models, their architectures, tokenization, and attention mechanisms will be displayed.
+ Large Language models, their architectures, tokenization, and attention mechanisms.
 """)
 
-#  Select Transformer
+#Selection
 model_name = st.selectbox(
-    "Choose a Transformer :",
-    ["bigscience/bloom", "openai/whisper-base", "facebook/wav2vec2-base-960h"]
+    "Select Large Language Model:",
+    ["gpt-j-6b", "opt-175b", "bigscience/bloom-176b"]
 )
 
 # Load Tokenizer & Model
 st.write(f"Loading model: `{model_name}`...")
-if "bloom" in model_name:
-    tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
-    model = AutoModel.from_pretrained(model_name)
-elif "whisper" in model_name:
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModel.from_pretrained(model_name)
-elif "wav2vec2" in model_name:
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModel.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModel.from_pretrained(model_name)
 
 # Display Model Details
 st.subheader("🛠 Model Details")
@@ -42,75 +35,57 @@ 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 = {
-    "bigscience/bloom": 176,
-    "openai/whisper-base": 74,
-    "facebook/wav2vec2-base-960h": 317
+    "gpt-j-6b": 6,
+    "opt-175b": 175,
+    "bigscience/bloom-176b": 176
 }
-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")
+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")
 st.plotly_chart(fig)
 
 # Tokenization Section
 st.subheader("📝 Tokenization Visualization")
 input_text = st.text_input("Enter Text:", "Hello, how are you?")
-
-if "whisper" in model_name:
-    st.write("Note: Whisper is an audio model and doesn't use text tokenization")
-    st.write("Instead, it processes raw audio waveforms")
-else:
-    tokens = tokenizer.tokenize(input_text)
-    st.write("Tokenized Output:", tokens)
+tokens = tokenizer.tokenize(input_text)
+st.write("Tokenized Output:", tokens)
 
 # Token Embeddings Visualization (Fixed PCA Projection)
 st.subheader("🧩 Token Embeddings Visualization")
 with torch.no_grad():
-    if "whisper" in model_name:
-        st.write("Note: Whisper uses a different embedding structure for audio features")
-        st.write("Cannot directly visualize token embeddings as with text models")
-    else:
-        inputs = tokenizer(input_text, return_tensors="pt")
-        outputs = model(**inputs)
-        if hasattr(outputs, "last_hidden_state"):
-            embeddings = outputs.last_hidden_state.squeeze(0).numpy()
-            # Ensure the number of tokens and embeddings match
-            n_tokens = min(len(tokens), embeddings.shape[0])
-            embeddings = embeddings[:n_tokens]  # Trim embeddings to match token count
-            tokens = tokens[:n_tokens]  # Trim tokens to match embeddings count
-            pca = PCA(n_components=2)
-            reduced_embeddings = pca.fit_transform(embeddings)
-            df_embeddings = pd.DataFrame(reduced_embeddings, columns=["PCA1", "PCA2"])
-            df_embeddings["Token"] = tokens
-            fig = px.scatter(df_embeddings, x="PCA1", y="PCA2", text="Token",
-                           title="Token Embeddings (PCA Projection)")
-            st.plotly_chart(fig)
+    inputs = tokenizer(input_text, return_tensors="pt")
+    outputs = model(**inputs)
+    if hasattr(outputs, "last_hidden_state"):
+        embeddings = outputs.last_hidden_state.squeeze(0).numpy()
+        # Ensure the number of tokens and embeddings match
+        n_tokens = min(len(tokens), embeddings.shape[0])
+        embeddings = embeddings[:n_tokens]  # Trim embeddings to match token count
+        tokens = tokens[:n_tokens]  # Trim tokens to match embeddings count
+        pca = PCA(n_components=2)
+        reduced_embeddings = pca.fit_transform(embeddings)
+        df_embeddings = pd.DataFrame(reduced_embeddings, columns=["PCA1", "PCA2"])
+        df_embeddings["Token"] = tokens
+        fig = px.scatter(df_embeddings, x="PCA1", y="PCA2", text="Token",
+                        title="Token Embeddings (PCA Projection)")
+        st.plotly_chart(fig)
 
-# Attention Visualization (for BERT & RoBERTa models)
-if "bloom" in model_name:
-    st.subheader("🔍 Attention Map")
-    with torch.no_grad():
-        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)
+# 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)
 
-# Text Generation Demo (for BLOOM)
-if "bloom" in model_name:
-    st.subheader("✍️ Text Generation & Token Probabilities")
-    generator = pipeline("text-generation", model=model_name, return_full_text=False)
-    generated_output = generator(input_text, max_length=50, return_tensors=True)
-    st.write("Generated Output:", generated_output[0]["generated_text"])
+# 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
-    model_gen = AutoModelForCausalLM.from_pretrained(model_name)
-    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)
+# Token Probability Visualization
+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)
 
-st.markdown("💡 *Explore more about Transformer models!*")
+st.markdown("💡 *Explore more about Large Language Models!*\n")