|
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:") |
|
|
|
|
|
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..."): |
|
|
|
max_length = max(50, int(word_count * compression_rate)) |
|
|
|
|
|
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'] |
|
|
|
|
|
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.") |