Deploy and tune basic models in Amazon SageMaker JumpStart with two lines of code

We are pleased to announce a simplified version of the Amazon SageMaker JumpStart SDK that simplifies building, training, and deploying basic models. The prediction code is also simplified. In this post, we demonstrate how you can use the simplified SageMaker Jumpstart SDK to get started using basic models in just a couple of lines of code.

For more information about the simplified SageMaker JumpStart SDK for deployment and training, see Low code implementation with JumpStartModel class and Low code tuning with JumpStartEstimator classrespectively.

Solution Overview

SageMaker JumpStart provides pre-trained open source models for a wide range of problem types to help you get started with machine learning (ML). You can progressively train and tune these models before deployment. JumpStart also provides solution templates that configure the infrastructure for common use cases and executable sample notebooks for ML with Amazon SageMaker. You can access pre-trained models, solution templates, and examples through the SageMaker JumpStart home page in Amazon SageMaker Studio or by using the SageMaker Python SDK.

To demonstrate the new features of the SageMaker JumpStart SDK, we show you how to use the pre-trained software. Flan T5 XL Hugging Face model for text generation for summary tasks. We also show how, with just a few lines of code, you can tune the Flan T5 XL model for summarization tasks. You can use any other template for text generation like call2, Hawkeither Mistral ai.

You can find the notebook for this solution using Flan T5 XL in the GitHub repository.

Implement and invoke the model.

Foundation models hosted in SageMaker JumpStart have model IDs. For the complete list of model IDs, see Integrated algorithms with pre-trained model table. For this post, we use the Flan T5 XL text generation model ID. We instantiate the model object and deploy it to a SageMaker endpoint by calling su deploy method. See the following code:

from sagemaker.jumpstart.model import JumpStartModel

# Replace with larger model if needed
pretrained_model = JumpStartModel(model_id="huggingface-text2text-flan-t5-base")
pretrained_predictor = pretrained_model.deploy()

We then invoke the model to create a summary of the provided text using the Flan T5 XL model. The new SDK interface makes it easy to invoke the model: you just need to pass the text to the predictor and it returns the model response as a Python dictionary.

text = """Summarize this content - Amazon Comprehend uses natural language processing (NLP) to extract insights about the content of documents. It develops insights by recognizing the entities, key phrases, language, sentiments, and other common elements in a document. Use Amazon Comprehend to create new products based on understanding the structure of documents. For example, using Amazon Comprehend you can search social networking feeds for mentions of products or scan an entire document repository for key phrases. 
You can access Amazon Comprehend document analysis capabilities using the Amazon Comprehend console or using the Amazon Comprehend APIs. You can run real-time analysis for small workloads or you can start asynchronous analysis jobs for large document sets. You can use the pre-trained models that Amazon Comprehend provides, or you can train your own custom models for classification and entity recognition. """
query_response = pretrained_predictor.predict(text)
print(query_response("generated_text"))

The following is the result of the summary task:

Understand how Amazon Comprehend works. Use Amazon Comprehend to analyze documents.

Tune and deploy the model

The SageMaker JumpStart SDK provides you with a new class, JumpStartEstimator, which simplifies adjustment. You can provide the location of the fitting data and optionally also pass validation data sets. After fine-tuning the model, use the implementation method of the Estimator object to deploy the fine-tuned model:

from sagemaker.jumpstart.estimator import JumpStartEstimator

estimator = JumpStartEstimator(
    model_id=model_id,
)
estimator.set_hyperparameters(instruction_tuned="True", epoch="3", max_input_length="1024")
estimator.fit({"training": train_data_location})
finetuned_predictor = estimator.deploy()

Customize new classes in SageMaker SDK

The new SDK simplifies the deployment and tuning of JumpStart models by setting many default parameters. You still have the option to override the default values ​​and customize the implementation and invocation as per your requirements. For example, you can customize the input payload format type, instance type, VPC configuration, and more for your environment and use case.

The following code shows how to override the instance type when deploying your model:

finetuned_predictor = estimator.deploy(instance_type="ml.g5.2xlarge")

The SageMaker JumpStart SDK deploy The feature will automatically select a default content type and serializer for you. If you want to change the format type of the input payload, you can use serializers and content_types objects to introspect the options available to you by passing the model_id of the model you are working with. In the following code, we set the payload input format as JSON by setting JSONSerializer as serializer and application/json as content_type:

from sagemaker import serializers
from sagemaker import content_types

serializer_options = serializers.retrieve_options(model_id=model_id, model_version=model_version)
content_type_options = content_types.retrieve_options(model_id=model_id, model_version=model_version)

pretrained_predictor.serializer = serializers.JSONSerializer()
pretrained_predictor.content_type="application/json"

You can then invoke the Flan T5 XL model for the summary task with a payload in JSON format. In the following code, we also pass inference parameters in the JSON payload to make the responses more precise:

from sagemaker import serializers

input_text= """Summarize this content - Amazon Comprehend uses natural language processing (NLP) to extract insights about the content of documents. It develops insights by recognizing the entities, key phrases, language, sentiments, and other common elements in a document. Use Amazon Comprehend to create new products based on understanding the structure of documents. For example, using Amazon Comprehend you can search social networking feeds for mentions of products or scan an entire document repository for key phrases.
You can access Amazon Comprehend document analysis capabilities using the Amazon Comprehend console or using the Amazon Comprehend APIs. You can run real-time analysis for small workloads or you can start asynchronous analysis jobs for large document sets. You can use the pre-trained models that Amazon Comprehend provides, or you can train your own custom models for classification and entity recognition. """

parameters = {
    "max_length": 600,
    "num_return_sequences": 1,
    "top_p": 0.01,
    "do_sample": False,
}

payload = {"text_inputs": input_text, **parameters} #JSON Input format

pretrained_predictor.serializer = serializers.JSONSerializer()
query_response = pretrained_predictor.predict(payload)
print(query_response("generated_texts")(0))

If you’re looking for more ways to customize inputs and other housing and tuning options, check out the documentation. JumpStartModel and JumpStart Estimator classes.

Conclusion

In this post, we show you how you can use the simplified SageMaker JumpStart SDK to build, train, and deploy basic and task-based models in just a few lines of code. We demonstrate the new classes like JumpStartModel and JumpStartEstimator using the Hugging Face Flan T5-XL model as an example. You can use any of the other basic SageMaker JumpStart models for use cases such as content writing, code generation, question answering, summarization, classification, information retrieval, and more. To see the full list of models available with SageMaker JumpStart, see Integrated algorithms with pre-trained model table. SageMaker JumpStart also supports specific task models for many popular problem types.

We hope the SageMaker JumpStart SDK’s simplified interface helps you get started quickly and delivers faster. We look forward to hearing how you use SageMaker’s simplified JumpStart SDK to create interesting apps!

About the authors

Evan Kravitz He is a software engineer at Amazon Web Services and works at SageMaker JumpStart. He is interested in the confluence of machine learning with cloud computing. Evan received his bachelor’s degree from Cornell University and his master’s degree from the University of California, Berkeley. In 2021, he presented a paper on adversarial neural networks at the ICLR conference. In his spare time, Evan enjoys cooking, traveling, and running around New York City.

Rachna Chadha is a Principal ai/ML Solutions Architect in Strategic Accounts at AWS. Rachna is an optimist who believes that the ethical and responsible use of ai can improve society in the future and generate economic and social prosperity. In her free time, Rachna enjoys spending time with her family, hiking, and listening to music.

Jonathan Guinea is a Senior Software Engineer at Amazon SageMaker JumpStart on AWS. He earned his master’s degree from Columbia University. His interests span machine learning, distributed systems and cloud computing, as well as democratizing the use of ai. Jonathan is originally from France and now lives in Brooklyn, New York.

Dr Ashish Khaitan is a Senior Applied Scientist with Amazon SageMaker integrated algorithms and helps develop machine learning algorithms. He earned his doctorate from the University of Illinois Urbana-Champaign. He is an active researcher in machine learning and statistical inference, and has published numerous papers at NeurIPS, ICML, ICLR, JMLR, ACL, and EMNLP conferences.