z00mP commited on
Commit
d953944
1 Parent(s): 0d13133

fix rerank bug

Browse files
Files changed (1) hide show
  1. backend/reranker.py +3 -4
backend/reranker.py CHANGED
@@ -7,7 +7,6 @@ def rerank_documents(ce_model_name, documents, query, top_k_rerank):
7
  for doc in documents:
8
  pairs.append((query, doc))
9
  ce_model = CrossEncoder(ce_model_name, max_length=512)
10
- scores = ce_model.predict(pairs)
11
- #sorted_pairs = [(s, x[1]) for s, x in sorted(zip(scores, pairs), key=lambda p: p[0], reverse = True)]
12
- reranked_docs = [x[1] for _, x in sorted(zip(scores, pairs), key=lambda p: p[0], reverse = True)]
13
- return reranked_docs[:top_k_rerank]
 
7
  for doc in documents:
8
  pairs.append((query, doc))
9
  ce_model = CrossEncoder(ce_model_name, max_length=512)
10
+ scores = ce_model.predict(pairs)
11
+ reranked_docs = [x[1] for _, x in sorted(zip(scores, pairs), key=lambda p: p[0], reverse = True)]
12
+ return reranked_docs[:top_k_rerank]