Update app.py

app.py

@@ -1,24 +1,34 @@
 import subprocess
 
-
 import streamlit as st
 from dotenv import load_dotenv
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.vectorstores import Chroma
 from langchain.embeddings import FastEmbedEmbeddings  # General embeddings from HuggingFace models.
 from langchain.memory import ConversationBufferMemory
-from langchain.chains import ConversationalRetrievalChain
+from langchain.callbacks.manager import CallbackManager
+from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
 from htmlTemplates import css, bot_template, user_template
 from langchain.llms import LlamaCpp  # For loading transformer models.
 from langchain.document_loaders import PyPDFLoader, TextLoader, JSONLoader, CSVLoader
 import tempfile 
 from langchain.chains import RetrievalQA
 from langchain.prompts import PromptTemplate
+from langchain import hub
 import os
 import glob
 
 
+# TEXT LOADERS
 def get_pdf_text(pdf_docs):
+    """
+    Purpose: A hypothetical loader for PDF files in Python.
+    Usage: Used to extract text or other information from PDF documents.
+    Load Function: A load_pdf function might be used to read and extract data from a PDF file.
+
+    input : pdf document path
+    returns : extracted text 
+    """
     temp_dir = tempfile.TemporaryDirectory() 
     temp_filepath = os.path.join(temp_dir.name, pdf_docs.name) 
 
@@ -31,6 +41,8 @@ def get_pdf_text(pdf_docs):
 
 
 def get_text_file(text_docs):
+    """
+    """
     temp_dir = tempfile.TemporaryDirectory()
     temp_filepath = os.path.join(temp_dir.name, text_docs.name)
     
@@ -70,8 +82,8 @@ def get_json_file(json_docs):
 
 def get_text_chunks(documents):
     text_splitter = RecursiveCharacterTextSplitter(
-        chunk_size=1000,  
-        chunk_overlap=200,  
+        chunk_size=512,  
+        chunk_overlap=50,  
         length_function=len  
     )
 
@@ -81,25 +93,24 @@ def get_text_chunks(documents):
 
 
 def get_vectorstore(text_chunks, embeddings):
-    # embeddings = FastEmbedEmbeddings( model_name= "BAAI/bge-small-en-v1.5", 
-    #                              cache_dir="./embedding_model/")
-     
+    
     vectorstore = Chroma.from_documents(documents= text_chunks, 
-                               embedding= embeddings,
+                               embedding= st.session_state.embeddings,
                                persist_directory= "./vectordb/")
+    # Document stored 
     return vectorstore  
 
 def get_conversation_chain(vectorstore):
-    # model_name_or_path = 'TheBloke/Llama-2-7B-chat-GGUF'
-    # model_basename = 'llama-2-7b-chat.Q2_K.gguf'
-    model_path = "./models/llama-2-13b-chat.Q4_K_S.gguf"
 
-    llm = LlamaCpp(model_path="./models/llama-2-13b-chat.Q4_K_S.gguf", 
-                   template = 0.4,
+    model_path = "models/llama-2-13b-chat.Q4_K_S.gguf"
+    callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])
+
+    llm = LlamaCpp(model_path= model_path, 
                    n_ctx=4000, 
-                   max_tokens=4000,
+                   max_tokens= 500,
                    n_gpu_layers = 50, 
                    n_batch = 512,
+                   callback_manager = callback_manager
                    verbose=True)
     
     memory = ConversationBufferMemory(
@@ -110,7 +121,8 @@ def get_conversation_chain(vectorstore):
     You are a Experience human Resource Manager. When the employee asks you a question, you will have to refer the company policy and respond in a professional way. Make sure to sound Empethetic while being professional and sound like a Human!
     Try to summarise the content and keep the answer to the point.
     If you don't know the answer, just say that you don't know, don't try to make up an answer.
-    When generating answer for the given question make sure to follow the example template!
+
+    Followe the template below
     Example: 
     Question : how many paid leaves do i have ?
     Answer : The number of paid leaves varies depending on the type of leave, like privilege leave you're entitled to a maximum of 21 days in a calendar year. Other leaves might have different entitlements. thanks for asking!
@@ -120,16 +132,22 @@ def get_conversation_chain(vectorstore):
     
     Question: {question}
     Answer:
+
+    Just answer to the point! 
     """
 
     rag_prompt_custom = PromptTemplate.from_template(template)
+
+    # prompt = hub.pull("rlm/rag-prompt")
     
     conversation_chain = RetrievalQA.from_chain_type(
         llm,
         retriever=vectorstore.as_retriever(),
         chain_type_kwargs={"prompt": rag_prompt_custom},
-        memory = memory
     )
+    conversation_chain.callback_manager = callback_manager
+    conversation_chain.memory = ConversationBufferMemory()
+
     return conversation_chain
 
 
@@ -154,6 +172,7 @@ def handle_userinput():
         if clear:
             st.session_state.conversation.clean()
         msg = st.session_state.conversation.run(prompt)
+        print(msg)
         st.session_state.messages.append({"role": "assistant", "content": msg})
         st.chat_message("assistant").write(msg)
 
@@ -182,7 +201,8 @@ def main():
     st.subheader("🚀 A HR powered by Generative AI")
     # user_question = st.text_input("Ask a question about your documents:")
 
-    st.session_state.embeddings = embeddings = FastEmbedEmbeddings( model_name= "BAAI/bge-small-en-v1.5", cache_dir="./embedding_model/")
+    st.session_state.embeddings = FastEmbedEmbeddings( model_name= "BAAI/bge-small-en-v1.5", 
+                                                                    cache_dir="./embedding_model/")
 
     if len(glob.glob("./vectordb/*.sqlite3")) > 0:
         
@@ -227,7 +247,6 @@ def main():
 
 
 if __name__ == '__main__':
-    # Define the command
     command = 'CMAKE_ARGS="-DLLAMA_CUBLAS=on" FORCE_CMAKE=1 pip install llama-cpp-python --no-cache-dir'
     
     # Run the command using subprocess
@@ -237,4 +256,5 @@ if __name__ == '__main__':
     except subprocess.CalledProcessError as e:
         print(f"Error: {e}")
         
-    main()
+    main()
+