Skip to content

Commit

Permalink
Fix xinference rerank issue. (#4499)
Browse files Browse the repository at this point in the history
### What problem does this PR solve?
#4495
### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
  • Loading branch information
KevinHuSh authored Jan 16, 2025
1 parent a75cda4 commit 3805621
Showing 1 changed file with 5 additions and 1 deletion.
6 changes: 5 additions & 1 deletion rag/llm/rerank_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -172,6 +172,10 @@ def __init__(self, key="xxxxxxx", model_name="", base_url=""):
def similarity(self, query: str, texts: list):
if len(texts) == 0:
return np.array([]), 0
pairs = [(query, truncate(t, 4096)) for t in texts]
token_count = 0
for _, t in pairs:
token_count += num_tokens_from_string(t)
data = {
"model": self.model_name,
"query": query,
Expand All @@ -183,7 +187,7 @@ def similarity(self, query: str, texts: list):
rank = np.zeros(len(texts), dtype=float)
for d in res["results"]:
rank[d["index"]] = d["relevance_score"]
return rank, res["meta"]["tokens"]["input_tokens"] + res["meta"]["tokens"]["output_tokens"]
return rank, token_count


class LocalAIRerank(Base):
Expand Down

0 comments on commit 3805621

Please sign in to comment.