|
import streamlit as st |
|
from transformers import pipeline |
|
import os |
|
|
|
|
|
os.environ['TRANSFORMERS_CACHE'] = os.getenv('HF_HOME', os.path.expanduser('~/.cache/huggingface/hub')) |
|
|
|
|
|
@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 |
|
|
|
|
|
st.title("Text Summarization with Pre-trained Models: BART, T5, Pegasus") |
|
|
|
|
|
with st.spinner("Loading models..."): |
|
bart_model, t5_model, pegasus_model = load_models() |
|
|
|
|
|
text_input = st.text_area("Enter text to summarize:") |
|
|
|
|
|
st.sidebar.header("Summary Length Settings") |
|
min_words = st.sidebar.slider("Minimum words in summary:", 10, 100, 50, step=5) |
|
max_words = st.sidebar.slider("Maximum words in summary:", min_words + 10, 300, 150, step=10) |
|
|
|
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..."): |
|
|
|
bart_summary = bart_model( |
|
text_input, |
|
max_length=max_words, |
|
min_length=min_words, |
|
num_beams=4, |
|
early_stopping=True |
|
)[0]['summary_text'] |
|
|
|
t5_summary = t5_model( |
|
text_input, |
|
max_length=max_words, |
|
min_length=min_words, |
|
num_beams=4, |
|
early_stopping=True |
|
)[0]['summary_text'] |
|
|
|
pegasus_summary = pegasus_model( |
|
text_input, |
|
max_length=max_words, |
|
min_length=min_words, |
|
num_beams=4, |
|
early_stopping=True |
|
)[0]['summary_text'] |
|
|
|
|
|
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.") |
|
|