import streamlit as st
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch
import nltk

# Download punkt for sentence tokenization
nltk.download('punkt')

# Load tokenizer and model from the Hugging Face Hub
tokenizer = AutoTokenizer.from_pretrained("your-huggingface-username/your-model-repo-name")
model = AutoModelForSeq2SeqLM.from_pretrained("your-huggingface-username/your-model-repo-name")

st.title("Dialogue Summarization with BART")

# Input dialogue
dialogue = st.text_area("Enter dialogue:", height=200)

if st.button("Summarize"):
    # Tokenize input
    inputs = tokenizer(dialogue, max_length=512, truncation=True, return_tensors="pt")

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

    # Display summary
    st.subheader("Summary:")
    st.write(summary)

st.markdown("---")
st.markdown("This app uses a fine-tuned BART model to summarize dialogues. The model was trained on the SAMSum dataset.")