Przeglądaj źródła

Fix xinference rerank issue. (#4499)

### What problem does this PR solve?
#4495
### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
tags/v0.16.0
Kevin Hu 9 miesięcy temu
rodzic
commit
3805621564
No account linked to committer's email address
1 zmienionych plików z 5 dodań i 1 usunięć
  1. 5
    1
      rag/llm/rerank_model.py

+ 5
- 1
rag/llm/rerank_model.py Wyświetl plik

@@ -172,6 +172,10 @@ class XInferenceRerank(Base):
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,
@@ -183,7 +187,7 @@ class XInferenceRerank(Base):
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):

Ładowanie…
Anuluj
Zapisz