Update app.py

app.py

@@ -131,9 +131,9 @@ def cross_encode():
 
 
 
-def display_as_table(model, top_k=2, score='score'):
+def display_as_table(model, score='score'):
     # Display the df with text and scores as a table
-    df = pd.DataFrame([(hit[score], passages[hit['corpus_id']]) for hit in model[0:top_k]], columns=['Score', 'Text'])
+    df = pd.DataFrame([(hit[score], passages[hit['corpus_id']]) for hit in model[0:2]], columns=['Score', 'Text'])
     df['Score'] = round(df['Score'], 2)
 
     return df
@@ -146,7 +146,7 @@ window_size = 3
 
 bi_encoder_type="multi-qa-mpnet-base-dot-v1"
 # This function will search all wikipedia articles for passages that answer the query
-def search_func(query, top_k=2):
+def search_func(query):
     global bi_encoder, cross_encoder
 
     st.subheader(f"Search Query: {query}")
@@ -162,7 +162,7 @@ def search_func(query, top_k=2):
     # Encode the query using the bi-encoder and find potentially relevant passages
     question_embedding = bi_encoder.encode(query, convert_to_tensor=True)
     question_embedding = question_embedding.cpu()
-    hits = util.semantic_search(question_embedding, corpus_embeddings, top_k=top_k, score_function=util.dot_score)
+    hits = util.semantic_search(question_embedding, corpus_embeddings, top_k=2, score_function=util.dot_score)
     hits = hits[0]  # Get the hits for the first query
 
     # Now, score all retrieved passages with the cross_encoder
@@ -175,18 +175,18 @@ def search_func(query, top_k=2):
 
     # Output of top hits from bi-encoder
     st.markdown("\n-------------------------\n")
-    st.subheader(f"Top-{top_k} Bi-Encoder Retrieval hits")
+    st.subheader(f"Top 2 Bi-Encoder Retrieval hits")
     hits = sorted(hits, key=lambda x: x['score'], reverse=True)
 
-    cross_df = display_as_table(hits, top_k)
+    cross_df = display_as_table(hits, )
     st.write(cross_df.to_html(index=False), unsafe_allow_html=True)
 
     # Output of top hits from cross encoder
     st.markdown("\n-------------------------\n")
-    st.subheader(f"Top-{top_k} Cross-Encoder Re-ranker hits")
+    st.subheader(f"Top-2 Cross-Encoder Re-ranker hits")
     hits = sorted(hits, key=lambda x: x['cross-score'], reverse=True)
 
-    rerank_df = display_as_table(hits, top_k, 'cross-score')
+    rerank_df = display_as_table(hits, 'cross-score')
     st.write(rerank_df.to_html(index=False), unsafe_allow_html=True)
 
 
@@ -239,7 +239,7 @@ if search:
 
         with st.spinner(
                 text="Embedding completed, searching for relevant text for given query and hits..."):
-            search_func(search_query, top_k=2)
+            search_func(search_query)
 
 st.markdown("""
             """)