Update retrieval.py

retrieval.py

@@ -2,9 +2,6 @@
 LLM chain retrieval
 """
 
-
-
-
 import json
 import gradio as gr
 
@@ -14,16 +11,6 @@ from langchain_huggingface import HuggingFaceEndpoint
 from langchain_core.prompts import PromptTemplate
 
 
-# Add system template for RAG application
-PROMPT_TEMPLATE = """
-You are an assistant for question-answering tasks. Use the following pieces of context to answer the question at the end. 
-If you don't know the answer, just say that you don't know, don't try to make up an answer. Keep the answer concise.
-Question: {question} 
-Context: {context} 
-Helpful Answer:
-"""
-
-
 # Initialize langchain LLM chain
 def initialize_llmchain(
     llm_model,
@@ -37,22 +24,11 @@ def initialize_llmchain(
     """Initialize Langchain LLM chain"""
 
     progress(0.1, desc="Initializing HF tokenizer...")
-    # 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 'Llama' in llm_model:
-    #     task = "conversational"
-    # else:
-    #     task = "text-generation"
-    # print(f"Task: {task}")
 
     llm = HuggingFaceEndpoint(
         repo_id=llm_model,
         task="text-generation",
-        #task="conversational",
         provider="hf-inference",
         temperature=temperature,
         max_new_tokens=max_tokens,
@@ -62,18 +38,20 @@ def initialize_llmchain(
 
     progress(0.75, desc="Defining buffer memory...")
     memory = ConversationBufferMemory(
-        memory_key="chat_history", output_key="answer", return_messages=True
+        memory_key="chat_history",
+        output_key="answer",
+        return_messages=True,
     )
-    # retriever=vector_db.as_retriever(search_type="similarity", search_kwargs={'k': 3})
-    retriever = vector_db.as_retriever()
+    retriever = vector_db.as_retriever(search_type="similarity", search_kwargs={'k': top_k})
 
     progress(0.8, desc="Defining retrieval chain...")
     with open('prompt_template.json', 'r') as file:
-    	system_prompt = json.load(file)
+        system_prompt = json.load(file)
     prompt_template = system_prompt["prompt"]
     rag_prompt = PromptTemplate(
         template=prompt_template, input_variables=["context", "question"]
     )
+
     qa_chain = ConversationalRetrievalChain.from_llm(
         llm,
         retriever=retriever,
@@ -81,17 +59,16 @@ def initialize_llmchain(
         memory=memory,
         combine_docs_chain_kwargs={"prompt": rag_prompt},
         return_source_documents=True,
-        # return_generated_question=False,
         verbose=False,
     )
-    progress(0.9, desc="Done!")
 
+    progress(0.9, desc="Done!")
     return qa_chain
 
 
+# Format chat history
 def format_chat_history(message, chat_history):
-    """Format chat history for llm chain"""
-
+    """Format chat history for LLM"""
     formatted_chat_history = []
     for user_message, bot_message in chat_history:
         formatted_chat_history.append(f"User: {user_message}")
@@ -99,27 +76,22 @@ def format_chat_history(message, chat_history):
     return formatted_chat_history
 
 
+# Invoke QA chain with history
 def invoke_qa_chain(qa_chain, message, history):
     """Invoke question-answering chain"""
-
     formatted_chat_history = format_chat_history(message, history)
-    # print("formatted_chat_history",formatted_chat_history)
 
-    # Generate response using QA chain
-    response = qa_chain.invoke(
-        {"question": message, "chat_history": formatted_chat_history}
-    )
+    response = qa_chain.invoke({
+        "question": message,
+        "chat_history": formatted_chat_history,
+    })
 
     response_sources = response["source_documents"]
-
     response_answer = response["answer"]
-    if response_answer.find("Helpful Answer:") != -1:
-        response_answer = response_answer.split("Helpful Answer:")[-1]
-
-    # Append user message and response to chat history
-    new_history = history + [(message, response_answer)]
 
-    # print ('chat response: ', response_answer)
-    # print('DB source', response_sources)
+    # Clean up if "Helpful Answer:" is included
+    if "Helpful Answer:" in response_answer:
+        response_answer = response_answer.split("Helpful Answer:")[-1].strip()
 
+    new_history = history + [(message, response_answer)]
     return qa_chain, new_history, response_sources