Fine-tuning and deploying a Hugging Face summarization model on SageMaker with your own scripts and 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.

For ease of use, we advise opening this notebook in an Amazon SageMaker notebook instance using the conda_pytorch_latest_p36 kernel, or in Amazon SageMaker Studio using a Python 3 (PyTorch 1.8 Python 3.6 CPU Optimized) kernel on a ml.t3.medium instance.

In this notebook, we will see how to fine-tune and deploy one of the 🤗 Transformers model for a summarization task on Amazon SageMaker with your own scripts and data.

In the first part "Preparing the dataset" we show how to load your own dataset to s3 into separated files for training, validation and testing. We will use the Women's E-Commerce Clothing Reviews dataset which contains e-commerce clothing reviews and review titles, but we also provide code to do it for your own custom dataset. In our case the text and summary columns are called review_text and title respectively, and the data is saved in s3 under the prefix DEMO-sagemaker-huggingface-summarization.

Afterwards, we walk you through how to create your own train and inference scripts to fine-tune and deploy a Hugging Face model on Amazon SageMaker.

Make sure that the latest version of SageMaker SDK is installed

Part 1: Preparing the dataset for Hugging Face on Amazon SageMaker

One way to prepare your dataset for training on Amazon SageMaker is to have your training, validation and test datasets saved separately. This enables to effectively decouple data preparation from training in an architecture and for example ensure that the same datasets can be reused by different models with the same split. In this example we download the Women's E-Commerce Clothing Reviews dataset and prepare it for Hugging Face using the datasets library. Any dataset containing text and something that could be considered a summary (e.g. titles) can work here.

We first import required packages and define the prefix where we will save the data:

We read the raw dataset directly from its source

This raw dataset has missing values in the columns that are interesting for us: "Review text" and "Title". So we drop rows with missing values in those 2 columns. Additionally, we reformat the column names to be lowercase and replace space by underscore.

The cleaned dataset should contain 19675 rows.

Now that we've cleaned the data from missing reviews and titles, we will split it into train, validation and test set using the load_dataset() functions from the datasets library.

We can inspect an example review:

Finally, we write the training, validation and test data frames to separate CSVs and upload them to S3.

Use the save_to_disk method to directly save your dataset to S3 in Hugging Face dataset format. The format is backed by the Apache Arrow format which enables processing of large datasets with zero-copy reads without any memory constraints for optimal speed and efficiency. You can use the load_to_disk method in your train script to directly load the dataset in the format it was saved.

Part 2: Fine-tune and deploy a Hugging Face model on Amazon SageMaker

Now that the data is ready and saved in s3, we will demonstrate how to fine-tune and deploy a Hugging Face model on Amazon SageMaker with your own scripts.

This notebook is built to run with any model checkpoint from the Model Hub as long as that model has a sequence-to-sequence version in the Transformers library. Here we picked the pegasus-xsum checkpoint.

Write the training script

To fine-tune a Hugging Face model with a custom dataset on Amazon SageMaker, we will write a training script to be used by the Amazon SageMaker Training Job.

The training script will need to do the following steps:

• Load a pretrained Tokenizer and Model
• Load and Tokenize datasets
• Define the Training Arguments
• Define a Trainer
• Train the model and save the checkpoint with the best performance on the validation set
• Evaluate the best checkpoint on the test set

These steps will be done in a train() function which uses a couple helper functions: tokenize() takes a batch, specified text and target columns, and tokenizes them with the Tokenizer loaded in memory, load_and_tokenize() which reads data from s3 and applies the tokenize() function, and compute_metrics() to compute ROUGE scores for evaluation.

The script uses AutoTokenizer and AutoModelForSeq2SeqLM which works with any 🤗 Transformers model for summarization. You might however want to change some hyperparameters depending on what works best for each model. Here we used adafactor as optimizer for Pegasus for example.

All computations will be running inside Amazon SageMaker Hugging Face training and inference containers, which we call using the SageMaker SDK

By default, the Trainer saves several checkpoints before selecting the best one. Once the best checkpoint is loaded in memory and saved, those remaining checkpoints are not needed anymore. They can be safely deleted (which we do in the last line of the train()) to liberate space in the SM_MODEL_DIR which content will be used later for creating a SageMaker Model and deploy it to an endpoint.

Fine-tuning the model on SageMaker

We first load a couple of libraries and objects, namely sagemaker and the HuggingFace SageMaker Estimator which will be used to launch a training job.

We define a few arguments to be sent to the training script which will be read by the parser.

Thanks to 🤗 Transformers' Trainer seamless integration with SageMaker Distributed Data Parallel, we can make use of instances with several GPU units to parallelize and speed up training, without any modification to our training script.

When defining the SageMaker Hugging Face Estimator we specify a training script and source directory (here only containing train.py, but it could contain any additional modules and a requirements.txt), as well as the instance type on which to run the Training Job.

We then launch the training job by specifying where to read the data from. 'train' will be loaded inside SM_CHANNEL_TRAIN, 'validation' inside SM_CHANNEL_VALIDATION and 'test' inside SM_CHANNEL_TEST, which will be the data directories inside the container running train.py.

With distributed training on a p3.16xlarge instance, the training should take around 6 hours for 5 epochs.

Bring your own inference script

Our friends at Hugging Face have made inference on SageMaker for transformers model simpler than ever thanks to the SageMaker Hugging Face Inference Toolkit. You can directly deploy the previously trained model by simply setting up the environment variable "HF_TASK":"summarization" following the instructions on the HuggingFace website selecting "Deploy" and then "Amazon SageMaker", without the need to write an inference script.

However, when needing specific post-processing, for example if for a same input you want to return several summaries based on different text generation parameters, bringing your own inference.py script might be useful, and relatively straightforward:

As we can see, the only requirements to writing such an inference script for Hugging Face on SageMaker is that the inference script shall contain the following template functions:

• model_fn() reading the content of what was saved at the end of the training job inside SM_MODEL_DIR, or from an existing model weights directory saved as a tar.gz in s3. We will use it to load the trained Model and associated Tokenizer
• input_fn() used here simply to format the data receives from a request made to the endpoint.
• predict_fn() calling the output of model_fn() (so here the model and tokenizer) to run inference on the output of input_fn().

Optionally a output_fn() can be created for inference formatting, using the output of predict_fn(), but we did not use it here.

Create and deploy a SageMaker Model to an endpoint and test it

This time we will import the SageMaker HuggingFaceModel object which will help us create a SageMaker Model and deploy it to an endpoint.

Again, we specify here the inference script that we wrote earlier, a source directory (here again containing only inference.py but could contain modules and a requirements.txt) and model_data specifying where to load the model weights from. Using huggingface_estimator.model_data directly points to the s3 location where the output of the huggingface_estimator (after training) was saved, but any s3 arn containing pre-trained weights compressed as a tar.gz could work.

Finally, we deploy the register model by specifying the instance type.

Once the model is deployed, you can test it directly: Feel free to change the parameters list to see different predictions

Lastly, please remember to delete the Amazon SageMaker endpoint to avoid charges.

Conclusion

In this notebook, we trained and deployed a Hugging Face model for Text Summarization with custom scripts and data on Amazon SageMaker. You can use this solution to train and deploy other pretrained models provided by Hugging Face. Sample notebooks are available on GitHub.

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.