SageMaker JumpStart Foundation Models - Fine-tuning text generation GPT-J 6B model on domain specific dataset

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 Built-in Algorithms! You can use SageMaker Built-in algorithms to solve many Machine Learning tasks through SageMaker Python SDK. You can also use these algorithms through one-click in SageMaker Studio via JumpStart.

In this demo notebook, we demonstrate how to use the SageMaker Python SDK for finetuning Foundation Models and deploying the trained model for inference. The Foundation models perform Text Generation task. It takes a text string as input and predicts next words in the sequence.

• How to fine-tune GPT-J 6B model on a domain specific dataset, and then run inference on the fine-tuned model. In particular, the example dataset we demonstrated is publicly available SEC filing of Amazon from year 2021 to 2022. The expectation is that after fine-tuning, the model should be able to generate insightful text in financial domain.

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
2. Select Text Generation Model GTP-J 6B
3. Finetune the pre-trained model on a custom dataset
   • Set Training parameters
   • Start Training
   • Deploy & run Inference on the fine-tuned model

1. Set Up

Before executing the notebook, there are some initial steps required for setup.

2. Select Text Generation Model GTP-J 6B

3. Fine-tune the pre-trained model on a custom dataset

Fine-tuning refers to the process of taking a pre-trained language model and retraining it for a different but related task using specific data. This approach is also known as transfer learning, which involves transferring the knowledge learned from one task to another. Large language models (LLMs) like GPT-J 6B are trained on massive amounts of unlabeled data and can be fine-tuned on domain domain datasets, making the model perform better on that specific domain.

We will use financial text from SEC filings to fine tune a LLM GPT-J 6B for financial applications.

• Input: A train and an optional validation directory. Each directory contains a CSV/JSON/TXT file.
  • For CSV/JSON files, the train or validation data is used from the column called 'text' or the first column if no column called 'text' is found.
  • The number of files under train and validation (if provided) should equal to one.
• Output: A trained model that can be deployed for inference. Below is an example of a TXT file for fine-tuning the Text Generation model. The TXT file is SEC filings of Amazon from year 2021 to 2022.

This report includes estimates, projections, statements relating to our
business plans, objectives, and expected operating results that are “forward-
looking statements” within the meaning of the Private Securities Litigation
Reform Act of 1995, Section 27A of the Securities Act of 1933, and Section 21E
of the Securities Exchange Act of 1934. Forward-looking statements may appear
throughout this report, including the following sections: “Business” (Part I,
Item 1 of this Form 10-K), “Risk Factors” (Part I, Item 1A of this Form 10-K),
and “Management’s Discussion and Analysis of Financial Condition and Results
of Operations” (Part II, Item 7 of this Form 10-K). These forward-looking
statements generally are identified by the words “believe,” “project,”
“expect,” “anticipate,” “estimate,” “intend,” “strategy,” “future,”
“opportunity,” “plan,” “may,” “should,” “will,” “would,” “will be,” “will
continue,” “will likely result,” and similar expressions. Forward-looking
statements are based on current expectations and assumptions that are subject
to risks and uncertainties that may cause actual results to differ materially.
We describe risks and uncertainties that could cause actual results and events
to differ materially in “Risk Factors,” “Management’s Discussion and Analysis
of Financial Condition and Results of Operations,” and “Quantitative and
Qualitative Disclosures about Market Risk” (Part II, Item 7A of this Form
10-K). Readers are cautioned not to place undue reliance on forward-looking
statements, which speak only as of the date they are made. We undertake no
obligation to update or revise publicly any forward-looking statements,
whether because of new information, future events, or otherwise.

GENERAL

Embracing Our Future ...

SEC filings data of Amazon is downloaded from publicly available EDGAR. Instruction of accessing the data is shown here.

Set Training parameters

Now that we are done with all the setup that is needed, we are ready to fine-tune our model. To begin, let us create a sageMaker.estimator.Estimator object. This estimator will launch the training job.

There are two kinds of parameters that need to be set for training.

The first one are the parameters for the training job. These include: Training data path. This is S3 folder in which the input data is stored

The second set of parameters are algorithm specific training hyper-parameters.

Start Training

We start by creating the estimator object with all the required assets and then launch the training job. Since default hyperparameter values are model-specific, inspect estimator.hyperparameters() to view default values for your selected model.

Deploy & run Inference on the fine-tuned model

A trained model does nothing on its own. We now want to use the model to perform inference.

Next, we query the finetuned model and print the predictions.

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.