Update app.py

app.py

@@ -1,22 +1,38 @@
 import streamlit as st
 import transformers
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
 
-from transformers import pipeline
+# Title
+st.markdown("<h1 style='text-align: center; color: black;'>Text Summarization App</h1>", unsafe_allow_html=True)
+st.markdown("---")
 
-st.title("Text Summarization App")
+# Model Selection
+model_choice = st.selectbox("Select a Summarization Model", ["BART", "T5", "PEGASUS"])
 
-model_choice = st.selectbox("Choose a Model", ["BART", "T5", "PEGASUS"])
+# Load model and tokenizer
+@st.cache_resource
+def load_model(model_name):
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+    return pipeline("summarization", model=model, tokenizer=tokenizer)
 
-text_input = st.text_area("Enter text to summarize:")
+# Map choices to model names or paths
+model_map = {
+    "BART": "facebook/bart-large-cnn",
+    "T5": "t5-small",
+    "PEGASUS": "google/pegasus-cnn_dailymail"
+}
 
+# Text Input
+text_input = st.text_area("Enter the text you want to summarize:", height=300)
+
+# Button to generate summary
 if st.button("Summarize"):
-    if model_choice == "BART":
-        summarizer = pipeline("summarization", model="BART")
-    elif model_choice == "T5":
-        summarizer = pipeline("summarization", model="T5")
+    if not text_input.strip():
+        st.warning("Please enter some text!")
     else:
-        summarizer = pipeline("summarization", model="PEGASUS")
-    
-    summary = summarizer(text_input, max_length=150, min_length=40, do_sample=False)
-    st.subheader("Summary:")
-    st.write(summary[0]['summary_text'])
+        summarizer = load_model(model_map[model_choice])
+        summary = summarizer(text_input, max_length=150, min_length=40, do_sample=False)
+        
+        st.markdown("### Summary:")
+        st.success(summary[0]['summary_text'])