import streamlit as st
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Load the model and tokenizer
model_path = '/content/drive/MyDrive/Distibert_stramlit_samsum'
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForSeq2SeqLM.from_pretrained(model_path)

# Define the summarization function
def generate_summary(text):
    inputs = tokenizer.encode("summarize: " + text, return_tensors="pt", max_length=1024, truncation=True)
    summary_ids = model.generate(
        inputs,
        max_length=150,
        min_length=50,
        length_penalty=2.0,
        num_beams=4,
        early_stopping=True
    )
    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
    return summary

# Streamlit app
st.title('Text Summarization with Custom Model')
st.write('Enter text to summarize:')

# Text input box
user_input = st.text_area("Text", "")

if st.button('Summarize'):
    if user_input:
        summary = generate_summary(user_input)
        st.write('Summary:')
        st.write(summary)
    else:
        st.write('Please enter some text to summarize.')