HuggingFace serves as a home to many popular open-source NLP models. Many of these models are effective as is, but often require some sort of training or fine-tuning to improve performance for your specific use-case. As the LLM implosion continues, we will take a step back in this article to revisit some of the core building blocks HuggingFace provides that simplify the training of NLP models.
Traditionally NLP models can be trained using vanilla PyTorch, TensorFlow/Keras, and other popular ML frameworks. While you can go this route, it does require a deeper understanding of the framework you are utilizing as well as more code to write the training loop. With HuggingFace’s Trainer class, there’s a simpler way to interact with the NLP Transformers models that you want to utilize.
Trainer is a class specifically optimized for Transformers models and also provides tight integration with other Transformers libraries such as Datasets and Evaluate. Trainer at a more advanced level also supports distributed training libraries and can be easily integrated with infrastructure platforms such as Amazon SageMaker.
In this example we’ll take a look at using the Trainer class locally to fine-tune the popular BERT model on the IMBD dataset for a Text Classification use-case(Large Movie Reviews Dataset Citation).
NOTE: This article assumes basic knowledge of Python and the domain of NLP. We will not get into any specific Machine Learning theory around model building or selection, this article is dedicated to understanding how we can fine-tune the existing pre-trained models available in the HuggingFace Model Hub.
- Setup
- Fine-Tuning BERT
- Additional Resources & Conclusion
For this example, we’ll be working in SageMaker Studio and utilize a conda_python3 kernel on a ml.g4dn.12xlarge instance. Note that you can use a smaller instance type, but this might impact the training speed depending on the number of CPUs/workers that are available.