Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import spaces | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| import torch | |
| from marker.convert import convert_single_pdf | |
| from marker.output import markdown_exists, save_markdown, get_markdown_filepath | |
| from marker.pdf.utils import find_filetype | |
| from marker.pdf.extract_text import get_length_of_text | |
| from marker.models import load_all_models | |
| from marker.settings import settings | |
| from marker.logger import configure_logging | |
| 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" | |
| def generate(paper_text): | |
| messages = [ | |
| {"role": "system", "content": "You are a pirate."}, | |
| {"role": "user", "content": paper_text} | |
| ] | |
| input_ids = tokenizer.apply_chat_template( | |
| messages, | |
| add_generation_prompt=True, | |
| return_tensors='pt' | |
| ).to(model.device) | |
| generated_ids = model.generate( | |
| input_ids=input_ids, | |
| max_new_tokens=256 | |
| ) | |
| generated_ids = [ | |
| output_ids[len(input_ids):] for input_ids, output_ids in zip(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() | |