logo
rerank
repo-copy-icon

copied

Browse Source

Add rerank op

Signed-off-by: shiyu22 <shiyu.chen@zilliz.com>
main
shiyu22 2 years ago
parent
commit
f5b2df224b
  1. 4
      __init__.py
  2. 1
      requirements.txt
  3. 18
      rerank.py

4
__init__.py

@ -0,0 +1,4 @@
from .rerank import ReRank
def rerank(*args, **kwargs):
return ReRank(*args, **kwargs)

1
requirements.txt

@ -0,0 +1 @@
sentence_transformers

18
rerank.py

@ -0,0 +1,18 @@
import numpy as np
from typing import List
from sentence_transformers import CrossEncoder
from towhee.operator import NNOperator
class ReRank(NNOperator):
def __init__(self, model_name: str = 'cross-encoder/ms-marco-MiniLM-L-12-v2'):
super().__init__()
self._model_name = model_name
self._model = CrossEncoder(self._model_name, max_length=1000)
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
Loading…
Cancel
Save