Update app.py

app.py

@@ -17,7 +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-v0.3"]
 list_llm_simple = [os.path.basename(llm) for llm in list_llm]
 
 # Load PDF document and create doc splits
@@ -64,32 +64,19 @@ def create_db(splits, collection_name, db_type):
     return vectordb
 
 # Initialize langchain LLM chain
-# 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",
@@ -101,11 +88,12 @@ 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
 
@@ -266,7 +254,7 @@ def demo():
         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: gr.State(prompt), 
+        set_prompt_btn.click(lambda prompt: prompt, 
                              inputs=prompt_input, 
                              outputs=initial_prompt)
         qachain_btn.click(initialize_LLM,