zoya23's picture
Update app.py
fec6ac6 verified
raw
history blame
1.27 kB
import streamlit as st
import transformers
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
# Title
st.markdown("<h1 style='text-align: center; color: black;'>Text Summarization App</h1>", unsafe_allow_html=True)
st.markdown("---")
# Model Selection
model_choice = st.selectbox("Select a Summarization 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)
# 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 not text_input.strip():
st.warning("Please enter some text!")
else:
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'])