fix xinference add rerank model bug

infiniflow · Oct 9, 2024 · ac67f83 · ac67f83
1 parent 7f44cf5
commit ac67f83
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/rag/llm/rerank_model.py b/rag/llm/rerank_model.py
@@ -15,6 +15,8 @@
 #
 import re
 import threading
+from urllib.parse import urljoin
+
 import requests
 from huggingface_hub import snapshot_download
 import os
@@ -154,8 +156,8 @@ def similarity(self, query: str, texts: list):
 
 class XInferenceRerank(Base):
     def __init__(self, key="xxxxxxx", model_name="", base_url=""):
-        if base_url.split("/")[-1] != "v1":
-            base_url = os.path.join(base_url, "v1")
+        if base_url.find("/v1") == -1:
+            base_url = urljoin(base_url, "/v1/rerank")
         self.model_name = model_name
         self.base_url = base_url
         self.headers = {