Update app.py

app.py

@@ -1,15 +1,15 @@
 import streamlit as st
 import tempfile
-import pandas as pd 
-from langchain import HuggingFacePipeline 
-from transformers import AutoTokenizer 
-from langchain.embeddings import HuggingFaceEmbeddings 
-from langchain.document_loaders.csv_loader import CSVLoader 
-from langchain.vectorstores import FAISS 
-from langchain.chains import RetrievalQA 
-import transformers 
-import torch 
-import textwrap 
+import pandas as pd
+from langchain import HuggingFacePipeline
+from transformers import AutoTokenizer
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.document_loaders.csv_loader import CSVLoader
+from langchain.vectorstores import FAISS
+from langchain.chains import RetrievalQA
+import transformers
+import torch
+import textwrap
 
 def main():
     st.set_page_config(page_title="Talk with BORROWER data")
@@ -19,32 +19,25 @@ def main():
     if st.button("Submit Query", type="primary"):
         DB_FAISS_PATH = "vectorstore/db_faiss"
 
-
         loader = CSVLoader(file_path="./borrower_data.csv", encoding="utf-8", csv_args={
                     'delimiter': ','})
         data = loader.load()
 
-        model = "daryl149/llama-2-7b-chat-hf" 
-        tokenizer = AutoTokenizer.from_pretrained(model) 
-        pipeline = transformers.pipeline("text-generation", #task 
-                                          model=model, 
-                                          tokenizer=tokenizer, 
-                                          torch_dtype=torch.bfloat16, 
-                                          trust_remote_code=True, 
-                                          device_map="auto",  
-                                          do_sample=True, 
-                                          top_k=1, 
-                                          num_return_sequences=1, 
-                                          eos_token_id=tokenizer.eos_token_id 
-        ) 
-
-        llm = HuggingFacePipeline(pipeline = pipeline, model_kwargs = {'temperature':0}) 
-        embeddings = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2') 
+        model = "daryl149/llama-2-7b-chat-hf"
+        tokenizer = AutoTokenizer.from_pretrained(model)
+        pipeline = transformers.pipeline("text-generation", model=model, tokenizer=tokenizer, torch_dtype=torch.bfloat16, trust_remote_code=True, device_map="auto", do_sample=True, top_k=1, num_return_sequences=1, eos_token_id=tokenizer.eos_token_id)
+
+        llm = HuggingFacePipeline(pipeline=pipeline, model_kwargs={'temperature': 0})
+        embeddings = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2')
         vectorstore = FAISS.from_documents(data, embeddings)
         vectorstore.save_local(DB_FAISS_PATH)
-        chain=retrievalQA.from_chain_type(llm=llm, chain_type = "stuff",return_source_documents=True, retriever=vectorstore.as_retriever()) 
-        result=chain(query)
+
+        # Load the saved vectorstore
+        vectorstore = FAISS.load_local(DB_FAISS_PATH, embeddings)
+
+        chain = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", return_source_documents=True, retriever=vectorstore.as_retriever())
+        result = chain(query)
         st.write(result['result'])
 
 if __name__ == '__main__':
-    main()
+    main()