feat(api): support wenxin bge-large and tao embedding model. (#7393)

This commit is contained in:
Chengyu Yan 2024-08-19 22:25:09 +08:00 committed by GitHub
parent 31f9977411
commit 1f944c6eeb
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 89 additions and 2 deletions

View File

@ -118,6 +118,9 @@ class _CommonWenxin:
'ernie-4.0-turbo-8k-preview': 'https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/ernie-4.0-turbo-8k-preview',
'yi_34b_chat': 'https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/yi_34b_chat',
'embedding-v1': 'https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/embeddings/embedding-v1',
'bge-large-en': 'https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/embeddings/bge_large_en',
'bge-large-zh': 'https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/embeddings/bge_large_zh',
'tao-8k': 'https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/embeddings/tao_8k',
}
function_calling_supports = [

View File

@ -0,0 +1,9 @@
model: bge-large-en
model_type: text-embedding
model_properties:
context_size: 512
max_chunks: 16
pricing:
input: '0.0005'
unit: '0.001'
currency: RMB

View File

@ -0,0 +1,9 @@
model: bge-large-zh
model_type: text-embedding
model_properties:
context_size: 512
max_chunks: 16
pricing:
input: '0.0005'
unit: '0.001'
currency: RMB

View File

@ -0,0 +1,9 @@
model: tao-8k
model_type: text-embedding
model_properties:
context_size: 8192
max_chunks: 1
pricing:
input: '0.0005'
unit: '0.001'
currency: RMB

View File

@ -5,7 +5,7 @@ from core.model_runtime.entities.text_embedding_entities import TextEmbeddingRes
from core.model_runtime.model_providers.wenxin.text_embedding.text_embedding import WenxinTextEmbeddingModel
def test_invoke_embedding_model():
def test_invoke_embedding_v1():
sleep(3)
model = WenxinTextEmbeddingModel()
@ -21,4 +21,61 @@ def test_invoke_embedding_model():
assert isinstance(response, TextEmbeddingResult)
assert len(response.embeddings) == 3
assert isinstance(response.embeddings[0], list)
assert isinstance(response.embeddings[0], list)
def test_invoke_embedding_bge_large_en():
sleep(3)
model = WenxinTextEmbeddingModel()
response = model.invoke(
model='bge-large-en',
credentials={
'api_key': os.environ.get('WENXIN_API_KEY'),
'secret_key': os.environ.get('WENXIN_SECRET_KEY')
},
texts=['hello', '你好', 'xxxxx'],
user="abc-123"
)
assert isinstance(response, TextEmbeddingResult)
assert len(response.embeddings) == 3
assert isinstance(response.embeddings[0], list)
def test_invoke_embedding_bge_large_zh():
sleep(3)
model = WenxinTextEmbeddingModel()
response = model.invoke(
model='bge-large-zh',
credentials={
'api_key': os.environ.get('WENXIN_API_KEY'),
'secret_key': os.environ.get('WENXIN_SECRET_KEY')
},
texts=['hello', '你好', 'xxxxx'],
user="abc-123"
)
assert isinstance(response, TextEmbeddingResult)
assert len(response.embeddings) == 3
assert isinstance(response.embeddings[0], list)
def test_invoke_embedding_tao_8k():
sleep(3)
model = WenxinTextEmbeddingModel()
response = model.invoke(
model='tao-8k',
credentials={
'api_key': os.environ.get('WENXIN_API_KEY'),
'secret_key': os.environ.get('WENXIN_SECRET_KEY')
},
texts=['hello', '你好', 'xxxxx'],
user="abc-123"
)
assert isinstance(response, TextEmbeddingResult)
assert len(response.embeddings) == 3
assert isinstance(response.embeddings[0], list)