Text_summarizer / app.py
JayAjmera's picture
Upload app.py
7bdd009 verified
from transformers import PegasusTokenizer, PegasusForConditionalGeneration
import streamlit as st
# Load the model and tokenizer
tokenizer = PegasusTokenizer.from_pretrained("google/pegasus-cnn_dailymail")
model = PegasusForConditionalGeneration.from_pretrained("google/pegasus-cnn_dailymail")
def summarize(text, model, tokenizer, max_length=100, min_length=30):
# Tokenize input text
inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
# Generate summary
summary_ids = model.generate(inputs['input_ids'], max_length=max_length, min_length=min_length,
length_penalty=2.0, num_beams=4, early_stopping=True)
# Decode and return summary
return tokenizer.decode(summary_ids[0], skip_special_tokens=True)
def main():
st.title("Text Summarization App")
st.write("Enter text below to get a summarized version using the Pegasus model.")
# Input text area
user_input = st.text_area("Enter your text here", height=300)
if st.button("Summarize"):
if user_input:
# Generate summary
summary = summarize(user_input,model,tokenizer)
st.subheader("Summary:")
st.write(summary)
else:
st.write("Please enter some text to summarize.")
if __name__ == "__main__":
main()