Update app.py

app.py

@@ -50,7 +50,7 @@ def construct_index(doc):
   max_chunk_overlap = 1
 
     # Set chunk overlap ratio 
-  chunk_overlap_ratio = 0.5
+  chunk_overlap_ratio = 0.2
 
     # Define prompt helper
   prompt_helper = PromptHelper(max_input_size, num_output, max_chunk_overlap, chunk_size_limit, chunk_overlap_ratio)
@@ -64,11 +64,31 @@ def construct_index(doc):
   ## Indexation process and saving in the disk
   index = GPTVectorStoreIndex.from_documents(doc, service_context=service_context)
 
-  # save index to disk
-  index.set_index_id("vector_index")
-
   return index
 
+
+def upload_doc(file):
+    # Open the PDF file in binary mode
+    with open(file.name, 'rb') as f:
+        # Initialize a PDF file reader object
+        pdf_reader = PdfReader(f)
+
+        # Initialize an empty string for storing the extracted text
+        text = ''
+
+        # Loop through the number of pages
+        for page in pdf_reader.pages:
+            # Add the text from each page to the text string
+            text += page.extract_text()
+        
+        index = construct_index([Document(text)])
+  
+        # Save index to Azure blob storage
+        file_name = os.path.basename(file.name)
+        index.storage_context.persist(f'gpt/storage_demo/{file_name}', fs=fs)
+
+        return ''
+
 def extract_text(file):
     # Open the PDF file in binary mode
     with open(file.name, 'rb') as f:
@@ -85,7 +105,8 @@ def extract_text(file):
 
         return text, os.path.basename(file.name)
   
-def ask_ai(doc, question):
+  
+def ask_ai_upload(doc, question):
   
   text, file_name = extract_text(doc)
   index = construct_index([Document(text)])
@@ -103,45 +124,82 @@ def ask_ai(doc, question):
   query_engine = index.as_query_engine(optimizer=SentenceEmbeddingOptimizer(percentile_cutoff=0.8)) 
   query = 'Your task is to  answer a question on the report loaded and give insights to an investment team in Infrastructure. Make your response as clear and precise as possible. The question is:' + str(question)
   response = query_engine.query(query)
+
+  return response.response
+
+def respond_document_upload(message, chat_history, doc):
+
+  bot_message = ask_ai_upload(doc, message)
+  chat_history.append((message, bot_message))
+  time.sleep(2)
+
+  return "", chat_history
+
+
+def ask_ai_choose(doc, question):
   
-  # Display the chunks retrieved to produce the response
-  sources = []
-  for node in response.source_nodes:
-    node_text_start= 'START: ' + node.node.text.strip().replace('\n', ' ')[:100]
-    node_text_end = 'END: ' + node.node.text.strip().replace('\n', ' ')[-100:]
-    sources.append((node_text_start, node_text_end))
+  # Rebuild storage context
+  name_doc = str(doc)+'.pdf'
+  storage_context = StorageContext.from_defaults(persist_dir=f'gpt/storage_demo/{name_doc}', fs=fs)
+
+  # Load index
+  index = load_index_from_storage(storage_context)
+  
+  # Define the query & the querying method
+  query_engine = index.as_query_engine(optimizer=SentenceEmbeddingOptimizer(percentile_cutoff=0.8)) 
+  query = 'Your task is to  answer a question on the report loaded and give insights to an investment team in Infrastructure. Make your response as clear and precise as possible. The question is:' + str(question)
+  response = query_engine.query(query)
 
   return response.response
 
+def respond_document_choose(message, chat_history, doc):
+
+  bot_message = ask_ai_choose(doc, message)
+  chat_history.append((message, bot_message))
+  time.sleep(2)
+
+  return "", chat_history
+    
+# Retrieve the documents name whose indexes are stored
+path_list = fs.ls('gpt/storage_demo')
+documents_list = [Path(path).name[:-4] for path in path_list]
+
+# Configure Gradio platform 
+
 header = """<center><b><p style=\"color: #E13C32; font-size: 36px;\">My Ardian Chatbot</p></b></center>
 <i><p style=\"font-size: 16px; color: grey;\">Please make sure to formulate clear and precise questions and to add contextual information when possible. This will help the tool produce the most relevant response. Adopt an iterative approach and ask for more details or explanations when necessary.</br><i/></p>"""
 
 footnote = "<p style=\"font-size: 16px; color: grey;\"> ⚠ The chatbot doesn't have a memory, it doesn't remember what it previously generated.</a></p>"
 
 theme = gr.themes.Base(
-  primary_hue="red",
-  secondary_hue="gray",
-  font=['FuturaTOT', '=', '36px']
+    primary_hue="red",
+    secondary_hue="gray",
+    font=['FuturaTOT', '=']
 )
 
 with gr.Blocks(theme=theme) as demo:
     gr.Markdown(header)
 
-    download_button = gr.inputs.File(label="Upload a PDF")
-    chatbot = gr.Chatbot()
-    question = gr.Textbox(label='Question', info="Please write your question here.")
-    clear = gr.Button("Clear")
-    
-    def respond(message, chat_history, doc):
-        bot_message = ask_ai(doc, message)
-        chat_history.append((message, bot_message))
-        time.sleep(2)
-        return "", chat_history
-
-    question.submit(respond, [question, chatbot, download_button], [question, chatbot])
-    
-    clear.click(lambda: None, None, chatbot, queue=False)
-    gr.Markdown(footnote)
+    with gr.Tab("Upload a document & ask a question 📥"):
+      upload_file = gr.inputs.File(label="Upload your PDF document")
+      upload_button = gr.Button("Save the document in Ardian Knowledge Library")
+      output = gr.Textbox(label='Output', visible=False)
+      chatbot = gr.Chatbot()
+      question = gr.Textbox(label='Question', info="Please write your question here.")
+      clear = gr.Button("Clear")
+
+      question.submit(respond_document_upload, [question, chatbot, upload_file], [question, chatbot])
+      upload_button.click(upload_doc, inputs=upload_file, outputs=output)
+      clear.click(lambda: None, None, chatbot, queue=False)
+        
+    with gr.Tab("Choose a document & ask a question 📚"):
+      list_button = gr.Dropdown(documents_list, multiselect=False, label="Document", info="Please select the report you want to ask questions on.")
+      chatbot = gr.Chatbot()
+      question = gr.Textbox(label='Question', info="Please write your question here.")
+      clear = gr.Button("Clear")
+  
+      question.submit(respond_document_choose, [question, chatbot, list_button], [question, chatbot])
+      clear.click(lambda: None, None, chatbot, queue=False)
 
 demo.launch(auth=(os.environ['username'],os.environ['password']))