Upload app.py

app.py

@@ -3,7 +3,6 @@ import spaces
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 import torch
 from threading import Thread
-import re
 
 from marker.convert import convert_single_pdf
 from marker.output import markdown_exists, save_markdown, get_markdown_filepath
@@ -38,6 +37,7 @@ model = AutoModelForCausalLM.from_pretrained(
     device_map="auto"
 )
 tokenizer = AutoTokenizer.from_pretrained(model_name)
+streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, decode_kwargs=dict(skip_special_tokens=True))
 
 # 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.
@@ -87,7 +87,8 @@ Please assign the paper a numerical rating on the following scale to indicate th
 2: fair
 1: poor
 
-**Justification for Soundness Rating:** Please provide specific reasons for your soundness score, explaining which aspects of the technical claims, methodology, or evidence support your rating.
+## Soundness Explanation
+Please provide specific reasons for your soundness score, explaining which aspects of the technical claims, methodology, or evidence support your rating.
 
 ## 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:
@@ -96,7 +97,8 @@ Please assign the paper a numerical rating on the following scale to indicate th
 2: fair
 1: poor
 
-**Justification for Presentation Rating:** Please explain your presentation score by addressing specific aspects of writing clarity, organization, figure quality, and how well the work is positioned within existing literature.
+## Presentation Explanation
+Please explain your presentation score by addressing specific aspects of writing clarity, organization, figure quality, and how well the work is positioned within existing literature.
 
 ## 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:
@@ -105,7 +107,8 @@ Please assign the paper a numerical rating on the following scale to indicate th
 2: fair
 1: poor
 
-**Justification for Contribution Rating:** Please justify your contribution score by explaining the significance of the research questions, the originality of the approach, and the potential impact on the field.
+## Contribution Explanation
+Please justify your contribution score by explaining the significance of the research questions, the originality of the approach, and the potential impact on the field.
 
 ## 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.
@@ -140,16 +143,10 @@ Please provide an "overall score" for this submission. Choose from the following
 8: accept, good paper
 10: strong accept, should be highlighted at the conference
 
-**Justification for Overall Rating:** Please provide a comprehensive justification for your overall rating, synthesizing the key strengths and weaknesses that led to this decision. Explain how you weighted different aspects of the paper in arriving at your final score.
+Please provide a comprehensive justification for your overall rating, synthesizing the key strengths and weaknesses that led to this decision. Explain how you weighted different aspects of the paper in arriving at your final score.
 
 """
 
-JUSTIFICATION_PROMPT = """Based on the review above, please provide detailed justifications for each numerical rating. For each rating section, explain WHY you gave that specific score:
-
-{missing_justifications}
-
-Please provide substantive explanations that reference specific aspects of the paper. Each justification should be 2-4 sentences explaining your reasoning."""
-
 # functions
 def create_messages(review_fields, paper_text):
     messages = [
@@ -158,100 +155,6 @@ def create_messages(review_fields, paper_text):
     ]
     return messages
 
-def extract_ratings(review_text):
-    """Extract numerical ratings from the review"""
-    ratings = {}
-    
-    # Pattern to find ratings
-    soundness_pattern = r'## Soundness\s*\n.*?(\d):\s*\w+'
-    presentation_pattern = r'## Presentation\s*\n.*?(\d):\s*\w+'
-    contribution_pattern = r'## Contribution\s*\n.*?(\d):\s*\w+'
-    overall_pattern = r'## Rating\s*\n.*?(\d+):\s*[^#]+'
-    
-    soundness_match = re.search(soundness_pattern, review_text, re.IGNORECASE | re.DOTALL)
-    if soundness_match:
-        ratings['soundness'] = soundness_match.group(1)
-    
-    presentation_match = re.search(presentation_pattern, review_text, re.IGNORECASE | re.DOTALL)
-    if presentation_match:
-        ratings['presentation'] = presentation_match.group(1)
-    
-    contribution_match = re.search(contribution_pattern, review_text, re.IGNORECASE | re.DOTALL)
-    if contribution_match:
-        ratings['contribution'] = contribution_match.group(1)
-    
-    overall_match = re.search(overall_pattern, review_text, re.IGNORECASE | re.DOTALL)
-    if overall_match:
-        ratings['overall'] = overall_match.group(1)
-    
-    return ratings
-
-def check_for_justifications(review_text):
-    """Check which justifications are missing"""
-    missing = []
-    
-    # Check for each justification
-    if "justification for soundness" not in review_text.lower() or \
-       not re.search(r'justification for soundness.*?:\s*.{20,}', review_text, re.IGNORECASE | re.DOTALL):
-        missing.append("Soundness")
-    
-    if "justification for presentation" not in review_text.lower() or \
-       not re.search(r'justification for presentation.*?:\s*.{20,}', review_text, re.IGNORECASE | re.DOTALL):
-        missing.append("Presentation")
-    
-    if "justification for contribution" not in review_text.lower() or \
-       not re.search(r'justification for contribution.*?:\s*.{20,}', review_text, re.IGNORECASE | re.DOTALL):
-        missing.append("Contribution")
-    
-    if "justification for overall rating" not in review_text.lower() or \
-       not re.search(r'justification for overall rating.*?:\s*.{20,}', review_text, re.IGNORECASE | re.DOTALL):
-        missing.append("Overall Rating")
-    
-    return missing
-
-def insert_justifications(original_review, justifications_text):
-    """Insert the generated justifications into the appropriate places in the review"""
-    review = original_review
-    
-    # Extract individual justifications from the response
-    justification_dict = {}
-    
-    # Parse justifications for each category
-    patterns = {
-        'soundness': r'(?:soundness|Soundness).*?justification.*?:(.*?)(?=\n\n|\n(?:Presentation|Contribution|Overall|$))',
-        'presentation': r'(?:presentation|Presentation).*?justification.*?:(.*?)(?=\n\n|\n(?:Contribution|Overall|$))',
-        'contribution': r'(?:contribution|Contribution).*?justification.*?:(.*?)(?=\n\n|\n(?:Overall|$))',
-        'overall': r'(?:overall rating|Overall Rating).*?justification.*?:(.*?)(?=\n\n|$)'
-    }
-    
-    for key, pattern in patterns.items():
-        match = re.search(pattern, justifications_text, re.IGNORECASE | re.DOTALL)
-        if match:
-            justification_dict[key] = match.group(1).strip()
-    
-    # Insert justifications into the review
-    if 'soundness' in justification_dict:
-        pattern = r'(## Soundness.*?\d:\s*\w+)\n'
-        replacement = f'\\1\n\n**Justification for Soundness Rating:** {justification_dict["soundness"]}\n'
-        review = re.sub(pattern, replacement, review, flags=re.DOTALL)
-    
-    if 'presentation' in justification_dict:
-        pattern = r'(## Presentation.*?\d:\s*\w+)\n'
-        replacement = f'\\1\n\n**Justification for Presentation Rating:** {justification_dict["presentation"]}\n'
-        review = re.sub(pattern, replacement, review, flags=re.DOTALL)
-    
-    if 'contribution' in justification_dict:
-        pattern = r'(## Contribution.*?\d:\s*\w+)\n'
-        replacement = f'\\1\n\n**Justification for Contribution Rating:** {justification_dict["contribution"]}\n'
-        review = re.sub(pattern, replacement, review, flags=re.DOTALL)
-    
-    if 'overall' in justification_dict:
-        pattern = r'(## Rating.*?\d+:\s*[^#]+)\n'
-        replacement = f'\\1\n\n**Justification for Overall Rating:** {justification_dict["overall"]}\n'
-        review = re.sub(pattern, replacement, review, flags=re.DOTALL)
-    
-    return review
-
 @spaces.GPU()
 def convert_file(filepath):
     full_text, images, out_metadata = convert_single_pdf(
@@ -278,6 +181,10 @@ def process_file(file):
     except spaces.zero.gradio.HTMLError as e:
         print(e)
         return "Error. GPU quota exceeded. Please return later."
+    # except gradio.exceptions.Error as e:
+    #     if 'GPU task aborted' in str(e):
+    #         print(e)
+    #         return 'GPU task aborted'
     except Exception as e:
         print(traceback.format_exc())
         print(f"Error converting {filepath}: {e}")
@@ -285,92 +192,33 @@ def process_file(file):
     return paper_text
 
 
-@spaces.GPU(duration=120)
+@spaces.GPU(duration=190)
 def generate(paper_text, review_template):
-    # First generate the main review
     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"Generating initial review...")
-    
-    # Generate initial review
-    outputs = model.generate(
-        input_ids,
-        max_new_tokens=4096,
-        do_sample=True,
-        temperature=0.6,
-        top_p=0.9,
-        pad_token_id=tokenizer.eos_token_id
-    )
-    
-    initial_review = tokenizer.decode(outputs[0][input_ids.shape[-1]:], skip_special_tokens=True)
-    
-    # Check for missing justifications
-    missing_justifications = check_for_justifications(initial_review)
-    
-    if missing_justifications:
-        print(f"Missing justifications for: {missing_justifications}")
-        print("Generating justifications...")
-        
-        # Extract ratings from the review
-        ratings = extract_ratings(initial_review)
-        
-        # Build the request for missing justifications
-        missing_text = ""
-        if "Soundness" in missing_justifications and 'soundness' in ratings:
-            missing_text += f"\n- Soundness (you rated it {ratings['soundness']}): Explain why you gave this soundness score."
-        if "Presentation" in missing_justifications and 'presentation' in ratings:
-            missing_text += f"\n- Presentation (you rated it {ratings['presentation']}): Explain why you gave this presentation score."
-        if "Contribution" in missing_justifications and 'contribution' in ratings:
-            missing_text += f"\n- Contribution (you rated it {ratings['contribution']}): Explain why you gave this contribution score."
-        if "Overall Rating" in missing_justifications and 'overall' in ratings:
-            missing_text += f"\n- Overall Rating (you rated it {ratings['overall']}): Explain why you gave this overall rating."
-        
-        # Create follow-up message
-        follow_up_messages = messages + [
-            {"role": "assistant", "content": initial_review},
-            {"role": "user", "content": JUSTIFICATION_PROMPT.format(missing_justifications=missing_text)}
-        ]
-        
-        follow_up_input_ids = tokenizer.apply_chat_template(
-            follow_up_messages,
-            add_generation_prompt=True,
-            return_tensors='pt'
-        ).to(model.device)
-        
-        # Generate justifications
-        justification_outputs = model.generate(
-            follow_up_input_ids,
-            max_new_tokens=1024,
-            do_sample=True,
-            temperature=0.6,
-            top_p=0.9,
-            pad_token_id=tokenizer.eos_token_id
-        )
-        
-        justifications = tokenizer.decode(
-            justification_outputs[0][follow_up_input_ids.shape[-1]:], 
-            skip_special_tokens=True
-        )
-        
-        # Combine the initial review with justifications
-        final_review = insert_justifications(initial_review, justifications)
-        
-        # Use streaming for the final output
-        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, decode_kwargs=dict(skip_special_tokens=True))
-        for chunk in final_review.split():
-            yield " ".join([chunk])
-        yield final_review
-    else:
-        # If justifications are already present, return the initial review
-        yield initial_review
+    print(f"input_ids shape: {input_ids.shape}")
+    generation_kwargs = dict(input_ids=input_ids, streamer=streamer, max_new_tokens=4096, do_sample=True, temperature=0.6, top_p=0.9)
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+    generated_text = ""
+    for new_text in streamer:
+        generated_text += new_text
+        yield generated_text.replace("<|eot_id|>", "")
+
 
 
 # ui
+#8C1B13 red
+#4D8093 blue
+#767676 med grey
+#EFECE3 light grey
+#DDDDDD silver below red
+#FFFDFA white
+
 title = """<h1 align="center">OpenReviewer</h1>
 <div align="center">Using <a href="https://huggingface.co/maxidl/Llama-OpenReviewer-8B" target="_blank"><code>Llama-OpenReviewer-8B</code></a> - Built with Llama</div>
 """
@@ -389,8 +237,6 @@ Take a look at the Review Template to properly interpret the generated review. Y
 
 To obtain more than one review, just generate again.
 
-**Note:** The system will automatically add justifications for all numerical ratings if they are not initially provided.
-
 **GPU quota:** If exceeded, either sign in with your HF account or come back later. Your quota has a half-life of 2 hours.
 
 """
@@ -414,5 +260,7 @@ with gr.Blocks(theme=theme) as demo:
     demo.title = "OpenReviewer"
 
 
+
+
 if __name__ == "__main__":
     demo.launch()