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

# Available backend options are: "jax", "tensorflow", "torch".
import os
os.environ["KERAS_BACKEND"] = "tensorflow"


preprocessor = keras_nlp.models.BartSeq2SeqLMPreprocessor.from_preset(
    "hf://Grey01/bart_billsum",
    encoder_sequence_length=512,
    decoder_sequence_length=128,
)

bart_billsum = keras_nlp.models.BartSeq2SeqLM.from_preset("hf://Grey01/bart_billsum", preprocessor=preprocessor)


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:")
if user_input:
    text = user_input
else:
    text = text

def generate_text(model, input_texts, max_length=500, print_time_taken=False):
    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(generated_summaries)