Update app.py

app.py

@@ -29,50 +29,44 @@ def set_custom_prompt():
                             input_variables=['context', 'question'])
     return prompt
 
-#Retrieval QA Chain
+# Retrieval QA Chain
 def retrieval_qa_chain(llm, prompt, db):
     qa_chain = RetrievalQA.from_chain_type(llm=llm,
-                                       chain_type='stuff',
-                                       retriever=db.as_retriever(search_kwargs={'k': 2}),
-                                       return_source_documents=True,
-                                       chain_type_kwargs={'prompt': prompt}
-                                       )
+                                           chain_type='stuff',
+                                           retriever=db.as_retriever(search_kwargs={'k': 2}),
+                                           return_source_documents=True,
+                                           chain_type_kwargs={'prompt': prompt}
+                                           )
     return qa_chain
 
-#Loading the model
+# Loading the model
 def load_llm():
     # Load the locally downloaded model here
     llm = CTransformers(
-        model = "TheBloke/Llama-2-7B-Chat-GGML",
+        model="TheBloke/Llama-2-7B-Chat-GGML",
         model_type="llama",
-        max_new_tokens = 512,
-        temperature = 0.5
+        max_new_tokens=512,
+        temperature=0.5
     )
     return llm
 
-#QA Model Function
+# QA Model Function
 def qa_bot():
     embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2",
                                        model_kwargs={'device': 'cpu'})
-    db = FAISS.load_local(DB_FAISS_PATH, embeddings)
+    db = FAISS.load_local(DB_FAISS_PATH, embeddings, allow_dangerous_deserialization=True)
     llm = load_llm()
     qa_prompt = set_custom_prompt()
     qa = retrieval_qa_chain(llm, qa_prompt, db)
-    db = faiss.load_local(DB_FAISS_PATH, embeddings, allow_dangerous_deserialization=True)
-    # 1. Verify file ownership
-    if not os.path.exists(DB_FAISS_PATH) or os.stat(DB_FAISS_PATH).st_uid != os.getuid():
-        raise RuntimeError("Vector store file may have been tampered with. Aborting.")
-
-
     return qa
 
-#output function
+# Output function
 def final_result(query):
     qa_result = qa_bot()
     response = qa_result({'query': query})
     return response
 
-#chainlit code
+# Chainlit code
 @cl.on_chat_start
 async def start():
     chain = qa_bot()
@@ -85,7 +79,7 @@ async def start():
 
 @cl.on_message
 async def main(message: cl.Message):
-    chain = cl.user_session.get("chain") 
+    chain = cl.user_session.get("chain")
     cb = cl.AsyncLangchainCallbackHandler(
         stream_final_answer=True, answer_prefix_tokens=["FINAL", "ANSWER"]
     )
@@ -101,7 +95,6 @@ async def main(message: cl.Message):
 
     await cl.Message(content=answer).send()
 
-
 def main():
     st.title("Medical Bot")
     st.text_input("Enter your query:", key="query")
@@ -113,4 +106,4 @@ def main():
         st.markdown(response) 
 
 if __name__ == "__main__":
-    main()
+    main()