app.py

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
from surya.settings import settings as surya_settings
import traceback


# marker
configure_logging()
MAX_PAGES = 20
MIN_LENGTH=200
settings.EXTRACT_IMAGES = False
settings.DEBUG = False
surya_settings.IN_STREAMLIT = True

model_refs = load_all_models()
metadata = {}

# prepare LLM
model_name = "maxidl/arena-test"
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Define prompts
SYSTEM_PROMPT_TEMPLATE = """You are an expert reviewer for AI conferences. You follow best practices and review papers according to the reviewer guidelines.

Reviewer guidelines:
1. Read the paper: It’s important to carefully read through the entire paper, and to look up any related work and citations that will help you comprehensively evaluate it. Be sure to give yourself sufficient time for this step.
2. While reading, consider the following:
    - Objective of the work: What is the goal of the paper? Is it to better address a known application or problem, draw attention to a new application or problem, or to introduce and/or explain a new theoretical finding? A combination of these? Different objectives will require different considerations as to potential value and impact.
    - Strong points: is the submission clear, technically correct, experimentally rigorous, reproducible, does it present novel findings (e.g. theoretically, algorithmically, etc.)?
    - Weak points: is it weak in any of the aspects listed in b.?
    - Be mindful of potential biases and try to be open-minded about the value and interest a paper can hold for the community, even if it may not be very interesting for you.
3. Answer four key questions for yourself, to make a recommendation to Accept or Reject:
    - What is the specific question and/or problem tackled by the paper?
    - Is the approach well motivated, including being well-placed in the literature?
    - Does the paper support the claims? This includes determining if results, whether theoretical or empirical, are correct and if they are scientifically rigorous.
    - What is the significance of the work? Does it contribute new knowledge and sufficient value to the community? Note, this does not necessarily require state-of-the-art results. Submissions bring value to the community when they convincingly demonstrate new, relevant, impactful knowledge (incl., empirical, theoretical, for practitioners, etc).
4. Write your review including the following information: 
    - Summarize what the paper claims to contribute. Be positive and constructive.
    - List strong and weak points of the paper. Be as comprehensive as possible.
    - Clearly state your initial recommendation (accept or reject) with one or two key reasons for this choice.
    - Provide supporting arguments for your recommendation.
    - Ask questions you would like answered by the authors to help you clarify your understanding of the paper and provide the additional evidence you need to be confident in your assessment.
    - Provide additional feedback with the aim to improve the paper. Make it clear that these points are here to help, and not necessarily part of your decision assessment.

Your write reviews in markdown format. Your reviews contain the following sections:

# Review

{review_fields}

Your response must only contain the review in markdown format with sections as defined above.
"""

USER_PROMPT_TEMPLATE = """Review the following paper:

{paper_text}
"""

# For now, use fixed review fields
REVIEW_FIELDS = """## Summary
Briefly summarize the paper and its contributions. This is not the place to critique the paper; the authors should generally agree with a well-written summary.

## Soundness
Please assign the paper a numerical rating on the following scale to indicate the soundness of the technical claims, experimental and research methodology and on whether the central claims of the paper are adequately supported with evidence. Choose from the following:
4: excellent
3: good
2: fair
1: poor

## Presentation
Please assign the paper a numerical rating on the following scale to indicate the quality of the presentation. This should take into account the writing style and clarity, as well as contextualization relative to prior work. Choose from the following:
4: excellent
3: good
2: fair
1: poor

## Contribution
Please assign the paper a numerical rating on the following scale to indicate the quality of the overall contribution this paper makes to the research area being studied. Are the questions being asked important? Does the paper bring a significant originality of ideas and/or execution? Are the results valuable to share with the broader ICLR community? Choose from the following:
4: excellent
3: good
2: fair
1: poor

## Strengths
A substantive assessment of the strengths of the paper, touching on each of the following dimensions: originality, quality, clarity, and significance. We encourage reviewers to be broad in their definitions of originality and significance. For example, originality may arise from a new definition or problem formulation, creative combinations of existing ideas, application to a new domain, or removing limitations from prior results.

## Weaknesses
A substantive assessment of the weaknesses of the paper. Focus on constructive and actionable insights on how the work could improve towards its stated goals. Be specific, avoid generic remarks. For example, if you believe the contribution lacks novelty, provide references and an explanation as evidence; if you believe experiments are insufficient, explain why and exactly what is missing, etc.

## Questions
Please list up and carefully describe any questions and suggestions for the authors. Think of the things where a response from the author can change your opinion, clarify a confusion or address a limitation. This is important for a productive rebuttal and discussion phase with the authors.

## Flag For Ethics Review
If there are ethical issues with this paper, please flag the paper for an ethics review and select area of expertise that would be most useful for the ethics reviewer to have. Please select all that apply. Choose from the following:
No ethics review needed.
Yes, Discrimination / bias / fairness concerns
Yes, Privacy, security and safety
Yes, Legal compliance (e.g., GDPR, copyright, terms of use)
Yes, Potentially harmful insights, methodologies and applications
Yes, Responsible research practice (e.g., human subjects, data release)
Yes, Research integrity issues (e.g., plagiarism, dual submission)
Yes, Unprofessional behaviors (e.g., unprofessional exchange between authors and reviewers)
Yes, Other reasons (please specify below)

## Details Of Ethics Concerns
Please provide details of your concerns.

## Rating
Please provide an "overall score" for this submission. Choose from the following:
1: strong reject
3: reject, not good enough
5: marginally below the acceptance threshold
6: marginally above the acceptance threshold
8: accept, good paper
10: strong accept, should be highlighted at the conference


"""

# functions
def create_messages(review_fields, paper_text):
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT_TEMPLATE.format(review_fields=review_fields)},
        {"role": "user", "content": USER_PROMPT_TEMPLATE.format(paper_text=paper_text)},
    ]
    return messages

@spaces.GPU(duration=60)
def convert_file(filepath):
    full_text, images, out_metadata = convert_single_pdf(
            filepath, model_refs, metadata=metadata, max_pages=MAX_PAGES
    )
    return full_text.strip()

def process_file(file):
    print(file.name)
    filepath = file.name
    try:
        if MIN_LENGTH:
            filetype = find_filetype(filepath)
            if filetype == "other":
                raise ValueError()

            length = get_length_of_text(filepath)
            if length < MIN_LENGTH:
                raise ValueError()
        paper_text = convert_file(filepath)
        if not len(paper_text) > MIN_LENGTH:
            raise ValueError()
    except spaces.zero.gradio.HTMLError as e:
        print(e)
        return "GPU quota exceeded"
    except Exception as e:
        print(traceback.format_exc())
        print(f"Error converting {filepath}: {e}")
        return "Error processing pdf"
    return paper_text


@spaces.GPU(duration=60)
def generate(paper_text, review_template):
    # messages = [
    #     {"role": "system", "content": "You are a pirate."},
    #     {"role": "user", "content": paper_text}
    # ]
    messages = create_messages(review_template, paper_text)
    input_ids = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors='pt'
    ).to(model.device)
    print(f"input_ids shape: {input_ids.shape}")
    generated_ids = model.generate(
        input_ids=input_ids,
        max_new_tokens=512
    )
    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"



# ui
css = """
#warning {background: red;}

.gradio-container {background-color: #FFFDFA !important;}
.topbar {background-color: #8C1B13 !important; margin-left: -15px !important; margin-right: -15px !important}
.title {background-color: #FFFDFA !important; color: white !important;}
"""
#4D8093 blue
#767676 med grey
#EFECE3 light grey
#DDDDDD silver below red
#FFFDFA white

title = "# OpenReviewer"
description = """Placeholder Description"""

theme = gr.themes.Default()
with gr.Blocks(theme=theme, css=css) as demo:
    topbar = gr.HTML("""<div>OpenReviewer</div>"""
    , elem_classes=['topbar'])
    title = gr.Markdown(title, elem_classes=["title"])
    description = gr.Markdown(description)
    instr = gr.Markdown("## Upload your paper in pdf format")
    file_input = gr.File(file_types=[".pdf"], file_count="single")
    paper_text_field= gr.Textbox(label="Paper Text", max_lines=20, autoscroll=False)
    review_template_field = gr.Textbox(label="Review Template", max_lines=20, autoscroll=False, value=REVIEW_FIELDS)
    generate_button = gr.Button("Generate Review", interactive=not paper_text_field)
    # generate_button = gr.Button("Generate Review")
    file_input.upload(process_file, file_input, paper_text_field)
    paper_text_field.change(lambda text: gr.update(interactive=True) if len(text) > 200 else gr.update(interactive=False), paper_text_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, [paper_text_field, review_template_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()