Update app.py

app.py

@@ -16,68 +16,59 @@ from utills import load_txt_documents, split_docs, load_uploaded_documents, retr
 from langchain.text_splitter import TokenTextSplitter, RecursiveCharacterTextSplitter
 from langchain_community.document_loaders.directory import DirectoryLoader
 
-def main():
-    st.set_page_config(page_title="Chat with multiple PDFs", page_icon=":books:")
-    st.header("Chat with multiple PDFs :books:")
-
-    if "pdf_docs" not in st.session_state:
-        st.session_state.pdf_docs = []
+script_dir = os.path.dirname(os.path.abspath(__file__))
+data_path = os.path.join(script_dir, "data/")
+model_path = os.path.join(script_dir, 'qwen2-0_5b-instruct-q4_0.gguf')
+store = {}
 
-    if "conversation_chain" not in st.session_state:
-        st.session_state.conversation_chain = None
+model_name = "sentence-transformers/all-mpnet-base-v2"
+model_kwargs = {'device': 'cpu'}
+encode_kwargs = {'normalize_embeddings': True}
 
-    with st.sidebar:
-        st.subheader("Your documents")
-        pdf_docs = st.file_uploader("Upload your PDFs here and click on 'Process'", accept_multiple_files=True)
-        if pdf_docs:
-            st.session_state.pdf_docs.extend(pdf_docs)
+hf = HuggingFaceEmbeddings(
+    model_name=model_name,
+    model_kwargs=model_kwargs,
+    encode_kwargs=encode_kwargs
+)
 
-        if st.button("Process"):
-            with st.spinner("Processing"):
-                raw_text = get_pdf_text(st.session_state.pdf_docs)
-                text_chunks = get_text_chunks(raw_text)
-                vectorstore = get_vectorstore(text_chunks)
-                st.session_state.conversation_chain = get_conversation_chain(vectorstore)
-                st.success("Documents processed and conversation chain created successfully.")
+def get_vectorstore(text_chunks):
+    model_name = "sentence-transformers/all-mpnet-base-v2"
+    model_kwargs = {'device': 'cpu'}
+    encode_kwargs = {'normalize_embeddings': True}
+    hf = HuggingFaceEmbeddings(
+        model_name=model_name,
+        model_kwargs=model_kwargs,
+        encode_kwargs=encode_kwargs
+    )
 
-    user_question = st.text_input("Ask a question about your documents:")
-    if user_question:
-        handle_userinput(st.session_state.conversation_chain, user_question)
+    vectorstore = Chroma.from_documents(documents=text_chunks, embedding=hf, persist_directory="docs/chroma/")
+    return vectorstore
 
 def get_pdf_text(pdf_docs):
-    text = ""
-    for pdf in pdf_docs:
-        pdf_reader = PdfReader(pdf)
-        for page in pdf_reader.pages:
-            text += page.extract_text()
-    return text    
+    document_loader = DirectoryLoader(pdf_docs)
+    return document_loader.load()
 
 def get_text_chunks(text):
     text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
-        chunk_size=600, chunk_overlap=50,
-        separators=["\n \n \n", "\n \n", "\n1" , "(?<=\. )", " ", ""],
+        separator="\n",
+        chunk_size=1000,
+        chunk_overlap=200,
+        length_function=len
     )
     chunks = text_splitter.split_text(text)
     return chunks
 
-def get_vectorstore(text_chunks):
-    model_name = "sentence-transformers/all-mpnet-base-v2"
-    model_kwargs = {'device': 'cpu'}
-    encode_kwargs = {'normalize_embeddings': True}
-    embeddings = HuggingFaceEmbeddings(
-        model_name=model_name,
-        model_kwargs=model_kwargs,
-        encode_kwargs=encode_kwargs
-    )
-    vectorstore = Chroma.from_texts(
-        texts=text_chunks, embedding=embeddings, persist_directory="docs/chroma/")
-    return vectorstore    
+def create_conversational_rag_chain(vectorstore):
+    
+    script_dir = os.path.dirname(os.path.abspath(__file__))
+    model_path = os.path.join(script_dir, 'qwen2-0_5b-instruct-q4_0.gguf')
+    
+    retriever = vectorstore.as_retriever(search_type='mmr', search_kwargs={"k": 7})
 
-def get_conversation_chain(vectorstore):
     callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])
-    script_dir = os.path.dirname(os.path.abspath(__file__))
+
     llm = llamacpp.LlamaCpp(
-        model_path = os.path.join(script_dir, 'qwen2-0_5b-instruct-q4_0.gguf'),
+        model_path=os.path.join(model_path),
         n_gpu_layers=1,
         temperature=0.1,
         top_p=0.9,
@@ -88,26 +79,14 @@ def get_conversation_chain(vectorstore):
         verbose=False,
     )
 
-    retriever = vectorstore.as_retriever(search_type="mmr", search_kwargs={"k": 7})
-
-    contextualize_q_system_prompt = """Given a context, chat history and the latest user question, formulate a standalone question
+    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}"),
-        ]
-    )
-    
-    history_aware_retriever = create_history_aware_retriever(llm, retriever, contextualize_q_prompt)
+    ha_retriever = create_history_aware_retriever(llm, retriever, contextualize_q_system_prompt)
 
-    qa_system_prompt = """Use the following pieces of retrieved context to answer the question{input}. \
-    Be informative but don't make too long answers, be polite and formal. \
-    If you don't know the answer, say "I don't know the answer." \
-    {context}"""
+    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}"""
 
     qa_prompt = ChatPromptTemplate.from_messages(
         [
@@ -116,12 +95,12 @@ def get_conversation_chain(vectorstore):
             ("human", "{input}"),
         ]
     )
-    
+
     question_answer_chain = create_stuff_documents_chain(llm, qa_prompt)
 
-    rag_chain = create_retrieval_chain(history_aware_retriever, question_answer_chain)
+    rag_chain = create_retrieval_chain(ha_retriever, question_answer_chain)
     msgs = StreamlitChatMessageHistory(key="special_app_key")
-    
+
     conversation_chain = RunnableWithMessageHistory(
         rag_chain,
         lambda session_id: msgs,
@@ -131,17 +110,63 @@ def get_conversation_chain(vectorstore):
     )
     return conversation_chain
 
-def handle_userinput(conversation_chain, prompt):
-    msgs = StreamlitChatMessageHistory(key="special_app_key")
-    st.chat_message("human").write(prompt)
-    input_dict = {"input": prompt, "chat_history": msgs.messages}
-    config = {"configurable": {"session_id": 1}}
+def main():
+    """Main function for the Streamlit app."""
+    # Initialize chat history if not already present in session state
+
+    documents = []
+    
+    script_dir = os.path.dirname(os.path.abspath(__file__))
+    data_path = os.path.join(script_dir, "data/")
+    if not os.path.exists(data_path):
+        st.error(f"Data path does not exist: {data_path}")
+        return
     
     try:
-        response = conversation_chain.invoke(input_dict, config)
-        st.chat_message("ai").write(response["answer"])
+        # Load documents from the data path
+        documents = load_txt_documents(data_path)
+        if not documents:
+            st.warning("No documents found in the data path.")
+        else:
+            # Split the documents into chunks
+            docs = split_docs(documents, 350, 40)
+            # Add your logic here to use `docs`
+            st.success("Documents loaded and processed successfully.")
     except Exception as e:
-        st.error(f"Error invoking conversation chain: {e}")
+        st.error(f"An error occurred while loading documents: {e}")
+
+    
+    
+    
+    documents = load_txt_documents(data_path)
+    docs = split_docs(documents, 350, 40)
+
+    vectorstore = get_vectorstore(docs)
+    
+    msgs = st.session_state.get("chat_history", StreamlitChatMessageHistory(key="special_app_key"))
+    chain_with_history = create_conversational_rag_chain(vectorstore)
+
+    st.title("Conversational RAG Chatbot")
+
+    if prompt := st.chat_input():
+        st.chat_message("human").write(prompt)
+
+        # Prepare the input dictionary with the correct keys
+        input_dict = {"input": prompt, "chat_history": msgs.messages}
+        config = {"configurable": {"session_id": "any"}}
+
+        # Process user input and handle response
+        response = chain_with_history.invoke(input_dict, config)
+        st.chat_message("ai").write(response["answer"])
+
+        # Display retrieved documents (if any and present in response)
+        if "docs" in response and response["documents"]:
+            for index, doc in enumerate(response["documents"]):
+                with st.expander(f"Document {index + 1}"):
+                    st.write(doc)
+
+    # Update chat history in session state
+    st.session_state["chat_history"] = msgs
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()