Kevin Hu commited on
Commit
3f95476
·
1 Parent(s): 1f55425

Fix xinference rerank issue. (#4499)

Browse files

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

- [x] Bug Fix (non-breaking change which fixes an issue)

Files changed (1) hide show
  1. rag/llm/rerank_model.py +5 -1
rag/llm/rerank_model.py CHANGED
@@ -172,6 +172,10 @@ class XInferenceRerank(Base):
172
  def similarity(self, query: str, texts: list):
173
  if len(texts) == 0:
174
  return np.array([]), 0
 
 
 
 
175
  data = {
176
  "model": self.model_name,
177
  "query": query,
@@ -183,7 +187,7 @@ class XInferenceRerank(Base):
183
  rank = np.zeros(len(texts), dtype=float)
184
  for d in res["results"]:
185
  rank[d["index"]] = d["relevance_score"]
186
- return rank, res["meta"]["tokens"]["input_tokens"] + res["meta"]["tokens"]["output_tokens"]
187
 
188
 
189
  class LocalAIRerank(Base):
 
172
  def similarity(self, query: str, texts: list):
173
  if len(texts) == 0:
174
  return np.array([]), 0
175
+ pairs = [(query, truncate(t, 4096)) for t in texts]
176
+ token_count = 0
177
+ for _, t in pairs:
178
+ token_count += num_tokens_from_string(t)
179
  data = {
180
  "model": self.model_name,
181
  "query": query,
 
187
  rank = np.zeros(len(texts), dtype=float)
188
  for d in res["results"]:
189
  rank[d["index"]] = d["relevance_score"]
190
+ return rank, token_count
191
 
192
 
193
  class LocalAIRerank(Base):