feat: add timer

app.py

@@ -1,6 +1,7 @@
 import streamlit as st
 from st_pages import Page, show_pages, add_page_title, Section 
 from lib.utils.model import get_model, get_similarities 
+from lib.utils.timer import timer 
 
 add_page_title()
 
@@ -39,8 +40,11 @@ if button:
 
     st.text(f'IRRA model loaded with {sum(p.numel() for p in model.parameters()) / 1e6:.0f}M parameters')
     
+    time = timer()
     with st.spinner('Computing and ranking similarities'):
-        similarities = get_similarities(caption, images, model).squeeze(0)
+        with timer() as t:
+            similarities = get_similarities(caption, images, model).squeeze(0)
+    elapsed = t()
 
     indices = similarities.argsort(descending=True).cpu().tolist()[:ranks]
     
@@ -61,6 +65,8 @@ if button:
         with c3:
             st.text(f'{similarities[idx].cpu():.2f}')
 
+    st.success(f'Done in {elapsed:.2f}s')
+
 with st.sidebar:
     st.title('IRRA Text-To-Image Retrival')
 

lib/utils/timer.py

@@ -0,0 +1,7 @@
+from time import perf_counter
+from contextlib import contextmanager
+
+@contextmanager
+def timer():
+    start = perf_counter()
+    yield lambda: perf_counter() - start