madanagrawal's picture
Update app.py
6a27b92 verified
raw
history blame
1.32 kB
import streamlit as st
from transformers import AutoTokenizer
from transformers import AutoModelForSeq2SeqLM
# Set up the Streamlit app
st.title("Text Summarization")
# Load the summarization model
@st.cache_resource
def load_summarizer():
return AutoModelForSeq2SeqLM.from_pretrained("madanagrawal/summarization_model")
# Load the tokenizer
@st.cache_resource
def load_tokenizer():
return AutoTokenizer.from_pretrained("madanagrawal/summarization_model")
tokenizer = load_tokenizer()
summarizer = load_summarizer()
# Create a text input for the user
text_input = st.text_area("Enter the text you want to summarize:", height=200)
# Create a slider for selecting the maximum length of the summary
max_length = st.slider("Select the maximum length of the summary:", min_value=50, max_value=500, value=150, step=10)
# Create a button to trigger the summarization
if st.button("Summarize"):
if text_input:
inputs = tokenizer(text_input, return_tensors="pt").input_ids
# Generate the summary
summary = summarizer.generate(inputs, max_new_tokens=max_length, do_sample=False)
# Display the summary
st.subheader("Summary:")
st.write(tokenizer.decode(summary[0], skip_special_tokens=True))
else:
st.warning("Please enter some text to summarize.")