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