feat: support xinference rerank model (#1466)

### What problem does this PR solve? support xinference rerank model #1455 ### Type of change - [x] New Feature (non-breaking change which adds functionality)
infiniflow · Jul 11, 2024 · 009e18f · 009e18f
1 parent 9c023b6
commit 009e18f
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 3 deletions.
diff --git a/api/db/init_data.py b/api/db/init_data.py
@@ -109,7 +109,7 @@ def init_superuser():
     "name": "Ollama",
     "logo": "",
     "tags": "LLM,TEXT EMBEDDING,SPEECH2TEXT,MODERATION",
-        "status": "1",
+    "status": "1",
 }, {
     "name": "Moonshot",
     "logo": "",
@@ -123,8 +123,8 @@ def init_superuser():
 }, {
     "name": "Xinference",
     "logo": "",
-    "tags": "LLM,TEXT EMBEDDING,SPEECH2TEXT,MODERATION",
-        "status": "1",
+    "tags": "LLM,TEXT EMBEDDING,SPEECH2TEXT,MODERATION,TEXT RE-RANK",
+    "status": "1",
 },{
     "name": "Youdao",
     "logo": "",

diff --git a/rag/llm/__init__.py b/rag/llm/__init__.py
@@ -68,4 +68,5 @@
     "BAAI": DefaultRerank,
     "Jina": JinaRerank,
     "Youdao": YoudaoRerank,
+    "Xinference": XInferenceRerank
 }
diff --git a/rag/llm/rerank_model.py b/rag/llm/rerank_model.py
@@ -136,4 +136,22 @@ 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
+        self.headers = {
+            "Content-Type": "application/json",
+            "accept": "application/json"
+        }
 
+    def similarity(self, query: str, texts: list):
+        data = {
+            "model":self.model_name,
+            "query":query,
+            "return_documents": "true",
+            "return_len": "true",
+            "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"]
diff --git a/web/src/pages/user-setting/setting-model/ollama-modal/index.tsx b/web/src/pages/user-setting/setting-model/ollama-modal/index.tsx
@@ -74,6 +74,7 @@ const OllamaModal = ({
           <Select placeholder={t('modelTypeMessage')}>
             <Option value="chat">chat</Option>
             <Option value="embedding">embedding</Option>
+            <Option value="rerank">rerank</Option>
           </Select>
         </Form.Item>
         <Form.Item<FieldType>