Update app.py

app.py

@@ -230,20 +230,18 @@ def re_rank_candidates(query, candidates, method):
             key=lambda x: x[1]['cross_score'] + x[1]['bi_score'],
             reverse=True
         )
-    data_dicts = [{'query': item[0], **item[1]} for item in filtered_sorted_result]
-    # Convert the list of dictionaries into a DataFrame
-    df = pd.DataFrame(data_dicts)
-    return df
+    return filtered_sorted_result
 
 col1, col2 = st.columns(2)
 if st.button('Generated Expansion'): 
     candidates = generate_query_expansion_candidates(query = user_query)
-    col1.subheader('Raw Candidates:')
-    df = re_rank_candidates(user_query, candidates, method='cross_encoder')
-    result = list(df['query'][:maxtags_sidebar])
-    st.write(result)
-
-    col2.subheader('Rerank By GMS:')
-    df = re_rank_candidates(user_query, candidates, method='gms')
-    result = df[['query', 'gms']][:maxtags_sidebar]
-    st.json(result.to_dict('list'))
+    with col1:
+        st.header('Raw Candidates:')
+        candidates_rerank = re_rank_candidates(user_query, candidates, method='cross_encoder')[:maxtags_sidebar]
+        result = [item[0] for item in candidates_rerank]
+        st.write(result)
+    with col2:
+        st.header('Rerank By GMS:')
+        candidates_rerank = re_rank_candidates(user_query, candidates, method='gms')[:maxtags_sidebar]
+        result = [{item[0]: item[1]['gms']} for item in candidates_rerank]
+        st.json(result)