AI21 Studio Python SDK Guide
This SDK enables interaction with all of AI21 Studio's models and services, including Jurassic-2 and Task-Specific APIs.
Set up
To get started with AI21's SDK, you'll need to install it first. To do so, simply run the following command:
$ pip install ai21
API Keys
Before you can start using the SDK, you'll need to obtain your API key from AI21 Studio. Once you have your API key, you can set it globally in your Python code like this:
import ai21
ai21.api_key = <YOUR_API_KEY>
Building a Prompt
To generate text using Jurassic-2 language models, you need to provide a prompt for the model to complete. We recommend following our Prompt Engineering best practices to achieve the best results.
Here are a couple of examples of prompts you can use to get started:
"Zero Shot" prompt:
prompt = "These are a few of my favorite"
"Few Shot" prompt:
prompt = """
Complete the following sentence.
Sentence: These are a few of my favorite
Completion: things. Cream-colored ponies and crisp apple strudels.
##
Complete the following sentence.
Sentence: These are a few of my favorite
Completion: rock bands: The Beatles, The Doors, Fleetwood Mac.
##
Complete the following sentence.
Sentence: These are a few of my favorite
Completion:
"""
Jurassic-2
Jurassic-2 with Python SDK Code Example
Below is a sample usage of the AI21 Python SDK to interact with Jurassic-2 models, demonstrating some of the frequently utilized parameters:
import ai21
ai21.api_key = <YOUR_API_KEY>
# J2 Grande
response_grande = ai21.Completion.execute(
model="j2-grande",
prompt="These are a few of my favorite",
numResults=1,
maxTokens=2,
temperature=0.4,
topKReturn=0,
topP=1,
stopSequences=["##"]
)
# J2 Grande Instruct
response_grande_instruct = ai21.Completion.execute(
model="j2-grande-instruct",
prompt="explain black holes to 8th graders",
numResults=1,
maxTokens=100,
temperature=0.7,
topKReturn=0,
topP=1,
stopSequences=["##"]
)
print(response_grande)
print(response_grande_instruct)
By customizing these 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-large, j2-grande, j2-jumbo, j2-grande-instruct, and j2-grande-jumbo.
Response
Here's an example of a response object from executing a j2-grande 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 as a float value..
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.
Task-Specific APIs
Paraphrase API with Python SDK Code Example
import ai21
ai21.api_key = <YOUR_API_KEY>
response = ai21.Paraphrase.execute(
text="You only live once, but if you do it right, once is enough.",
style="general"
)
print(response.summaries[0].text)
Note: This API will return a list of up to 10 suggested paraphrases with every request. Read more here.
Grammatical Error Correction (GEC) API with Python SDK Code Example
import ai21
ai21.api_key = <YOUR_API_KEY>
response = ai21.GEC.execute(
text="jazzz is a great stile off music",
)
print(response.corrections[0].suggestion)
print(response.corrections[0].startIndex)
print(response.corrections[0].endIndex)
print(response.corrections[0]['originalText'])
print(response.corrections[0]['correctionType'])
Note: This API can return 7 different types of grammatical error corrections: Grammar, Missing Word, Punctuation, Spelling, Word Repetition, and Wrong Word. Read more here.
Text Improvements API with Python SDK Code Example
import ai21
ai21.api_key = <YOUR_API_KEY>
response = ai21.Improvements.execute(
text="Affiliated with the profession of project management, I have ameliorated myself with a different set of hard skills as well as soft skills",
types=["fluency", "vocabulary/specificity"]
)
print(response.improvements[0].originalText)
print(response.improvements[0].suggestions)
print(response.improvements[0].suggestions[0])
print(response.improvements[0].improvementType)
print(response['improvements'][1].startIndex)
print(response['improvements'][1].endIndex)
Note: This API can return as many as 5 different text improvement types: fluency, vocabulary/specificity, vocabulary/variety, clarity/short-sentences, and clarity/conciseness. Read more here.
Summarize API with Python SDK Code Example
import ai21
ai21.api_key = <YOUR_API_KEY>
response = ai21.Summarize.execute(
source="https://en.wikipedia.org/wiki/Koala",
sourceType="URL"
)
print(response.summary)
Note: sourceType can be either "URL" or "TEXT". Read more here.
Text Segmentation API with Python SDK Code Example
import ai21
ai21.api_key = <YOUR_API_KEY>
response = ai21.Segmentation.execute(
source="https://www.ai21.com/blog/summarizing-legal-documents-for-different-personas-using-ai21-studio",
sourceType="URL"
)
print(response.segments[0].segmentText)
print(response.segments[0].segmentType)
Note: sourceType can be either "URL" or "TEXT". Read more here.