Update app.py

app.py

@@ -14,16 +14,7 @@ from langchain.llms import HuggingFaceHub
 from pathlib import Path
 import chromadb
 
-from transformers import AutoTokenizer
-import transformers
-import torch
-import tqdm 
-import accelerate
-
-# Update list of LLM models
-list_llm = ["mistralai/Mistral-7B-Instruct-v0.2"]
-list_llm_simple = [os.path.basename(llm) for llm in list_llm]
-
+# Load PDF document and create doc splits
 def load_doc(list_file_path, chunk_size, chunk_overlap):
     loaders = [PyPDFLoader(x) for x in list_file_path]
     pages = []
@@ -35,6 +26,7 @@ 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()
     new_client = chromadb.EphemeralClient()
@@ -46,6 +38,7 @@ def create_db(splits, collection_name):
     )
     return vectordb
 
+# Initialize langchain LLM chain
 def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, progress=gr.Progress()):
     llm = HuggingFaceHub(
         repo_id=llm_model,
@@ -60,7 +53,7 @@ def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, pr
     qa_chain = ConversationalRetrievalChain.from_llm(
         llm,
         retriever=retriever,
-        chain_type="stuff", 
+        chain_type="stuff",
         memory=memory,
         return_source_documents=True,
         verbose=False,
@@ -68,18 +61,20 @@ def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, pr
     progress(0.9, desc="Done!")
     return qa_chain
 
-def initialize_database(list_file_obj, chunk_size, chunk_overlap, llm_temperature, max_tokens, top_k, progress=gr.Progress()):
+# Initialize database and LLM chain
+def initialize_demo(list_file_obj, chunk_size, chunk_overlap, progress=gr.Progress()):
     list_file_path = [x.name for x in list_file_obj if x is not None]
     collection_name = Path(list_file_path[0]).stem.replace(" ", "-")[:50]
     doc_splits = load_doc(list_file_path, chunk_size, chunk_overlap)
     vector_db = create_db(doc_splits, collection_name)
     qa_chain = initialize_llmchain(
-        list_llm[0], 
-        llm_temperature, 
-        max_tokens, 
-        top_k, 
-        vector_db, 
-        progress)
+        "mistralai/Mistral-7B-Instruct-v0.2",
+        0.7,
+        1024,
+        3,
+        vector_db,
+        progress
+    )
     return vector_db, collection_name, qa_chain, "Complete!"
 
 def format_chat_history(message, chat_history):
@@ -105,13 +100,6 @@ def conversation(qa_chain, message, history):
     new_history = history + [(message, response_answer)]
     return qa_chain, gr.update(value=""), new_history, response_source1, response_source1_page, response_source2, response_source2_page, response_source3, response_source3_page
 
-def upload_file(file_obj):
-    list_file_path = []
-    for idx, file in enumerate(file_obj):
-        file_path = file_obj.name
-        list_file_path.append(file_path)
-    return list_file_path
-
 def demo():
     with gr.Blocks(theme="base") as demo:
         vector_db = gr.State()
@@ -125,21 +113,25 @@ def demo():
         When generating answers, it takes past questions into account (via conversational memory), and includes document references for clarity purposes.</i>
         <br><b>Warning:</b> This space uses the free CPU Basic hardware from Hugging Face. Some steps and LLM models used below (free inference endpoints) can take some time to generate an output.<br>
         """)
-        with gr.Tab("Chatbot"):
-            with gr.Row():
-                document = gr.Files(height=100, file_count="multiple", file_types=["pdf"], interactive=True, label="Upload your PDF documents (single or multiple)")
-                db_btn = gr.Button("Generate vector database...")
-            with gr.Accordion("Advanced options - Document text splitter", open=False):
-                with gr.Row():
-                    slider_chunk_size = gr.Slider(minimum=100, maximum=1000, value=600, step=20, label="Chunk size", info="Chunk size", interactive=True)
-                with gr.Row():
-                    slider_chunk_overlap = gr.Slider(minimum=10, maximum=200, value=40, step=10, label="Chunk overlap", info="Chunk overlap", interactive=True)
-            with gr.Row():
-                db_progress = gr.Textbox(label="Vector database initialization", value="None")
-            with gr.Row():
-                llm_btn = gr.Radio(list_llm_simple, label="LLM models", value=list_llm_simple[0], type="index", info="Choose your LLM model")
-            with gr.Accordion("Advanced options - LLM model", open=False):
-                with gr.Row():
-                    slider_temperature = gr.Slider(minimum=0.0, maximum=1.0, value=0.7, step=0.1, label="Temperature", info="Model temperature", interactive=True)
-                with gr.Row():
-                    slider_maxtokens = gr
+
+        document = gr.Files(height=100, file_count="multiple", file_types=["pdf"], interactive=True, label="Upload your PDF documents (single or multiple)")
+        slider_chunk_size = gr.Slider(minimum=100, maximum=1000, value=600, step=20, label="Chunk size", info="Chunk size", interactive=True)
+        slider_chunk_overlap = gr.Slider(minimum=10, maximum=200, value=40, step=10, label="Chunk overlap", info="Chunk overlap", interactive=True)
+        db_progress = gr.Textbox(label="Vector database initialization", value="None")
+
+        # Initialize vector database and LLM chain in the background
+        vector_db, collection_name, qa_chain, status = initialize_demo([document], slider_chunk_size, slider_chunk_overlap, db_progress)
+
+        chatbot = gr.Chatbot(height=300)
+        msg = gr.Textbox(placeholder="Type message", container=True)
+        submit_btn = gr.Button("Submit")
+        clear_btn = gr.ClearButton([msg, chatbot])
+
+        msg.submit(conversation, inputs=[qa_chain, msg, chatbot], outputs=[qa_chain, msg, chatbot], queue=False)
+        submit_btn.click(conversation, inputs=[qa_chain, msg, chatbot], outputs=[qa_chain, msg, chatbot], queue=False)
+        clear_btn.click(lambda:[None,"",0,"",0,"",0], inputs=None, outputs=[chatbot], queue=False)
+
+    demo.queue().launch(debug=True)
+
+if __name__ == "__main__":
+    demo()