import os import logging import openai import tiktoken import re import anthropic import cohere import google.generativeai as genai import time from file_utils import read_file from openai import OpenAI class Paper: def __init__(self, arxiv_id, tex_file): self.arxiv_id = arxiv_id self.tex_file = tex_file class PaperProcessor: MAX_TOKENS = 127192 encoding = tiktoken.encoding_for_model("gpt-4") def __init__(self, prompt_dir, model, openai_api_key, claude_api_key, gemini_api_key, commandr_api_key): self.prompt_dir = prompt_dir self.model = model self.openai_api_key = os.environ.get('OPENAI_API_KEY') self.claude_api_key = os.environ.get('ANTHROPIC_API_KEY') self.gemini_api_key = os.environ.get('GEMINI_API_KEY') self.commandr_api_key = os.environ.get('COMMANDR_API_KEY') def count_tokens(self, text): return len(self.encoding.encode(text)) def truncate_content(self, content): token_count = self.count_tokens(content) logging.debug(f"Token count before truncation: {token_count}") if token_count > self.MAX_TOKENS: tokens = self.encoding.encode(content) truncated_tokens = tokens[:self.MAX_TOKENS] truncated_content = self.encoding.decode(truncated_tokens) logging.debug(f"Content truncated. Token count after truncation: {self.count_tokens(truncated_content)}") return truncated_content return content def prepare_base_prompt(self, paper): logging.debug(f"Preparing base prompt for paper: {paper.arxiv_id}") logging.debug(f"Paper content: {paper.tex_file[:500]}") # Log the first 500 characters return paper.tex_file def call_model(self, prompt, model_type): system_role_file_path = os.path.join(self.prompt_dir, "systemrole.txt") if not os.path.exists(system_role_file_path): logging.error(f"System role file not found: {system_role_file_path}") return None system_role = read_file(system_role_file_path) logging.debug(f"Token count of full prompt: {self.count_tokens(prompt)}") logging.info(f"Sending the following prompt to {model_type}: {prompt}") try: if model_type == 'gpt-4-turbo-2024-04-09': client = OpenAI(api_key=self.openai_api_key) messages = [{"role": "system", "content": system_role}, {"role": "user", "content": prompt}] completion = client.chat.completions.create( model="gpt-4-turbo-2024-04-09", messages=messages, temperature=1 ) print(completion) return completion.choices[0].message.content.strip() elif model_type == 'gpt-4o': client = OpenAI(api_key=self.openai_api_key) messages = [{"role": "system", "content": system_role}, {"role": "user", "content": prompt}] completion = client.chat.completions.create( model="gpt-4o", messages=messages, temperature=1 ) print(completion) return completion.choices[0].message.content.strip() elif model_type == 'claude-3-opus-20240229': client = anthropic.Anthropic(api_key=self.claude_api_key) response = client.messages.create( model='claude-3-opus-20240229', max_tokens=4096, system=system_role, temperature=0.5, messages=[{"role": "user", "content": prompt}] ) print(response) return response.content[0].text elif model_type == 'command-r-plus': co = cohere.Client(self.commandr_api_key) response = co.chat( model="command-r-plus", message=prompt, preamble=system_role ) print(response) return response.text elif model_type == 'gemini-pro': genai.configure(api_key=self.gemini_api_key) model = genai.GenerativeModel('gemini-pro') response = model.generate_content(prompt) print(response) return response.candidates[0].content.parts[0].text except Exception as e: logging.error(f"Exception occurred: {e}") print(e) return None def is_content_appropriate(self, content): try: response = openai.moderations.create(input=content) return not response["results"][0]["flagged"] except Exception as e: logging.error(f"Exception occurred while checking content appropriateness: {e}") return True # In case of an error, default to content being appropriate def get_prompt_files(self, prompt_dir): return [f for f in os.listdir(prompt_dir) if f.endswith('.txt') and f.startswith('question')] def process_paper(self, paper): openai.api_key = self.openai_api_key start_time = time.time() base_prompt = self.prepare_base_prompt(paper) print("BASE PROMPT:", base_prompt) if base_prompt is None: return "Error: Base prompt could not be prepared." moderation_response = openai.moderations.create(input=base_prompt) if moderation_response.results[0].flagged: return ["Desk Rejected", "The paper contains inappropriate or harmful content."] review_output = [] previous_responses = [] header = ['Summary:', 'Soundness:', 'Presentation:', 'Contribution:', 'Strengths:', 'Weaknesses:', 'Questions:', 'Flag For Ethics Review:', 'Rating:', 'Confidence:', 'Code Of Conduct:'] for i in range(1, 12): question_file = os.path.join(self.prompt_dir, f"question{i}.txt") question_text = read_file(question_file) if i == 1: prompt = f"{question_text}\n\n####\n{base_prompt}\n####" else: prompt = f"\nHere is your review so far:\n{' '.join(previous_responses)}\n\nHere are your reviewer instructions. Please answer the following question:\n{question_text}" truncated_prompt = self.truncate_content(prompt) logging.info(f"Processing prompt for question {i}") response = self.call_model(truncated_prompt, self.model) if response is None: response = "N/A" if i in [2, 3, 4, 10]: number_match = re.search(r'\b\d+\b', response) if number_match: number = int(number_match.group(0)) response = '5/5' if number > 5 else number_match.group(0) + '/5' elif i == 9: number_match = re.search(r'\b\d+\b', response) if number_match: response = number_match.group(0) + '/10' response_with_header = f"{header[i-1]} {response}" review_output.append(response_with_header) previous_responses.append(response) end_time = time.time() elapsed_time = end_time - start_time print(f"Time taken to process paper: {elapsed_time:.2f} seconds") return review_output