import streamlit as st
import tensorflow as tf
from tensorflow import keras
import keras_nlp
import PyPDF2
import docx2txt
import huggingface_hub

import os
os.environ['KERAS_BACKEND'] = 'tensorflow'
os.environ['HF_USERNAME'] = 'Grey01'



from huggingface_hub import whoami

# Get the username
hf_username = whoami()['name']

# Define the model repository
model_repo = f'{hf_username}/bart_billsum'

# Load the Keras NLP model
bart_billsum = keras_nlp.models.BartSeq2SeqLM.from_preset(model_repo)

st.title("SummarizeIt")

# File uploader
uploaded_file = st.file_uploader("Choose a file", type=["pdf", "txt", "docx"])

# Text extraction
text = ""
if uploaded_file is not None:
    if uploaded_file.type == "application/pdf":
        pdf_reader = PyPDF2.PdfReader(uploaded_file)
        for page in pdf_reader.pages:
            text += page.extract_text()
    elif uploaded_file.type == "text/plain":
        text = uploaded_file.read().decode("utf-8")
    elif uploaded_file.type == "application/vnd.openxmlformats-officedocument.wordprocessingml.document":
        text = docx2txt.process(uploaded_file)
# Text input for direct text entry
user_input = st.text_area("Or paste your text here:")
text = user_input if user_input else text  # Prioritize user input over file

def generate_text(model, input_texts, max_length=200, print_time_taken=False):
    # Convert input_texts to a list if it's a Dataset
    if isinstance(input_texts, datasets.Dataset):
        input_texts = input_texts.to_list()
    chunks = [input_texts[i:i+512] for i in range(0, len(input_texts), 512)]
    #initialize an empty list to store summaries
    summaries = []
    # generate summaries for each chunk
    for chunk in chunks:
        # Assuming your model's generate method can handle a batch of inputs
        summary = model.generate(input_texts, max_length=max_length)
        summaries.append(summary)
    return summary

generated_summaries = generate_text(
    bart_billsum,
    text,  # Pass the list of documents directly
)
st.subheader("Generated Summary:")
st.write(summary)