Add suuport for AWS Bedrock Cohere embedding (#3444)

This commit is contained in:
kerlion 2024-04-16 13:22:38 +08:00 committed by GitHub
parent 5e02a83b53
commit 200010be19
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 58 additions and 15 deletions

View File

@ -1 +1,3 @@
- amazon.titan-embed-text-v1
- cohere.embed-english-v3
- cohere.embed-multilingual-v3

View File

@ -0,0 +1,8 @@
model: cohere.embed-english-v3
model_type: text-embedding
model_properties:
context_size: 512
pricing:
input: '0.1'
unit: '0.000001'
currency: USD

View File

@ -0,0 +1,8 @@
model: cohere.embed-multilingual-v3
model_type: text-embedding
model_properties:
context_size: 512
pricing:
input: '0.1'
unit: '0.000001'
currency: USD

View File

@ -1,4 +1,5 @@
import json
import logging
import time
from typing import Optional
@ -24,6 +25,7 @@ from core.model_runtime.errors.invoke import (
)
from core.model_runtime.model_providers.__base.text_embedding_model import TextEmbeddingModel
logger = logging.getLogger(__name__)
class BedrockTextEmbeddingModel(TextEmbeddingModel):
@ -53,17 +55,19 @@ class BedrockTextEmbeddingModel(TextEmbeddingModel):
embeddings = []
token_usage = 0
model_prefix = model.split('.')[0]
if model_prefix == "amazon":
for text in texts:
body = {
"inputText": text,
}
response_body = self._invoke_bedrock_embedding(model, bedrock_runtime, body)
embeddings.extend([response_body.get('embedding')])
token_usage += response_body.get('inputTextTokenCount')
result = TextEmbeddingResult(
if model_prefix == "amazon" :
for text in texts:
body = {
"inputText": text,
}
response_body = self._invoke_bedrock_embedding(model, bedrock_runtime, body)
embeddings.extend([response_body.get('embedding')])
token_usage += response_body.get('inputTextTokenCount')
logger.warning(f'Total Tokens: {token_usage}')
result = TextEmbeddingResult(
model=model,
embeddings=embeddings,
usage=self._calc_response_usage(
@ -71,11 +75,32 @@ class BedrockTextEmbeddingModel(TextEmbeddingModel):
credentials=credentials,
tokens=token_usage
)
)
else:
raise ValueError(f"Got unknown model prefix {model_prefix} when handling block response")
return result
)
return result
if model_prefix == "cohere" :
input_type = 'search_document' if len(texts) > 1 else 'search_query'
for text in texts:
body = {
"texts": [text],
"input_type": input_type,
}
response_body = self._invoke_bedrock_embedding(model, bedrock_runtime, body)
embeddings.extend(response_body.get('embeddings'))
token_usage += len(text)
result = TextEmbeddingResult(
model=model,
embeddings=embeddings,
usage=self._calc_response_usage(
model=model,
credentials=credentials,
tokens=token_usage
)
)
return result
#others
raise ValueError(f"Got unknown model prefix {model_prefix} when handling block response")
def get_num_tokens(self, model: str, credentials: dict, texts: list[str]) -> int: