From e43a511c070560f77e764711c67ad0f449d38830 Mon Sep 17 00:00:00 2001 From: shiyu22 Date: Tue, 20 Jun 2023 16:21:39 +0800 Subject: [PATCH] Update output Signed-off-by: shiyu22 --- rerank.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/rerank.py b/rerank.py index 3907113..2e3be93 100644 --- a/rerank.py +++ b/rerank.py @@ -13,6 +13,6 @@ class ReRank(NNOperator): def __call__(self, query: str, docs: List): scores = self._model.predict([(query, doc) for doc in docs]) re_ids = sorted(range(len(scores)), key=lambda k: scores[k], reverse=True) - print(re_ids, docs) re_docs = [docs[i] for i in re_ids] - return re_docs \ No newline at end of file + scores.sort(reverse=True) + return re_docs, scores \ No newline at end of file