AWS SageMaker procedures
These procedures leverage the Amazon SageMaker API.
Here is a list of all available Aws SageMaker procedures:
name | description |
---|---|
apoc.ml.sagemaker.custom(body, $config) |
To create a customizable SageMaker API call |
apoc.ml.sagemaker.embedding(texts, $config) |
To create an API call to generate embedding |
apoc.ml.sagemaker.chat(messages, $config) |
To create a Chat Completion API call |
apoc.ml.sagemaker.completion(prompt, $config) |
To create a Text Completion API call |
All the procedures, leverage the apoc.ml.sagemaker.custom
procedures,
and support the same config parameter, but unlike the custom
one,
they have some different default parameters and endpoint name.
Moreover, the return data is consistent with the called API,
instead of returning a generic Object
as a result
Config
name | type | default | description |
---|---|---|---|
keyId |
String |
null |
The AWS key ID. We can also evaluate it via |
secretKey |
String |
null |
The AWS secret access key. We can also evaluate it via |
region |
String |
us-east-1 |
The AWS region |
endpoint |
String |
see below |
The AWS endpoint. |
method |
String |
|
The HTTP Method |
headers |
Map<String, Object> |
|
The HTTP Header |
endpointName |
String |
see below |
The SageMaker Endpoint name |
The endpoint
config takes precedence over the model
one.
In case of all procedures, the default endpoint
is "https://runtime.sagemaker.<regionConfigValue>.amazonaws.com/endpoints/<endpointNameConfigValue>/invocations"
.
So, with the default region
config, i.e. "us-east-1", the default endpoint is "https://runtime.sagemaker.us-east-1.amazonaws.com/endpoints/<endpointNameConfigValue>/invocations"
.
If we use the ml.sagemaker.custom
procedure, either the <endpointNameConfigValue>
part or the endpoint
url must be configured ,
while with the sagemaker.chat
, sagemaker.completion
, sagemaker.embedding
ones,
has a default value of "Endpoint-Distilbart-xsum-1-1-1", "Endpoint-GPT-2-1", "Endpoint-Jina-Embeddings-v2-Base-en-1" respectively.
Since SageMaker allows you to choose the endpoint to use, this parameter will most likely always need to be set |
Authentication settings
To authenticate to SageMaker services, we can set in the apoc.conf
file the following entries.
apoc.aws.key.id=<AWS Key ID>
apoc.aws.secret.key=<AWS Secret Access Key>
Alternatively we can set them as $config
parameters, i.e.: {keyId: '<AWS Key ID>', secretKey:'<AWS Secret Access Key>'}
.
Or also, we can put an Authorization header, by using the header
parameter,
i.e. {header: {Authorization: 'AWS4-HMAC-SHA256 <CredentialAndSignature..>', …other entries…} }
.
Note that the default Content-Type: application/json
and the Accept: */*
header entries,
are always passed to the http request, unless overridden via the config header
.
In the following examples,
we assume that we set Key id and Secret Access Key via apoc.conf
.
Usage Examples
Chat Completion API
This procedure apoc.ml.sagemaker.chat
takes a list of maps of chat exchanges between assistant and user (with optional system context), and will return the next message in the flow.
Additional configuration can be passed to the API, the default endpoint name is Endpoint-Distilbart-xsum-1-1-1
.
CALL apoc.ml.sagemaker.chat([{role: "admin", content: "Test answer"}],
{endpointName: 'Endpoint-Distilbart-xsum-1-1-1', region: 'us-east-1'}
)
value |
---|
{ "summary_text": " The ABritainThe,A" } |
Text Completion API
This procedure apoc.ml.sagemaker.completion
can continue/complete a given text.
Additional configuration can be passed to the API, the default endpoint name is Endpoint-GPT-2-1
.
CALL apoc.ml.sagemaker.completion('Test answer API',
{endpointName: 'Endpoint-GPT-2-1-0', headers: {`Content-Type`: "application/x-text"}, region: "eu-central-1"}
)
value |
---|
{ "generated_text": "Test answer API The previous answer in this question was: <summary> - this is in fact an answer to this question, and has a description of the "language" it was written in. Please test in your language first," } |
Embedding API
This procedure apoc.ml.sagemaker.embedding
can take a list of text strings, and will return one row per string, with the embedding data as a 768 element vector.
Additional configuration can be passed to the API, the default endpoint name is Endpoint-Jina-Embeddings-v2-Base-en-1
.
CALL apoc.ml.sagemaker.embedding(['Some Text'], {region: 'eu-central-1'})
index | text | embedding |
---|---|---|
0 |
"Some Text" |
[-0.0065358975, -7.9563365E-4, …. -0.010693862, -0.005087272] |
Custom AWS API Call
Via the apoc.ml.sagemaker.custom
we can create a customizable SageMaker API Request, by choosing the HTTP Method, the endpoint, the region and the additional headers.
Useful both for invoke an endpoint,
in the case the response is incompatible with the previous procedures, and to use any other SageMaker API.
For example, we can execute a POST
request to call the ListEndpoints API,
we can execute:
CALL apoc.ml.sagemaker.custom({SortBy: "Name"},
{
endpoint: "https://api.sagemaker.us-east-1.amazonaws.com/",
headers: {`X-Amz-Target`: "SageMaker.ListEndpoints", `Content-Type`: "application/x-amz-json-1.1"},
method: "POST"
})