SageMaker JumpStart Foundation Models - HuggingFace Text2Text Instruction Fine-Tuning

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 use the SageMaker Python SDK to fine-tune a Text2Text model. Such a model takes prompting text as input and generates text as output. The prompt can include a task description in natural language. Accordingly, the model can be used for a variety of NLP tasks (e.g., text summarization, question answering, etc.).

We will fine-tune a pre-trained FLAN T5 model from Hugging Face. While pre-trained FLAN T5 models can be used "as is" for many tasks, fine-tuning can improve model performance on a particular task or language domain. As an example, we will fine-tune the model for a task that was not used for pre-training. After fine-tuning we will deploy two inference endpoints, one with a pre-trained and one with a fine-tuned model. We will then run the same inference query against both endpoints and compare results.

This notebook was tested on ml.t3.medium Amazon SageMaker Notebook instance with conda_python3 kernel.

In this notebook:

1. Setting up
2. Fine-tuning a model
3. Deploying inference endpoints
4. Running inference queries
5. Cleaning up resources

1. Setting up

We begin by installing and upgrading necessary packages. Restart the kernel after executing the cell below for the first time.

We will use the following variables throughout the notebook. In particular, we select FLAN T5 model size and select training and inference instance types. We also obtain execution role associated with the current notebook instance.

2. Fine-tuning a model

FLAN T5 models were pre-trained on a variety of tasks. In this demo, we fine-tune a model for a new task. In this task, given a piece of text, the model is asked to generate questions that are relevant to the text, but cannot be answered based on provided information. Examples are given in the inference section of this notebook.

2.1. Preparing training data

We will use a subset of SQuAD2.0 for supervised fine-tuning. This dataset contains questions posed by human annotators on a set of Wikipedia articles. In addition to questions with answers, SQuAD2.0 contains about 50k unanswerable questions. Such questions are plausible, but cannot be directly answered from the articles' content. We only use unanswerable questions for our task.

Citation: @article{rajpurkar2018know, title={Know what you don't know: Unanswerable questions for SQuAD}, author={Rajpurkar, Pranav and Jia, Robin and Liang, Percy}, journal={arXiv preprint arXiv:1806.03822}, year={2018} }

License: Creative Commons Attribution-ShareAlike License (CC BY-SA 4.0)

The Text2Text generation model can be fine-tuned on any text data provided that the data is in the expected format. The data must include a training and an optional validation parts. The best model is selected according to the validation loss, calculated at the end of each epoch. If a validation set is not given, an (adjustable) percentage of the training data is automatically split and used for validation.

The training data must be formatted in JSON lines (.jsonl) format, where each line is a dictionary representing a single data sample. All training data must be in a single folder, however it can be saved in multiple jsonl files. The .jsonl file extension is mandatory. The training folder can also contain a template.json file describing the input and output formats.

If no template file is given, the following default template will be used:

{
    "prompt": "{prompt}",
    "completion": "{completion}"
}

In this case, the data in the JSON lines entries must include prompt and completion fields.

In this demo, we are going to use a custom template (see below).

We will use simplified JumpStart SDK that defaults various input paramters.

2.2. Starting training

We are now ready to start the training job. This can take a while to complete, from 20 minutes to several hours, depending on the model size, amount of data, and so on (e.g., it can take a few hours for the xl model, 40k examples and 3 epochs).

Performance metrics such as training and validation loss can be accessed through CloudWatch during training. We can also fetch the most recent snapshot of metrics as follows.

3. Deploying inference endpoints

Remainder of the notebook should be executed once the training job is successfully completed. We will inference endpoints of pretrained and finetuned models and run the same request against the two endpoints and compare the results. Note that each endpoint deployment can take a few minutes.

4. Running inference queries

As the name suggests, a Text2Text model such as FLAN T5 receives a piece of text as input, and generates text as output. The input text will contain the description of the task. In this demo, our task is to generate questions given a piece of text. The questions must be relevant to the text, but the text should contain no answer. Such a task could arise when automating gathering additional information, or identifying gaps in technical documentation.

5. Cleaning up resources

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.