Update app.py

app.py

@@ -1,38 +1,49 @@
 import streamlit as st
-from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
 import matplotlib.pyplot as plt
 import seaborn as sns
+import numpy as np
 
-# App Title and Styling
-st.set_page_config(page_title="Transflower 🌸", layout="centered")
-st.markdown("<h1 style='text-align: center; color: #D16BA5;'>Transflower 🌸</h1>", unsafe_allow_html=True)
-st.markdown("<p style='text-align: center; color: #8E44AD;'>A girly & elegant app to visualize Transformer models</p>", unsafe_allow_html=True)
+# Page setup
+st.set_page_config(page_title="Transflower 🌸", page_icon="🌼", layout="centered")
+
+st.markdown(
+    "<h1 style='text-align: center; color: pink;'>🌸 Transflower 🌸</h1>"
+    "<p style='text-align: center; color: gray;'>A girly and cute app to visualize Transformer magic</p>",
+    unsafe_allow_html=True,
+)
 
 # Load model and tokenizer
-model_name = "sshleifer/distilbart-cnn-12-6"
+model_name = "t5-small"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSeq2SeqLM.from_pretrained(model_name, output_attentions=True)
 
-# Text Input
-input_text = st.text_area("🌼 Enter text to summarize or visualize:", height=150)
-
-# When user clicks the button
-if st.button("✨ Visualize Transformer Magic ✨") and input_text:
-    inputs = tokenizer(input_text, return_tensors="pt")
-    with torch.no_grad():
-        output = model.generate(**inputs, output_attentions=True, return_dict_in_generate=True)
-    
-    decoded = tokenizer.decode(output.sequences[0], skip_special_tokens=True)
-    st.success("🌸 Summary:")
-    st.markdown(f"`{decoded}`")
-
-    st.markdown("🌷 Attention Visualization (Encoder Self-Attention)")
-    
-    # Extract attentions
-    with torch.no_grad():
-        outputs = model(**inputs, output_attentions=True)
-        attentions = outputs.encoder_attentions[0][0]
-
-    fig, ax = plt.subplots(figsize=(10, 6))
-    sns.heatmap(attentions.mean(
+# Input area
+user_input = st.text_area("🌼 Enter text to summarize or visualize:", height=200)
+
+if st.button("✨ Visualize Transformer Magic ✨"):
+    if not user_input.strip():
+        st.warning("Please enter some text to visualize.")
+    else:
+        # Prepare input
+        input_ids = tokenizer.encode("summarize: " + user_input, return_tensors="pt", max_length=512, truncation=True)
+        
+        # Forward pass with attentions
+        with torch.no_grad():
+            outputs = model.generate(input_ids, output_attentions=True, return_dict_in_generate=True, output_scores=True)
+            decoded = tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
+
+        st.subheader("🌸 Summary:")
+        st.success(decoded)
+
+        # Visualization
+        st.subheader("💖 Attention Heatmap:")
+        fig, ax = plt.subplots(figsize=(10, 5))
+
+        # Get decoder self-attention from the last layer
+        attention_data = outputs.attentions[-1]  # List of attention tensors from each layer
+        avg_attention = attention_data[0].mean(dim=0).squeeze().detach().numpy()  # mean over heads
+
+        sns.heatmap(avg_attention, cmap="coolwarm", ax=ax)
+        st.pyplot(fig)