Spaces:
Sleeping
Sleeping
import streamlit as st | |
from transformers import AutoTokenizer, BartForConditionalGeneration | |
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.") | |