This commit is contained in:
yumoqing 2025-06-21 10:27:22 +08:00
parent 388ece4b18
commit bffea04c1e
2 changed files with 5 additions and 4 deletions

View File

@ -45,8 +45,9 @@ classs BaseReranker:
def rerank(self, query, docs, top_n=5, sys_prompt="", task=""): def rerank(self, query, docs, top_n=5, sys_prompt="", task=""):
pairs = self.build_pairs(query, docs, sys_prompt=sys_prompt, task=task) pairs = self.build_pairs(query, docs, sys_prompt=sys_prompt, task=task)
inputs = self.process_inputs(pairs) with torch.no_grad():
scores = self.compute_logits(inputs) inputs = self.process_inputs(pairs)
scores = self.compute_logits(inputs)
data = [] data = []
for i, s in enumerate(scores): for i, s in enumerate(scores):
d = { d = {
@ -68,4 +69,4 @@ classs BaseReranker:
"total_tokens": 0 "total_tokens": 0
} }
} }
return ret

View File

@ -75,7 +75,7 @@ async def rerank(request, params_kw, *params, **kw):
return arr return arr
def main(): def main():
parser = argparse.ArgumentParser(prog="Embedding") parser = argparse.ArgumentParser(prog="Rerank")
parser.add_argument('-w', '--workdir') parser.add_argument('-w', '--workdir')
parser.add_argument('-p', '--port') parser.add_argument('-p', '--port')
parser.add_argument('model_path') parser.add_argument('model_path')