Update app.py

app.py

@@ -4,12 +4,11 @@ import asyncio
 from helper_functions import *
 from rank_bm25 import BM25L
 import nest_asyncio
+import time
 nest_asyncio.apply()
 from aiogoogletrans import Translator
 import pprint
 
-
-
 # Initialize the translator
 translator = Translator()
 
@@ -53,6 +52,7 @@ async def encode_document(document: str):
     return semantic_model.encode(document, convert_to_tensor=True)
         
 async def predict(query):
+    start_time = time.time()
     normalized_query_list = (
         [normalizer.clean_text(query)]
     )
@@ -73,8 +73,10 @@ async def predict(query):
         return {"error": str(e)}
     except Exception as e:
         return {"error": f"An error occurred while making the request: {e}"}
-    
+    request_end_time = time.time()
+    request_time = request_end_time - start_time
     # Translate product representations to English
+    normalization_start_time = time.time()
     tasks = []
     for product in request_json:
         try:
@@ -90,28 +92,40 @@ async def predict(query):
         except:
             return {"error": "something wrong with the normalization step or some products are not defined correctly\nmake sure the products are in a dictionary format with fields ['name', 'brandName', 'providerName', 'categoryName'] existant."}
     
+    normalization_end_time = time.time()
+    normalization_time = normalization_end_time - normalization_time
+    
     try:
         # cateogorize products
+        categorize_start_time = time.time()
         predicted_categories = categorizer.predict(tasks)
         for idx, product in enumerate(request_json):
             product["Inferred Category"] = category_map[predicted_categories[0][idx][0]][0]
+        categorize_end_time = time.time()
+        categorize_time = categorize_end_time - categorize_start_time
     except Exception as e: 
         return {"error": f"An error occurred while categorizing products: {e}"}
-    
+
     try:
+        translation_start_time = time.time()
         representation_list = await translate_bulk(tasks)
     except Exception as e:
         representation_list = tasks
         print(f"An error occurred while translating: {e}")
+    translation_time = time.time() - translation_start_time
     
     try:
         # Tokenize representations for keyword search
+        tokenization_start_time = time.time()
         corpus = [set(representation.split(" ")) for representation in representation_list]
         keyword_search = BM25L(corpus)
+        tokenization_end_time = time.time()
+        tokenization_time = tokenization_end_time - tokenization_start_time
     except Exception as e:
         return {"error": f"An error occurred while tokenizing representations: {e}"}
     
     # Encode representations for semantic search
+    encode_start_time = time.time()
     try:
         embeddings = await asyncio.gather(
             *[encode_document(document) for document in representation_list]
@@ -122,14 +136,19 @@ async def predict(query):
             representation_list, convert_to_tensor=True
         )
         print(f"An error occurred while encoding documents: {e}")
-    
+    encode_end_time = time.time()
+    encode_time = encode_end_time - encode_start_time
     
     try:
         # Calculate interrelations between products
+        calculate_interrelations_start_time = time.time()
         calculate_interrelations(request_json, doc_embeddings)
+        calculate_interrelations_end_time = time.time()
+        calculate_interrelations_time = calculate_interrelations_end_time - calculate_interrelations_start_time
     
         # Perform hybrid search for each query
         # this will result in a dictionary of re-ranked search results for each query
+        process_time = time.time()
         for query in normalized_query_list:
             keyword_scores = check_validity(query, keyword_search)
             semantic_scores = semantic_search(query, doc_embeddings)
@@ -137,7 +156,15 @@ async def predict(query):
             is_cheapest(query, request_json)
             results[query] = rerank_results(request_json, hybrid_scores)
     
-        return print_results(results)
+        process_end_time = time.time()
+        process_time_taken = process_end_time - process_time
+        time_taken = time.time() - start_time
+        return {"results": results, "time_taken": time_taken, 
+                "request_time": request_time, "normalization_time": normalization_time,
+                "translation_time": translation_time, "categorize_time": categorize_time, 
+                "tokenization_time": tokenization_time, "encode_time": encode_time, 
+                "calculate_interrelations_time": calculate_interrelations_time,
+                "process_time": process_time_taken}
     
     except Exception as e:
         error_message = f"An error occurred during processing: {e}"
@@ -150,4 +177,4 @@ app = gr.Interface(
     title = "Re-Ranker"
     )
 
-app.launch()
+app.launch()