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()