fix bugs of rerank model with xinference

api/apps/llm_app.py

@@ -165,6 +165,17 @@ def add_llm():
         except Exception as e:
             msg += f"\nFail to access model({llm['llm_name']})." + str(
                 e)
+    elif llm["model_type"] == LLMType.RERANK:
+        mdl = RerankModel[factory](
+            key=None, model_name=llm["llm_name"], base_url=llm["api_base"]
+        )
+        try:
+            arr, tc = mdl.similarity("Hello~ Ragflower!", ["Hi, there!"])
+            if len(arr) == 0 or tc == 0:
+                raise Exception("Not known.")
+        except Exception as e:
+            msg += f"\nFail to access model({llm['llm_name']})." + str(
+                e)
     else:
         # TODO: check other type of models
         pass

rag/llm/rerank_model.py

@@ -136,22 +136,24 @@ def similarity(self, query: str, texts: list):
             else: res.extend(scores)
         return np.array(res), token_count
 
+
 class XInferenceRerank(Base):
-    def __init__(self,model_name="",base_url=""):
-        self.model_name=model_name
-        self.base_url=base_url
+    def __init__(self, key="xxxxxxx", model_name="", base_url=""):
+        self.model_name = model_name
+        self.base_url = base_url
         self.headers = {
             "Content-Type": "application/json",
             "accept": "application/json"
         }
 
     def similarity(self, query: str, texts: list):
         data = {
-            "model":self.model_name,
-            "query":query,
+            "model": self.model_name,
+            "query": query,
             "return_documents": "true",
             "return_len": "true",
-            "documents":texts
+            "documents": texts
         }
         res = requests.post(self.base_url, headers=self.headers, json=data).json()
-        return np.array([d["relevance_score"] for d in res["results"]]),res["tokens"]["input_tokens"]+res["tokens"]["output_tokens"]
+        return np.array([d["relevance_score"] for d in res["results"]]), res["tokens"]["input_tokens"] + res["tokens"][
+            "output_tokens"]