Jupyter AI with the SageMaker endpoint

This demo showcases the IPython magics Jupyter AI provides out-of-the-box for Amazon SageMaker.

First, make sure that you’ve set your AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY environment variables before starting JupyterLab as follows:

os.environ['AWS_ACCESS_KEY_ID'] = <your_aws_access_key_id>
os.environ['AWS_SECRET_ACCESS_KEY'] = <your_aws_secret_access_key>

If you prefer to set these keys interactively in this notebook, then use the following code:

# NOTE: Enter the AWS access key id and the AWS secret access key when prompted by the code below

import getpass

# Enter your keys
access_key = getpass.getpass('Enter your AWS ACCESS KEY ID: ')
secret_access_key = getpass.getpass('Enter your AWS SECRET ACCESS KEY: ')

# Set the environment variable without displaying the full key
os.environ['AWS_ACCESS_KEY_ID'] = access_key
os.environ['AWS_SECRET_ACCESS_KEY'] = secret_access_key

Note: You may also set these keys directly using the %env magic command, but the key value may be echoed in the cell output. If you prefer to use %env, be sure to not share the notebook with people you don’t trust, as this may leak your API keys.

Then, load the IPython extension:

%load_ext jupyter_ai

Jupyter AI supports language models hosted on SageMaker endpoints that use JSON APIs. Authenticate with AWS via the boto3 SDK and have the credentials stored in the default profile. Guidance on how to do this can be found in the `boto3 documentation <https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html>`__.

You will need to deploy a model in SageMaker, then provide it as your model name (as sagemaker-endpoint:my-model-name). See the documentation on how to deploy a JumpStart model.

All SageMaker endpoint requests require you to specify the --region-name, --request-schema, and --response-path options.

The --region-name parameter is set to the AWS region code, such as us-east-1 or eu-west-1.

The --request-schema parameter is the JSON object the endpoint expects as input, with the prompt being substituted into any value that matches the string literal "<prompt>". For example, the request schema {"text_inputs":"<prompt>"} will submit a JSON object with the prompt stored under the text_inputs key.

The --response-path option is a JSONPath string that retrieves the language model’s output from the endpoint’s JSON response. For example, if your endpoint returns an object with the schema {"generated_texts":["<output>"]}, its response path is generated_texts.[0].

%%ai sagemaker-endpoint:jumpstart-dft-hf-text2text-flan-t5-xl --region-name=us-east-1 --request-schema={"text_inputs":"<prompt>"} --response-path=generated_texts.[0]
Generate a question about weather today.

What is the weather like in the US today?