Update app.py

app.py

@@ -63,29 +63,38 @@ for output in outputs:
 search_query = sorted(query2outputs)[0]
 
 
+def preferences_from_hits(list_of_hits):
+    docid2id = {}
+    id2doc = {}
+    preferences = []
+
+    for result in list_of_hits:
+        for doc in result:
+            if doc["docid"] not in docid2id:
+                id = len(doc["docid"])
+                docid2id[doc["docid"]] = id
+                id2doc[id] = doc
+        preferences.append([docid2id[doc["rank"]] for doc in result])
+
+    #  = {v: k for k, v in docid2id.items()}
+    return np.array(preferences), id2doc
+
+
 def aggregate(list_of_hits):
     import numpy as np
     from permsc import KemenyOptimalAggregator, sum_kendall_tau, ranks_from_preferences
     from permsc import BordaRankAggregator
 
-    preferences = []
-    for result in list_of_hits:
-        preferences.append([doc["rank"] - 1 for doc in result])
-        print([doc["docid"] for doc in result])
-
-    preferences = np.array(preferences)
+    preferences, id2doc = preferences_from_hits(list_of_hits)
     y_optimal = KemenyOptimalAggregator().aggregate(preferences)
     # y_optimal = BordaRankAggregator().aggregate(preferences)
 
-    rank2doc = {}
-    for doc in list_of_hits[0]:
-        rank2doc[doc["rank"] - 1] = doc
-
     print("preference:")
     print(preferences)
     print("preferences shape: ", preferences.shape)
     print("y_optimal: ", y_optimal)
-    return [rank2doc[rank] for rank in y_optimal]
+
+    return [id2doc[id] for id in y_optimal]
 
 aggregated_ranking = aggregate(query2outputs[search_query])