Update app.py

app.py

@@ -31,7 +31,8 @@ temperature = 0.1
 max_tokens = 6000
 top_k = 3
 
-def load_doc(list_file_path, chunk_size, chunk_overlap):
+
+def load_doc(list_file_path):
     # Processing for one document only
     # loader = PyPDFLoader(file_path)
     # pages = loader.load()
@@ -46,6 +47,8 @@ def load_doc(list_file_path, chunk_size, chunk_overlap):
     doc_splits = text_splitter.split_documents(pages)
     return doc_splits
 
+
+
 # Create vector database
 def create_db(splits, collection_name):
     embedding = HuggingFaceEmbeddings()
@@ -67,12 +70,7 @@ def load_db():
         embedding_function=embedding)
     return vectordb
 
-
-
-
-
-#list_file_obj = document
-
+    
 # Initialize database
 def initialize_database(list_file_obj):
     # Create list of documents (when valid)
@@ -94,7 +92,7 @@ def initialize_database(list_file_obj):
     # print('list_file_path: ', list_file_path)
     print('Collection name: ', collection_name)
     # Load document and create splits
-    doc_splits = load_doc(list_file_path, chunk_size, chunk_overlap)
+    doc_splits = load_doc(list_file_path)
     # Create or load vector database
     # global vector_db
     vector_db = create_db(doc_splits, collection_name)
@@ -121,14 +119,6 @@ def initialize_LLM(vector_db):
     return qa_chain
 
 
-def format_chat_history(message, chat_history):
-    formatted_chat_history = []
-    for user_message, bot_message in chat_history:
-        formatted_chat_history.append(f"User: {user_message}")
-        formatted_chat_history.append(f"Assistant: {bot_message}")
-    return formatted_chat_history
-
-
 def conversation(qa_chain, message, history):
     formatted_chat_history = format_chat_history(message, history)
     #print("formatted_chat_history",formatted_chat_history)
@@ -153,43 +143,47 @@ def conversation(qa_chain, message, history):
     new_history = history + [(message, response_answer)]
     # return gr.update(value=""), new_history, response_sources[0], response_sources[1] 
     return qa_chain, new_history, response_source1, response_source1_page, response_source2, response_source2_page, response_source3, response_source3_page
-    
-#document = os.listdir(list_file_obj)
-#qa_chain = 
-
 
 def demo():
-    with gr.Blocks(theme='base') as demo:
+    with gr.Block() as demo:
         vector_db = gr.State()
         qa_chain = gr.State()
-        collection_name = gr.State()   
-        
-        vector_db, collection_name = initialize_database(list_file_obj)
+        collection_name = gr.State()
+
         chatbot = gr.Chatbot(height=300)
-        with gr.Accordion('References', open=True):
-            with gr.Row():
-                doc_source1 = gr.Textbox(label="Reference 1", lines=2, container=True, scale=20)
-                source1_page = gr.Number(label="Page", scale=1)
-            with gr.Row():
-                doc_source2 = gr.Textbox(label="Reference 2", lines=2, container=True, scale=20)
-                source2_page = gr.Number(label="Page", scale=1)
-            with gr.Row():
-                doc_source3 = gr.Textbox(label="Reference 3", lines=2, container=True, scale=20)
-                source3_page = gr.Number(label="Page", scale=1)
+        with gr.Accordion("Advanced - Document references", open=False):
         with gr.Row():
-            msg = gr.Textbox(placeholder = 'Ask your question', container = True)
+            doc_source1 = gr.Textbox(label="Reference 1", lines=2, container=True, scale=20)
+            source1_page = gr.Number(label="Page", scale=1)
         with gr.Row():
-            qa_chain_button = gr.Button('Start Chatbot')
-            submit_btn = gr.Button('Submit')
-            clear_button = gr.ClearButton([msg, chatbot])
-
-        qa_chain_button.click(initialize_LLM, \
+            doc_source2 = gr.Textbox(label="Reference 2", lines=2, container=True, scale=20)
+            source2_page = gr.Number(label="Page", scale=1)
+        with gr.Row():
+            doc_source3 = gr.Textbox(label="Reference 3", lines=2, container=True, scale=20)
+            source3_page = gr.Number(label="Page", scale=1)
+        with gr.Row():
+            msg = gr.Textbox(placeholder="Type message", container=True)
+        with gr.Row():
+            db_btn = gr.Button('Initialize database')
+            qachain_btn = gr.Button('Start chatbot')
+            submit_btn = gr.Button("Submit")
+            clear_btn = gr.ClearButton([msg, chatbot])
+
+        document = list_file_obj
+            
+        #upload_btn.upload(upload_file, inputs=[upload_btn], outputs=[document])
+        db_btn.click(initialize_database, \
+            inputs=[document], \
+            outputs=[vector_db, collection_name])
+        
+        qachain_btn.click(initialize_LLM, \
             inputs=[vector_db], \
-            outputs=[qa_chain]).then(lambda:[None,"",0,"",0,"",0], \
+            outputs=[qa_chain]).then(lambda:[0], \
             inputs=None, \
             outputs=[chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], \
-            queue=False)
-        
+            queue=False)    
+    
+        # Chatbot events
         msg.submit(conversation, \
             inputs=[qa_chain, msg, chatbot], \
             outputs=[qa_chain, msg, chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], \
@@ -198,12 +192,9 @@ def demo():
             inputs=[qa_chain, msg, chatbot], \
             outputs=[qa_chain, msg, chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], \
             queue=False)
-        clear_button.click(lambda:[None,"",0,"",0,"",0], \
+        clear_btn.click(lambda:[None,"",0,"",0,"",0], \
             inputs=None, \
             outputs=[chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], \
             queue=False)
     demo.queue().launch(debug=True)
 
-if __name__ == "__main__":
-    demo()
-