Playing with GPT-2 text generation in Python

February 21, 2020

OpenAI finally released the full 1.5B parameter version of GPT-2 back in November, after spending most of 2019 saying it was “too dangerous” to release. Meanwhile someone built AI Dungeon on top of it and half the internet is playing a text adventure game powered by a neural network. So I figured it was time to poke at it myself.

The easiest way to run GPT-2 today is the transformers library from Hugging Face. No TensorFlow checkpoints, no cloning the OpenAI repo:

pip install transformers torch

Then generating text is a few lines:

from transformers import GPT2LMHeadModel, GPT2Tokenizer
import torch

tokenizer = GPT2Tokenizer.from_pretrained("gpt2-medium")
model = GPT2LMHeadModel.from_pretrained("gpt2-medium")

prompt = "Docker Swarm is dead because"
input_ids = tokenizer.encode(prompt, return_tensors="pt")

output = model.generate(
    input_ids,
    max_length=100,
    do_sample=True,
    top_k=50,
    top_p=0.95,
)

print(tokenizer.decode(output[0], skip_special_tokens=True))

A few notes from a couple of evenings with it:

The gpt2-medium (355M) model runs fine on CPU on my laptop, just slowly — a few seconds per generation. The full gpt2-xl (1.5B) wants about 6GB of RAM and is painful without a GPU.

Sampling parameters matter a lot. Greedy decoding produces repetitive garbage (“the the the”). top_k + top_p (nucleus sampling) is what makes the output feel coherent.

And “coherent” is doing a lot of work in that sentence. The text is grammatical and stays on topic for a paragraph or two, then confidently drifts into nonsense. It will invent APIs, quotes, and facts with total confidence. Fun toy, but I wouldn’t put it anywhere near production.

Still, between this, BERT now running inside Google Search, and Google’s new Meena chatbot paper from a few weeks ago, it feels like NLP is where the interesting stuff is happening right now. The Hugging Face library making it a pip install away is a big part of that.

Next I want to try fine-tuning the small model on my own commit messages and see what horror it produces.