Update app.py

app.py

@@ -55,48 +55,53 @@ def text_to_chunks(texts, word_length=150, start_page=1):
 
 class SemanticSearch:
     
-    def __init__(self):
-        self.use = hub.load('https://tfhub.dev/google/universal-sentence-encoder/4')
+    def __init__(self, embedding_method='use'):
+        self.embedding_method = embedding_method
+        if embedding_method == 'use':
+            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, openAI_key=None, batch=1000, n_neighbors=5):
         self.data = data
-        self.embeddings = self.get_text_embedding(data, batch=batch)
+        self.embeddings = self.get_text_embedding(data, openAI_key, 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):
-        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, openAI_key=None, batch=1000):
         embeddings = []
+        
+        if self.embedding_method == 'openai':
+            if openAI_key is None:
+                raise ValueError("OpenAI key is required when using OpenAI embeddings.")
+            openai.api_key = openAI_key
+
         for i in range(0, len(texts), batch):
             text_batch = texts[i:(i+batch)]
-            emb_batch = self.use(text_batch)
+            
+            if self.embedding_method == 'use':
+                emb_batch = self.use(text_batch)
+            elif self.embedding_method == 'openai':
+                emb_batch = [openai.Embed.extract(prompt=text)["embeddings"] for text in text_batch]
+                emb_batch = np.vstack(emb_batch)
+            
             embeddings.append(emb_batch)
+
         embeddings = np.vstack(embeddings)
         return embeddings
 
 
 
-def load_recommender(path, start_page=1):
+
+def load_recommender(path, openAI_key=None, start_page=1):
     global recommender
     texts = pdf_to_text(path, start_page=start_page)
     chunks = text_to_chunks(texts, start_page=start_page)
-    recommender.fit(chunks)
+    recommender.fit(chunks, openAI_key=openAI_key)
     return 'Corpus Loaded.'
 
+
 def generate_text(openAI_key,prompt, engine="text-davinci-003"):
     openai.api_key = openAI_key
     completions = openai.Completion.create(
@@ -110,7 +115,7 @@ def generate_text(openAI_key,prompt, engine="text-davinci-003"):
     message = completions.choices[0].text
     return message
 
-def generate_answer1(question,openAI_key):
+def generate_answer(question,openAI_key):
     topn_chunks = recommender(question)
     prompt = ""
     prompt += 'search results:\n\n'
@@ -130,56 +135,6 @@ def generate_answer1(question,openAI_key):
     answer = generate_text(openAI_key, prompt,"text-davinci-003")
     return answer
 
-def generate_answer(question, openAI_key):
-    topn_chunks = recommender(question)
-    
-    max_tokens = 4096  # Maximum tokens allowed for text-davinci-003
-    completion_tokens = 512  # Tokens reserved for the completion
-    tokenizer = OpenAITokenizer()
-    max_prompt_tokens = max_tokens - completion_tokens
-
-    # Split search results into groups based on token count
-    search_results_groups = []
-    current_group = []
-    current_group_tokens = 0
-
-    for c in topn_chunks:
-        c_tokens = len(tokenizer.tokenize(c))
-        if current_group_tokens + c_tokens <= max_prompt_tokens:
-            current_group.append(c)
-            current_group_tokens += c_tokens
-        else:
-            search_results_groups.append(current_group)
-            current_group = [c]
-            current_group_tokens = c_tokens
-
-    if current_group:
-        search_results_groups.append(current_group)
-
-    # Generate response for each group of search results
-    responses = []
-    for search_results in search_results_groups:
-        prompt = 'search results:\n\n'
-        for c in search_results:
-            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: "
-
-        response = generate_text(openAI_key, prompt, "text-davinci-003")
-        responses.append(response)
-
-    # Combine and clean up the responses
-    final_answer = " ".join(responses).strip()
-
-    return final_answer
-
 
 def question_answer(url, file, question,openAI_key):
     if openAI_key.strip()=='':
@@ -193,14 +148,10 @@ def question_answer(url, file, question,openAI_key):
     if url.strip() != '':
         glob_url = url
         download_pdf(glob_url, 'corpus.pdf')
-        load_recommender('corpus.pdf')
+        load_recommender('corpus.pdf', openAI_key=openAI_key)
 
     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)
+        load_recommender(file_name, openAI_key=openAI_key)
 
     if question.strip() == '':
         return '[ERROR]: Question field is empty'
@@ -208,7 +159,9 @@ def question_answer(url, file, question,openAI_key):
     return generate_answer(question,openAI_key)
 
 
-recommender = SemanticSearch()
+#recommender = SemanticSearch()
+recommender = SemanticSearch(embedding_method='openai')
+
 
 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."""