Display a list of retrieved and reranked results

app.py

@@ -5,6 +5,7 @@
 
 import gradio as gr
 import os
+import pandas as pd
 
 from reranker.reranker import CrossEncReranker
 from retriever.es_retriever import ESRetriever
@@ -31,7 +32,16 @@ def retrieve_and_rerank(question: str, example: str):
         retrieval_results = es_retriever_client.retrieve(query)
         reranker_results = cross_enc_reranker.rerank(query, retrieval_results)
 
-        return reranker_results[0].generate_string(), reranker_results[0].text
+        law_refs = [i.generate_string() for i in reranker_results]
+        law_texts = [i.text for i in reranker_results]
+
+        df = pd.DataFrame({
+            'Rank': range(1, len(law_refs)+1),
+            'Reference': law_refs,
+            'Text': law_texts
+        })
+
+        return reranker_results[0].generate_string(), reranker_results[0].text, df
     except:
         return "-", "(Result Not Found)"
 
@@ -51,7 +61,11 @@ demo = gr.Interface(
             ]
         )
     ],
-    outputs=["label", "text"])
+    outputs=[
+        "label",
+        "text",
+        "dataframe"
+    ])
 
 if __name__ == "__main__":
     demo.launch()

requirements.txt

@@ -2,3 +2,4 @@ gradio~=3.28.1
 numpy>1.21.4
 requests==2.26.0
 sentence-transformers
+pandas~=1.5.3