Преглед изворни кода

Fix: bug #991 (#1013)

### 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>
tags/v0.8.0
Wang Baoling пре 1 година
родитељ
комит
c58a1c48eb
No account linked to committer's email address
2 измењених фајлова са 15 додато и 1 уклоњено
  1. 1
    1
      api/db/init_data.py
  2. 14
    0
      rag/llm/rerank_model.py

+ 1
- 1
api/db/init_data.py Прегледај датотеку

@@ -386,7 +386,7 @@ def init_llm_factory():
"fid": factory_infos[7]["name"],
"llm_name": "maidalun1020/bce-reranker-base_v1",
"tags": "RE-RANK, 8K",
"max_tokens": 8196,
"max_tokens": 512,
"model_type": LLMType.RERANK.value
},
# ------------------------ DeepSeek -----------------------

+ 14
- 0
rag/llm/rerank_model.py Прегледај датотеку

@@ -113,4 +113,18 @@ class YoudaoRerank(DefaultRerank):
YoudaoRerank._model = RerankerModel(
model_name_or_path=model_name.replace(
"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…
Откажи
Сачувај