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.")