Add dimensions parameter support for bedrock titan embedding v2 model #3136
Conversation
zane-neo
requested review from
b4sjoo,
dhrubo-os,
jngz-es,
model-collapse,
rbhavna,
ylwu-amzn,
Zhangxunmt,
austintlee,
HenryL27,
sam-herman and
xinyual
as code owners
zane-neo
temporarily deployed
to
ml-commons-cicd-env
GitHub Actions
Inactive
zane-neo
temporarily deployed
to
ml-commons-cicd-env
GitHub Actions
Inactive
zane-neo
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Failure
Signed-off-by: zane-neo <[email protected]>
Signed-off-by: zane-neo <[email protected]>
zane-neo
force-pushed
the
support-embedding-dimension-bedrock
branch
from
2d4c496 to
ab2d736
Compare
zane-neo
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Failure
zane-neo
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Failure
zane-neo
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Failure
zane-neo
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Failure
Signed-off-by: zane-neo <[email protected]>
zane-neo
temporarily deployed
to
ml-commons-cicd-env
GitHub Actions
Inactive
zane-neo
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Failure
| import org.opensearch.ml.common.dataset.TextDocsInputDataSet; | ||
| import org.opensearch.ml.common.dataset.remote.RemoteInferenceInputDataSet; | ||
| import org.opensearch.ml.common.input.MLInput; | ||
|
|
||
| import lombok.extern.slf4j.Slf4j; | ||
|
|
||
| @Slf4j |
There was a problem hiding this comment.
can we add @log4j instead to keep the consistency?
| TextDocsInputDataSet inputData = (TextDocsInputDataSet) mlInput.getInputDataset(); | ||
| // Amazon Titan Text Embeddings V2 model: https://docs.aws.amazon.com/bedrock/latest/userguide/titan-embedding-models.html | ||
| // Default dimension is 1024 | ||
| int dimensions = Optional.ofNullable(connectorParams.get("dimensions")).map(x -> NumberUtils.toInt(x, 1024)).orElse(1024); |
There was a problem hiding this comment.
Is BedrockEmbeddingPreProcessFunction only for titan models? What about other bedrock models?
There was a problem hiding this comment.
I also notice we have multimodal embedding model: https://docs.aws.amazon.com/bedrock/latest/userguide/titan-multiemb-models.html. Is the multimodal model using MultiModalConnectorPreProcessFunction, if so, can we also add dimension for this?
There was a problem hiding this comment.
Is BedrockEmbeddingPreProcessFunction only for titan models? What about other bedrock models?
Currently this is a titan specific process function, but if in the future bedrock has another similar(similar means they have same request body, e.g. inputText & dimensions etc) text embedding model, then this can be reused to that one.
BedRock integrated with multiple models like anthropic, cohere, meta etc, currently we only have titan, cohere pre process functions in our code, from cohere API: https://docs.cohere.com/reference/embed I didn't see they support dimension parameter so it's fine.
There was a problem hiding this comment.
In that case do we need to set the default dimension? If dimension is provided then we will add otherwise not?
zane-neo
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Failure
| Map inferenceResult = predictTextEmbeddingModelIgnoreFunctionName(modelId, mlInput); | ||
| String errorMsg = String | ||
| .format(Locale.ROOT, "Failing test case name: %s, inference result: %s", testCaseName, gson.toJson(inferenceResult)); | ||
| assertTrue(errorMsg, inferenceResult.containsKey("inference_results")); |
There was a problem hiding this comment.
Why it is showing error here?
|
After adding more details of the issue, this PR is not necessary and please check the new PR: #3262 |
Description
Functionusage by adding a new interfacePreProcessFunctionas there isn't cases that the Function accepts different type of parameters. Currently it only acceptsMLInputandMap<String, String>dimensionsparameter introduced in titan embedding v2 model.Related Issues
#3093
Check List
--signoff.By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.