github/DB-GPT
1
0
Fork 0
mirror of https://github.com/csunny/DB-GPT.git synced 2025-09-05 11:01:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(model): Support siliconflow rerank models (#2188)

Browse Source
This commit is contained in:
Fangyin Cheng
2024-12-11 18:36:44 +08:00
committed by GitHub
parent 23aedea092
commit abab4e3e65
5 changed files with 96 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
dbgpt/model/adapter/embeddings_loader.py
View File

@@ -109,6 +109,18 @@ class EmbeddingLoader:
if proxy_param.proxy_backend:
openapi_param["model_name"] = proxy_param.proxy_backend
return OpenAPIRerankEmbeddings(**openapi_param)
elif model_name in ["rerank_proxy_silicon_flow"]:
from dbgpt.rag.embedding.rerank import SiliconFlowRerankEmbeddings
proxy_param = cast(ProxyEmbeddingParameters, param)
openapi_param = {}
if proxy_param.proxy_server_url:
openapi_param["api_url"] = proxy_param.proxy_server_url
if proxy_param.proxy_api_key:
openapi_param["api_key"] = proxy_param.proxy_api_key
if proxy_param.proxy_backend:
openapi_param["model_name"] = proxy_param.proxy_backend
return SiliconFlowRerankEmbeddings(**openapi_param)
else:
from dbgpt.rag.embedding.rerank import CrossEncoderRerankEmbeddings