feat: support jina-clip-v1 embedding model (#5146)

This commit is contained in:
sino 2024-06-13 16:31:18 +08:00 committed by GitHub
parent 790543131a
commit 8210637bc5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 27 additions and 8 deletions

View File

@ -0,0 +1,9 @@
model: jina-clip-v1
model_type: text-embedding
model_properties:
context_size: 8192
max_chunks: 2048
pricing:
input: '0.001'
unit: '0.001'
currency: USD

View File

@ -52,16 +52,21 @@ class JinaTextEmbeddingModel(TextEmbeddingModel):
'Content-Type': 'application/json' 'Content-Type': 'application/json'
} }
def transform_jina_input_text(model, text):
if model == 'jina-clip-v1':
return {"text": text}
return text
data = { data = {
'model': model, 'model': model,
'input': texts 'input': [transform_jina_input_text(model, text) for text in texts]
} }
try: try:
response = post(url, headers=headers, data=dumps(data)) response = post(url, headers=headers, data=dumps(data))
except Exception as e: except Exception as e:
raise InvokeConnectionError(str(e)) raise InvokeConnectionError(str(e))
if response.status_code != 200: if response.status_code != 200:
try: try:
resp = response.json() resp = response.json()
@ -75,16 +80,19 @@ class JinaTextEmbeddingModel(TextEmbeddingModel):
else: else:
raise InvokeBadRequestError(msg) raise InvokeBadRequestError(msg)
except JSONDecodeError as e: except JSONDecodeError as e:
raise InvokeServerUnavailableError(f"Failed to convert response to json: {e} with text: {response.text}") raise InvokeServerUnavailableError(
f"Failed to convert response to json: {e} with text: {response.text}")
try: try:
resp = response.json() resp = response.json()
embeddings = resp['data'] embeddings = resp['data']
usage = resp['usage'] usage = resp['usage']
except Exception as e: except Exception as e:
raise InvokeServerUnavailableError(f"Failed to convert response to json: {e} with text: {response.text}") raise InvokeServerUnavailableError(
f"Failed to convert response to json: {e} with text: {response.text}")
usage = self._calc_response_usage(model=model, credentials=credentials, tokens=usage['total_tokens']) usage = self._calc_response_usage(
model=model, credentials=credentials, tokens=usage['total_tokens'])
result = TextEmbeddingResult( result = TextEmbeddingResult(
model=model, model=model,
@ -122,7 +130,8 @@ class JinaTextEmbeddingModel(TextEmbeddingModel):
try: try:
self._invoke(model=model, credentials=credentials, texts=['ping']) self._invoke(model=model, credentials=credentials, texts=['ping'])
except Exception as e: except Exception as e:
raise CredentialsValidateFailedError(f'Credentials validation failed: {e}') raise CredentialsValidateFailedError(
f'Credentials validation failed: {e}')
@property @property
def _invoke_error_mapping(self) -> dict[type[InvokeError], list[type[Exception]]]: def _invoke_error_mapping(self) -> dict[type[InvokeError], list[type[Exception]]]:
@ -144,7 +153,7 @@ class JinaTextEmbeddingModel(TextEmbeddingModel):
InvokeBadRequestError InvokeBadRequestError
] ]
} }
def _calc_response_usage(self, model: str, credentials: dict, tokens: int) -> EmbeddingUsage: def _calc_response_usage(self, model: str, credentials: dict, tokens: int) -> EmbeddingUsage:
""" """
Calculate response usage Calculate response usage
@ -185,7 +194,8 @@ class JinaTextEmbeddingModel(TextEmbeddingModel):
model_type=ModelType.TEXT_EMBEDDING, model_type=ModelType.TEXT_EMBEDDING,
fetch_from=FetchFrom.CUSTOMIZABLE_MODEL, fetch_from=FetchFrom.CUSTOMIZABLE_MODEL,
model_properties={ model_properties={
ModelPropertyKey.CONTEXT_SIZE: int(credentials.get('context_size')) ModelPropertyKey.CONTEXT_SIZE: int(
credentials.get('context_size'))
} }
) )