Upload 5 files

app.py

@@ -0,0 +1,61 @@
+
+import streamlit as st
+from streaming import StreamHandler
+import utils
+from langchain.callbacks.manager import CallbackManager
+from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler 
+from retriever import retriever
+from langchain.chains import RetrievalQA
+from langchain.llms import LlamaCpp
+from dotenv import load_dotenv
+
+class CustomDataChatbot:
+    def __init__(self):
+        # Initialize session state variables, including messages
+        st.session_state.messages = []
+
+    @st.spinner('Analyzing documents..')
+    def setup_qa_chain(self):
+        # Setup memory for contextual conversation        
+        # memory = ConversationBufferMemory(
+        #     memory_key='chat_history',
+        #     return_messages=True
+        # )
+
+        callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])
+        # Setup LLM and QA chain
+        llm = LlamaCpp(model_path="./models/openhermes-2.5-neural-chat-7b-v3-1-7b.Q5_K_M.gguf",
+            temperature=0.34,
+            max_tokens=4000,
+            n_ctx=4096,
+            top_p=1,
+            callback_manager=callback_manager,
+            verbose=True)
+
+        # qa_chain = ConversationalRetrievalChain.from_llm(llm, retriever=retriever(), memory=memory, verbose=True)
+
+
+        return RetrievalQA.from_chain_type( llm, retriever= retriever())
+
+    @utils.enable_chat_history
+    def main(self):
+        load_dotenv()
+        st.set_page_config(page_title="ChatPDF", page_icon="📄")
+        st.header('Chat with your documents')
+        st.write('Has access to custom documents and can respond to user queries by referring to the content within those documents')
+        st.write('[![view source code ](https://img.shields.io/badge/view_source_code-gray?logo=github)](https://github.com/shashankdeshpande/langchain-chatbot/blob/master/pages/4_%F0%9F%93%84_chat_with_your_documents.py)')
+
+        user_query = st.chat_input(placeholder="Ask me anything!")
+
+        if user_query:
+            qa_chain = self.setup_qa_chain()
+            utils.display_msg(user_query, 'user')
+
+            with st.chat_message("assistant"):
+                st_cb = StreamHandler(st.empty())
+                response = qa_chain.run(user_query, callbacks=[st_cb])
+                st.session_state.messages.append({"role": "assistant", "content": response})
+
+if __name__ == "__main__":
+    obj = CustomDataChatbot()
+    obj.main()

retriever.py

@@ -0,0 +1,38 @@
+import pickle
+from langchain.retrievers import EnsembleRetriever
+from langchain.vectorstores import FAISS
+from langchain.embeddings import HuggingFaceEmbeddings
+
+from langchain.embeddings import HuggingFaceEmbeddings
+from transformers import AutoModel
+
+
+
+def retriever():
+
+    # Embeddings 
+    # Defign our Embedding Model
+
+    model_name = "jinaai/jina-embeddings-v2-base-en"
+    model_kwargs = {'device': 'cpu'}
+    encode_kwargs = {'normalize_embeddings': False, } 
+
+    model = AutoModel.from_pretrained( model_name, trust_remote_code=True) 
+
+    embeddings = HuggingFaceEmbeddings( model_name=model_name, 
+                                    model_kwargs=model_kwargs, 
+                                    encode_kwargs=encode_kwargs)
+    
+
+    #to read bm25 object
+    with open('./bm25', 'rb') as file:
+        bm25_retriever = pickle.load(file)
+
+    bm25_retriever.k = 2
+
+    # Load FAISS 
+    faiss_vectorstore = FAISS.load_local("./Vector_DB/", embeddings)
+    faiss_retriever = faiss_vectorstore.as_retriever(search_kwargs={"k": 1})
+
+    # initialize the ensemble retriever
+    return  EnsembleRetriever( retrievers=[bm25_retriever, faiss_retriever], weights=[0.5, 0.5] )

streaming.py

@@ -0,0 +1,11 @@
+from langchain.callbacks.base import BaseCallbackHandler
+
+class StreamHandler(BaseCallbackHandler):
+    
+    def __init__(self, container, initial_text=""):
+        self.container = container
+        self.text = initial_text
+
+    def on_llm_new_token(self, token: str, **kwargs):
+        self.text += token
+        self.container.markdown(self.text)

utils.py

@@ -0,0 +1,39 @@
+import os
+import random
+import streamlit as st
+
+#decorator
+def enable_chat_history(func):
+    if os.environ.get("OPENAI_API_KEY"):
+
+        # to clear chat history after swtching chatbot
+        current_page = func.__qualname__
+        if "current_page" not in st.session_state:
+            st.session_state["current_page"] = current_page
+        if st.session_state["current_page"] != current_page:
+            try:
+                st.cache_resource.clear()
+                del st.session_state["current_page"]
+                del st.session_state["messages"]
+            except:
+                pass
+
+        # to show chat history on ui
+        if "messages" not in st.session_state:
+            st.session_state["messages"] = [{"role": "assistant", "content": "How can I help you?"}]
+        for msg in st.session_state["messages"]:
+            st.chat_message(msg["role"]).write(msg["content"])
+
+    def execute(*args, **kwargs):
+        func(*args, **kwargs)
+    return execute
+
+def display_msg(msg, author):
+    """Method to display message on the UI
+
+    Args:
+        msg (str): message to display
+        author (str): author of the message -user/assistant
+    """
+    st.session_state.messages.append({"role": author, "content": msg})
+    st.chat_message(author).write(msg)