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


model_name = "maxidl/arena-test"

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_name)


title = "# Placeholder Title"
steps = """Placeholder Description"""
# steps = """1. Converts uploaded pdf file to markdown. You can edit the intermediate markdown output.\n2. Generates a review for the paper"""

def process_file(file):
    return "Processed file"


@spaces.GPU(duration=60)
def generate(paper_text):
    messages = [
        {"role": "system", "content": "You are a pirate."},
        {"role": "user", "content": paper_text}
    ]
    model_inputs = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True
    ).to(model.device)

    generated_ids = model.generate(
        **model_inputs,
        max_new_tokens=256
    )
    generated_ids = [
        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
    ]
    
    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
    return response
    # return "Success"



with gr.Blocks() as demo:
    title = gr.Markdown(title)
    steps = gr.Markdown(steps)
    instr = gr.Markdown("## Upload your paper as a pdf file")
    file_input = gr.File(file_types=[".pdf"], file_count="single")
    markdown_field = gr.Textbox(label="Markdown", max_lines=20, autoscroll=False)
    # generate_button = gr.Button("Generate Review", interactive=not markdown_field)
    generate_button = gr.Button("Generate Review")
    file_input.upload(process_file, file_input, markdown_field)
    # markdown_field.change(lambda text: gr.update(interactive=True) if len(text) > 1000 else gr.update(interactive=False), markdown_field, generate_button)

    review_field = gr.Markdown(label="Review")
    # generate_button.click(fn=lambda: gr.update(interactive=False), inputs=None, outputs=generate_button).then(generate, markdown_field, review_field).then(fn=lambda: gr.update(interactive=True), inputs=None, outputs=generate_button)
    generate_button.click(fn=lambda: gr.update(interactive=False), inputs=None, outputs=generate_button).then(generate, markdown_field, review_field).then(fn=lambda: gr.update(interactive=True), inputs=None, outputs=generate_button)
    demo.title = "Paper Review Generator"



if __name__ == "__main__":
    demo.launch()