Intial commit

README.md

@@ -0,0 +1,14 @@
+---
+title: TMLR Paper Reviewer
+emoji: 📝
+colorFrom: blue
+colorTo: red
+sdk: gradio
+sdk_version: 4.19.2
+app_file: run_claude.py
+pinned: false
+---
+
+# TMLR Paper Reviewer
+
+This tool helps generate high-quality reviews for papers submitted to the Transactions on Machine Learning Research (TMLR).

requirements.txt

@@ -0,0 +1,6 @@
+gradio
+anthropic
+openai
+PyPDF2
+tiktoken
+pydantic

run_claude.py

@@ -0,0 +1,261 @@
+import anthropic
+import base64
+import pandas as pd
+import requests
+import os
+import numpy as np
+from openai import OpenAI
+import io
+import tiktoken
+import PyPDF2
+
+import prompts
+
+from typing import List, Literal
+from pydantic import BaseModel
+
+import time
+
+import gradio as gr
+
+ANTHROPIC_API_KEY = os.environ.get("ANTHROPIC_API_KEY")
+OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
+
+
+def ask_claude(
+    query: str, 
+    pdf_path: str = None, 
+    use_cache: bool = False, 
+    system: str = None,
+    max_tokens: int = 1024,
+    model: str = "claude-3-5-sonnet-20241022"
+) -> str:
+    """
+    Unified function to query Claude API with various options.
+    
+    Args:
+        query: Question/prompt for Claude
+        pdf_path: Optional path to PDF file (local or URL)
+        use_cache: Whether to enable prompt caching
+        system: Optional system prompt
+        max_tokens: Maximum tokens in response (default 1024)
+        model: Claude model to use (default claude-3-5-sonnet)
+        
+    Returns:
+        Claude's response as a string
+    """
+    client = anthropic.Anthropic(api_key=ANTHROPIC_API_KEY)
+    
+    # Handle PDF if provided
+    content = query
+    betas = []
+    if pdf_path:
+        # Get PDF content
+        if pdf_path.startswith(('http://', 'https://')):
+            response = requests.get(pdf_path)
+            binary_data = response.content
+        else:
+            with open(pdf_path, "rb") as pdf_file:
+                binary_data = pdf_file.read()
+        
+        pdf_data = base64.standard_b64encode(binary_data).decode("utf-8")
+        content = [
+            {
+                "type": "document",
+                "source": {
+                    "type": "base64",
+                    "media_type": "application/pdf",
+                    "data": pdf_data
+                }
+            },
+            {
+                "type": "text",
+                "text": query
+            }
+        ]
+        betas.append("pdfs-2024-09-25")
+    
+    # Add prompt caching if requested
+    if use_cache:
+        betas.append("prompt-caching-2024-07-31")
+    
+    # Prepare API call kwargs
+    kwargs = {
+        "model": model,
+        "max_tokens": max_tokens,
+        "messages": [{"role": "user", "content": content}]
+    }
+    
+    # Add optional parameters if provided
+    if system:
+        kwargs["system"] = system
+    if betas:
+        kwargs["betas"] = betas
+        
+    message = client.beta.messages.create(**kwargs)
+    return message.content[0].text
+
+class Point(BaseModel):
+    content: str
+    importance: Literal["critical", "minor"]
+
+
+class Review(BaseModel):
+    contributions: str
+    strengths: List[Point]
+    weaknesses: List[Point]
+    requested_changes: List[Point]
+    impact_concerns: str
+
+
+importance_mapping = {"critical": 2, "minor": 1}
+
+client = OpenAI(api_key=OPENAI_API_KEY)
+
+model_name = "gpt-4o-2024-08-06"
+
+def format_gpt(prompt):
+    chat_completion = client.beta.chat.completions.parse(
+        messages=[
+            {
+                "role": "user",
+                "content": prompt,
+            }
+        ],
+        model='gpt-4o',
+        response_format=Review,
+    )
+    return chat_completion.choices[0].message.parsed.model_dump()
+
+def parse_final(parsed, max_strengths=3, max_weaknesses=5, max_requested_changes=5):
+    new_parsed = {}
+    new_parsed["contributions"] = parsed["contributions"]
+    new_parsed["impact_concerns"] = parsed["impact_concerns"]
+    new_parsed["strengths"] = "\n".join(
+        [f'- {point["content"]}' for point in parsed["strengths"][:max_strengths]]
+    )
+    new_parsed["weaknesses"] = "\n".join(
+        [f'- {point["content"]}' for point in parsed["weaknesses"][:max_weaknesses]]
+    )
+    request_changes_sorted = sorted(
+        parsed["requested_changes"],
+        key=lambda x: importance_mapping[x["importance"]],
+        reverse=True,
+    )
+    new_parsed["requested_changes"] = "\n".join(
+        [
+            f"- {point['content']}"
+            for point in request_changes_sorted[:max_requested_changes]
+        ]
+    )
+    return new_parsed
+
+def process(file_content, progress=gr.Progress()):
+    # Create a list to store log messages
+    log_messages = []
+    
+    def log(msg):
+        print(msg)
+        log_messages.append(msg)
+        return "\n".join(log_messages)
+
+    if not os.path.exists("cache"):
+        os.makedirs("cache")
+    pdf_path = f"cache/{time.time()}.pdf"
+    with open(pdf_path, "wb") as f:
+        f.write(file_content)
+
+    progress(0, desc="Starting review process...")
+    log("Starting review process...")
+
+    all_reviews = []
+    for i in range(3):
+        progress((i + 1) / 3, desc=f"Generating review {i+1}/3")
+        log(f"Generating review {i+1}/3...")
+        all_reviews.append(ask_claude(prompts.review_prompt, pdf_path=pdf_path))
+    all_reviews_string = "\n\n".join([f"Review {i+1}:\n{review}" for i, review in enumerate(all_reviews)])
+    
+    progress(0.4, desc="Combining reviews...")
+    log("Combining reviews...")
+    combined_review = ask_claude(prompts.combine_prompt.format(all_reviews_string=all_reviews_string,
+                                                             review_format=prompts.review_format), pdf_path=pdf_path)
+
+    progress(0.6, desc="Defending paper...")
+    log("Defending paper...")
+    rebuttal = ask_claude(prompts.defend_prompt.format(combined_review=combined_review), pdf_path=pdf_path)
+
+    progress(0.8, desc="Revising review...")
+    log("Revising review...")
+    revised_review = ask_claude(prompts.revise_prompt.format(review_format=prompts.review_format, combined_review=combined_review, defended_paper=rebuttal), pdf_path=pdf_path)
+    log("Humanizing review...")
+    humanized_review = ask_claude(prompts.human_style.format(review=revised_review), pdf_path=pdf_path)
+
+    progress(0.9, desc="Formatting review...")
+    log("Formatting review...")
+    formatted_review = parse_final(format_gpt(prompts.formatting_prompt.format(review=humanized_review)))
+
+    log("Finished!")
+
+    contributions, strengths, weaknesses, requested_changes, impact_concerns = (
+                formatted_review["contributions"],
+                formatted_review["strengths"],
+                formatted_review["weaknesses"],
+                formatted_review["requested_changes"],
+                formatted_review["impact_concerns"],
+            )
+    contributions = f"# Contributions\n\n{contributions}"
+    strengths = f"# Strengths\n\n{strengths}"
+    weaknesses = f"# Weaknesses\n\n{weaknesses}"
+    requested_changes = f"# Requested Changes\n\n{requested_changes}"
+    impact_concerns = f"# Impact Concerns\n\n{impact_concerns}"
+    return (
+        contributions,
+        strengths,
+        weaknesses,
+        requested_changes,
+        impact_concerns,
+        "\n".join(log_messages),  # Return the log messages
+    )
+
+def gradio_interface():
+    with gr.Blocks() as demo:
+        gr.Markdown("# TMLR Reviewer")
+        gr.Markdown("This tool helps you generate high-quality reviews for the Transactions on Machine Learning Research (TMLR).")
+        
+        with gr.Row():
+            upload_component = gr.File(label="Upload PDF", type="binary")
+        
+        with gr.Row():
+            submit_btn = gr.Button("Generate Review")
+            
+        # Add progress log display
+        with gr.Row():
+            progress_log = gr.Textbox(label="Progress Log", interactive=False, lines=10)
+        
+        with gr.Row():
+            output_component_contributions = gr.Markdown(label="Contributions")
+            output_component_strengths = gr.Markdown(label="Strengths")
+            output_component_weaknesses = gr.Markdown(label="Weaknesses")
+            output_component_requested_changes = gr.Markdown(label="Requested Changes")
+            output_component_impact_concerns = gr.Markdown(label="Impact Concerns")
+
+        submit_btn.click(
+            fn=process,
+            inputs=upload_component,
+            outputs=[
+                output_component_contributions,
+                output_component_strengths,
+                output_component_weaknesses,
+                output_component_requested_changes,
+                output_component_impact_concerns,
+                progress_log,  # Add progress_log to outputs
+            ]
+        )
+        
+    demo.queue()
+    return demo
+
+if __name__ == "__main__":
+    demo = gradio_interface()
+    demo.launch(share=False)
+