Update app.py

app.py

@@ -3,10 +3,10 @@ from transformers import AutoTokenizer, AutoModel, AutoModelForSequenceClassific
 import torch
 import plotly.express as px
 import numpy as np
-from utils import visualize_attention, list_supported_models
+from sklearn.decomposition import PCA
+from utils import visualize_attention, list_supported_models, plot_token_embeddings
 
 st.set_page_config(page_title="Transformer Visualizer", layout="wide")
-
 st.title("🧠 Transformer Visualizer")
 st.markdown("Explore how Transformer models process and understand language.")
 
@@ -19,27 +19,40 @@ if st.button("Run"):
     st.info(f"Loading model: `{model_name}`...")
     
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    
     if task == "Text Classification":
         model = AutoModelForSequenceClassification.from_pretrained(model_name, output_attentions=True)
     else:
         model = AutoModel.from_pretrained(model_name, output_attentions=True)
-    
-    inputs = tokenizer(text_input, return_tensors="pt")
+
+    inputs = tokenizer(text_input, return_tensors="pt", return_token_type_ids=False)
     outputs = model(**inputs)
     attentions = outputs.attentions
 
     st.success("Model inference complete!")
 
+    # Tokenization Visualization
+    st.subheader("🔠 Tokenization")
+    tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
+    token_ids = inputs["input_ids"][0].tolist()
+    st.write(list(zip(tokens, token_ids)))
+
+    # Token Embeddings Visualization
+    st.subheader("🌐 Token Embedding Space (PCA)")
+    with torch.no_grad():
+        hidden_states = model.base_model.embeddings.word_embeddings(inputs["input_ids"]).squeeze(0)
+    fig_embed = plot_token_embeddings(hidden_states, tokens)
+    st.plotly_chart(fig_embed, use_container_width=True)
+
+    # Attention Visualization
     if attentions:
-        st.subheader("Attention Visualization")
+        st.subheader("👁️ Attention Visualization")
         fig = visualize_attention(attentions, tokenizer, inputs)
         st.plotly_chart(fig, use_container_width=True)
     else:
         st.warning("This model does not return attention weights.")
 
     if task == "Text Classification":
-        st.subheader("Prediction")
+        st.subheader("✅ Prediction")
         pipe = pipeline("text-classification", model=model, tokenizer=tokenizer)
         prediction = pipe(text_input)
         st.write(prediction)