Update app.py

app.py

@@ -46,6 +46,7 @@ def load_scraped_web_info():
 
 
 
+
 @st.cache_resource
 def load_embedding_model():
     embedding_model = HuggingFaceInstructEmbeddings(model_name='hkunlp/instructor-base',
@@ -57,6 +58,31 @@ def load_faiss_index():
     vector_database = FAISS.load_local("faiss_index", embedding_model)
     return vector_database
 
+@st.cache_resource
+def load_llm_model():
+    # llm = HuggingFacePipeline.from_model_id(model_id= 'lmsys/fastchat-t5-3b-v1.0', 
+    #                                         task= 'text2text-generation',
+    #                                         model_kwargs={ "device_map": "auto",
+    #                                                     "load_in_8bit": True,"max_length": 256, "temperature": 0,
+    #                                                     "repetition_penalty": 1.5})
+    
+    
+    llm = HuggingFacePipeline.from_model_id(model_id= 'lmsys/fastchat-t5-3b-v1.0', 
+                                        task= 'text2text-generation',
+                                        
+                                        model_kwargs={ "max_length": 256, "temperature": 0,
+                                                      "torch_dtype":torch.float32,
+                                                    "repetition_penalty": 1.3})
+
+
+    return llm
+
+
+def load_retriever(llm, db):
+    qa_retriever = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff",
+                            retriever=db.as_retriever())
+
+    return qa_retriever
 
 #--------------
 
@@ -65,19 +91,36 @@ def load_faiss_index():
 load_scraped_web_info()
 embedding_model = load_embedding_model()
 vector_database = load_faiss_index()
-print("load done")
+llm_model = load_llm_model()
+qa_retriever = load_retriever(llm= llm_model, db= vector_database)
 
 
+print("all load done")
+
 
 
 query_input = st.text_input(label= 'your question')
+
+
 def retrieve_document(query_input):
     related_doc = vector_database.similarity_search(query_input)
     return related_doc
 
-output = st.text_area(label = "Here is the relevant documents",
+def retrieve_answer(query_input):
+    answer = qa_retriever.run(query_input)
+    return answer
+
+output_1 = st.text_area(label = "Here is the relevant documents",
                       value = retrieve_document(query_input))
 
 
-faiss_retriever = vector_database.as_retriever()
-print("Succesfully had FAISS as retriever")
+output_2 = st.text_area(label = "Here is the answer",
+                      value = retrieve_answer(query_input))
+
+
+
+# faiss_retriever = vector_database.as_retriever()
+# print("Succesfully had FAISS as retriever")
+
+
+