### What problem does this PR solve? issue #991 ### Type of change - [x] Bug Fix (non-breaking change which fixes an issue) --------- Co-authored-by: KevinHuSh <kevinhu.sh@gmail.com>
This commit is contained in:
parent
fefe7124a1
commit
c58a1c48eb
@ -386,7 +386,7 @@ def init_llm_factory():
|
|||||||
"fid": factory_infos[7]["name"],
|
"fid": factory_infos[7]["name"],
|
||||||
"llm_name": "maidalun1020/bce-reranker-base_v1",
|
"llm_name": "maidalun1020/bce-reranker-base_v1",
|
||||||
"tags": "RE-RANK, 8K",
|
"tags": "RE-RANK, 8K",
|
||||||
"max_tokens": 8196,
|
"max_tokens": 512,
|
||||||
"model_type": LLMType.RERANK.value
|
"model_type": LLMType.RERANK.value
|
||||||
},
|
},
|
||||||
# ------------------------ DeepSeek -----------------------
|
# ------------------------ DeepSeek -----------------------
|
||||||
|
|||||||
@ -113,4 +113,18 @@ class YoudaoRerank(DefaultRerank):
|
|||||||
YoudaoRerank._model = RerankerModel(
|
YoudaoRerank._model = RerankerModel(
|
||||||
model_name_or_path=model_name.replace(
|
model_name_or_path=model_name.replace(
|
||||||
"maidalun1020", "InfiniFlow"))
|
"maidalun1020", "InfiniFlow"))
|
||||||
|
|
||||||
|
def similarity(self, query: str, texts: list):
|
||||||
|
pairs = [(query,truncate(t, self._model.max_length)) for t in texts]
|
||||||
|
token_count = 0
|
||||||
|
for _, t in pairs:
|
||||||
|
token_count += num_tokens_from_string(t)
|
||||||
|
batch_size = 32
|
||||||
|
res = []
|
||||||
|
for i in range(0, len(pairs), batch_size):
|
||||||
|
scores = self._model.compute_score(pairs[i:i + batch_size], max_length=self._model.max_length)
|
||||||
|
scores = sigmoid(np.array(scores)).tolist()
|
||||||
|
res.extend(scores)
|
||||||
|
return np.array(res), token_count
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user