Python SDK - with Amazon SageMaker

AI21 Studio Python SDK with Amazon SageMaker Guide

This guide covers how to use the AI21 Studio Python SDK with SageMaker integration for interaction with Jurassic-2 models.

Set up

To get started with AI21's SDK with SageMaker integration, you'll need to install it first. To do so, simply run the following command:

$ pip install -U "ai21[AWS]"

Using AI21 Studio Python SDK with SageMaker

To use the AI21 Studio Python SDK with SageMaker, first, create an Amazon SageMaker instance and configure it. Then, install the AI21 Studio Python SDK with SageMaker integration using the command mentioned above. After setting up the environment, you can follow the same steps as described in the AI21 Studio Python SDK Guide for Jurassic-2 models.

Example: Using AI21 Studio Python SDK with SageMaker for Jurassic-2

Below is a sample usage of the AI21 Python SDK with SageMaker integration to interact with Jurassic-2 models:

import ai21

# J2 Mid
response_mid = ai21.Completion.execute(
  destination=ai21.SageMakerDestination("<your_endpoint_name>"),
  prompt="explain black holes to 8th graders",
  numResults=1,
  maxTokens=100,
  temperature=0.7
)

print(response_mid)

By customizing the request parameters, you can control the content and style of the generated text. For a full list of available options, check out our Complete API page.

All Jurassic-2 models can be interacted with using the same ai21.Completion.execute() function, including j2-light, j2-mid, and j2-ultra.

Response

Here's an example of a response object from executing a j2-mid model:

{
   "id":"94078cb6-687e-4262-ef8f-1d7c2b0dbd2b",
   "prompt":{
      "text":"These are a few of my favorite",
      "tokens":[
         {
            "generatedToken":{
               "token":"▁These▁are",
               "logprob":-8.824776649475098,
               "raw_logprob":-8.824776649475098
            },
            "topTokens":"None",
            "textRange":{
               "start":0,
               "end":9
            }
         },
         {
            "generatedToken":{
               "token":"▁a▁few",
               "logprob":-4.798709869384766,
               "raw_logprob":-4.798709869384766
            },
            "topTokens":"None",
            "textRange":{
               "start":9,
               "end":15
            }
         },
         {
            "generatedToken":{
               "token":"▁of▁my▁favorite",
               "logprob":-1.0864331722259521,
               "raw_logprob":-1.0864331722259521
            },
            "topTokens":"None",
            "textRange":{
               "start":15,
               "end":30
            }
         }
      ]
   },
   "completions":[
      {
         "data":{
            "text":" things –",
            "tokens":[
               {
                  "generatedToken":{
                     "token":"▁things",
                     "logprob":-0.0003219324571546167,
                     "raw_logprob":-0.47372230887413025
                  },
                  "topTokens":"None",
                  "textRange":{
                     "start":0,
                     "end":7
                  }
               },
               {
                  "generatedToken":{
                     "token":"▁–",
                     "logprob":-7.797079563140869,
                     "raw_logprob":-4.319167613983154
                  },
                  "topTokens":"None",
                  "textRange":{
                     "start":7,
                     "end":9
                  }
               }
            ]
         },
         "finishReason":{
            "reason":"length",
            "length":2
         }
      }
   ]
}

The response is a nested data structure containing information about the processed request, prompt, and completions. At the top level, the response has the following fields:

ID

A unique string id for the processed request. Repeated identical requests receive different IDs.

prompt

The prompt includes the raw text, the tokens with their log probabilities, and the top-K alternative tokens at each position, if requested. It has two nested fields:

text (string)
tokens (list of TokenData)

completions

A list of completions, including raw text, tokens, and log probabilities. The number of completions corresponds to the requested numResults. Each completion has two nested fields:

data, which contains the text (string) and tokens (list of TokenData) for the completion.
finishReason, a nested data structure describing the reason generation was terminated for this completion.

TokenData

The TokenData object provides detailed information about each token in both the prompt and the completions. It includes the following fields:

generatedToken:

The generatedToken field consists of two nested fields:

token: The string representation of the token.
logprob: The predicted log probability of the token after applying the sampling parameters as a float value.
raw_logprob: The raw predicted log probability of the token as a float value. For the indifferent values (namely, temperature=1, topP=1) we get raw_logprob=logprob.

topTokens

The topTokens field is a list of the top K alternative tokens for this position, sorted by probability, according to the topKReturn request parameter. If topKReturn is set to 0, this field will be null.

Each token in the list includes:

token: The string representation of the alternative token.
logprob: The predicted log probability of the alternative token as a float value.

textRange

The textRange field indicates the start and end offsets of the token in the decoded text string:

start: The starting index of the token in the decoded text string.
end: The ending index of the token in the decoded text string.

Additional Resources

For a detailed example of using Jurassic-2 Mid on SageMaker through Model Packages, you can refer to this notebook on AI21 Labs' SageMaker GitHub repository.

By using the AI21 Studio Python SDK with SageMaker integration, you can seamlessly utilize AI21 Studio Jurassic-2 models in your SageMaker environment, allowing for streamlined development and deployment of your machine learning solutions. Note that the SDK version for SageMaker supports Jurassic-2 models but not the Task Specific Models. Additionally, there's no need to use an AI21 Studio API key when using the SDK version for SageMaker.