Update app.py

app.py

@@ -55,8 +55,15 @@ 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')
+class SemanticSearch:
+    
+    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
     
     
@@ -79,15 +86,29 @@ class SemanticSearch:
             return neighbors
     
     
-    def get_text_embedding(self, texts, batch=1000):
+    '''def get_text_embedding(self, texts, batch=1000):
         embeddings = []
         for i in range(0, len(texts), batch):
             text_batch = texts[i:(i+batch)]
             emb_batch = self.use(text_batch)
             embeddings.append(emb_batch)
         embeddings = np.vstack(embeddings)
-        return embeddings
-
+        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):
@@ -163,13 +184,16 @@ 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."""
 
-with gr.Blocks() as demo:
+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."""
 
+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')
@@ -179,10 +203,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()