Browse Source

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)
tags/v0.9.0
zhuhao 1 year ago
parent
commit
009e18f094
No account linked to committer's email address

+ 3
- 3
api/db/init_data.py View File

@@ -109,7 +109,7 @@ factory_infos = [{
"name": "Ollama",
"logo": "",
"tags": "LLM,TEXT EMBEDDING,SPEECH2TEXT,MODERATION",
"status": "1",
"status": "1",
}, {
"name": "Moonshot",
"logo": "",
@@ -123,8 +123,8 @@ factory_infos = [{
}, {
"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": "",

+ 1
- 0
rag/llm/__init__.py View File

@@ -68,4 +68,5 @@ RerankModel = {
"BAAI": DefaultRerank,
"Jina": JinaRerank,
"Youdao": YoudaoRerank,
"Xinference": XInferenceRerank
}

+ 18
- 0
rag/llm/rerank_model.py View File

@@ -136,4 +136,22 @@ class YoudaoRerank(DefaultRerank):
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"]

+ 1
- 0
web/src/pages/user-setting/setting-model/ollama-modal/index.tsx View File

@@ -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>

Loading…
Cancel
Save