Dhanush S Gowda
Update app.py
d769310 verified
raw
history blame
2.58 kB
import streamlit as st
from transformers import pipeline
import os
# Set Hugging Face cache directory
os.environ['TRANSFORMERS_CACHE'] = os.getenv('HF_HOME', os.path.expanduser('~/.cache/huggingface/hub'))
# Function to load all three models
@st.cache_resource
def load_models():
bart_summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
t5_summarizer = pipeline("summarization", model="t5-large")
pegasus_summarizer = pipeline("summarization", model="google/pegasus-cnn_dailymail")
return bart_summarizer, t5_summarizer, pegasus_summarizer
# Streamlit app layout
st.title("Text Summarization with Pre-trained Models: BART, T5, Pegasus")
# Load models
with st.spinner("Loading models..."):
bart_model, t5_model, pegasus_model = load_models()
# Input text
text_input = st.text_area("Enter text to summarize:")
# Compression rate slider
compression_rate = st.slider(
"Summary Compression Rate",
min_value=0.1,
max_value=0.5,
value=0.3,
step=0.05,
help="Adjust how much of the original text to keep in the summary"
)
if text_input:
word_count = len(text_input.split())
st.write(f"**Input Word Count:** {word_count}")
if st.button("Generate Summaries"):
with st.spinner("Generating summaries..."):
# Calculate dynamic max length based on compression rate
max_length = max(50, int(word_count * compression_rate))
# Generate summaries
bart_summary = bart_model(
text_input, max_length=max_length, min_length=30, num_beams=4, early_stopping=True
)[0]['summary_text']
t5_summary = t5_model(
text_input, max_length=max_length, min_length=30, num_beams=4, early_stopping=True
)[0]['summary_text']
pegasus_summary = pegasus_model(
text_input, max_length=max_length, min_length=30, num_beams=4, early_stopping=True
)[0]['summary_text']
# Display summaries
st.subheader("BART Summary")
st.write(bart_summary)
st.write(f"**Word Count:** {len(bart_summary.split())}")
st.subheader("T5 Summary")
st.write(t5_summary)
st.write(f"**Word Count:** {len(t5_summary.split())}")
st.subheader("Pegasus Summary")
st.write(pegasus_summary)
st.write(f"**Word Count:** {len(pegasus_summary.split())}")
else:
st.warning("Please enter text to summarize.")