import gradio as gr
import PyPDF2
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import spaces


# Initialize the model and tokenizer
device = "cuda" if torch.cuda.is_available() else "cpu"
model = AutoModelForCausalLM.from_pretrained(
    "pints-ai/1.5-Pints-2k-v0.1",
    device_map=device,
    attn_implementation="flash_attention_2"
)
tokenizer = AutoTokenizer.from_pretrained("pints-ai/1.5-Pints-2k-v0.1")

def extract_pdf_content(file):
    pdf_reader = PyPDF2.PdfReader(file)
    content = ""
    for page in pdf_reader.pages:
        content += page.extract_text() + "\n"
    return content

@spaces.GPU
def chat(message, history, pdf_content):
    messages = [
        {"role": "system", "content": f"You are an AI assistant that follows instruction extremely well. Use the following information from the uploaded PDF as context: {pdf_content}"},
        {"role": "user", "content": message}
    ]

    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )

    input = tokenizer([text], return_tensors="pt").to(device)

    generated_ids = model.generate(
        input.input_ids,
        max_new_tokens=512
    )

    input_length = len(input.input_ids[0])
    response = tokenizer.decode(generated_ids[0][input_length:])

    return response

def process_pdf(file):
    if file is None:
        return "Please upload a PDF file."
    content = extract_pdf_content(file)
    return f"PDF content extracted. Length: {len(content)} characters."

with gr.Blocks() as demo:
    pdf_content = gr.State("")

    with gr.Row():
        pdf_upload = gr.File(label="Upload PDF", file_types=[".pdf"])
        pdf_output = gr.Textbox(label="PDF Processing Output")

    chatbot = gr.Chatbot()
    msg = gr.Textbox()

    pdf_upload.upload(process_pdf, pdf_upload, pdf_output)
    pdf_upload.upload(lambda file: extract_pdf_content(file), pdf_upload, pdf_content)

    msg.submit(chat, [msg, chatbot, pdf_content], chatbot)

demo.launch()