Update app.py

app.py

@@ -4,6 +4,7 @@ import numpy as np
 
 import streamlit as st
 from pathlib import Path
+from collections import defaultdict
 
 import sys
 path_root = Path("./")
@@ -83,6 +84,22 @@ def preferences_from_hits(list_of_hits):
     return np.array(preferences), id2doc
 
 
+def load_qrels(name):
+    import ir_datasets
+    if name == "dl19":
+        ds_name = "msmarco-passage/trec-dl-2019/judged"
+    elif name == "dl20":
+        ds_name = "msmarco-passage/trec-dl-2020/judged"
+    else:
+        raise ValueError(name)
+
+    dataset = ir_datasets.load(ds_name)
+    qrels = defaultdict(dict)
+    for qrel in dataset.qrels_iter():
+        qrels[qrel.query_id][qrel.doc_id] = qrel.relevance
+    return qrels
+
+
 def aggregate(list_of_hits):
     import numpy as np
     from permsc import KemenyOptimalAggregator, sum_kendall_tau, ranks_from_preferences
@@ -92,15 +109,16 @@ def aggregate(list_of_hits):
     y_optimal = KemenyOptimalAggregator().aggregate(preferences)
     # y_optimal = BordaRankAggregator().aggregate(preferences)
 
-    print("-------------------------------------")
-    print("preference:")
-    print(preferences)
-    print("preferences shape: ", preferences.shape)
-    print("y_optimal: ", y_optimal)
+    # print("-------------------------------------")
+    # print("preference:")
+    # print(preferences)
+    # print("preferences shape: ", preferences.shape)
+    # print("y_optimal: ", y_optimal)
 
     return [id2doc[id] for id in y_optimal]
 
 aggregated_ranking = aggregate(query2outputs[search_query])
+qrels = load_qrels("dl19")
 
 if search_query or button_clicked:
 
@@ -112,17 +130,22 @@ if search_query or button_clicked:
     st.write(
         f'<p align=\"right\" style=\"color:grey;\"> Before aggregation for query [{search_query}] ms</p>', unsafe_allow_html=True)
 
+    qid = {result["qid"] for result in search_results}
+    assert len(qid) == 1
+    qid = list(qid)[0]
+
     for i, result in enumerate(search_results):
         result_id = result["docid"]
         contents = result["content"]
 
+        style = "style=\"color:grey;\"" if qrels[qid].get(result_id, 0) else  "style=\"color:red;\""
         # 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}'
+        output = f'<div class="row" {style}> <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)
+                f'<div class="row" {style}>{contents}</div>', unsafe_allow_html=True)
 
         except:
             pass