refine xinference (#2521)

### What problem does this PR solve? #1588 ### Type of change - [x] Refactoring
infiniflow · Sep 20, 2024 · 4a6a2a0 · 4a6a2a0
1 parent 9bbef82
commit 4a6a2a0
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 0 deletions.
diff --git a/rag/llm/cv_model.py b/rag/llm/cv_model.py
@@ -449,6 +449,8 @@ def __init__(self, key, model_name, base_url, lang="Chinese"):
 
 class XinferenceCV(Base):
     def __init__(self, key, model_name="", lang="Chinese", base_url=""):
+        if base_url.split("/")[-1] != "v1":
+            base_url = os.path.join(base_url, "v1")
         self.client = OpenAI(api_key="xxx", base_url=base_url)
         self.model_name = model_name
         self.lang = lang

diff --git a/rag/llm/embedding_model.py b/rag/llm/embedding_model.py
@@ -268,6 +268,8 @@ def encode_queries(self, text: str):
 
 class XinferenceEmbed(Base):
     def __init__(self, key, model_name="", base_url=""):
+        if base_url.split("/")[-1] != "v1":
+            base_url = os.path.join(base_url, "v1")
         self.client = OpenAI(api_key="xxx", base_url=base_url)
         self.model_name = model_name
 

diff --git a/rag/llm/rerank_model.py b/rag/llm/rerank_model.py
@@ -140,6 +140,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")
         self.model_name = model_name
         self.base_url = base_url
         self.headers = {

diff --git a/rag/llm/sequence2txt_model.py b/rag/llm/sequence2txt_model.py
@@ -93,6 +93,8 @@ def __init__(self, key, model_name, lang="Chinese", **kwargs):
 
 class XinferenceSeq2txt(Base):
     def __init__(self, key, model_name="", base_url=""):
+        if base_url.split("/")[-1] != "v1":
+            base_url = os.path.join(base_url, "v1")
         self.client = OpenAI(api_key="xxx", base_url=base_url)
         self.model_name = model_name