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 thetext
(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 getraw_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.
Updated about 1 month ago