added claude support + annotations display

app.py

@@ -1,8 +1,6 @@
 import streamlit as st
 import pandas as pd
 import ast
-import openai
-from openai import OpenAI
 import PyPDF2
 from PyPDF2 import PdfReader
 import openreview
@@ -11,9 +9,27 @@ import urllib.request
 import requests
 import os
 import pickle
-my_api_key = os.environ["OPENAI_API_KEY"]
-client_openai = OpenAI(api_key=my_api_key)
+import openai
+from openai import OpenAI
+import anthropic
+from anthropic import Anthropic
 
+openai_api_key = os.environ["OPENAI_API_KEY"]
+client_openai = OpenAI(api_key=openai_api_key)
+
+anthropic_api_key = os.environ["ANTHROPIC_API_KEY"]
+client_anthropic = Anthropic(api_key=anthropic_api_key)
+
+# load manual (human) annotations
+def load_annotations():
+    path = './annotations_8_26.csv'
+    annotations = pd.read_csv(path)
+    annotations.columns = annotations.columns.str.replace(' ', '_').str.lower()
+
+    # 36 0, 15 1
+    annotations['annotation_label'] = annotations['annotation_label'].astype(int)
+    
+    return annotations
 
 # Check if the password matches the env variable
 def check_password():
@@ -100,7 +116,7 @@ def parse_openreview_id(submissions, paper_id, reviewer_id):
     return review_content, pdf, pdf_text
 
 
-def create_feedback(review, pdf_text, agent_prompt):
+def create_feedback(review, pdf_text, agent_prompt, model):
 
     messages = [{
         "role": "user",
@@ -127,28 +143,33 @@ def create_feedback(review, pdf_text, agent_prompt):
       }
       ]}]
 
-    headers = {
-    "Content-Type": "application/json",
-    "Authorization": f"Bearer {my_api_key}"
-    }
-    payload = {
-    "model": "gpt-4o",
-    "messages": messages,
-    "max_tokens": 1000,
-    "temperature": 0.2
-    }
+    if model == 'gpt':
+        headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {openai_api_key}"
+        }
+        payload = {
+        "model": "gpt-4o",
+        "messages": messages,
+        "max_tokens": 1000,
+        "temperature": 0.2
+        }
+
+        try:
+            response = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=payload)
+            feedback = response.json()["choices"][0]["message"]["content"]
+        except Exception as e:
+            print(f"An unexpected error occurred: {e}")
+            feedback = "an error occured"
 
-    try:
-        response = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=payload)
-        feedback = response.json()["choices"][0]["message"]["content"]
-    except Exception as e:
-        print(f"An unexpected error occurred: {e}")
-        feedback = "an error occured"
+    else:
+        message = client_anthropic.messages.create(max_tokens=1024, messages=messages, model="claude-3-5-sonnet-20240620")
+        feedback = message.content[0].text
 
     return feedback
 
 
-def critic(review, feedback, pdf_text, critic_prompt):
+def critic(review, feedback, pdf_text, critic_prompt, model):
     
     messages = [{
         "role": "user",
@@ -183,23 +204,28 @@ def critic(review, feedback, pdf_text, critic_prompt):
       }
       ]}]
 
-    headers = {
-    "Content-Type": "application/json",
-    "Authorization": f"Bearer {my_api_key}"
-    }
-    payload = {
-    "model": "gpt-4o",
-    "messages": messages,
-    "max_tokens": 1000,
-    "temperature": 0.2
-    }
+    if model == 'gpt':
+        headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {openai_api_key}"
+        }
+        payload = {
+        "model": "gpt-4o",
+        "messages": messages,
+        "max_tokens": 1000,
+        "temperature": 0.2
+        }
+
+        try:
+            response = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=payload)
+            revised_feedback = response.json()["choices"][0]["message"]["content"]
+        except Exception as e:
+            print(f"An unexpected error occurred: {e}")
+            revised_feedback = "an error occured"
 
-    try:
-        response = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=payload)
-        revised_feedback = response.json()["choices"][0]["message"]["content"]
-    except Exception as e:
-        print(f"An unexpected error occurred: {e}")
-        revised_feedback = "an error occured"
+    else:
+        message = client_anthropic.messages.create(max_tokens=1024, messages=messages, model="claude-3-5-sonnet-20240620")
+        revised_feedback = message.content[0].text
 
     return revised_feedback
 
@@ -385,6 +411,7 @@ if option == "OpenReview paper ID":
     paper_id = paper_id.strip()
     reviewer_id = st.text_input("Enter reviewer ID (e.g. gNxe):")
     reviewer_id = reviewer_id.strip()
+    model = st.text_input("Which model? ('GPT' or 'Claude')")
     if paper_id and reviewer_id:
         upload_file = False
         user_input = True
@@ -395,6 +422,7 @@ if option == "OpenReview paper ID":
 elif option == "Upload PDF with Review":
     user_text = st.text_area("Enter review:")
     uploaded_file = st.file_uploader("Upload PDF", type="pdf")
+    model = st.text_input("Which model? ('GPT' or 'Claude')")
     if user_text and uploaded_file:
         upload_file = True
         user_input = True
@@ -406,34 +434,68 @@ elif option == "Upload PDF with Review":
         # st.write(f"Uploaded file: {uploaded_file.name}")
 
 if user_input:
+    annotations = load_annotations()
+    
     # Run your pipeline to generate the dataframe based on user input
-    feedback = create_feedback(review, pdf_text, agent_prompt)
+    feedback = create_feedback(review, pdf_text, agent_prompt, model)
 
-    revised_feedback = critic(review, feedback, pdf_text, critic_prompt)
+    revised_feedback = critic(review, feedback, pdf_text, critic_prompt, model)
 
     st.title(f'Review feedback')
 
-    # Create three columns
-    col1, col2, col3 = st.columns(3)
-
-    # Column 1: Display review fields
-    with col1:
-        st.subheader("Review")
-        st.write(json.dumps(review))
-
-        if not upload_file:
-            pdf_url = f"https://openreview.net/pdf?id={paper_id}"
-            st.markdown(f"[Click here to view PDF]({pdf_url})")
-
-    # Column 2: Display feedback
-    with col2:
-        st.subheader('Feedback')
-        st.write(feedback)
+    if paper_id in annotations['submission_id']:
+        # Create four columns
+        col1, col2, col3, col4 = st.columns(4)
+    
+        # Column 1: Display review fields
+        with col1:
+            st.subheader("Review")
+            st.write(json.dumps(review))
+    
+            if not upload_file:
+                pdf_url = f"https://openreview.net/pdf?id={paper_id}"
+                st.markdown(f"[Click here to view PDF]({pdf_url})")
+    
+        # Column 2: Display human annotations
+        with col2:
+            st.subheader('Human annotations')
+            if annotations['annotation_label'] == 1:
+                st.write(annotations['annotation_sentences'])
+            else:
+                st.write("No human feedback on this review")
+        
+        # Column 3: Display feedback
+        with col3:
+            st.subheader('Feedback')
+            st.write(feedback)
+    
+        # Column 4: Display revised feedback (from critic)
+        with col4:
+            st.subheader('Revised Feedback')
+            st.write(revised_feedback)
+    else:
 
-    # Column 3: Display revised feedback (from critic)
-    with col3:
-        st.subheader('Revised Feedback')
-        st.write(revised_feedback)
+        # Create three columns
+        col1, col2, col3 = st.columns(3)
+    
+        # Column 1: Display review fields
+        with col1:
+            st.subheader("Review")
+            st.write(json.dumps(review))
+    
+            if not upload_file:
+                pdf_url = f"https://openreview.net/pdf?id={paper_id}"
+                st.markdown(f"[Click here to view PDF]({pdf_url})")
+    
+        # Column 2: Display feedback
+        with col2:
+            st.subheader('Feedback')
+            st.write(feedback)
+    
+        # Column 3: Display revised feedback (from critic)
+        with col3:
+            st.subheader('Revised Feedback')
+            st.write(revised_feedback)
 
 else:
     st.title('Please enter OpenReview ID or upload PDF with review to generate feedback')