continue icon indicating copy to clipboard operation
continue copied to clipboard

feat: Add SageMaker Embeddings

Open kbreese-x opened this issue 1 year ago • 1 comments

Description

This pull request adds SageMaker as an embeddings model provider. The new provider is primarily designed to work with SageMaker endpoints deployed using the Huggingface TEI containers, but can be customized to work with other endpoints too through the use of requestOptions.extraBodyProperties. Documentation on how to set up the new provider has been added to both the Embedding Models and Sagemaker Provider pages. The config_schema has been updated for both vscode and intellij to include sagemaker as a new provider, as well as providing some default region options for the most commonly used AWS regions.

Checklist

  • [x] The base branch of this PR is dev, rather than main
  • [x] The relevant docs, if any, have been updated or created

Screenshots

[ For visual changes, include screenshots. ]

Testing

Requires: AWS Sagemaker endpoint deployed with TEI container (example here) Requires: an AWS profile called "sagemaker" or a default profile with permissions to invoke endpoints

Add the following configuration: "embeddingsProvider": { "provider": "sagemaker", "model": "embedding-test-endpoint", "maxBatchSize": 12, "profile": "default", "region": "us-west-2", "requestOptions": { // this is an example for an endpoint deployed using TEI "extraBodyProperties": { "parameters": { "truncate": true } } }

To learn more about the parameters you can send to TEI backed endpoints, see the API reference here

kbreese-x avatar Oct 04 '24 20:10 kbreese-x