Build a Genre-Based Story Generator | Pranav Vadrevu | Towards Data Science

After discovering time travel, the Earth’s inhabitants now live in futuristic cities, which are controlled by the government, for the duration of a decade. The government plans to send two elite teams of scientists to the city, in order to investigate the origin of these machines and discover the existence of the “God”.
Wouldn’t it be fun to generate stories for your favorite genre? That’s what we’ll be doing today. We’ll learn how to build a story generator that creates stories based on genres, just like the one that created the sci-fi story above (the user-provided input is bold in the story above).
You can also use this post as a launching pad to develop your own text generator. For example, you can generate headlines of topics like Tech, Science, and Politics. Or generate lyrics of your favorite artists.
As an avid movies and TV-Shows fan, I loved the idea of a story-generator that would generate stories based on genres, input prompts, or even titles. After learning about GPT-2, I wanted to bring that idea to life. That’s what led me to build this model.
Intended use: To have fun and test this idea of fusion storytelling where we can generate stories by blending our creativity (by providing a prompt) with the model’s creativity (by generating the rest of story using the prompt).
Before taking a deep dive into building the generator, let’s first give story generation a shot. Check out my story generator at this Huggingface link or run the cells in this Colab notebook to generate stories. The model input format is of the form:
<BOS> <genre> Optional prompt...
e.g. <BOS> <sci_fi> After discovering time travel,
where genre belongs to: superhero, sci_fi, action, drama, horror, thriller
The model will generate its story using this prompt. There’s also a more intuitive way to generate stories: a web-app using my model in action (keep in mind this app’s generation is slower than the Huggingface link).
Now that you’re done experimenting with the model, let’s explore the idea: We’re fine-tuning the OpenAI GPT-2 model on a dataset that contains movie plots of different genres. This walkthrough follows a Three-Act structure:
Act I: What is GPT-2?
Act II: Fine-Tuning Time...
Act III: Generation Time!
If you’re familiar with the key ideas of GPT-2, you can fast-forward to the next act. Otherwise, get a quick refresher on GPT-2 below.
Recently, we saw the humongous hype behind text-generation model GPT-3 and its dazzling performance for tasks like generating code using zero-shot learning. GPT-2 is the predecessor to GPT-3.
Wait, why are we using GPT-2 and not GPT-3? Well, it’s in Beta. And...
Just like its hype, GPT-3’s size is also gigantic (rumored to require 350 GB of RAM). While still bulky, GPT-2 takes up much less space (largest variant takes up 6GB disk space). And it even comes in different sizes. So, we can use it on many devices (including iPhones).
If you’re looking for a detailed overview of GPT-2, you can get it straight from the horse’s mouth.
But, If you just want a quick summary of GPT-2, here’s a distilled rundown:
GPT-2 is a text-generation language model using a Decoder-only Transformer (a variant of Transformers architecture). If this seems like mumbo-jumbo, just know that Transformer is a state-of-the-art architecture for NLP models.
It was pre-trained (with 40GB of text data) on the task of predicting the next word (more formally, token) at each time-step given preceding text as input.
We can fine-tune (further train) pre-trained models like GPT-2 on a chosen dataset to tailor their performance towards that dataset. To fine-tune and use the GPT-2 pre-trained model, we will use the huggingface/transformers library. It does all the heavy-lifting for us.
For this idea, I created the dataset files by cleaning, transforming and combining the Kaggle Wikipedia Movie Plots dataset along with scraped superhero comic book plots from Wikipedia.
The training file has over 30,000 stories. Each line in the file is a story of this format:
<BOS> <genre> Story goes here... <EOS>
The genres: superhero, sci_fi, horror, action, drama, thriller
To create your own dataset for another task, like generating research paper abstracts based on subjects, each example’s format can be similar like this:
<BOS> <subjects> Abstract here... <EOS>
The subjects: physics, chemistry, computer_science, etc.
I recommend that you train the model on Google Colab (Set runtime to GPU). The link to the colab notebook for fine-tuning our model is here. You can create a copy of this notebook and modify it for your own dataset as well.
Here’s the 6_genre_clean_training_data.txt (training file) and 6_genre_eval_data.txt (evaluation file) that you need to upload to your Colab notebook’s environment before running the code.
I will go through some of the core code from the Colab notebook in this post.
!pip install transformers torch
Above, we are installing the core libraries. We need the transformers library to load the pre-trained GPT-2 checkpoint and fine-tune it. We need torch because it is used for training.
Note: I took the code from this transformers example file and modified it.
In the code snippet above, we specify the model’s arguments (ModelArguments), data arguments (DataTrainingArguments), and training arguments (TrainingArguments). Let’s quickly go through the key arguments for these parameters. To properly check out all parameters, look here.
model_name_or_path: Used to specify the model name or its path. The model is then downloaded to our cache directory. Check here for all available models. Here, we specify the model_name_or_path as gpt2. We also have other options like gpt2-medium or gpt2-xl.
model_type: We are specifying that we want a gpt2 model. This is different from the above parameter because, we only specify the model type, not the name (name refers to gpt2-xl, gpt2-medium, etc.).
train_data_file: We provide the training file.
eval_data_file: We provide a data file to evaluate the model (using perplexity). Perplexity measures how likely our language model generates text from the eval_data_file. The lower the perplexity, the better.
line_by_line: Set to true since each line in our file is a new story.
block_size: Shorten each training example to only have maximum of block_size tokens.
Below, n refers to the value of these parameters.
output_dir: Where to save the final model checkpoint after fine-tuning.
do_train, do_eval: Set to true since we are training and evaluating.
logging_steps: Log the model’s loss after every n optimization steps.
per_device_train_batch_size: Each training optimization step involves n number of training examples.
num_train_epochs: Number of full passes through the training dataset.
save_steps: Save intermediate checkpoint after every n optimization steps (Recommended if Colab keeps disconnecting after couple hours).
save_total_limit: Number of intermediate checkpoints to store at any point.
Now, it’s time to load the model, its configuration and the tokenizer.
The Classes AutoConfig, AutoTokenizer, GPT2LMHeadModel load the respective configuration (GPT2Config), tokenizer (GPT2Tokenizer), and model (GPT2LMHeadModel) based on the model_name.
The GPT2Tokenizer object tokenizes text (converts text into lists of tokens) and encodes these tokens into numbers. Note that tokens can be words, or even subwords (GPT-2 uses Byte Pair Encoding to create tokens). Below is an example of tokenization (without encoding).
After this, we have to encode the tokens into numbers since computers only deal with numbers.
We have GPT2Config object is used to load up the configuration of the GPT-2 model based on the model name.
Finally, we have GPT2LMHeadModel object that loads the model based on the model name and GPT2Config we have.
Some of you might be wondering: what exactly is an “LMHeadModel”? Simply put, for text-generation, we need a language model (a model that assigns probability distribution to tokens in vocabulary). GPT2LMHeadModel is a language model (it assigns scores to each token in the vocabulary). So, we can use this model for text-generation.
Next main step is adding special tokens to integrate the model with our data.
The datasets have some special tokens, and we need to specify to the model what these tokens are. We can pass in a special_tokens_dict which can have some keys like “bos_token”. To see what keys are allowed, visit this link.
bos_token (“<BOS>”) is the token that appears at the start of each story.
eos_token (“<EOS>”) is the token that appears at the end of each story.
pad_token (“<PAD>”) refers to the padding token to pad shorter output to a fixed length.
However, we have some extra special tokens (our genre tokens in this case) that don’t have their own keys in our special_tokens_dict. These extra tokens can be placed as a list under one collective key “additional_special_tokens”.
Replace the genre tokens with your own categories if you’re working with your own dataset.
We’re ready for training, so let’s create the Trainer object.
The Trainer object can be used for training, and evaluating our models.
To create the Trainer object, we specify the model, and datasets along with training arguments (as shown above). The data_collator parameter is used for batching training and evaluation examples.
We then call the Trainer’s train method and save the model to our output directory after training. It’s time to relax and let the machine train the model.
After training the model, download the model checkpoint folder from Colab onto your computer. You can deploy it as a Huggingface community model, develop a web app, or even develop a mobile app using the model.
Here’s the exciting part! Run this cell below in your Colab notebook if you’re itching to generate stories!
Here, we are using the TextGenerationPipeline object that simplifies text-generation with our model since it abstracts away from input pre-preprocessing and output post-processing steps.
text_generator = TextGenerationPipeline(model=model, tokenizer=tokenizer)
The text_generator is an object of the TextGenerationPipeline class. How are we able to use this object as a function? This class has a __call__ method that allows its objects to behave like functions (this method is called when using the object as a function). Here, we use this object like a function to generate text by providing the input_prompt. That’s all there is to it.
We still haven’t explained some text-generation parameters like top_p. Stick after the end-credits if you’d like to learn about these parameters.
To recap, we took a dataset of movie plots of various genres and fed it to the GPT2LMHeadModel to fine-tune our model to generate genre-specific stories. Using these ideas, you can also create other datasets to generate text based on those datasets as well.
To test my movie story generator, click here (or use the web app here).
A huge thank you to the Huggingface Team for their transformers library and detailed examples, and this Article by Raymond Cheng for helping me create my model.
First, let’s learn what the output of the GPT2LMHeadModel is:
The model produces logits for all the possible tokens in its vocabulary when predicting the next token. To simplify, think of these logits as scores for each token. Tokens that have higher scores (logits) mean that they have higher probability to be an appropriate next token. We then apply a softmax operation on the logits of all tokens. We now get a softmax score of each token that is between 0 and 1. The softmax scores of all tokens sum up to 1. These softmax scores can be considered as probabilities (although they aren’t) of a token being an appropriate next token given some previous text.
Below is a basic outlin of the parameters:
max_length [int]: Specifies the maximum number of tokens to generate.
do_sample [bool]: Specifies whether to use sampling (such as top_p or use greedy search). Greedy search is selecting most probable word at each time step (not recommended as text can get repetitive).
repetition_penalty [float]: Specifies penalty for repetition. Increase the repetition_penalty parameter to reduce repetition.
temperature [float]: Used to increase or decrease reliance on high-logits tokens. Increasing temperature value reduces the chance of only few tokens having very high softmax scores.
top_p [float]: Specifies to consider only the tokens whose sum of probability (formally, softmax) scores does not exceed the value of top_p
top_k [int]: Tells us to only consider top_k number of tokens (ranked by their softmax scores) when generating a text.
And that’s it.
Some rights reserved