timer added

app.py

@@ -1,5 +1,6 @@
 import os
 import string
+import time
 from typing import Any, Dict, List, Tuple, Union
 
 import chromadb
@@ -8,7 +9,6 @@ import openai
 import pandas as pd
 import requests
 import streamlit as st
-import time
 from datasets import load_dataset
 from langchain.document_loaders import TextLoader
 from langchain.embeddings.sentence_transformer import SentenceTransformerEmbeddings
@@ -70,11 +70,14 @@ if option == "YSA":
     dataset = load_dataset(
         "eagle0504/youthless-homeless-shelter-web-scrape-dataset-qa-formatted"
     )
-    end_t  = time.time()
+    end_t = time.time()
     st.success(f"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"Database loaded. | Time: {end_t - begin_t} sec")
     initial_input = "Tell me about Larkin"
 
 
@@ -99,11 +102,14 @@ collection = client.create_collection(combined_string)
 # Embed and store the first N supports for this demo
 with st.spinner("Loading, please be patient with us ... 🙏"):
     L = len(dataset["train"]["questions"])
+    begin_t = time.time()
     collection.add(
         ids=[str(i) for i in range(0, L)],  # IDs are just strings
         documents=dataset["train"]["questions"],  # Enter questions here
         metadatas=[{"type": "support"} for _ in range(0, L)],
     )
+    end_t = time.time()
+    st.success(f"Add to VectorDB. | Time: {end_t - begin_t} sec")
 
 
 # React to user input
@@ -115,7 +121,10 @@ if prompt := st.chat_input(initial_input):
         st.session_state.messages.append({"role": "user", "content": prompt})
 
         question = prompt
+        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")
         idx = results["ids"][0]
         idx = [int(i) for i in idx]
         ref = pd.DataFrame(
@@ -141,7 +150,10 @@ if prompt := st.chat_input(initial_input):
 
         if option == "YSA":
             try:
+                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")
             except:
                 st.warning("Sorry, the inference endpoint is temporarily down. 😔")
                 llm_response = "NA."
@@ -158,6 +170,7 @@ if prompt := st.chat_input(initial_input):
         # add ai judge as additional rating
         if run_ai_judge == "Yes":
             independent_ai_judge_score = []
+            begin_t = time.time()
             for i in range(final_ref.shape[0]):
                 this_content = final_ref["answers"][i]
                 if len(this_content) > 3:
@@ -171,6 +184,9 @@ 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")
+
         engineered_prompt = f"""
             Based on the context: {ref_from_db_search}
 
@@ -179,7 +195,10 @@ if prompt := st.chat_input(initial_input):
             Answer the question directly (don't say "based on the context, ...")
         """
 
+        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")
         response = answer
 
     # Display assistant response in chat message container