demo-pegasus / app.py
fatihfauzan26's picture
Create app.py
7bd9b65 verified
import streamlit as st
from transformers import PegasusForConditionalGeneration, PegasusTokenizer, pipeline
# Load the model and tokenizer
model = PegasusForConditionalGeneration.from_pretrained("fatihfauzan26/PEGASUS_liputan6")
tokenizer = PegasusTokenizer.from_pretrained("google/pegasus-cnn_dailymail")
# Initialize the summarization pipeline
summarizer = pipeline("summarization", model=model, tokenizer=tokenizer)
# Streamlit interface
st.title("Summarization App using PEGASUS")
# Input article for summarization
sample_article = st.text_area('Enter the article you want to summarize', height=300)
if sample_article:
# Generate summary
input_ids = tokenizer.encode(sample_article, return_tensors='pt')
summary_ids = model.generate(input_ids,
min_length=30,
max_length=128,
num_beams=8,
repetition_penalty=2.0,
length_penalty=0.8,
early_stopping=True,
no_repeat_ngram_size=2,
use_cache=True,
do_sample=True,
temperature=1.2,
top_k=50,
top_p=0.95)
# Decode the summary
summary_text = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
# Display results
st.subheader("Summary")
st.write(summary_text)