Browse Source
Update output
Signed-off-by: shiyu22 <shiyu.chen@zilliz.com>
main
1 changed files with
2 additions and
2 deletions
-
rerank.py
|
|
@ -13,6 +13,6 @@ class ReRank(NNOperator): |
|
|
def __call__(self, query: str, docs: List): |
|
|
def __call__(self, query: str, docs: List): |
|
|
scores = self._model.predict([(query, doc) for doc in docs]) |
|
|
scores = self._model.predict([(query, doc) for doc in docs]) |
|
|
re_ids = sorted(range(len(scores)), key=lambda k: scores[k], reverse=True) |
|
|
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] |
|
|
re_docs = [docs[i] for i in re_ids] |
|
|
return re_docs |
|
|
|
|
|
|
|
|
scores.sort(reverse=True) |
|
|
|
|
|
return re_docs, scores |
