|
import streamlit as st |
|
from transformers import AutoTokenizer, BartForConditionalGeneration |
|
|
|
@st.cache_resource |
|
def load_model(): |
|
summarizer = BartForConditionalGeneration.from_pretrained("sshleifer/distilbart-cnn-12-6") |
|
tokenizer_sum = AutoTokenizer.from_pretrained("sshleifer/distilbart-cnn-12-6") |
|
return summarizer, tokenizer_sum |
|
|
|
summarizer, tokenizer_sum = load_model() |
|
|
|
def generate_summary(text, length): |
|
inputs = tokenizer_sum([text], max_length=1024, return_tensors="pt") |
|
summary_ids = summarizer.generate(inputs["input_ids"], num_beams=2, min_length=1, max_length=length) |
|
out = tokenizer_sum.batch_decode(summary_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] |
|
st.write(out) |
|
|
|
st.title('Summarizer') |
|
st.write('Submit a news article in the field below, and the Bart-based model with provide a summary.') |
|
|
|
length = st.slider('Maximum length of summary', value = 50, min_value = 15, max_value = 150, step = 1) |
|
user_input = st.text_area("Enter your text:") |
|
if st.button("Send a review for processing"): |
|
if user_input: |
|
generate_summary(user_input, length) |
|
else: |
|
st.warning("Please enter some text before processing.") |
|
|