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 either before starting JupyterLab or using the %env magic command within JupyterLab.

Then, load the IPython extension:

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.

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].