AI21 Jurassic-1 foundation model is now available on Amazon SageMaker

Today we are excited to announce that AI21 Jurassic-1 (J1) foundation models are available for customers using Amazon SageMaker. Jurassic-1 models are highly versatile, capable of both human-like text generation, as well as solving complex tasks such as question answering, text classification, and many others. You can easily try out this model and use it with Amazon SageMaker JumpStart. JumpStart is the machine learning (ML) hub of SageMaker that provides access to foundation models in addition to built-in algorithms and end-to-end solution templates to help you quickly get started with ML.

In this post, we walk through how to use the Jurassic-1 Grande model in SageMaker.

Foundation models in SageMaker

JumpStart provides access to a range of models from popular model hubs including Hugging Face, PyTorch Hub, and TensorFlow Hub, which you can use within your ML development workflow in SageMaker. Recent advances in ML have given rise to a new class of models known as foundation models, which are typically trained on billions of parameters and are adaptable to a wide category of use cases, such as text summarization, generating digital art, and language translation. Because these models are expensive to train, customers want to use existing pre-trained foundation models and fine-tune them as needed, rather than train these models themselves. SageMaker provides a curated list of models that you can choose from on the SageMaker console.

You can now find foundation models from different model providers within JumpStart, enabling you to get started with foundation models quickly. You can find foundation models based on different tasks or model providers, and easily review model characteristics and usage terms. You can also try out these models using a test UI widget. When you want to use a foundation model at scale, you can do so easily without leaving SageMaker by using pre-built notebooks from model providers. Because the models are hosted and deployed on AWS, you can rest assured that your data, whether used for evaluating or using the model at scale, is never shared with third parties.

Jurassic-1 foundation model

Jurassic-1 is the first generation in a series of large language models trained and made widely accessible by AI21 Labs. For a complete description of Jurassic-1, including benchmarks and quantitative comparisons with other models, refer to the following technical paper. All J1 models were trained on a massive corpus of English text, making them highly versatile general purpose text-generators, capable of composing human-like text and solving complex tasks such as question answering, text classification, and many others. J1 can be applied to virtually any language task by crafting a suitable prompt containing a description of the task and a few examples, a process commonly known as prompt engineering. Popular use cases include generating marketing copy, powering chatbots, and assisting creative writing.

“We are building world-class foundation models for text and want to help our customers innovate with the latest Jurassic-1 models. Amazon SageMaker offers the deepest and broadest set of ML services, and we’re excited to collaborate with Amazon SageMaker so that customers will be able to use these foundation models on SageMaker within their development environment. Now customers can rapidly innovate, lower time-to-value, and drive efficiency in their businesses.”

-Ori Goshen, co-CEO of AI21 Labs.

Walkthrough

Let’s take you on a tour to test the J1-Grande model in SageMaker. You can try out the experience in three simple steps:

Choose the Jurassic-1 model on the SageMaker console.
Evaluate the model using a test widget.
Use a notebook associated with the foundation model to deploy it in your environment.

Let’s expand each step in detail.

Choose the Jurassic-1 model on the SageMaker console

The first step is to login to the AWS Management Console for Amazon SageMaker and request access to the list of foundation models from the foundation model category under JumpStart here:

After your account is allow listed, you can see a list of models on this page. You can quickly search for the Jurassic-1 Grande model from the same view.

Evaluate the Jurassic-1 Grande model with a test widget

On the Jurassic-1 Grande listing, choose View Model. You will see a description of the model and the tasks that you can perform. Read through the EULA for the model before proceeding.

Let’s first try out the model for text summarization. Choose Try out model.

You’re taken to the page in a separate browser tab where you can give sample prompts to the J1-Grande model and view the output.

The following example generates a summary about a restaurant based on reviews.

Note that foundation models and their output are from the model provider, and AWS is not responsible for the content or accuracy therein.

The model output may vary depending on the settings and the prompt. You can generate text from the model using simple instructions, but by providing the model with more examples in the prompt, just as a human would, it can produce completions that are more aligned with your intentions. The best way to guide the model is to provide several examples of input/output pairs in the prompt. This establishes a pattern for the model to mimic. Then add the input for a query example and let the model complete it with an appropriate generation.

After you have played with the model, it’s time to use the notebook and deploy it as an endpoint in your environment.

Deploy the foundation model from a notebook

Go back to the model listing shown earlier and choose View notebook. You should see the Jurassic-1 Grande Jupyter notebook with the walkthrough to deploy the model.

Let’s use this notebook from Amazon SageMaker Studio. Open Studio and pull in the notebook using the Git repo URL https://github.com/AI21Labs/SageMaker.git.

The notebook example uses both the Boto3 SDK and the AI21 SDK to deploy and interact with the endpoint.

Note that this example uses an ml.g5.12xlarge instance. If your default limit for your AWS account is 0, you need to request a limit increase for this GPU instance.

Let’s create the endpoint using SageMaker inference. First we set the necessary variables, then we deploy the model from the model package:

model_name = "j1-grande"

content_type = "application/json"

real_time_inference_instance_type = (
    "ml.g5.12xlarge"
)

# create a deployable model from the model package.
model = ModelPackage(
    role=role, model_package_arn=model_package_arn, sagemaker_session=sagemaker_session
)

# Deploy the model
predictor = model.deploy(1, real_time_inference_instance_type, endpoint_name=model_name, 
                         model_data_download_timeout=3600,
                         container_startup_health_check_timeout=600,
                        )

After the endpoint is deployed, you can run inference queries against the model.

You can think of Jurassic-1 Grande as a smart auto-completion algorithm: it’s very good at latching on to hints and patterns expressed in plain English, and generating text that follows the same patterns. After the model is deployed, you can interact with the deployed endpoint using the following code snippet:

response = ai21.Completion.execute(sm_endpoint="j1-grande",
                                   prompt="To be or",
                                   maxTokens=4,
                                   temperature=0,
                                   numResults=1)

print(response['completions'][0]['data']['text'])

The notebook also contains a walkthrough on how you can run inference queries with the AI21 SDK.

The following video walks through the workflow.

Clean up

After you have tested the endpoint, make sure you delete the SageMaker inference endpoint and delete the model to avoid incurring charges.

Conclusion

In this post, we showed you how you can test and use AI21’s Jurassic Grande model using Amazon SageMaker. Request access, try out the foundation model in SageMaker today and let us know your feedback!

About the authors

Karthik Bharathy is the product leader for the Amazon SageMaker team with over a decade of product management, product strategy, execution, and launch experience.

Tomer Asida is an algo team lead at AI21 Labs. As an algo team lead, Tomer heads the algorithm development efforts of our developer platform Ai21 Studio including Jurassic-1 models and associated APIs.

AWS Machine Learning Blog