Update app.py

app.py

@@ -1,45 +1,38 @@
 import streamlit as st
-from PIL import Image
-from ultralytics import YOLO
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
-
-st.set_page_config(page_title="Animal Detection App", layout="centered")
-
-# Load YOLOv8 model
-@st.cache_resource
-def load_model():
-    return YOLO("yolov8s.pt")
-
-model = load_model()
-
-st.title("🐾 Animal Detection App")
-st.write("Upload an image and let the YOLOv8 model detect animals!")
-
-uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
-
-if uploaded_file is not None:
-    image = Image.open(uploaded_file)
-    st.image(image, caption="Uploaded Image", use_container_width=True)
+import matplotlib.pyplot as plt
+import seaborn as sns
+
+# 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)
+
+# Load model and tokenizer
+model_name = "sshleifer/distilbart-cnn-12-6"
+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)
     
-    with st.spinner("Detecting..."):
-        results = model.predict(image)
-
-        for r in results:
-            rendered_img = r.plot()  # Draws boxes and labels on the image
-            st.image(rendered_img, caption="Detected Image", use_container_width=True)
-
+    decoded = tokenizer.decode(output.sequences[0], skip_special_tokens=True)
+    st.success("🌸 Summary:")
+    st.markdown(f"`{decoded}`")
 
-        result_img = Image.fromarray(results[0].plot()[:, :, ::-1])
-        st.image(result_img, caption="Detected Animals", use_column_width=True)
-
-        # Filter animal predictions
-        animal_labels = ["cat", "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "bird"]
-        names = model.names
-        detections = results[0].boxes.data.cpu().numpy()
+    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]
 
-        st.subheader("Detections:")
-        for det in detections:
-            class_id = int(det[5])
-            label = names[class_id]
-            if label in animal_labels:
-                st.markdown(f"- **{label}** (Confidence: {det[4]:.2f})")
+    fig, ax = plt.subplots(figsize=(10, 6))
+    sns.heatmap(attentions.mean(