Michael Bernovskiy commited on
Commit
488bee9
1 Parent(s): 4dc9474

reranker fixed

Browse files
Files changed (1) hide show
  1. backend/reranker.py +1 -1
backend/reranker.py CHANGED
@@ -7,4 +7,4 @@ reranker = FlagReranker('BAAI/bge-reranker-large',
7
  def rerank(query: str, documents: [str], k: int) -> [str]:
8
  scores = reranker.compute_score([(query, document) for document in documents])
9
  sorted_docs = sorted(zip(scores, documents), key=lambda x: x[0], reverse=True)
10
- return sorted_docs[:k][1]
 
7
  def rerank(query: str, documents: [str], k: int) -> [str]:
8
  scores = reranker.compute_score([(query, document) for document in documents])
9
  sorted_docs = sorted(zip(scores, documents), key=lambda x: x[0], reverse=True)
10
+ return [sorted_doc[1] for sorted_doc in sorted_docs[:k]]