SageMaker JumpStart Foundation Models - GPT-J, GPT-Neo Few-shot learning

This notebook's CI test result for us-west-2 is as follows. CI test results in other regions can be found at the end of the notebook.

Welcome to Amazon SageMaker JumpStart! You can use Sagemaker JumpStart to solve many Machine Learning tasks through one-click in SageMaker Studio, or through SageMaker Python SDK.

In this demo notebook, we demonstrate how to use the SageMaker Python SDK for deploying Foundation Models as an endpoint and use them for various NLP tasks. The Foundation models perform Text Generation. It takes a text string as input and predicts a sequence of next few words. These models can, for example, fill in incomplete text or paraphrase.

Here, we show how to use the state-of-the-art pre-trained GPT-J-6B models from Hugging Face for Text Generation in the following tasks. You can use GPT-J-6B model with few shot learning for many NLP tasks, without fine-tuning the model. There are other text generation models like Bloom 7b1, Bloom 3b, GPT-NEO-2b7, and etc.

The few shot leanring allows a machine learning model to make predictions with only a couple of examples. There is no need to train a new model as models like GPT-J, Bloom have a significant amount of parameters such that they can easily adapt to many contexts without being re-trained. Below are tasks that text generation can solve using few shot learning method.

• Text summarization
• Code generation
• Name entity recognition (NER)
• Question answering
• Grammar and spelling correction
• Product description and generalization
• Sentence / sentiment classification
• Chatbot and conversational AI
• Tweet generation
• Machine translation
• Pharapharsing
• Intent classification

For domain adaption finetuning of text generation models such as GPT-J 6B, please check the notebook domain-adaption-finetuning-gpt-j-6b.ipynb in the same directory.

1. Set Up
2. Select a model
3. Retrieve Artifacts & Deploy an Endpoint
4. Query endpoint and parse response
5. Advanced features: How to use varisous parameters to control the generated text
6. Advanced features: How to use prompts engineering to solve different tasks
7. Clean up the endpoint

Note: This notebook was tested on ml.t3.medium instance in Amazon SageMaker Studio with Python 3 (Data Science) kernel and in Amazon SageMaker Notebook instance with conda_python3 kernel.

1. Set Up

Before executing the notebook, there are some initial steps required for set up. This notebook requires ipywidgets.

Permissions and environment variables

To host on Amazon SageMaker, we need to set up and authenticate the use of AWS services. Here, we use the execution role associated with the current notebook as the AWS account role with SageMaker access.

2. Select a pre-trained model

You can continue with the default model, or can choose a different model from the dropdown generated upon running the next cell. A complete list of SageMaker pre-trained models can also be accessed at Sagemaker pre-trained Models.

[Optional] Select a different Sagemaker pre-trained model. Here, we download the model_manifest file from the Built-In Algorithms s3 bucket, filter-out all the Text Generation models and select a model for inference.

The notebook in the following sections uses GPT-J-6B as an example. You are welcome to try Bloom 7b1, Bloom 3b, GPT-NEO-2b7 and many others yourself. Please modify model_id based on the dropdown list shown as below.

Choose a model for Inference

3. Retrieve Artifacts & Deploy an Endpoint

Using SageMaker, we can perform inference on the pre-trained model, even without fine-tuning it first on a new dataset. We start by retrieving the deploy_image_uri, deploy_source_uri, and model_uri for the pre-trained model. To host the pre-trained model, we create an instance of sagemaker.model.Model and deploy it. This may take a few minutes.

4. Query endpoint and parse response

Input to the endpoint is any string of text formatted as json and encoded in utf-8 format. Output of the endpoint is a json with generated text.

Below, we put in some example input text. You can put in any text and the model predicts next words in the sequence. Longer sequences of text can be generated by calling the model repeatedly.

5. Advanced features: How to use various advanced parameters to control the generated text

This model also supports many advanced parameters while performing inference. They include:

• max_length: Model generates text until the output length (which includes the input context length) reaches max_length. If specified, it must be a positive integer.
• num_return_sequences: Number of output sequences returned. If specified, it must be a positive integer.
• num_beams: Number of beams used in the greedy search. If specified, it must be integer greater than or equal to num_return_sequences.
• no_repeat_ngram_size: Model ensures that a sequence of words of no_repeat_ngram_size is not repeated in the output sequence. If specified, it must be a positive integer greater than 1.
• temperature: Controls the randomness in the output. Higher temperature results in output sequence with low-probability words and lower temperature results in output sequence with high-probability words. If temperature -> 0, it results in greedy decoding. If specified, it must be a positive float.
• early_stopping: If True, text generation is finished when all beam hypotheses reach the end of stence token. If specified, it must be boolean.
• do_sample: If True, sample the next word as per the likelyhood. If specified, it must be boolean.
• top_k: In each step of text generation, sample from only the top_k most likely words. If specified, it must be a positive integer.
• top_p: In each step of text generation, sample from the smallest possible set of words with cumulative probability top_p. If specified, it must be a float between 0 and 1.
• seed: Fix the randomized state for reproducibility. If specified, it must be an integer.

We may specify any subset of the parameters mentioned above while invoking an endpoint. Next, we show an example of how to invoke endpoint with these arguments

6. Advanced features: How to use prompts engineering to solve different few shot learning NLP tasks

Note . The notebook in the following sections uses GPT-J-6B as an example. You are welcome to try Bloom 7b1, Bloom 3b, GPT-NEO-2b7 and many others yourself. Please modify model_id in above section.

6.1. Summarization

Define the text article you want to summarize.

6.2. Code generation

How about SQL code?

6.3. Name entity recognition (NER)

6.4. Question answering

6.5. Grammar and spelling correction

6.6. Product description and generalization

6.7. Sentence / Sentiment Classification

Define the sentence you want to classifiy and the corresponded options.

6.8. Chatbot and Conversational AI

Define the sentence and the language you want to translate the sentence to.

6.9. Tweet generation

Define the sentence, pronoun, and options you want to reason.

6.10. Machine translation

6.11. Paraphrasing

6.12. Intent classification

7. Clean up the endpoint

For domain adaption finetuning of text generation models such as GPT-J 6B, please check the notebook domain-adaption-finetuning-gpt-j-6b.ipynb in the same directory.

Notebook CI Test Results

This notebook was tested in multiple regions. The test results are as follows, except for us-west-2 which is shown at the top of the notebook.