Update app.py

app.py

@@ -53,6 +53,7 @@ from langchain.document_loaders.base import BaseLoader
 from langchain.schema import Document
 import gradio as gr
 import tempfile
+import timeit
 
 FILE_LOADER_MAPPING = {
     "csv": (CSVLoader, {"encoding": "utf-8"}),
@@ -259,7 +260,7 @@ def create_chain(llm, prompt, db):
     # )
     chain = RetrievalQA.from_chain_type(llm=llm,
                                        chain_type='stuff',
-                                       retriever=vectordb.as_retriever(search_kwargs={'k': 3}),
+                                       retriever=db.as_retriever(search_kwargs={'k': 3}),
                                        return_source_documents=True,
                                        chain_type_kwargs={'prompt': prompt}
                                        )
@@ -373,12 +374,16 @@ def main():
             if query:
                 # Load model, set prompts, create vector database, and retrieve answer
                 try:
+                    start = timeit.default_timer()
                     llm = load_model()
                     prompt = set_custom_prompt()
                     CONDENSE_QUESTION_PROMPT = set_custom_prompt_condense()
                     db = create_vector_database(loaded_documents)
                     # st.write(f"db: {db}") 
                     result, sources = retrieve_bot_answer(query,loaded_documents)
+                    end = timeit.default_timer()
+                    st.write("Elapsed time:")
+                    st.write(end - start)
                     # st.write(f"response: {response}") 
                     # Display bot response
                     st.write("Bot Response:")