Spaces:
Running
Running
| from PyPDF2 import PdfReader | |
| import openai | |
| from .prompt import BASE_POINTS, READING_PROMT_V2 | |
| from .paper import Paper | |
| from .model_interface import OpenAIModel | |
| # Setting the API key to use the OpenAI API | |
| class PaperReader: | |
| """ | |
| A class for summarizing research papers using the OpenAI API. | |
| Attributes: | |
| openai_key (str): The API key to use the OpenAI API. | |
| token_length (int): The length of text to send to the API at a time. | |
| model (str): The GPT model to use for summarization. | |
| points_to_focus (str): The key points to focus on while summarizing. | |
| verbose (bool): A flag to enable/disable verbose logging. | |
| """ | |
| def __init__(self, openai_key, token_length=4000, model="gpt-3.5-turbo", | |
| points_to_focus=BASE_POINTS, verbose=False): | |
| # Setting the API key to use the OpenAI API | |
| openai.api_key = openai_key | |
| # Initializing prompts for the conversation | |
| self.init_prompt = READING_PROMT_V2.format(points_to_focus) | |
| self.summary_prompt = 'You are a researcher helper bot. Now you need to read the summaries of a research paper.' | |
| self.messages = [] # Initializing the conversation messages | |
| self.summary_msg = [] # Initializing the summary messages | |
| self.token_len = token_length # Setting the token length to use | |
| self.keep_round = 2 # Rounds of previous dialogues to keep in conversation | |
| self.model = model # Setting the GPT model to use | |
| self.verbose = verbose # Flag to enable/disable verbose logging | |
| self.model = OpenAIModel(api_key=openai_key, model=model) | |
| def drop_conversation(self, msg): | |
| # This method is used to drop previous messages from the conversation and keep only recent ones | |
| if len(msg) >= (self.keep_round + 1) * 2 + 1: | |
| new_msg = [msg[0]] | |
| for i in range(3, len(msg)): | |
| new_msg.append(msg[i]) | |
| return new_msg | |
| else: | |
| return msg | |
| def send_msg(self, msg): | |
| return self.model.send_msg(msg) | |
| def _chat(self, message): | |
| # This method is used to send a message and get a response from the OpenAI API | |
| # Adding the user message to the conversation messages | |
| self.messages.append({"role": "user", "content": message}) | |
| # Sending the messages to the API and getting the response | |
| response = self.send_msg(self.messages) | |
| # Adding the system response to the conversation messages | |
| self.messages.append({"role": "system", "content": response}) | |
| # Dropping previous conversation messages to keep the conversation history short | |
| self.messages = self.drop_conversation(self.messages) | |
| # Returning the system response | |
| return response | |
| def summarize(self, paper: Paper): | |
| # This method is used to summarize a given research paper | |
| # Adding the initial prompt to the conversation messages | |
| self.messages = [ | |
| {"role": "system", "content": self.init_prompt}, | |
| ] | |
| # Adding the summary prompt to the summary messages | |
| self.summary_msg = [{"role": "system", "content": self.summary_prompt}] | |
| # Reading and summarizing each part of the research paper | |
| for (page_idx, part_idx, text) in paper.iter_pages(): | |
| print('page: {}, part: {}'.format(page_idx, part_idx)) | |
| # Sending the text to the API and getting the response | |
| summary = self._chat('now I send you page {}, part {}：{}'.format(page_idx, part_idx, text)) | |
| # Logging the summary if verbose logging is enabled | |
| if self.verbose: | |
| print(summary) | |
| # Adding the summary of the part to the summary messages | |
| self.summary_msg.append({"role": "user", "content": '{}'.format(summary)}) | |
| # Adding a prompt for the user to summarize the whole paper to the summary messages | |
| self.summary_msg.append({"role": "user", "content": 'Now please make a summary of the whole paper'}) | |
| # Sending the summary messages to the API and getting the response | |
| result = self.send_msg(self.summary_msg) | |
| # Returning the summary of the whole paper | |
| return result | |
| def read_pdf_and_summarize(self, pdf_path): | |
| # This method is used to read a research paper from a PDF file and summarize it | |
| # Creating a PdfReader object to read the PDF file | |
| pdf_reader = PdfReader(pdf_path) | |
| paper = Paper(pdf_reader) | |
| # Summarizing the full text of the research paper and returning the summary | |
| print('reading pdf finished') | |
| summary = self.summarize(paper) | |
| return summary | |
| def get_summary_of_each_part(self): | |
| # This method is used to get the summary of each part of the research paper | |
| return self.summary_msg | |
| def question(self, question): | |
| # This method is used to ask a question after summarizing a paper | |
| # Adding the question to the summary messages | |
| self.summary_msg.append({"role": "user", "content": question}) | |
| # Sending the summary messages to the API and getting the response | |
| response = self.send_msg(self.summary_msg) | |
| # Adding the system response to the summary messages | |
| self.summary_msg.append({"role": "system", "content": response}) | |
| # Returning the system response | |
| return response | |