fix: flashrank ordering (#46)

AnswerDotAI · Nov 12, 2024 · 89c76f2 · 89c76f2
1 parent 69f09b7
commit 89c76f2
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/rerankers/models/flashrank_ranker.py b/rerankers/models/flashrank_ranker.py
@@ -45,12 +45,13 @@ def rank(
 
         ranked_results = [
             Result(
-                document=docs[idx],
+                document=docs[result["id"]],  # Returns reranked documents.
                 score=result["score"],
                 rank=idx + 1,
             )
             for idx, result in enumerate(flashrank_results)
         ]
+
         return RankedResults(results=ranked_results, query=query, has_scores=True)
 
     def score(self, query: str, doc: str) -> float: