Add timer

backend/semantic_search.py

@@ -6,6 +6,8 @@ import gradio as gr
 from sentence_transformers import SentenceTransformer
 from FlagEmbedding import FlagReranker
 
+from gradio_app.utils.time_decorator import timeit
+
 db = lancedb.connect(".lancedb")
 
 TABLE = db.open_table(os.getenv("TABLE_NAME"))
@@ -30,6 +32,7 @@ def rerank(query: str, documents: typing.List[str], k: int):
     return [documents[i] for i in best_indices]
 
 
+@timeit
 def retrieve(query, k):
     query_vec = retriever.encode(query)
     try:

utils/time_decorator.py

@@ -0,0 +1,14 @@
+from functools import wraps
+import time
+
+
+def timeit(func):
+    @wraps(func)
+    def timeit_wrapper(*args, **kwargs):
+        start_time = time.perf_counter()
+        result = func(*args, **kwargs)
+        end_time = time.perf_counter()
+        total_time = end_time - start_time
+        print(f'Function {func.__name__}{args} {kwargs} Took {total_time:.4f} seconds')
+        return result
+    return timeit_wrapper