import streamlit as st
from transformers import BartTokenizer, BartForConditionalGeneration, pipeline
import zipfile
import os
import nltk

# Download NLTK data
nltk.download('punkt')
from nltk.tokenize import sent_tokenize

# Define the path to the saved model zip file
zip_model_path = 'bart_model-20240724T051306Z-001.zip'

# Define the directory to extract the model
model_dir = './model'

# Unzip the model
with zipfile.ZipFile(zip_model_path, 'r') as zip_ref:
    zip_ref.extractall(model_dir)

# After unzipping, the model should be in a specific directory, check the directory structure
model_path = os.path.join(model_dir, 'Bart')

# Print out the model_path for debugging
print("Model Path:", model_path)

# Verify that the directory exists and contains the necessary files
if not os.path.exists(model_path):
    st.error(f"Model directory {model_path} does not exist or is incorrect.")
else:
    # Load the tokenizer and model
    tokenizer = BartTokenizer.from_pretrained('facebook/bart-large-cnn')
    model = BartForConditionalGeneration.from_pretrained(model_path)

    # Create a summarization pipeline
    summarizer = pipeline("summarization", model=model, tokenizer=tokenizer)

    # Set the title for the Streamlit app
    st.title("BART Summary Generator")

    # Text input for the user
    text = st.text_area("Enter your text: ")

    def generate_summary(input_text):
        # Perform summarization
        summary = summarizer(input_text, max_length=200, min_length=40, do_sample=False)
        return summary[0]['summary_text']

    if st.button("Generate"):
        if text:
            generated_summary = generate_summary(text)
            # Display the generated summary
            st.subheader("Generated Summary")
            st.write(generated_summary)
        else:
            st.warning("Please enter some text to generate a summary.")