Update app.py

app.py

@@ -1,3 +1,9 @@
+"""
+This module provides functions for working with PDF files and URLs. It uses the urllib.request library
+to download files from URLs, and the fitz library to extract text from PDF files. And GPT3 modules to generate
+text completions.
+"""
+
 import urllib.request
 import fitz
 import re
@@ -7,18 +13,55 @@ import openai
 import gradio as gr
 import os
 from sklearn.neighbors import NearestNeighbors
+from typing import Optional, Union, IO
+
+
+def download_pdf(url: str, output_path: str) -> None:
+    """
+    Downloads a PDF file from the given URL and saves it to the specified output path.
+
+    Args:
+    url (str): The URL of the PDF file to be downloaded.
+    output_path (str): The file path where the downloaded PDF file will be saved.
 
-def download_pdf(url, output_path):
+    Returns:
+    None
+    """
     urllib.request.urlretrieve(url, output_path)
 
 
-def preprocess(text):
+def preprocess(text: str) -> str:
+    """
+    Preprocesses the given text by replacing newline characters with spaces and removing extra whitespaces.
+
+    Args:
+    text (str): The input text to be preprocessed.
+
+    Returns:
+    str: The preprocessed text with newline characters replaced by spaces and extra whitespaces removed.
+
+    Example:
+        >>> preprocess("Hello\\n   world!")
+        'Hello world!'
+    """
     text = text.replace('\n', ' ')
-    text = re.sub('\s+', ' ', text)
+    text = re.sub(r'\s+', ' ', text)
     return text
 
 
-def pdf_to_text(path, start_page=1, end_page=None):
+def pdf_to_text(path: str, start_page: int = 1, end_page: Optional[int] = None) -> list[str]:
+    """
+    Converts a PDF file to a list of text strings.
+
+    Args:
+        path (str): The path to the PDF file.
+        start_page (int): The page number to start extracting text from (default is 1).
+        end_page (int): Page number to stop extracting text at (default is None, which means extract text from all ]
+        pages)
+
+    Returns:
+        list: A list of text strings extracted from the PDF file.
+    """
     doc = fitz.open(path)
     total_pages = doc.page_count
 
@@ -36,16 +79,33 @@ def pdf_to_text(path, start_page=1, end_page=None):
     return text_list
 
 
-def text_to_chunks(texts, word_length=150, start_page=1):
+def text_to_chunks(texts: list[str], word_length: int = 150, start_page: int = 1) -> list[str]:
+    """
+    Splits a list of texts into chunks of specified length and formats them as strings.
+
+    Args:
+    - texts: A list of strings to be split into chunks.
+    - word_length: An integer representing the maximum number of words in each chunk. Default is 150.
+    - start_page: An integer representing the starting page number. Default is 1.
+
+    Returns:
+    - A list of formatted string chunks, where each chunk contains a page number, enclosed in square brackets,
+    followed by the chunk of text enclosed in double quotes.
+
+    Example:
+    >>> texts = ['This is a sample text for testing the function.', 'It should split the text into chunks of 5 words.']
+    >>> text_to_chunks(texts, word_length=5, start_page=3)
+    ['[3] "This is a sample text for"', '[3] "testing the function. It should"',
+        '[4] "split the text into chunks of"','[4] "5 words."']
+    """
     text_toks = [t.split(' ') for t in texts]
-    page_nums = []
     chunks = []
-    
+
     for idx, words in enumerate(text_toks):
         for i in range(0, len(words), word_length):
             chunk = words[i:i+word_length]
             if (i+word_length) > len(words) and (len(chunk) < word_length) and (
-                len(text_toks) != (idx+1)):
+                    len(text_toks) != (idx+1)):
                 text_toks[idx+1] = chunk + text_toks[idx+1]
                 continue
             chunk = ' '.join(chunk).strip()
@@ -54,33 +114,78 @@ def text_to_chunks(texts, word_length=150, start_page=1):
     return chunks
 
 
-class SemanticSearch:
-    
-    def __init__(self):
-        self.use = hub.load('https://tfhub.dev/google/universal-sentence-encoder/4')
+class SemanticSearch(object):
+    """
+    This class provides functionality for semantic search.
+    """
+
+    def __init__(self) -> None:
+        """
+        Initializes an instance of the class.
+
+        Attributes:
+        -----------
+        use : tensorflow_hub.KerasLayer
+            A pre-trained Universal Sentence Encoder model from TensorFlow Hub.
+        fitted : bool
+            A flag indicating whether the model has been fitted to data or not.
+        """
+        self.use = hub.load(
+            'https://tfhub.dev/google/universal-sentence-encoder/4')
         self.fitted = False
-    
-    
-    def fit(self, data, batch=1000, n_neighbors=5):
+
+    def fit(self, data: list[str], batch: int = 1000, n_neighbors: int = 5) -> None:
+        """
+        Fits the nearest neighbor model to the given data.
+
+        Args:
+            data (list[str]): A list of strings to fit the model on.
+            batch (int): The batch size to use when computing text embeddings. Defaults to 1000.
+            n_neighbors (int): The number of nearest neighbors to find for each query. Defaults to 5.
+
+        Returns:
+            None
+        """
         self.data = data
         self.embeddings = self.get_text_embedding(data, batch=batch)
         n_neighbors = min(n_neighbors, len(self.embeddings))
         self.nn = NearestNeighbors(n_neighbors=n_neighbors)
         self.nn.fit(self.embeddings)
         self.fitted = True
-    
-    
-    def __call__(self, text, return_data=True):
+
+    def __call__(self, text: str, return_data: bool = True) -> Union[list[str], np.ndarray]:
+        """
+        Finds nearest neighbors to a given text in the embedding space.
+
+        Args:
+            text (str): The input text to find nearest neighbors for.
+            return_data (bool): Whether to return the actual data points corresponding to the nearest neighbors.
+                If False, returns only the indices of the nearest neighbors. Defaults to True.
+
+        Returns:
+            Union[List[str], np.ndarray]: If return_data is True, returns a list of strings representing the
+            nearest neighbors. If return_data is False, returns a numpy array of shape (n_neighbors,)
+            containing the indices of the nearest neighbors.
+        """
         inp_emb = self.use([text])
         neighbors = self.nn.kneighbors(inp_emb, return_distance=False)[0]
-        
+
         if return_data:
             return [self.data[i] for i in neighbors]
         else:
             return neighbors
-    
-    
-    def get_text_embedding(self, texts, batch=1000):
+
+    def get_text_embedding(self, texts: list[str], batch: int = 1000) -> np.ndarray:
+        """
+        Generates embeddings for a list of texts using the Universal Sentence Encoder.
+
+        Args:
+            texts (List[str]): A list of strings to generate embeddings for.
+            batch (int): The batch size to use when generating embeddings. Defaults to 1000.
+
+        Returns:
+            np.ndarray: An array of shape (n_texts, embedding_size) containing the embeddings for each text.
+        """
         embeddings = []
         for i in range(0, len(texts), batch):
             text_batch = texts[i:(i+batch)]
@@ -90,16 +195,47 @@ class SemanticSearch:
         return embeddings
 
 
+def load_recommender(path: str, start_page: int = 1) -> str:
+    """
+    Loads embeddings from file if available, otherwise generates embeddings and saves them to file.
+
+    Args:
+    path (str): The path of the PDF file.
+    start_page (int): The page number to start generating embeddings from. Default is 1.
 
-def load_recommender(path, start_page=1):
+    Returns:
+    str: A message indicating whether embeddings were loaded from file or generated and saved to file.
+    """
     global recommender
+    pdf_file = os.path.basename(path)
+    embeddings_file = f"{pdf_file}_{start_page}.npy"
+
+    if os.path.isfile(embeddings_file):
+        embeddings = np.load(embeddings_file)
+        recommender.embeddings = embeddings
+        recommender.fitted = True
+        return "Embeddings loaded from file"
+
     texts = pdf_to_text(path, start_page=start_page)
     chunks = text_to_chunks(texts, start_page=start_page)
     recommender.fit(chunks)
+    np.save(embeddings_file, recommender.embeddings)
     return 'Corpus Loaded.'
 
-def generate_text(openAI_key,prompt, engine="text-davinci-003"):
-    openai.api_key = openAI_key
+
+def generate_text(openai_key: str, prompt: str, engine: str = "text-davinci-003") -> str:
+    """
+    Generates text using OpenAI's GPT-3 language model.
+
+    Parameters:
+    openai_key (str): The API key for accessing OpenAI's API.
+    prompt (str): The starting text prompt to generate the text from.
+    engine (str): The ID of the language model to use. Defaults to "text-davinci-003".
+
+    Returns:
+    str: The generated text based on the given prompt.
+    """
+    openai.api_key = openai_key
     completions = openai.Completion.create(
         engine=engine,
         prompt=prompt,
@@ -112,60 +248,100 @@ def generate_text(openAI_key,prompt, engine="text-davinci-003"):
     return message
 
 
-def generate_answer(question,openAI_key):
+def generate_text2(openai_key: str, prompt: str, engine: str = "gpt-3.5-turbo-0301") -> str:
+    """
+    Generates text using OpenAI's GPT-3 language model.
+
+    Args:
+        openai_key (str): The API key for accessing OpenAI's GPT-3 language model.
+        prompt (str): The user's prompt to generate a response to.
+        engine (str, optional): The name of the GPT-3 engine to use. Defaults to "gpt-3.5-turbo-0301".
+
+    Returns:
+        str: The generated text response from the GPT-3 language model.
+    """
+    openai.api_key = openai_key
+    messages = [{'role': 'system', 'content': 'You are a helpful assistant.'},
+                {'role': 'user', 'content': prompt}]
+
+    completions = openai.ChatCompletion.create(
+        model=engine,
+        messages=messages,
+        max_tokens=512,
+        n=1,
+        stop=None,
+        temperature=0.7,
+    )
+    message = completions.choices[0].message['content']
+    return message
+
+
+def generate_answer(question: str, openai_key: str) -> str:
+    """
+    Generates an answer to the given question using OpenAI's GPT-3 language model.
+
+    Args:
+        question (str): The question to answer.
+        openai_key (str): The API key for accessing OpenAI's GPT-3 API.
+
+    Returns:
+        str: The generated answer to the question.
+    """
     topn_chunks = recommender(question)
     prompt = ""
     prompt += 'search results:\n\n'
     for c in topn_chunks:
         prompt += c + '\n\n'
-        
+
     prompt += "Instructions: Compose a comprehensive reply to the query using the search results given. "\
-              "Cite each reference using [ Page Number] notation (every result has this number at the beginning). "\
-              "Citation should be done at the end of each sentence. If the search results mention multiple subjects "\
-              "with the same name, create separate answers for each. Only include information found in the results and "\
-              "don't add any additional information. Make sure the answer is correct and don't output false content. "\
-              "If the text does not relate to the query, simply state 'Text Not Found in PDF'. Ignore outlier "\
-              "search results which has nothing to do with the question. Only answer what is asked. The "\
-              "answer should be short and concise. Answer step-by-step. \n\nQuery: {question}\nAnswer: "
-    
+        "Cite each reference using [ Page Number] notation (every result has this number at the beginning). "\
+        "Citation should be done at the end of each sentence. If the search results mention multiple subjects "\
+        "with the same name, create separate answers for each. Only include information found in the results and "\
+        "don't add any additional information. Make sure the answer is correct and don't output false content. "\
+        "If the text does not relate to the query, simply state 'Text Not Found in PDF'. Ignore outlier "\
+        "search results which has nothing to do with the question. Only answer what is asked. The "\
+        "answer should be short and concise. Answer step-by-step. \n\nQuery: {question}\nAnswer: "
+
     prompt += f"Query: {question}\nAnswer:"
-    answer = generate_text(openAI_key, prompt,"text-davinci-003")
+    answer = generate_text(openai_key, prompt, "text-davinci-003")
     return answer
 
 
-def question_answer(url, file, question,openAI_key):
-    if openAI_key.strip()=='':
+def question_answer(url: str, file: IO[str], question: str, openai_key: str) -> str:
+    """
+    Generates an answer to a given question using OpenAI's GPT-3 model.
+
+    Parameters:
+    -----------
+    url : str
+        The URL of a webpage to extract text from. If provided, the text will be saved as a PDF and used
+        as input for the model.
+    file : file-like object
+        A file object containing a PDF document to use as input for the model. If provided, the text will
+        be extracted from the PDF and used as input for the model.
+    question : str
+        The question to generate an answer for.
+    openai_key : str
+        An API key for accessing OpenAI's GPT-3 model.
+
+    Returns:
+    --------
+    str
+        The generated answer to the given question.
+
+    Raises:
+    -------
+    ValueError
+        If both `url` and `file` are empty or if both are provided.
+        If `question` is empty.
+    """
+    if openai_key.strip() == '':
         return '[ERROR]: Please enter you Open AI Key. Get your key here : https://platform.openai.com/account/api-keys'
     if url.strip() == '' and file == None:
-        return '[ERROR]: Both URL and PDF is empty. Provide atleast one.'
-    
-    if url.strip() != '' and file != None:
-        return '[ERROR]: Both URL and PDF is provided. Please provide only one (eiter URL or PDF).'
-
-    if url.strip() != '':
-        glob_url = url
-        download_pdf(glob_url, 'corpus.pdf')
-        load_recommender('corpus.pdf')
-
-    else:
-        old_file_name = file.name
-        file_name = file.name
-        file_name = file_name[:-12] + file_name[-4:]
-        os.rename(old_file_name, file_name)
-        load_recommender(file_name)
-
-    if question.strip() == '':
-        return '[ERROR]: Question field is empty'
+        return '[ERROR]: Both URL and PDF is empty. Provide at least one.'
 
-    return generate_answer(question,openAI_key)
-def question_answer(url, file, question,openAI_key):
-    if openAI_key.strip()=='':
-        return '[ERROR]: Please enter you Open AI Key. Get your key here : https://platform.openai.com/account/api-keys'
-    if url.strip() == '' and file == None:
-        return '[ERROR]: Both URL and PDF is empty. Provide atleast one.'
-    
     if url.strip() != '' and file != None:
-        return '[ERROR]: Both URL and PDF is provided. Please provide only one (eiter URL or PDF).'
+        return '[ERROR]: Both URL and PDF is provided. Please provide only one (either URL or PDF).'
 
     if url.strip() != '':
         glob_url = url
@@ -182,13 +358,22 @@ def question_answer(url, file, question,openAI_key):
     if question.strip() == '':
         return '[ERROR]: Question field is empty'
 
-    return generate_answer(question,openAI_key)
+    return generate_answer(question, openai_key)
 
 
 recommender = SemanticSearch()
 
 title = 'PDF GPT'
-description = """ PDF GPT allows you to chat with your PDF file using Universal Sentence Encoder and Open AI. It gives hallucination free response than other tools as the embeddings are better than OpenAI. The returned response can even cite the page number in square brackets([]) where the information is located, adding credibility to the responses and helping to locate pertinent information quickly."""
+description = """ What is PDF GPT ?
+1. The problem is that Open AI has a 4K token limit and cannot take an entire PDF file as input. Additionally,
+it sometimes returns irrelevant responses due to poor embeddings. ChatGPT cannot directly talk to external data.
+The solution is PDF GPT, which allows you to chat with an uploaded PDF file using GPT functionalities.
+The application breaks the document into smaller chunks and generates embeddings using a powerful Deep Averaging
+Network Encoder. A semantic search is performed on your query, and the top relevant chunks are used to generate a
+response.
+2. The returned response can even cite the page number in square brackets([]) where the information is located,
+adding credibility to the responses and helping to locate pertinent information quickly. The Responses are much
+better than the naive responses by Open AI."""
 
 with gr.Blocks() as demo:
 
@@ -196,13 +381,18 @@ with gr.Blocks() as demo:
     gr.Markdown(description)
 
     with gr.Row():
-        
+
         with gr.Group():
-            gr.Markdown(f'<p style="text-align:center">Get your Open AI API key <a href="https://platform.openai.com/account/api-keys">here</a></p>')
-            openAI_key=gr.Textbox(label='Enter your OpenAI API key here')
+            gr.Markdown(
+                '<p style="text-align:center">'
+                'Get your Open AI API key <a href="https://platform.openai.com/account/api-keys">here</a>'
+                '</p>'
+            )
+            openAI_key = gr.Textbox(label='Enter your OpenAI API key here')
             url = gr.Textbox(label='Enter PDF URL here')
             gr.Markdown("<center><h4>OR<h4></center>")
-            file = gr.File(label='Upload your PDF/ Research Paper / Book here', file_types=['.pdf'])
+            file = gr.File(
+                label='Upload your PDF/ Research Paper / Book here', file_types=['.pdf'])
             question = gr.Textbox(label='Enter your question here')
             btn = gr.Button(value='Submit')
             btn.style(full_width=True)
@@ -210,6 +400,7 @@ with gr.Blocks() as demo:
         with gr.Group():
             answer = gr.Textbox(label='The answer to your question is :')
 
-        btn.click(question_answer, inputs=[url, file, question,openAI_key], outputs=[answer])
-#openai.api_key = os.getenv('Your_Key_Here') 
-demo.launch()
+        btn.click(question_answer, inputs=[
+            url, file, question, openAI_key], outputs=[answer])
+
+demo.launch()