Update app.py

app.py

@@ -20,13 +20,6 @@ import torch
 import tqdm 
 import accelerate
 
-#set parameters
-slider_chunk_size = 4096
-slider_chunk_overlap = 256
-slider_temperature = 0.1
-slider_maxtokens = 2048
-slider_topk = 3
-llm_model = "mistralai/Mistral-7B-Instruct-v0.2"
 
 
 # default_persist_directory = './chroma_HF/'
@@ -79,13 +72,63 @@ def load_db():
 
 
 # Initialize langchain LLM chain
-def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db):
-    llm = HuggingFaceHub(repo_id=llm_model, model_kwargs={"temperature": 
-                                                          temperature, "max_new_tokens": 
-                                                          max_tokens, "top_k": top_k, 
-                                                          "load_in_8bit": True})
-
+def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, progress=gr.Progress()):
+    progress(0.1, desc="Initializing HF tokenizer...")
+    # HuggingFacePipeline uses local model
+    # Note: it will download model locally...
+    # tokenizer=AutoTokenizer.from_pretrained(llm_model)
+    # progress(0.5, desc="Initializing HF pipeline...")
+    # pipeline=transformers.pipeline(
+    #     "text-generation",
+    #     model=llm_model,
+    #     tokenizer=tokenizer,
+    #     torch_dtype=torch.bfloat16,
+    #     trust_remote_code=True,
+    #     device_map="auto",
+    #     # max_length=1024,
+    #     max_new_tokens=max_tokens,
+    #     do_sample=True,
+    #     top_k=top_k,
+    #     num_return_sequences=1,
+    #     eos_token_id=tokenizer.eos_token_id
+    #     )
+    # llm = HuggingFacePipeline(pipeline=pipeline, model_kwargs={'temperature': temperature})
+    
+    # HuggingFaceHub uses HF inference endpoints
+    progress(0.5, desc="Initializing HF Hub...")
+    # Use of trust_remote_code as model_kwargs
+    # Warning: langchain issue
+    # URL: https://github.com/langchain-ai/langchain/issues/6080
+    if llm_model == "mistralai/Mixtral-8x7B-Instruct-v0.1":
+        llm = HuggingFaceHub(
+            repo_id=llm_model, 
+            model_kwargs={"temperature": temperature, "max_new_tokens": max_tokens, "top_k": top_k, "load_in_8bit": True}
+        )
+    elif llm_model == "microsoft/phi-2":
+        raise gr.Error("phi-2 model requires 'trust_remote_code=True', currently not supported by langchain HuggingFaceHub...")
+        llm = HuggingFaceHub(
+            repo_id=llm_model, 
+            model_kwargs={"temperature": temperature, "max_new_tokens": max_tokens, "top_k": top_k, "trust_remote_code": True, "torch_dtype": "auto"}
+        )
+    elif llm_model == "TinyLlama/TinyLlama-1.1B-Chat-v1.0":
+        llm = HuggingFaceHub(
+            repo_id=llm_model, 
+            model_kwargs={"temperature": temperature, "max_new_tokens": 250, "top_k": top_k}
+        )
+    elif llm_model == "meta-llama/Llama-2-7b-chat-hf":
+        raise gr.Error("Llama-2-7b-chat-hf model requires a Pro subscription...")
+        llm = HuggingFaceHub(
+            repo_id=llm_model, 
+            model_kwargs={"temperature": temperature, "max_new_tokens": max_tokens, "top_k": top_k}
+        )
+    else:
+        llm = HuggingFaceHub(
+            repo_id=llm_model, 
+            # model_kwargs={"temperature": temperature, "max_new_tokens": max_tokens, "top_k": top_k, "trust_remote_code": True, "torch_dtype": "auto"}
+            model_kwargs={"temperature": temperature, "max_new_tokens": max_tokens, "top_k": top_k}
+        )
     
+    progress(0.75, desc="Defining buffer memory...")
     memory = ConversationBufferMemory(
         memory_key="chat_history",
         output_key='answer',
@@ -93,6 +136,7 @@ def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db):
     )
     # retriever=vector_db.as_retriever(search_type="similarity", search_kwargs={'k': 3})
     retriever=vector_db.as_retriever()
+    progress(0.8, desc="Defining retrieval chain...")
     qa_chain = ConversationalRetrievalChain.from_llm(
         llm,
         retriever=retriever,
@@ -103,14 +147,16 @@ def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db):
         #return_generated_question=False,
         verbose=False,
     )
+    progress(0.9, desc="Done!")
     return qa_chain
 
 
 # Initialize database
-def initialize_database(list_file_obj, chunk_size, chunk_overlap):
+def initialize_database(list_file_obj, chunk_size, chunk_overlap, progress=gr.Progress()):
     # Create list of documents (when valid)
     list_file_path = [x.name for x in list_file_obj if x is not None]
     # Create collection_name for vector database
+    progress(0.1, desc="Creating collection name...")
     collection_name = Path(list_file_path[0]).stem
     # Fix potential issues from naming convention
     ## Remove space
@@ -124,20 +170,23 @@ def initialize_database(list_file_obj, chunk_size, chunk_overlap):
         collection_name[-1] = 'Z'
     # print('list_file_path: ', list_file_path)
     print('Collection name: ', collection_name)
+    progress(0.25, desc="Loading document...")
     # Load document and create splits
     doc_splits = load_doc(list_file_path, chunk_size, chunk_overlap)
     # Create or load vector database
+    progress(0.5, desc="Generating vector database...")
     # global vector_db
     vector_db = create_db(doc_splits, collection_name)
-    return vector_db, collection_name
+    progress(0.9, desc="Done!")
+    return vector_db, collection_name, "Complete!"
 
 
-def initialize_LLM(llm_option, llm_temperature, max_tokens, top_k, vector_db):
+def initialize_LLM(llm_option, llm_temperature, max_tokens, top_k, vector_db, progress=gr.Progress()):
     # print("llm_option",llm_option)
     llm_name = list_llm[llm_option]
     print("llm_name: ",llm_name)
     qa_chain = initialize_llmchain(llm_name, llm_temperature, max_tokens, top_k, vector_db, progress)
-    return qa_chain
+    return qa_chain, "Complete!"
 
 
 def format_chat_history(message, chat_history):
@@ -171,7 +220,7 @@ def conversation(qa_chain, message, history):
     # Append user message and response to chat 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
+    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):
@@ -189,35 +238,71 @@ def demo():
         vector_db = gr.State()
         qa_chain = gr.State()
         collection_name = gr.State()
-
-        document = gr.Files(value = '/home/user/app/pdfs/Annual-Report-2022-2023-English_1.pdf',visible=False,height=100, file_count="multiple", file_types=["pdf"])
-        chatbot = gr.Chatbot(height=300)
-        db_btn = gr.Radio(["ChromaDB"], label="Vector database type", value = "ChromaDB", type="index", info="Choose your vector database", visible=False)
-        with gr.Accordion("Document references", open=False):
+        
+        gr.Markdown(
+        """<center><h2>PDF-based chatbot (powered by LangChain and open-source LLMs)</center></h2>
+        <h3>Ask any questions about your PDF documents, along with follow-ups</h3>
+        <b>Note:</b> This AI assistant performs retrieval-augmented generation from your PDF documents. \
+        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("Step 1 - Document pre-processing"):
+            with gr.Row():
+                document = gr.Files(height=100, file_count="multiple", file_types=["pdf"], interactive=True, label="Upload your PDF documents (single or multiple)")
+                # upload_btn = gr.UploadButton("Loading document...", height=100, file_count="multiple", file_types=["pdf"], scale=1)
+            with gr.Row():
+                db_btn = gr.Radio(["ChromaDB"], label="Vector database type", value = "ChromaDB", type="index", info="Choose your 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():
-                doc_source1 = gr.Textbox(label="Reference 1", lines=2, container=True, scale=20)
-                source1_page = gr.Number(label="Page", scale=1)
+                db_progress = gr.Textbox(label="Vector database initialization", value="None")
+            with gr.Row():
+                db_btn = gr.Button("Generate vector database...")
+            
+        with gr.Tab("Step 2 - QA chain initialization"):
+            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.Slider(minimum = 224, maximum = 4096, value=1024, step=32, label="Max Tokens", info="Model max tokens", interactive=True)
+                with gr.Row():
+                    slider_topk = gr.Slider(minimum = 1, maximum = 10, value=3, step=1, label="top-k samples", info="Model top-k samples", interactive=True)
+            with gr.Row():
+                llm_progress = gr.Textbox(value="None",label="QA chain initialization")
+            with gr.Row():
+                qachain_btn = gr.Button("Initialize question-answering chain...")
+
+        with gr.Tab("Step 3 - Conversation with chatbot"):
+            chatbot = gr.Chatbot(height=300)
+            with gr.Accordion("Advanced - Document references", open=False):
+                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.Row():
-                doc_source2 = gr.Textbox(label="Reference 2", lines=2, container=True, scale=20)
-                source2_page = gr.Number(label="Page", scale=1)
+                msg = gr.Textbox(placeholder="Type message", container=True)
             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("Generate vector database...")
-            qachain_btn = gr.Button("Initialize question-answering chain...")
-            submit_btn = gr.Button("Submit")
-            clear_btn = gr.ClearButton([msg, chatbot])
+                submit_btn = gr.Button("Submit")
+                clear_btn = gr.ClearButton([msg, chatbot])
             
         # Preprocessing events
         #upload_btn.upload(upload_file, inputs=[upload_btn], outputs=[document])
         db_btn.click(initialize_database, \
             inputs=[document, slider_chunk_size, slider_chunk_overlap], \
-            outputs=[vector_db, collection_name])
+            outputs=[vector_db, collection_name, db_progress])
         qachain_btn.click(initialize_LLM, \
-            inputs=[llm_model, slider_temperature, slider_maxtokens, slider_topk, vector_db], \
+            inputs=[llm_btn, slider_temperature, slider_maxtokens, slider_topk, vector_db], \
             outputs=[qa_chain, llm_progress]).then(lambda:[None,"",0,"",0,"",0], \
             inputs=None, \
             outputs=[chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], \