From 3805621564bcbffeef83f06238a28703d05af950 Mon Sep 17 00:00:00 2001
From: Kevin Hu <kevinhu.sh@gmail.com>
Date: Thu, 16 Jan 2025 11:35:51 +0800
Subject: [PATCH] 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)
---
 rag/llm/rerank_model.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/rag/llm/rerank_model.py b/rag/llm/rerank_model.py
index c937e7798..406faf2f9 100644
--- a/rag/llm/rerank_model.py
+++ b/rag/llm/rerank_model.py
@@ -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):