continue
continue copied to clipboard
feat: Add SageMaker Embeddings
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 thanmain - [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