import streamlit as st
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Load the tokenizer and model
@st.cache_resource
def load_model():
    model_path = "./bart_samsum"  # Local directory where the model is stored
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
    return tokenizer, model

tokenizer, model = load_model()

st.title("Summarization with BART")

# Text input
dialogue = st.text_area("Enter the dialogue to summarize:", height=200)

# Summarize button
if st.button("Summarize"):
    inputs = tokenizer(dialogue, max_length=512, truncation=True, return_tensors="pt")
    summary_ids = model.generate(inputs["input_ids"], max_length=128, min_length=30, length_penalty=2.0, num_beams=4, early_stopping=True)
    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
    st.write("### Summary")
    st.write(summary)