Update app.py

app.py

@@ -1,6 +1,5 @@
 import gradio as gr
 import os
-import getpass
 from pathlib import Path
 import re
 from unidecode import unidecode
@@ -10,7 +9,7 @@ from langchain_community.document_loaders import PyPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import Chroma
 from langchain.chains import ConversationalRetrievalChain
-from langchain_community.embeddings import HuggingFaceEmbeddings 
+from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.llms import HuggingFacePipeline
 from langchain.chains import ConversationChain
 from langchain.memory import ConversationBufferMemory
@@ -18,9 +17,7 @@ from langchain_community.llms import HuggingFaceEndpoint
 import torch
 api_token = os.getenv("HF_TOKEN")
 
-
-
-list_llm = ["meta-llama/Meta-Llama-3-8B-Instruct", "mistralai/Mistral-7B-Instruct-v0.3"]  
+list_llm = ["meta-llama/Meta-Llama-3-8B-Instruct", "mistralai/Mistral-7B-Instruct"]
 list_llm_simple = [os.path.basename(llm) for llm in list_llm]
 
 # Load PDF document and create doc splits
@@ -66,39 +63,20 @@ def create_db(splits, collection_name, db_type):
     
     return vectordb
 
-# Load vector database
-def load_db():
-    embedding = HuggingFaceEmbeddings()
-    vectordb = Chroma(
-        embedding_function=embedding)
-    return vectordb
-
 # Initialize langchain LLM chain
-def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, progress=gr.Progress()):
+def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, initial_prompt, progress=gr.Progress()):
     progress(0.1, desc="Initializing HF tokenizer...")
     
     progress(0.5, desc="Initializing HF Hub...")
 
-    
-    if llm_model == "meta-llama/Meta-Llama-3-8B-Instruct":
-        llm = HuggingFaceEndpoint(
-            repo_id=llm_model,
-            huggingfacehub_api_token=api_token,
-            temperature=temperature,
-            max_new_tokens=max_tokens,
-            top_k=top_k,
-        )
-   
-    else:
-        
-        llm = HuggingFaceEndpoint(
-            repo_id=llm_model,
-            huggingfacehub_api_token=api_token,
-            temperature=temperature,
-            max_new_tokens=max_tokens,
-            top_k=top_k,        
-        )
-    
+    llm = HuggingFaceEndpoint(
+        repo_id=llm_model,
+        huggingfacehub_api_token=api_token,
+        temperature=temperature,
+        max_new_tokens=max_tokens,
+        top_k=top_k,
+    )
+
     progress(0.75, desc="Defining buffer memory...")
     memory = ConversationBufferMemory(
         memory_key="chat_history",
@@ -110,18 +88,19 @@ 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,
     )
+    qa_chain({"question": initial_prompt})  # Initialize with the initial prompt
     progress(0.9, desc="Done!")
     return qa_chain
 
 # Generate collection name for vector database
 def create_collection_name(filepath):
     collection_name = Path(filepath).stem
-    collection_name = collection_name.replace(" ", "-") 
+    collection_name = collection_name.replace(" ", "-")
     collection_name = unidecode(collection_name)
     collection_name = re.sub('[^A-Za-z0-9]+', '-', collection_name)
     collection_name = collection_name[:50]
@@ -147,10 +126,10 @@ def initialize_database(list_file_obj, chunk_size, chunk_overlap, db_type, progr
     progress(0.9, desc="Done!")
     return vector_db, collection_name, "Complete!"
 
-def initialize_LLM(llm_option, llm_temperature, max_tokens, top_k, vector_db, progress=gr.Progress()):
+def initialize_LLM(llm_option, llm_temperature, max_tokens, top_k, vector_db, initial_prompt, progress=gr.Progress()):
     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)
+    qa_chain = initialize_llmchain(llm_name, llm_temperature, max_tokens, top_k, vector_db, initial_prompt, progress)
     return qa_chain, "Complete!"
 
 def format_chat_history(message, chat_history):
@@ -162,7 +141,6 @@ def format_chat_history(message, chat_history):
 
 def conversation(qa_chain, message, history):
     formatted_chat_history = format_chat_history(message, history)
-   
     response = qa_chain({"question": message, "chat_history": formatted_chat_history})
     response_answer = response["answer"]
     if "Helpful Answer:" in response_answer:
@@ -178,6 +156,13 @@ 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 conversation_no_doc(llm, message, history):
+    formatted_chat_history = format_chat_history(message, history)
+    response = llm({"question": message, "chat_history": formatted_chat_history})
+    response_answer = response["answer"]
+    new_history = history + [(message, response_answer)]
+    return llm, gr.update(value=""), new_history
+
 def upload_file(file_obj):
     list_file_path = []
     for file in file_obj:
@@ -189,7 +174,9 @@ def demo():
         vector_db = gr.State()
         qa_chain = gr.State()
         collection_name = gr.State()
-        
+        initial_prompt = gr.State()
+        llm_no_doc = gr.State()
+
         gr.Markdown(
         """<center><h2>PDF-based chatbot</center></h2>
         <h3>Ask any questions about your PDF documents</h3>""")
@@ -197,7 +184,7 @@ def demo():
         """<b>Note:</b> Esta é a lucIAna, primeira Versão da IA para seus PDF documentos. 
         Este chatbot leva em consideração perguntas anteriores ao gerar respostas (por meio de memória conversacional) e inclui referências a documentos para fins de clareza.
         """)
-        
+
         with gr.Tab("Step 1 - Upload PDF"):
             with gr.Row():
                 document = gr.Files(height=100, file_count="multiple", file_types=["pdf"], interactive=True, label="Upload your PDF documents (single or multiple)")
@@ -215,7 +202,13 @@ def demo():
             with gr.Row():
                 db_btn = gr.Button("Generate vector database")
         
-        with gr.Tab("Step 3 - Initialize QA chain"):
+        with gr.Tab("Step 3 - Set Initial Prompt"):
+            with gr.Row():
+                prompt_input = gr.Textbox(label="Initial Prompt", lines=5, value="Você é um advogado sênior, onde seu papel é analisar e trazer as informações sem inventar, dando a sua melhor opinião sempre trazendo contexto e referência. Aprenda o que é jurisprudência.")
+            with gr.Row():
+                set_prompt_btn = gr.Button("Set Prompt")
+
+        with gr.Tab("Step 4 - Initialize QA chain"):
             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")
@@ -231,7 +224,7 @@ def demo():
             with gr.Row():
                 qachain_btn = gr.Button("Initialize Question Answering chain")
 
-        with gr.Tab("Step 4 - Chatbot"):
+        with gr.Tab("Step 5 - Chatbot with document"):
             chatbot = gr.Chatbot(height=300)
             with gr.Accordion("Advanced - Document references", open=False):
                 with gr.Row():
@@ -248,19 +241,30 @@ def demo():
             with gr.Row():
                 submit_btn = gr.Button("Submit message")
                 clear_btn = gr.ClearButton([msg, chatbot], value="Clear conversation")
-            
+
+        with gr.Tab("Step 6 - Chatbot without document"):
+            chatbot_no_doc = gr.Chatbot(height=300)
+            with gr.Row():
+                msg_no_doc = gr.Textbox(placeholder="Type message to chat with lucIAna", container=True)
+            with gr.Row():
+                submit_btn_no_doc = gr.Button("Submit message")
+                clear_btn_no_doc = gr.ClearButton([msg_no_doc, chatbot_no_doc], value="Clear conversation")
+
         # Preprocessing events
         db_btn.click(initialize_database, 
             inputs=[document, slider_chunk_size, slider_chunk_overlap, db_type_radio], 
             outputs=[vector_db, collection_name, db_progress])
+        set_prompt_btn.click(lambda prompt: prompt, 
+                             inputs=prompt_input, 
+                             outputs=initial_prompt)
         qachain_btn.click(initialize_LLM, 
-            inputs=[llm_btn, slider_temperature, slider_maxtokens, slider_topk, vector_db], 
+            inputs=[llm_btn, slider_temperature, slider_maxtokens, slider_topk, vector_db, initial_prompt], 
             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], 
             queue=False)
 
-        # Chatbot events
+        # Chatbot events with document
         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], 
@@ -273,7 +277,18 @@ def demo():
             inputs=None, 
             outputs=[chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], 
             queue=False)
-    demo.queue().launch(debug=True)
+
+        # Chatbot events without document
+        submit_btn_no_doc.click(conversation_no_doc, 
+            inputs=[llm_no_doc, msg_no_doc, chatbot_no_doc], 
+            outputs=[llm_no_doc, msg_no_doc, chatbot_no_doc], 
+            queue=False)
+        clear_btn_no_doc.click(lambda:[None,""], 
+            inputs=None, 
+            outputs=[chatbot_no_doc, msg_no_doc], 
+            queue=False)
+
+    demo.queue().launch(debug=True, share=True)
 
 if __name__ == "__main__":
-    demo()
+    demo()