import streamlit as st
from transformers import BartForConditionalGeneration, BartTokenizer

st.set_page_config(page_title="BART Text Summarization", layout="centered")

@st.cache_resource
def load_model():
    model = BartForConditionalGeneration.from_pretrained("Arjun9/bart_samsum")
    tokenizer = BartTokenizer.from_pretrained("Arjun9/bart_samsum")
    return model, tokenizer

model, tokenizer = load_model()

def main():
    st.title("Meeting summarization")

    # Get user input
    input_text = st.text_area("Enter text to summarize", height=200)

    if st.button("Summarize"):
        # Tokenize the input text
        inputs = tokenizer(input_text, return_tensors="pt", truncation=True)

        # Generate summary
        summary_ids = model.generate(inputs["input_ids"], num_beams=4, max_length=100, early_stopping=True)
        summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)

        # Display the summary
        st.write(f"Summary: {summary}")

if __name__ == "__main__":
    main()