use time button added

app.py

@@ -56,6 +56,7 @@ special_threshold = st.sidebar.number_input(
     value=0.2,
     placeholder="Type a number...",
 )
+user_timer = st.sidebar.selectbox("Shall we time each step?", ("No", "Yes"))
 st.sidebar.success(
     "The 'distances' score indicates the proximity of your question to our database questions (lower is better). The 'ai_judge' ranks the similarity between user's question and database answers independently (higher is better)."
 )
@@ -71,13 +72,15 @@ if option == "YSA":
         "eagle0504/ysa-web-scrape-dataset-qa-formatted-small-version"
     )
     end_t = time.time()
-    st.success(f"{option} Database loaded. | Time: {end_t - begin_t} sec")
+    if user_timer == "Yes":
+        st.success(f"{option} Database loaded. | Time: {end_t - begin_t} sec")
     initial_input = "Tell me about YSA"
 else:
     begin_t = time.time()
     dataset = load_dataset("eagle0504/larkin-web-scrape-dataset-qa-formatted")
     end_t = time.time()
-    st.success(f"{option} Database loaded. | Time: {end_t - begin_t} sec")
+    if user_timer == "Yes":
+        st.success(f"{option} Database loaded. | Time: {end_t - begin_t} sec")
     initial_input = "Tell me about Larkin"
 
 
@@ -109,7 +112,8 @@ with st.spinner("Loading, please be patient with us ... 🙏"):
         metadatas=[{"type": "support"} for _ in range(0, L)],
     )
     end_t = time.time()
-    st.success(f"Add to VectorDB. | Time: {end_t - begin_t} sec")
+    if user_timer == "Yes":
+        st.success(f"Add to VectorDB. | Time: {end_t - begin_t} sec")
 
 
 # React to user input
@@ -124,7 +128,8 @@ if prompt := st.chat_input(initial_input):
         begin_t = time.time()
         results = collection.query(query_texts=question, n_results=5)
         end_t = time.time()
-        st.success(f"Query answser. | Time: {end_t - begin_t} sec")
+        if user_timer == "Yes":
+            st.success(f"Query answser. | Time: {end_t - begin_t} sec")
         idx = results["ids"][0]
         idx = [int(i) for i in idx]
         ref = pd.DataFrame(
@@ -138,7 +143,8 @@ if prompt := st.chat_input(initial_input):
         # special_threshold = st.sidebar.slider('How old are you?', 0, 0.6, 0.1) # 0.3
         filtered_ref = ref[ref["distances"] < special_threshold]
         if filtered_ref.shape[0] > 0:
-            st.success("There are highly relevant information in our database.")
+            if user_timer == "Yes":
+                st.success("There are highly relevant information in our database.")
             ref_from_db_search = filtered_ref["answers"].str.cat(sep=" ")
             final_ref = filtered_ref
         else:
@@ -153,7 +159,8 @@ if prompt := st.chat_input(initial_input):
                 begin_t = time.time()
                 llm_response = llama2_7b_ysa(question)
                 end_t = time.time()
-                st.success(f"Running LLM. | Time: {end_t - begin_t} sec")
+                if user_timer == "Yes":
+                    st.success(f"Running LLM. | Time: {end_t - begin_t} sec")
             except:
                 st.warning("Sorry, the inference endpoint is temporarily down. 😔")
                 llm_response = "NA."
@@ -185,7 +192,8 @@ if prompt := st.chat_input(initial_input):
             final_ref["ai_judge"] = independent_ai_judge_score
 
             end_t = time.time()
-            st.success(f"Using AI Judge. | Time: {end_t - begin_t} sec")
+            if user_timer == "Yes":
+                st.success(f"Using AI Judge. | Time: {end_t - begin_t} sec")
 
         engineered_prompt = f"""
             Based on the context: {ref_from_db_search}
@@ -198,7 +206,8 @@ if prompt := st.chat_input(initial_input):
         begin_t = time.time()
         answer = call_chatgpt(engineered_prompt)
         end_t = time.time()
-        st.success(f"Final API Call. | Time: {end_t - begin_t} sec")
+        if user_timer == "Yes":
+            st.success(f"Final API Call. | Time: {end_t - begin_t} sec")
         response = answer
 
     # Display assistant response in chat message container