Update app.py

app.py

@@ -1,68 +1,83 @@
 import streamlit as st
-from langchain.document_loaders import PyPDFLoader
-from langchain.indexes import VectorstoreIndexCreator
-from langchain.chains import RetrievalQA, RetrievalQAWithSourcesChain
-from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
+from langchain_community.document_loaders.pdf import PyPDFDirectoryLoader
+from langchain.text_splitter import CharacterTextSplitter
+from langchain_community.embeddings import HuggingFaceInstructEmbeddings
+from langchain_community.vectorstores import FAISS
+from langchain.chains import RetrievalQA
+from langchain.memory import ConversationBufferMemory
+from langchain_community.llms import HuggingFaceHub
 
+def make_vectorstore(embeddings):
+    loader = PyPDFDirectoryLoader("data") 
+    documents = loader.load()
+    text_splitter = CharacterTextSplitter(chunk_size=200, chunk_overlap=0)
+    texts = text_splitter.split_documents(documents)
+    docsearch = FAISS.from_documents(texts, embeddings)
     
-def respond_to_question(question, model, tokenizer):
+    return docsearch
 
-    prompt = [{'role': 'user', 'content': question}]
-    inputs = tokenizer.apply_chat_template(
-        prompt,
-        add_generation_prompt=True,
-        return_tensors='pt'
-    )
+def get_conversation(vectorstore, model):
+    
+    memory = ConversationBufferMemory(memory_key="messages", return_messages=True)
 
-    tokens = model.generate(
-        inputs.to(model.device),
-        max_new_tokens=1024,
-        temperature=0.8,
-        do_sample=True
-    )
+    conversation_chain = RetrievalQA.from_llm(
+        llm=model,
+        retriever=vectorstore.as_retriever(),
+        memory=memory)
+    
+    return conversation_chain
 
-    print(tokenizer.decode(tokens[0], skip_special_tokens=False))
-    return tokenizer.decode(tokens[0], skip_special_tokens=False)
-    # prompt = "write me a python function that prints the fibonacci sequence"
-    # messages = [
-    #     {
-    #         "role": "system",
-    #         "content": "You are a friendly chatbot who can code",
-    #     },
-    #     {"role": "user", "content": prompt},
-    # ]
-    # prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    # outputs = pipe(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
-    # print(outputs[0]["generated_text"].split("<|assistant|>")[1])
-    # return outputs[0]["generated_text"].split("<|assistant|>")[1]
+def get_response(conversation_chain, query):
+    # get the response
+    response = conversation_chain.invoke(query)
+    response = response["result"]
+    answer = response.split('\nHelpful Answer: ')[1]
+    return answer
 
 def main():
-    st.title("LangChain Demo")
-
-    tokenizer = AutoTokenizer.from_pretrained('stabilityai/stablelm-zephyr-3b')
-    model = AutoModelForCausalLM.from_pretrained(
-            'stabilityai/stablelm-zephyr-3b',
-            trust_remote_code=True,
-            device_map="auto"
-        )
-
-    if 'messages' not in st.session_state:
+    st.title("Chat LLM")
+    
+    print("Downloading Embeddings Model")
+    with st.spinner('Downloading Embeddings Model...'):
+        embeddings = HuggingFaceInstructEmbeddings(model_name="hkunlp/instructor-base", model_kwargs = {'device': 'cpu'})
+    
+    print("Loading LLM from HuggingFace")
+    with st.spinner('Loading LLM from HuggingFace...'):
+        llm = HuggingFaceHub(repo_id="HuggingFaceH4/zephyr-7b-beta", model_kwargs={"temperature":0.7, "max_new_tokens":512, "top_p":0.95, "top_k":50},)
+    
+    # multiple pdfs uploader in the side bar
+    st.sidebar.title("Upload PDFs")
+    uploaded_files = st.sidebar.file_uploader("Upload PDFs", accept_multiple_files=True)
+    if uploaded_files:
+        for file in uploaded_files:
+            with open(f"data/{file.name}", "wb") as f:
+                f.write(file.getbuffer())
+        st.sidebar.success("PDFs uploaded successfully")
+    else:
+        st.sidebar.warning("Please upload PDFs")
+    # add a clear chat button which will clear the session state and the conversation history
+    
+    if "messages" not in st.session_state:
         st.session_state.messages = []
-        
-    prompt = st.text_input("Enter your question here:")
-
+    
     for message in st.session_state.messages:
-        st.chat_message(message['role']).markdown(message['text'])
-
-    if prompt:
-        st.session_state.messages.append({'role': 'user', 'text': prompt})
-        st.chat_message("user").markdown(prompt)
-        model_response = respond_to_question(prompt, model, tokenizer)
-        st.session_state.messages.append({'role': 'Assistant', 'text': model_response})
-        st.chat_message("system").markdown(model_response)
+        if message["role"] == "user":
+            st.chat_message("user").markdown(message["content"])
+        else:
+            st.chat_message("bot").markdown(message["content"])
+    
+    with st.spinner('making a vectorstore database...'):
+        vectorstore = make_vectorstore(embeddings)
+    with st.spinner('making a conversation chain...'):
+        conversation_chain = get_conversation(vectorstore, llm)
+    
+    user_prompt = st.chat_input("ask a question", key="user")
+    if user_prompt:
+        st.chat_message("user").markdown(user_prompt)
+        st.session_state.messages.append({"role": "user", "content": user_prompt})
+        response = get_response(conversation_chain, user_prompt)
+        st.chat_message("bot").markdown(response)
+        st.session_state.messages.append({"role": "bot", "content": response})
         
 if __name__ == "__main__":
     main()