jyapayne/dify
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: xinference rerank model support (#1615)

Browse source
This commit is contained in:
takatost 2023-11-25 03:56:00 +08:00 • committed by GitHub
commit 0e627c920f
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
9 changed files with 215 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

8
api/core/model_providers/providers/xinference_provider.py
View file

@ -2,11 +2,13 @@ import json
from typing import Type
import requests
from xinference_client.client.restful.restful_client import Client
from core.helper import encrypter
from core.model_providers.models.embedding.xinference_embedding import XinferenceEmbedding
from core.model_providers.models.entity.model_params import KwargRule, ModelKwargsRules, ModelType, ModelMode
from core.model_providers.models.llm.xinference_model import XinferenceModel
from core.model_providers.models.reranking.xinference_reranking import XinferenceReranking
from core.model_providers.providers.base import BaseModelProvider, CredentialsValidateFailedError
from core.model_providers.models.base import BaseProviderModel
@ -40,6 +42,8 @@ class XinferenceProvider(BaseModelProvider):
model_class = XinferenceModel
elif model_type == ModelType.EMBEDDINGS:
model_class = XinferenceEmbedding
elif model_type == ModelType.RERANKING:
model_class = XinferenceReranking
else:
raise NotImplementedError
@ -113,6 +117,10 @@ class XinferenceProvider(BaseModelProvider):
)
embedding.embed_query("ping")
elif model_type == ModelType.RERANKING:
rerank_client = Client(credential_kwargs['server_url'])
model = rerank_client.get_model(credential_kwargs['model_uid'])
model.rerank(query="ping", documents=["ping", "pong"], top_n=2)
except Exception as ex:
raise CredentialsValidateFailedError(str(ex))