Update app.py

app.py

@@ -16,59 +16,61 @@ 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
 
-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 = {}
-
-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
-)
-
-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
-    )
+def main():
+    st.set_page_config(page_title="Conversational RAG Chatbot", page_icon=":robot:")
+    st.title("Conversational RAG Chatbot")
 
-    vectorstore = Chroma.from_documents(documents=text_chunks, embedding=hf, persist_directory="docs/chroma/")
-    return vectorstore
+    if "documents" not in st.session_state:
+        st.session_state.documents = []
 
-def get_pdf_text(pdf_docs):
-    document_loader = DirectoryLoader(pdf_docs)
-    return document_loader.load()
+    if "conversation_chain" not in st.session_state:
+        st.session_state.conversation_chain = None
 
-def get_text_chunks(text):
-    text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
-        separator="\n",
-        chunk_size=1000,
-        chunk_overlap=200,
-        length_function=len
-    )
-    chunks = text_splitter.split_text(text)
-    return chunks
+    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:
+        documents = load_txt_documents(data_path)
+        if not documents:
+            st.warning("No documents found in the data path.")
+        else:
+            st.session_state.documents = documents
+            docs = split_docs(documents, 350, 40)
+            vectorstore = retriever_from_chroma(docs, HuggingFaceEmbeddings(), "mmr", 7)
+            st.session_state.conversation_chain = create_conversational_rag_chain(vectorstore)
+            st.success("Documents loaded and processed successfully.")
+    except Exception as e:
+        st.error(f"An error occurred while loading documents: {e}")
+
+    if prompt := st.text_input("Enter your question:"):
+        msgs = st.session_state.get("chat_history", StreamlitChatMessageHistory(key="special_app_key"))
+        st.chat_message("human").write(prompt)
+
+        input_dict = {"input": prompt, "chat_history": msgs.messages}
+        config = {"configurable": {"session_id": "any"}}
+
+        response = st.session_state.conversation_chain.invoke(input_dict, config)
+        st.chat_message("ai").write(response["answer"])
+
+        if "docs" in response and response["documents"]:
+            for index, doc in enumerate(response["documents"]):
+                with st.expander(f"Document {index + 1}"):
+                    st.write(doc)
+
+        st.session_state["chat_history"] = msgs
 
 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})
 
     callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])
 
     llm = llamacpp.LlamaCpp(
-        model_path=os.path.join(model_path),
+        model_path=model_path,
         n_gpu_layers=1,
         temperature=0.1,
         top_p=0.9,
@@ -84,7 +86,7 @@ def create_conversational_rag_chain(vectorstore):
     which can be understood without the chat history. Do NOT answer the question,
     just reformulate it if needed and otherwise return it as is."""
 
-    ha_retriever = create_history_aware_retriever(llm, retriever, contextualize_q_system_prompt)
+    ha_retriever = history_aware_retriever(llm, vectorstore.as_retriever(), contextualize_q_system_prompt)
 
     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}"""
 
@@ -110,63 +112,5 @@ def create_conversational_rag_chain(vectorstore):
     )
     return conversation_chain
 
-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:
-        # 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"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__":
     main()