change to retrieval method

app.py

@@ -181,11 +181,9 @@ def load_llama2_llamaCpp():
 
 def set_custom_prompt():
     custom_prompt_template = """ Use the following pieces of information to answer the user's question.
-    If you don't know the answer, please just say that you don't know the answer, don't try to make up
-    an answer.
+    If you don't know the answer, don't try to make up an answer.
 
     Context : {context}
-    chat_history : {chat_history}
     Question : {question}
 
     Only returns the helpful answer below and nothing else.
@@ -193,7 +191,7 @@ def set_custom_prompt():
     """
     prompt = PromptTemplate(template=custom_prompt_template, input_variables=['context',
                                                                               'question',
-                                                                              'chat_history'])
+                                                                              ])
     return prompt
 
 
@@ -228,10 +226,10 @@ def main():
 
     llm = load_llama2_llamaCpp()
     qa_prompt = set_custom_prompt()
-    memory = ConversationBufferWindowMemory(k = 0, return_messages=True,  input_key= 'question', output_key='answer', memory_key="chat_history")
+    #memory = ConversationBufferWindowMemory(k = 0, return_messages=True,  input_key= 'question', output_key='answer', memory_key="chat_history")
     #memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
-    doc_chain = load_qa_chain(llm, chain_type="stuff", prompt = qa_prompt)
-    question_generator = LLMChain(llm=llm, prompt=CONDENSE_QUESTION_PROMPT)
+    #doc_chain = load_qa_chain(llm, chain_type="stuff", prompt = qa_prompt)
+    #question_generator = LLMChain(llm=llm, prompt=CONDENSE_QUESTION_PROMPT)
     embeddings = load_embeddings()
 
     
@@ -243,15 +241,27 @@ def main():
         for page in pdf_reader.pages:
             text += page.extract_text()
         db = FAISS.from_texts(text, embeddings)
-        qa_chain = ConversationalRetrievalChain(
-            retriever =db.as_retriever(search_kwargs={'k':2}), 
-            question_generator=question_generator,
-            #condense_question_prompt=CONDENSE_QUESTION_PROMPT,
-            combine_docs_chain=doc_chain,
-            return_source_documents=True,
+
+        memory = ConversationBufferMemory(memory_key="chat_history", 
+                                  return_messages=True, 
+                                  input_key="query", 
+                                  output_key="result")
+        qa_chain = RetrievalQA.from_chain_type(
+            llm = llm,
+            chain_type = "stuff",
+            retriever = db.as_retriever(search_kwargs = {'k':3}), 
+            return_source_documents = True,
             memory = memory,
-            #get_chat_history=lambda h :h
-        )
+            chain_type_kwargs = {"prompt":qa_prompt}) 
+        # qa_chain = ConversationalRetrievalChain(
+        #     retriever =db.as_retriever(search_kwargs={'k':2}), 
+        #     question_generator=question_generator,
+        #     #condense_question_prompt=CONDENSE_QUESTION_PROMPT,
+        #     combine_docs_chain=doc_chain,
+        #     return_source_documents=True,
+        #     memory = memory,
+        #     #get_chat_history=lambda h :h
+        # )
 
         for message in st.session_state.messages:
             with st.chat_message(message["role"]):
@@ -267,20 +277,20 @@ def main():
 
             start = time.time()
 
-            response = qa_chain({'question': query})
+            response = qa_chain({'query': query})
 
             # url_list = set([i.metadata['source']  for i in response['source_documents']])
             #print(f"condensed quesion : {question_generator.run({'chat_history': response['chat_history'], 'question' : query})}")
             
             with st.chat_message("assistant"):
-                st.markdown(response['answer'])
+                st.markdown(response['result'])
             
             end = time.time()
             st.write("Respone time:",int(end-start),"sec")
             print(response)
                     
             # Add assistant response to chat history
-            st.session_state.messages.append({"role": "assistant", "content": response['answer']})
+            st.session_state.messages.append({"role": "assistant", "content": response['result']})
 
             # with st.expander("See the related documents"):
             #     for count, url in enumerate(url_list):