Create app.py

app.py

@@ -0,0 +1,85 @@
+import time
+import json
+
+import streamlit as st
+from pathlib import Path
+
+import sys
+path_root = Path("./")
+sys.path.append(str(path_root))
+
+
+st.set_page_config(page_title="PSC Runtime",
+                   page_icon='🌸', layout="centered")
+
+# cola, colb, colc = st.columns([5, 4, 5])
+
+# colaa, colbb, colcc = st.columns([1, 8, 1])
+# with colbb:
+#     runtime = st.select_slider(
+#         'Select a runtime type',
+#         options=['PyTorch', 'ONNX Runtime'])
+#     st.write('Now using: ', runtime)
+
+
+# colaa, colbb, colcc = st.columns([1, 8, 1])
+# with colbb:
+#     encoder = st.select_slider(
+#         'Select a query encoder',
+#         options=['uniCOIL', 'SPLADE++ Ensemble Distil', 'SPLADE++ Self Distil'])
+#     st.write('Now Running Encoder: ', encoder)
+
+# if runtime == 'PyTorch':
+#     runtime = 'pytorch'
+#     runtime_index = 1
+# else:
+#     runtime = 'onnx'
+#     runtime_index = 0
+
+
+col1, col2 = st.columns([9, 1])
+with col1:
+    search_query = st.text_input(label="search query", placeholder="Search")
+
+with col2:
+    st.write('#')
+    button_clicked = st.button("🔎")
+
+
+import torch
+fn = ""
+object = torch.load(fn)
+outputs = [x[2] for x in object]
+query2outputs = {}
+for output in outputs:
+    all_queries = {x['query'] for x in output}
+    assert len(all_queries) == 1
+    query = list(all_queries)[0]
+    query2outputs[query] = [x['hits'] for x in output]
+
+search_query = sorted(query2outputs)[0]
+
+if search_query or button_clicked:
+    
+    num_results = None
+    t_0 = time.time()
+    search_results = query2outputs[search_query]
+
+    st.write(
+        f'<p align=\"right\" style=\"color:grey;\"> Before aggregation for query [{search_query}] ms</p>', unsafe_allow_html=True)
+
+    for i, result in enumerate(search_results):
+        result_id = result["docid"]
+        contents = result["content"]
+
+        # output = f'<div class="row"> <b>Rank</b>: {i+1} | <b>Document ID</b>: {result_id} | <b>Score</b>:{result_score:.2f}</div>'
+        output = f'<div class="row"> <b>Rank</b>: {i+1} | <b>Document ID</b>: {result_id}'
+
+        try:
+            st.write(output, unsafe_allow_html=True)
+            st.write(
+                f'<div class="row">{contents}</div>', unsafe_allow_html=True)
+
+        except:
+            pass
+        st.write('---')