Update app.py

app.py

@@ -53,55 +53,70 @@ def text_to_chunks(texts, word_length=150, start_page=1):
             chunks.append(chunk)
     return chunks
 
+    
 class SemanticSearch:
     
-    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')
+    def __init__(self, embedder='openai'):
+        if embedder == 'openai':
+            self.embedder = openai.Engine("davinci")
+        elif embedder == 'use':
+            self.embedder = hub.load('https://tfhub.dev/google/universal-sentence-encoder/4')
+        else:
+            raise ValueError("Invalid embedder. Must be either 'openai' or 'use'.")
         self.fitted = False
-
-    def fit(self, data, openAI_key=None, batch=1000, n_neighbors=5):
+    
+    
+    def fit(self, data, batch=1000, n_neighbors=5):
         self.data = data
-        self.embeddings = self.get_text_embedding(data, openAI_key, batch=batch)
+        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 get_text_embedding(self, texts, openAI_key=None, batch=1000):
-        embeddings = []
+    
+    
+    def __call__(self, text, return_data=True):
+        inp_emb = self.use([text])
+        neighbors = self.nn.kneighbors(inp_emb, return_distance=False)[0]
         
-        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
-
+        if return_data:
+            return [self.data[i] for i in neighbors]
+        else:
+            return neighbors
+    
+    
+    '''def get_text_embedding(self, texts, batch=1000):
+        embeddings = []
         for i in range(0, len(texts), batch):
             text_batch = texts[i:(i+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)
-            
+            emb_batch = self.use(text_batch)
             embeddings.append(emb_batch)
-
         embeddings = np.vstack(embeddings)
-        return embeddings
-
-
-
-
-def load_recommender(path, openAI_key=None, start_page=1):
+        return embeddings'''
+	def get_text_embedding(self, texts):
+		embeddings = []
+		if isinstance(self.embedder, openai.Engine):
+			for text in texts:
+				response = self.embedder.search(
+					documents=texts,
+					query=text,
+					max_rerank=1
+				)
+				embeddings.append(response["data"][0]["score"])
+		elif isinstance(self.embedder, hub.Module):
+			embeddings = self.embedder(texts)
+		else:
+			raise ValueError("Invalid embedder.")
+		return np.array(embeddings)
+
+
+def load_recommender(path, 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, openAI_key=openAI_key)
+    recommender.fit(chunks)
     return 'Corpus Loaded.'
 
-
 def generate_text(openAI_key,prompt, engine="text-davinci-003"):
     openai.api_key = openAI_key
     completions = openai.Completion.create(
@@ -148,10 +163,14 @@ def question_answer(url, file, question,openAI_key):
     if url.strip() != '':
         glob_url = url
         download_pdf(glob_url, 'corpus.pdf')
-        load_recommender('corpus.pdf', openAI_key=openAI_key)
+        load_recommender('corpus.pdf')
 
     else:
-        load_recommender(file_name, openAI_key=openAI_key)
+        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'
@@ -159,20 +178,21 @@ def question_answer(url, file, question,openAI_key):
     return generate_answer(question,openAI_key)
 
 
-#recommender = SemanticSearch()
-recommender = SemanticSearch(embedding_method='openai')
+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."""
 
 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."""
 
-with gr.Blocks() as demo:
+embedder = gr.Dropdown(['openai', 'use'], label='Select Embedder')
+recommender = SemanticSearch(embedder=embedder)
 
+with gr.Blocks() as demo:
     gr.Markdown(f'<center><h1>{title}</h1></center>')
     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')
@@ -182,10 +202,7 @@ with gr.Blocks() as demo:
             question = gr.Textbox(label='Enter your question here')
             btn = gr.Button(value='Submit')
             btn.style(full_width=True)
-
         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()