ML horoscope generation pipeline as a REST API using GPT, Transformers, Fast API and GCP (Part 1)
This project implements a complete cloud-hosted NLP pipeline for generating (hopefully) well-written horoscopes using Python 3 for the machine learning, data processing and API serving. It uses a variety of open source libraries such as Huggingface’s Transformers, PyTorch and FastAPI. It also has a small demo website implemented in HTML, CSS and JS that can be used to preview generated horoscopes for each of the zodiac signs. If that seems interesting to you, please read on !
If you’ve already read this or you’re more interested in the model wrapper, the API and hosting it on Google Cloud, click below to read the second part.
NLP And Its Applications
Recently, NLP (Natural Language Processing) research has made leaps of progress and its applications become more and more impressive by the day.
NLP is a field of linguistics and machine learning focused on understanding everything related to human language. The aim of NLP tasks is not only to understand single words individually, but to be able to understand the context of those words.
Here are a few applications that have made a lot of progress lately:
- Classifying whole sentences: Getting the sentiment of a review, detecting if an email is spam, determining if a sentence is grammatically correct or whether two sentences are logically related or not
- Classifying each word in a sentence: Identifying the grammatical components of a sentence or the named entities
- Generating text content: Completing a prompt with auto-generated text, filling in the blanks in a text with masked words (this is the one we’re looking at this time)
- Extracting an answer from a text: Given a question and a context, extracting the answer to the question based on the information provided in the context
- Generating a new sentence from an input text: Translating a text into another language, summarizing a text
NLP isn’t limited to written text though. It also tackles complex challenges in speech recognition and computer vision, such as generating a transcript of an audio sample or a description of an image.
I’ve always wanted to implement and explore one of those applications on my own, but I’ve always found the task quite daunting. After all, there is a lot to know and understand about machine learning in general, how it leads to NLP and how to actually use it in the real world.
After witnessing the absolutely stunning applications of the newly released GPT3 and its open-source brethren (you can check out some examples below) and gaining some experience with machine learning applications, I decided to take the dive.
Project Summary
For the said application, I turned towards horoscopes. Not because I believe in the influence of star alignment on our personalities and life decisions, but because I was curious to see if an educated machine-generated version of said horoscopes would be easily distinguishable from the human-written ones you would find in newspapers, magazines and online.
For the horoscope generation, I imagined the following pipeline:
- Scrape horoscopes from websites known to publish horoscopes daily, and store them as a CSV for later use or use other scraped horoscopes
- Use a pre-trained open-source version of GPT3 (GPTNeo) and fine-tune it to generate horoscopes when given a particular zodiac sign
- Create a model wrapper that would handle data pre-processing, horoscope post-processing and deliver predictions from the horoscope model
- Create a REST API using FastAPI that would handle prediction requests for the various zodiac signs
- Containerize the API and the model using Docker and host it for public predictions on the Google Cloud Platform
- Create a small website to test out the horoscope generation easily and scrape some rust off my web developing skills
To check out all the project’s code or run the API yourself, you can check out the GitHub link below. If you want a more detailed breakdown, read on 😊!
Gathering the data from existing horoscopes
Finding scraped horoscopes wasn’t that easy. There are scripts that exist that were written a while ago, and most websites nowadays host their horoscopes for one week or one month, but we need a substantial amount of data for every zodiac for the training. I compounded different datasets I could find, because it seemed less confusing and time-consuming than writing scraping scripts and finding websites that store their horoscopes for more than a month back.
Once all the data was gathered in a single file, the pre-processing consisted in transforming it in a CSV with every line following this format:
{zodiac_sign}, {horoscope}
{zodiac_sign}, {horoscope}
…
{zodiac_sign}, {horoscope}
It makes sense to structure it like that, because GPT3 and other generative transformers function really well when given a keyword or a start to a sentence in order to generate the rest. This would make it easier for our model to generate a horoscope when given the beginning of the sentence, which is the zodiac sign followed by a comma. For example, giving the horoscope model the input:
“Aries, “
It can generate the following short horoscope (after post-processing):
You are not averse to taking risks, in fact it is one of the things that makes you feel most alive. Cosmic activity in your fellow Fire sign of Sagittarius means the odds are very much on your side at the moment, so go for it! You need to make a bold decision.
If you wish to use the horoscope dataset yourself, it’s freely available on the GitHub project mentioned above.
Fine-tuning GPT Neo on Google Colab
For fine-tuning GPT Neo, I’ve opted to use Google Colab because it fits my needs perfectly, and speeds up the process tenfold.
Google Colab is a free cloud service that gives you access to a Jupyter Notebook (for writing Python code) on a virtual machine hosted by Google, and is one of the favorite platforms for programmers that want to train models because it gives you access to a GPU.
Using a GPU for training is essential for bigger models, because it’s significantly faster than training it on a CPU, which means that comparing model performance and fine-tuning the important parameters is a lot faster.
If you’ve never used Google Colab before, check out the tutorial below.
The whole Colab notebook used for the training can be found in the link below, and what follows is the explanation of the more important parts of the code. Bear in mind that I often fiddle around with this code, and it’s not the cleanest, so if you do find it hard to understand, feel free to reach out.
For obtaining a pre-trained version of the model (built with PyTorch), we are going to use Huggingface’s Transformers library. Transformers provides general-purpose architectures (BERT, GPT-2, GPTNeo, …) for Natural Language Understanding (NLU) and Natural Language Generation (NLG) with a lot of pre-trained models in 100+ languages and deep interoperability between Jax, PyTorch and TensorFlow. It also offers a neat interface for simplifying training and transforming input and output from and to human-readable text.
We are going to use EleutherAI’s GPT Neo for this project which is conveniently hosted on the Huggingface platform. In the colab notebook, we fetch the model and its associated transformer directly from their repositories in the following way:
The lines above download the GPT Neo model and transformer from the repository, and set the tokens necessary for the fine-tuning of the model.
Then, the scraped horoscopes are read from the CSV file and the transformer is used to encode the words in the horoscopes in a machine-readable format.
We then create a dataset using the previously processed horoscopes and set the parameters and arguments used for the training. Most of them have self-explanatory names, and the rest are the defaults used for training GPT Neo. With these parameters, the training takes a few dozen hours to reach really accurate horoscopes, and the model is saved every 20000 training steps. Using Google Colab, I had to restart the training a few times from the last checkpoint because the amount of time you can use a GPU for free is limited.
And that’s it! Pretty easy, right? The next part will cover the model wrapper, input pre-processing and output post-processing, API serving and GCP hosting. To head to the second part, click here to open the Medium story.
