Teach an AI Model to Write like Shakespeare — For Free
Use SageMaker Studio Lab to train your own NLP models
What is this about?
Training machine learning models has becoming increasingly resource intensive. Today’s state-of-the-art Natural Language Processing (NLP) models, for example, are almost exclusively neural networks that require GPU power to train them.
Fortunately, the number of freely accessible development environments with GPU power is also growing. In December 2021, AWS announced SageMaker Studio Lab. Studio Lab is a web-based IDE based on the popular Jupyter Lab environment. And, best of all, it comes with free GPU power.
In this tutorial we will harness this GPU to teach an NLP model to write like Shakespeare. Doing so has become ridiculously easy, thanks to example scripts developed Hugging Face. In fact, it will take us only one command and less than five minutes to train this NLP model.
All the code for this tutorial can be found in this Github repo.
Getting started with Studio Lab
To get started we need to sign up for an account. We can do so on the website, it’s a straightforward process:
After filling in the form it may take between a few hours and a few days until your account i created. Once you receive the email that your account is ready you can log in and get started:
In this menu we can choose whether we want to start a project with CPU or GPU. Because modern NLP models are based on Deep Learning (i.e. they are based on neural networks) we should use GPU for our project. After we selected GPU we can start the runtime by clicking the button with the play symbol. This should only take a few seconds. Then, finally, we can open the Studio Lab environment by pressing the “Open Project” button.
How to teach an AI model
As mentioned, Studio Lab is based on Jupyter Lab and if you are familiar with Jupyter Lab then this environment should look familiar. On the left we have the file browser and to the right we have a launcher where we can launch the notebook that we will use for training the NLP model:
Teaching an AI model to generate text in a certain style requires a few things: (1) A pre-trained model that already knows how to generate text (2) A training set, in our case Shakespeare text (3) A training script.
The model
At the point of writing, Hugging Face provides more than 3,800 models in its Model Hub that can be used for text generation:
In this tutorial we will use distilgpt2, which is one of the smaller models, but has a decent performance considering its size.
The dataset
Once again we revert back to Hugging Face website, as they provide a dataset that is perfect for our purpose, tiny_shakespeare:
This dataset contains 40,000 lines of Shakespeare text and is available under the MIT License.
The training script
Lastly we need a training script. This is where we prepare the text so it can be used to teach an AI model & where we do the actual training. Normally this would take the lion’s share of our time, but we are, once again, lucky: Hugging Face provides a fully functional training script we can use. All we need to do is to point the script to the model and dataset we want to use.
To learn more about these readily available training scripts we can check out the documentation. Here we learn that, in order to use the script, we need to install the transformers library from source. How to do that is explained in detail in the documentation and is straightforward:
Kicking off the training
Now that we have all ingredients for training our model we can kick it off with just one command:
Training our model takes about 3–4 minutes in Studio Lab to complete:
The model will be saved in the folder we specified in the training command (in my case it’s called shakespeare-clm but you can name it as you like).
Testing our model
At this point we can download the model and use it anywhere we like. But while we are in Studio Lab, we should do a quick test to see if the training worked.
To do so we will first use the NLP model before we taught it to write like Shakespeare and compare its output to our Shakespeare model.
Untrained text generation model
Generating text with these NLP models is straightforward thanks to the Hugging Face Pipeline API. Using this API we can produce some text with the original model, i.e. before we trained it:
This text looks like “normal” text and not really in the style of Shakespeare. Let’s try our trained Shakespeare model.
Shakespeare text generation model
Similar to above we can use the Pipeline API to generate text with our model. We just need to explicitly tell the API where it can find our model:
The text generated with our model is strikingly different from the first model. This looks much more like Shakespeare text.
Conclusion
With one command and within just a few minutes we taught an AI model to write like Shakespeare. The resulting model can be downloaded and used anywhere we like.
I hope you found this tutorial useful – please reach out in the comments if you have any questions. Also, please go ahead and train AI models on other texts – I’d love to hear about the text generation models you are building.
