Update app.py

app.py

@@ -1,84 +1,60 @@
 import streamlit as st
-from llmware.prompts import Prompt
-import io, os, re
-import PyPDF2
+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 ConversationalRetrievalChain
+from langchain_community.llms import HuggingFaceHub
+from langchain.memory import ConversationBufferMemory
 
-def register_gguf_model():
+def make_vectorstore(embeddings):
+    # use glob to find all the pdf files in the data folder in the base directory
+    loader = PyPDFDirectoryLoader("data")   
 
-    prompter = Prompt()
-    your_model_name = "llama"
-    hf_repo_name = "TheBloke/Llama-2-7B-Chat-GGUF"
-    model_file = "llama-2-7b-chat.Q3_K_M.gguf"
-    print("registering models")
-    prompter.model_catalog.register_gguf_model(your_model_name,hf_repo_name, model_file, prompt_wrapper="open_chat")
-    your_model_name = "open_gpt4"
-    hf_repo_name = "TheBloke/Open_Gpt4_8x7B-GGUF"
-    model_file = "open_gpt4_8x7b.Q3_K_M.gguf"
-    prompter.model_catalog.register_gguf_model(your_model_name,hf_repo_name, model_file, prompt_wrapper="open_chat")
-    your_model_name = "phi2"
-    hf_repo_name = "TheBloke/phi-2-GGUF"
-    model_file = "phi-2.Q3_K_M.gguf"
-    prompter.model_catalog.register_gguf_model(your_model_name,hf_repo_name, model_file, prompt_wrapper="open_chat")
-    your_model_name = "mistral"
-    hf_repo_name = "TheBloke/Mistral-7B-Instruct-v0.2-GGUF"
-    model_file = "mistral-7b-instruct-v0.2.Q3_K_M.gguf"
-    prompter.model_catalog.register_gguf_model(your_model_name,hf_repo_name, model_file, prompt_wrapper="open_chat")
-    return prompter
+    # load the documents
+    documents = loader.load()
 
-def main():
-    st.title("BetterZila RAG Enabled LLM")
-    with st.spinner("Registering Models for use..."):
-        prompter = register_gguf_model()
-    
-    data_path = "data/"
-    
-    st.sidebar.subheader("Select Model")
-    model_name = st.sidebar.selectbox("Select Model", ["llama", "open_gpt4", "phi2", "mistral"])
-    with st.spinner("Loading Model..."):
-        prompter.load_model(model_name)
-    st.success("Model Loaded!")
+    # split the documents into chunks of 1400 characters with 0 overlap
+    text_splitter = CharacterTextSplitter(chunk_size=1400, chunk_overlap=0)
+
+    # split the documents into chunks of 1400 characters with 0 overlap
+    texts = text_splitter.split_documents(documents)
+
+    # create a vector store from the documents
+    docsearch = FAISS.from_documents(texts, embeddings)
     
-    queries = ['Can you give me an example from history where the enemy was crushed totally from the book?', "What's the point of making myself less accessible?", "Can you tell me the story of Queen Elizabeth I from this 48 laws of power book?"]
+    return docsearch
+
+def get_conversation(vectorstore):
+
+    # create a memory object to store the conversation history
+    memory = ConversationBufferMemory(memory_key="chat_history",return_messages=True,)
+
+    conversation_chain = ConversationalRetrievalChain.from_chain_type(
+        llm=HuggingFaceHub(repo_id="google/flan-t5-xxl", model_kwargs={"temperature":0.5, "max_length":512}, huggingfacehub_api_token = st.secrets["hf_token"]),
+        chain_type="stuff", 
+        retriever=vectorstore.as_retriever(),
+        memory=memory)
     
-    st.subheader("Query")
+    return conversation_chain
 
-    with st.spinner("Loading PDF file..."):
-        for file in os.listdir(data_path):
-            if file.endswith(".pdf"):
-                print("Found PDF file: ", file)
-                pdf_file = file
-                break
-    print("loading Source...")
-    source = prompter.add_source_document(data_path, pdf_file, query=None)
+def get_response(conversation_chain, query):
+    # get the response
+    response = conversation_chain.run(query)
+    return response
 
+def main():
+    st.title("BetterZila RAG Enabled LLM")
+    embeddings = HuggingFaceInstructEmbeddings(model_name="google/t5-v1_1-xl", model_kwargs = {'device': 'cpu'})
+    vectorstore = make_vectorstore(embeddings)
+    conversation_chain = get_conversation(vectorstore)
+    queries = ["Can you give me an example from history where the enemy was crushed totally from the book?", "What's the point of making myself less accessible?", "Can you tell me the story of Queen Elizabeth I from this 48 laws of power book?"]
     for query in queries:
         st.subheader(f"Query: {query}")
-        with st.spinner("Generating response..."):
-            responses = prompter.prompt_with_source(query, prompt_name="just_the_facts", temperature=0.3)
-            
-            for r, response in enumerate(responses):
-                st.write(query)
-                st.write(re.sub("[\n]", " ", response["llm_response"]).strip())
-
+        response = get_response(conversation_chain, query)
+        st.write(query)
+        st.write(response["llm_response"])
     st.success("Responses generated!")
-
-    # for query in queries:
-    #     st.subheader(f"Query: {query}")
-    #     with st.spinner("Generating response..."):
-    #         for file in os.listdir(data_path):
-    #             if file.endswith(".pdf"):
-    #                 print("Found PDF file: ", file)
-    #                 print("loading Source...")
-    #                 source = prompter.add_source_document(data_path, file, query=None)
-    #                 print("generating response...")
-    #                 responses = prompter.prompt_with_source(query, prompt_name="just_the_facts", temperature=0.3)
-    #                 print("response generated!")
-    #                 for r, response in enumerate(responses):
-    #                     print(query, ":", re.sub("[\n]"," ", response["llm_response"]).strip())
-    #                 prompter.clear_source_materials()
-    #                 st.write(query)
-    #                 st.write(re.sub("[\n]"," ", response["llm_response"]).strip())
-    #     st.success("Response generated!")
-        
+    
 if __name__ == "__main__":
-    main()
+    main()