Update app.py

app.py

@@ -10,14 +10,15 @@ st.set_page_config(page_title="Transflower 🌸", page_icon="🌼", layout="cent
 
 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 = "t5-small"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name, output_attentions=True)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+model.eval()
 
 # Input area
 user_input = st.text_area("🌼 Enter text to summarize or visualize:", height=200)
@@ -26,24 +27,25 @@ 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
+        # Encode input
+        inputs = tokenizer("summarize: " + user_input, return_tensors="pt", truncation=True)
+
+        # Forward pass manually to get attention
         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)
+            encoder_outputs = model.encoder(**inputs, output_attentions=True, return_dict=True)
+            attention = encoder_outputs.attentions[-1][0].mean(dim=0).detach().numpy()
 
-        st.subheader("🌸 Summary:")
-        st.success(decoded)
+        # Generate summary
+        with torch.no_grad():
+            summary_ids = model.generate(inputs["input_ids"], max_length=50)
+            summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
 
-        # Visualization
-        st.subheader("💖 Attention Heatmap:")
-        fig, ax = plt.subplots(figsize=(10, 5))
+        st.subheader("🌸 Summary:")
+        st.success(summary)
 
-        # 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
+        st.subheader("💖 Encoder Attention Heatmap:")
 
-        sns.heatmap(avg_attention, cmap="coolwarm", ax=ax)
+        fig, ax = plt.subplots(figsize=(10, 6))
+        sns.heatmap(attention, cmap="YlGnBu", ax=ax)
+        ax.set_title("Encoder Self-Attention Heatmap 💫")
         st.pyplot(fig)