fix-rerank (castorini#174)

leungjch · Apr 8, 2021 · 13b308e · 13b308e
1 parent b9b8254
commit 13b308e
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 1 deletion.
diff --git a/README.md b/README.md
@@ -67,7 +67,6 @@ for i in range(0, 10):
 
 # Finally, rerank:
 reranked = reranker.rerank(query, texts)
-reranked.sort(key=lambda x: x.score, reverse=True)
 
 # Print out reranked results:
 for i in range(0, 10):

diff --git a/pygaggle/rerank/transformer.py b/pygaggle/rerank/transformer.py
@@ -71,6 +71,8 @@ def rerank(self, query: Query, texts: List[Text]) -> List[Text]:
             batch_log_probs = batch_scores[:, 1].tolist()
             for doc, score in zip(batch.documents, batch_log_probs):
                 doc.score = score
+
+        texts.sort(key=lambda x: x.score, reverse=True)
         return texts
 
 
@@ -122,6 +124,9 @@ def rerank(self, query: Query, texts: List[Text]) -> List[Text]:
 
         for text in texts:
             text.score = scores[text.metadata['docid']]
+
+        texts.sort(key=lambda x: x.score, reverse=True)
+
         return texts
 
 
@@ -168,6 +173,9 @@ def rerank(self, query: Query, texts: List[Text]) -> List[Text]:
             for text in texts:
                 if text.score != max_score:
                     text.score = max_score - 10000
+
+        texts.sort(key=lambda x: x.score, reverse=True)
+
         return texts
 
 
@@ -210,6 +218,9 @@ def rerank(self, query: Query, texts: List[Text]) -> List[Text]:
                                 output, 1)[0, -1].item()
             else:
                 text.score = output.item()
+
+        texts.sort(key=lambda x: x.score, reverse=True)
+
         return texts
 
 
@@ -241,4 +252,7 @@ def rerank(self, query: Query, texts: List[Text]) -> List[Text]:
             smax_val, smax_idx = start_scores.max(0)
             emax_val, emax_idx = end_scores.max(0)
             text.score = max(smax_val.item(), emax_val.item())
+
+        texts.sort(key=lambda x: x.score, reverse=True)
+
         return texts