ranjangoel/GPT-PDF

1from PyPDF2 import PdfReader2import openai3from .prompt import BASE_POINTS, READING_PROMT_V24from .paper import Paper5from .model_interface import OpenAIModel6 7 8# Setting the API key to use the OpenAI API9class PaperReader:10 11    """12    A class for summarizing research papers using the OpenAI API.13 14    Attributes:15        openai_key (str): The API key to use the OpenAI API.16        token_length (int): The length of text to send to the API at a time.17        model (str): The GPT model to use for summarization.18        points_to_focus (str): The key points to focus on while summarizing.19        verbose (bool): A flag to enable/disable verbose logging.20 21    """22 23    def __init__(self, openai_key, token_length=4000, model="gpt-3.5-turbo",24                 points_to_focus=BASE_POINTS, verbose=False):25 26        # Setting the API key to use the OpenAI API27        openai.api_key = openai_key28 29        # Initializing prompts for the conversation30        self.init_prompt = READING_PROMT_V2.format(points_to_focus)31 32        self.summary_prompt = 'You are a researcher helper bot. Now you need to read the summaries of a research paper.'33        self.messages = []  # Initializing the conversation messages34        self.summary_msg = []  # Initializing the summary messages35        self.token_len = token_length  # Setting the token length to use36        self.keep_round = 2  # Rounds of previous dialogues to keep in conversation37        self.model = model  # Setting the GPT model to use38        self.verbose = verbose  # Flag to enable/disable verbose logging39        self.model = OpenAIModel(api_key=openai_key, model=model)40 41    def drop_conversation(self, msg):42        # This method is used to drop previous messages from the conversation and keep only recent ones43        if len(msg) >= (self.keep_round + 1) * 2 + 1:44            new_msg = [msg[0]]45            for i in range(3, len(msg)):46                new_msg.append(msg[i])47            return new_msg48        else:49            return msg50 51    def send_msg(self, msg):52        return self.model.send_msg(msg)53 54    def _chat(self, message):55        # This method is used to send a message and get a response from the OpenAI API56 57        # Adding the user message to the conversation messages58        self.messages.append({"role": "user", "content": message})59        # Sending the messages to the API and getting the response60        response = self.send_msg(self.messages)61        # Adding the system response to the conversation messages62        self.messages.append({"role": "system", "content": response})63        # Dropping previous conversation messages to keep the conversation history short64        self.messages = self.drop_conversation(self.messages)65        # Returning the system response66        return response67 68    def summarize(self, paper: Paper):69        # This method is used to summarize a given research paper70 71        # Adding the initial prompt to the conversation messages72        self.messages = [73            {"role": "system", "content": self.init_prompt},74        ]75        # Adding the summary prompt to the summary messages76        self.summary_msg = [{"role": "system", "content": self.summary_prompt}]77 78         # Reading and summarizing each part of the research paper79        for (page_idx, part_idx, text) in paper.iter_pages():80            print('page: {}, part: {}'.format(page_idx, part_idx))81            # Sending the text to the API and getting the response82            summary = self._chat('now I send you page {}, part {}：{}'.format(page_idx, part_idx, text))83            # Logging the summary if verbose logging is enabled84            if self.verbose:85                print(summary)86            # Adding the summary of the part to the summary messages87            self.summary_msg.append({"role": "user", "content": '{}'.format(summary)})88 89        # Adding a prompt for the user to summarize the whole paper to the summary messages90        self.summary_msg.append({"role": "user", "content": 'Now please make a summary of the whole paper'})91        # Sending the summary messages to the API and getting the response92        result = self.send_msg(self.summary_msg)93        # Returning the summary of the whole paper94        return result95 96    def read_pdf_and_summarize(self, pdf_path):97        # This method is used to read a research paper from a PDF file and summarize it98        99        # Creating a PdfReader object to read the PDF file100        print(pdf_path)101        pdf_reader = PdfReader(pdf_path)102        pdf_data = ""103        for page in pdf_reader.pages:104            file_data += page.extract_text()105        print(file_data)106        107        paper = Paper(pdf_reader)108        # Summarizing the full text of the research paper and returning the summary109        print('reading pdf finished')110        summary = self.summarize(paper)111        return summary112 113    def get_summary_of_each_part(self):114        # This method is used to get the summary of each part of the research paper115        return self.summary_msg116 117    def question(self, question):118        # This method is used to ask a question after summarizing a paper119 120        # Adding the question to the summary messages121        self.summary_msg.append({"role": "user", "content": question})122        # Sending the summary messages to the API and getting the response123        response = self.send_msg(self.summary_msg)124        # Adding the system response to the summary messages125        self.summary_msg.append({"role": "system", "content": response})126        # Returning the system response127        return response128