Update app.py

app.py

@@ -20,7 +20,7 @@ from HTML_templates import css, bot_template, user_template
 
 
 
-def get_vectorstore(text_chunks):
+def retriever_from_chroma(docs, search_type, k):
     model_name = "sentence-transformers/all-mpnet-base-v2"
     model_kwargs = {'device': 'cpu'}
     encode_kwargs = {'normalize_embeddings': True}
@@ -34,7 +34,9 @@ def get_vectorstore(text_chunks):
         os.makedirs(vectorstore_path)
     vectorstore = Chroma.from_documents(
         documents=text_chunks, embedding=embeddings, persist_directory="docs/chroma/")
-    return vectorstore    
+    retriever = vectordb.as_retriever(search_type=search_type, search_kwargs={"k": k})
+    return retriever
+        
 
 
 
@@ -57,16 +59,16 @@ for filename in os.listdir(data_path):
 
 
 
-docs = split_docs(documents, 350, 40)
+docs = split_docs(documents, 250, 20)
 
-vectorstore = get_vectorstore(docs)
+retriever = retriever_from_chroma(docs,'mmr',7)
 
 
 
 
 
 
-def main(vectorstore):
+def main(retriever):
 
     st.set_page_config(page_title="Chat with multiple PDFs",
                        page_icon=":books:")
@@ -78,34 +80,28 @@ def main(vectorstore):
         st.session_state.chat_history = None
 
     st.header("Chat with multiple PDFs :books:")
+
+    with st.chat_message("Assistant"):
+        st.write("Hello my name is Robert, how can i help you? ")
     user_question = st.text_input("Ask a question about your documents:")
+    with st.chat_message("User"):
+        st.write(user_question)
     if user_question:
         handle_userinput(user_question,vectorstore)
 
 
-def handle_userinput(user_question,vectorstore):
-    Rag_chain = create_conversational_rag_chain(vectorstore)
-    msgs = StreamlitChatMessageHistory(key="special_app_key")
-    response = Rag_chain({'question': user_question})
-    st.session_state.chat_history = response['chat_history']
-
-    for i, message in enumerate(st.session_state.chat_history):
-        if i % 2 == 0:
-            st.write(user_template.replace(
-                "{{MSG}}", message.content), unsafe_allow_html=True)
-        else:
-            st.write(bot_template.replace(
-                "{{MSG}}", message.content), unsafe_allow_html=True)
-
-    if 'retrieved_documents' in response:
-        st.subheader("Retrieved Documents")
-        for doc in response['source_documents']:
-            st.write(f"Document: {doc.metadata['source']}")
-            st.write(doc.page_content)       
-
+def handle_userinput(user_question,retriever):
+    docs = retriever.invoke(question)
+    
+    doc_txt = [doc.page_content for doc in docs]
+    
+    Rag_chain = create_conversational_rag_chain(retriever)
+    response = rag_chain.invoke({"context": doc_txt, "question": question})
+    with st.chat_message("Assistant"):
+        st.write(response)
 
 
-def create_conversational_rag_chain(vectorstore):
+def create_conversational_rag_chain(retriever):
     
     model_path = ('qwen2-0_5b-instruct-q4_0.gguf')
 
@@ -119,47 +115,19 @@ def create_conversational_rag_chain(vectorstore):
         n_ctx=22000,
         max_tokens=200,
         repeat_penalty=1.7,
-        callback_manager=callback_manager,
+        # callback_manager=callback_manager,
         verbose=False,
     )
 
-    contextualize_q_system_prompt = """Given a context, chat history and the latest user question
-    which maybe reference context in the chat history, formulate a standalone question
-    which can be understood without the chat history. Do NOT answer the question,
-    just reformulate it if needed and otherwise return it as is."""
-
-    contextualize_q_prompt = ChatPromptTemplate.from_messages(
-        [
-            ("system", contextualize_q_system_prompt),
-            MessagesPlaceholder("chat_history"),
-            ("human", "{input}"),
-        ]
-    )
+    template = """Answer the question based only on the following context:
+    {context}
 
-    ha_retriever = create_history_aware_retriever(llm, vectorstore.as_retriever(), contextualize_q_prompt)
+    Question: {question}
+    """
+    prompt = ChatPromptTemplate.from_template(template)
 
-    qa_system_prompt = """You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. Be as informative as possible, be polite and formal.\n{context}"""
+    rag_chain = prompt | llm | StrOutputParser()
 
-    qa_prompt = ChatPromptTemplate.from_messages(
-        [
-            ("system", qa_system_prompt),
-            MessagesPlaceholder("chat_history"),
-            ("human", "{input}"),
-        ]
-    )
-
-    question_answer_chain = create_stuff_documents_chain(llm, qa_prompt)
-    rag_chain = create_retrieval_chain(ha_retriever, question_answer_chain)
-    msgs = StreamlitChatMessageHistory(key="special_app_key")
-
-    conversational_rag_chain = RunnableWithMessageHistory(
-        rag_chain,
-        lambda session_id: msgs,
-        input_messages_key="input",
-        history_messages_key="chat_history",
-        output_messages_key="answer",
-        return_source_documents=True
-    )
 
     return rag_chain