3 changed files with 23 additions and 0 deletions
@ -0,0 +1,4 @@ |
|||
from .rerank import ReRank |
|||
|
|||
def rerank(*args, **kwargs): |
|||
return ReRank(*args, **kwargs) |
@ -0,0 +1 @@ |
|||
sentence_transformers |
@ -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…
Reference in new issue