Skip to content

Commit

Permalink
feat: extend api params for Jina Embeddings V3 (#8657)
Browse files Browse the repository at this point in the history
  • Loading branch information
DresAaron authored Sep 23, 2024
1 parent 03fdf5e commit 3618a97
Show file tree
Hide file tree
Showing 2 changed files with 76 additions and 9 deletions.
43 changes: 43 additions & 0 deletions api/core/model_runtime/model_providers/jina/jina.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -67,3 +67,46 @@ model_credential_schema:
required: false
type: text-input
default: '8192'
- variable: task
label:
zh_Hans: 下游任务
en_US: Downstream task
placeholder:
zh_Hans: 选择将使用向量模型的下游任务。模型将返回针对该任务优化的向量。
en_US: Select the downstream task for which the embeddings will be used. The model will return the optimized embeddings for that task.
required: false
type: select
options:
- value: retrieval.query
label:
en_US: retrieval.query
- value: retrieval.passage
label:
en_US: retrieval.passage
- value: separation
label:
en_US: separation
- value: classification
label:
en_US: classification
- value: text-matching
label:
en_US: text-matching
- variable: dimensions
label:
zh_Hans: 输出维度
en_US: Output dimensions
placeholder:
zh_Hans: 输入您的输出维度
en_US: Enter output dimensions
required: false
type: text-input
- variable: late_chunking
label:
zh_Hans: 后期分块
en_US: Late chunking
placeholder:
zh_Hans: 应用后期分块技术来利用模型的长上下文功能来生成上下文块向量化。
en_US: Apply the late chunking technique to leverage the model's long-context capabilities for generating contextual chunk embeddings.
required: false
type: switch
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,38 @@ class JinaTextEmbeddingModel(TextEmbeddingModel):

api_base: str = "https://api.jina.ai/v1"

def _to_payload(self, model: str, texts: list[str], credentials: dict) -> dict:
"""
Parse model credentials
:param model: model name
:param credentials: model credentials
:param texts: texts to embed
:return: parsed credentials
"""

def transform_jina_input_text(model, text):
if model == "jina-clip-v1":
return {"text": text}
return text

data = {"model": model, "input": [transform_jina_input_text(model, text) for text in texts]}

task = credentials.get("task")
dimensions = credentials.get("dimensions")
late_chunking = credentials.get("late_chunking")

if task is not None:
data["task"] = task

if dimensions is not None:
data["dimensions"] = int(dimensions)

if late_chunking is not None:
data["late_chunking"] = late_chunking

return data

def _invoke(
self, model: str, credentials: dict, texts: list[str], user: Optional[str] = None
) -> TextEmbeddingResult:
Expand All @@ -49,15 +81,7 @@ def _invoke(
url = base_url + "/embeddings"
headers = {"Authorization": "Bearer " + api_key, "Content-Type": "application/json"}

def transform_jina_input_text(model, text):
if model == "jina-clip-v1":
return {"text": text}
return text

data = {"model": model, "input": [transform_jina_input_text(model, text) for text in texts]}

if model == "jina-embeddings-v3":
data["task"] = "text-matching"
data = self._to_payload(model=model, texts=texts, credentials=credentials)

try:
response = post(url, headers=headers, data=dumps(data))
Expand Down

0 comments on commit 3618a97

Please sign in to comment.