!pip install transformers
from transformers import GPT2LMHeadModel, GPT2Tokenizer, set_seed
import torch
import torch.nn.functional as F
import numpy as np
# Load model and tokenizer
model = GPT2LMHeadModel.from_pretrained('gpt2')
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')Decoding strategies
Decoding from GPT2
This tutorial investigates how to use GPT2 (the forerunner of GPT3) to generate text. There are a number of ways to do this that trade-off the realism of the text against the amount of variation.
At every stage, GPT2 takes an input string and returns a probability for each of the possible subsequent tokens. We can choose what to do with these probability. We could always greedily choose the most likely next token, or we could draw a sample randomly according to the probabilities. There are also intermediate strategies such as top-k sampling and nucleus sampling, that have some controlled randomness.
We’ll also investigate beam search – the idea is that rather than greedily take the next best token at each stage, we maintain a set of hypotheses (beams)as we add each subsequent token and return the most likely overall hypothesis. This is not necessarily the same result we get from greedily choosing the next token.
First, let’s investigate the token themselves. The code below prints out the vocabulary size and shows 20 random tokens.
np.random.seed(1)
print("Number of tokens in dictionary = %d"%(tokenizer.vocab_size))
for i in range(20):
index = np.random.randint(tokenizer.vocab_size)
print("Token: %d "%(index)+tokenizer.decode(torch.tensor(index), skip_special_tokens=True))Sampling
Each time we run GPT2 it will take in a set of tokens, and return a probability over each of the possible next tokens. The simplest thing we could do is to just draw a sample from this probability distribution each time.
def sample_next_token(input_tokens, model, tokenizer):
# Run model to get prediction over next output
outputs = model(input_ids = input_tokens['input_ids'], attention_mask = input_tokens['attention_mask'])
# Find prediction
prob_over_tokens = F.softmax(outputs.logits, dim=-1).detach().numpy()[0,-1]
# TODO Draw a random token according to the probabilities
# next_token should be an array with an sole integer in it (as below)
# Use: https://numpy.org/doc/stable/reference/random/generated/numpy.random.choice.html
# Replace this line
next_token = [5000]
# Append token to sentence
output_tokens = input_tokens
output_tokens["input_ids"] = torch.cat((output_tokens['input_ids'],torch.tensor([next_token])),dim=1)
output_tokens['attention_mask'] = torch.cat((output_tokens['attention_mask'],torch.tensor([[1]])),dim=1)
output_tokens['last_token_prob'] = prob_over_tokens[next_token]
return output_tokens
# Expected output:
# "The best thing about Bath is that they don't even change or shrink anymore."
set_seed(0)
input_txt = "The best thing about Bath is"
input_tokens = tokenizer(input_txt, return_tensors='pt')
for i in range(10):
input_tokens = sample_next_token(input_tokens, model, tokenizer)
print(tokenizer.decode(input_tokens["input_ids"][0], skip_special_tokens=True))
# TODO Modify the code below by changing the number of tokens generated and the initial sentence
# to get a feel for how well this works. Since I didn't reset the seed, it will give a different
# answer every time that you run it.
# TODO Experiment with changing this line:
input_txt = "The best thing about Bath is"
input_tokens = tokenizer(input_txt, return_tensors='pt')
# TODO Experiment with changing this line:
for i in range(10):
input_tokens = sample_next_token(input_tokens, model, tokenizer)
print(tokenizer.decode(input_tokens["input_ids"][0], skip_special_tokens=True))Greedy token selection
You probably (correctly) got the impression that the text from pure sampling of the probability model can be kind of random. How about if we choose most likely token at each step?
def get_best_next_token(input_tokens, model, tokenizer):
# Run model to get prediction over next output
outputs = model(input_ids = input_tokens['input_ids'], attention_mask = input_tokens['attention_mask'])
# Find prediction
prob_over_tokens = F.softmax(outputs.logits, dim=-1).detach().numpy()[0,-1]
# TODO -- find the token index with the maximum probability
# It should be returns as a list (i.e., put squared brackets around it)
# Use https://numpy.org/doc/stable/reference/generated/numpy.argmax.html
# Replace this line
next_token = [5000]
# Append token to sentence
output_tokens = input_tokens
output_tokens["input_ids"] = torch.cat((output_tokens['input_ids'],torch.tensor([next_token])),dim=1)
output_tokens['attention_mask'] = torch.cat((output_tokens['attention_mask'],torch.tensor([[1]])),dim=1)
output_tokens['last_token_prob'] = prob_over_tokens[next_token]
return output_tokens
# Expected output:
# The best thing about Bath is that it's a place where you can go to
set_seed(0)
input_txt = "The best thing about Bath is"
input_tokens = tokenizer(input_txt, return_tensors='pt')
for i in range(10):
input_tokens = get_best_next_token(input_tokens, model, tokenizer)
print(tokenizer.decode(input_tokens["input_ids"][0], skip_special_tokens=True))
# TODO Modify the code below by changing the number of tokens generated and the initial sentence
# to get a feel for how well this works.
# TODO Experiment with changing this line:
input_txt = "The best thing about Bath is"
input_tokens = tokenizer(input_txt, return_tensors='pt')
# TODO Experiment with changing this line:
for i in range(10):
input_tokens = get_best_next_token(input_tokens, model, tokenizer)
print(tokenizer.decode(input_tokens["input_ids"][0], skip_special_tokens=True))Top-K sampling
You probably noticed that the greedy strategy produces quite realistic text, but it’s kind of boring. It produces generic answers. Also, if this was a chatbot, then we wouldn’t necessarily want it to produce the same answer to a question each time.
Top-K sampling is a compromise strategy that samples randomly from the top K most probable tokens. We could just choose them with a uniform distribution, or (as here) we could sample them according to their original probabilities.
def get_top_k_token(input_tokens, model, tokenizer, k=20):
# Run model to get prediction over next output
outputs = model(input_ids = input_tokens['input_ids'], attention_mask = input_tokens['attention_mask'])
# Find prediction
prob_over_tokens = F.softmax(outputs.logits, dim=-1).detach().numpy()[0,-1]
# Draw a sample from the top K most likely tokens.
# Take copy of the probabilities and sort from largest to smallest (use np.sort)
# TODO -- replace this line
sorted_prob_over_tokens = prob_over_tokens
# Find the probability at the k'th position
# TODO -- replace this line
kth_prob_value = 0.0
# Set all probabilities below this value to zero
prob_over_tokens[prob_over_tokens < kth_prob_value] = 0
# Expected output:
# The best thing about Bath is that you get to see all the beautiful faces of
set_seed(0)
input_txt = "The best thing about Bath is"
input_tokens = tokenizer(input_txt, return_tensors='pt')
for i in range(10):
input_tokens = get_top_k_token(input_tokens, model, tokenizer, k=10)
print(tokenizer.decode(input_tokens["input_ids"][0], skip_special_tokens=True))
# TODO
# Experiment with different values of k
# If you set it to a lower number (say 3) the text will be less random
# If you set it to a higher number (say 5000) the text will be more random
set_seed(0)
input_txt = "The best thing about Bath is"
input_tokens = tokenizer(input_txt, return_tensors='pt')
for i in range(10):
input_tokens = get_top_k_token(input_tokens, model, tokenizer, k=10)
print(tokenizer.decode(input_tokens["input_ids"][0], skip_special_tokens=True))Nucleus Sampling
Top-K sampling has the disadvantage that sometimes there are only a few plausible next tokens, and sometimes there are a lot. How do we adapt to this situation? One way is to sample from a fixed proportion of the probability mass. That is we order the tokens in terms of probability and cut off the possibility of sampling when the cumulative sum is greater than a threshold.
This way, we adapt the number of possible tokens that we can choose.
def get_nucleus_sampling_token(input_tokens, model, tokenizer, thresh=0.25):
# Run model to get prediction over next output
outputs = model(input_ids = input_tokens['input_ids'], attention_mask = input_tokens['attention_mask'])
# Find prediction
prob_over_tokens = F.softmax(outputs.logits, dim=-1).detach().numpy()[0,-1]
# Find the most likely tokens that make up the first (thresh) of the probability
# TODO -- sort the probabilities in decreasing order
# Replace this line
sorted_probs_decreasing = prob_over_tokens
# TODO -- compute the cumulative sum of these probabilities
# Replace this line
cum_sum_probs = sorted_probs_decreasing
# Find index where that the cumulative sum is greater than the threshold
thresh_index = np.argmax(cum_sum_probs>thresh)
print("Choosing from %d tokens"%(thresh_index))
# TODO: Find the probability value to threshold
# Replace this line:
thresh_prob = 0.5
# Set any probabilities less than this to zero
prob_over_tokens[prob_over_tokens < thresh_prob] = 0
# Expected output:
# The best thing about Bath is that it's not a city that has been around
set_seed(0)
input_txt = "The best thing about Bath is"
input_tokens = tokenizer(input_txt, return_tensors='pt')
for i in range(10):
input_tokens = get_nucleus_sampling_token(input_tokens, model, tokenizer, thresh = 0.2)
print(tokenizer.decode(input_tokens["input_ids"][0], skip_special_tokens=True))
# TODO -- experiment with setting the threshold probability to larger or smaller values
input_txt = "The best thing about Bath is"
input_tokens = tokenizer(input_txt, return_tensors='pt')
for i in range(10):
input_tokens = get_nucleus_sampling_token(input_tokens, model, tokenizer, thresh = 0.2)
print(tokenizer.decode(input_tokens["input_ids"][0], skip_special_tokens=True))Beam Search
All of the methods we’ve seen so far choose the tokens one by one. But this isn’t necessarily sensible. Even greedily choosing the best token doesn’t necessarily retrieve the sequence with the highest probability. It might be that the most likely token only has very unlikely tokens following it.
Beam search maintains hypotheses about the best possible continuation. It starts with the top continuations. Then for each of those, it finds the top K continuations, giving hypotheses. Then it retains just the top of these so that the number of hypotheses stays the same.
# This routine returns the k'th most likely next token.
# If k =0 then it returns the most likely token, if k=1 it returns the next most likely and so on
# We will need this for beam search
def get_kth_most_likely_token(input_tokens, model, tokenizer, k):
# Run model to get prediction over next output
outputs = model(input_ids = input_tokens['input_ids'], attention_mask = input_tokens['attention_mask'])
# Find prediction
prob_over_tokens = F.softmax(outputs.logits, dim=-1).detach().numpy()[0,-1]
# Find the k'th most likely token
# TODO Sort the probabilities from largest to smallest
# Replace this line:
sorted_prob_over_tokens = prob_over_tokens
# TODO Find the k'th sorted probability
# Replace this line
kth_prob_value = prob_over_tokens[0]
# Find position of this token.
next_token = np.where(prob_over_tokens == kth_prob_value)[0]
# Append token to sentence
output_tokens = input_tokens
output_tokens["input_ids"] = torch.cat((output_tokens['input_ids'],torch.tensor([next_token])),dim=1)
output_tokens['attention_mask'] = torch.cat((output_tokens['attention_mask'],torch.tensor([[1]])),dim=1)
output_tokens['last_token_prob'] = prob_over_tokens[next_token]
output_tokens['log_prob'] = output_tokens['log_prob'] + np.log(prob_over_tokens[next_token])
return output_tokens
# We can test this code and see that if we choose the 2nd most likely (K=1) token each time
# then we get much better generation results than if we choose the 2001st most likely token
# Expected output:
# The best thing about Bath is the way you get the most bang outta the
set_seed(0)
input_txt = "The best thing about Bath is"
input_tokens = tokenizer(input_txt, return_tensors='pt')
input_tokens['log_prob'] = 0.0
for i in range(10):
input_tokens = get_kth_most_likely_token(input_tokens, model, tokenizer, k=1)
print(tokenizer.decode(input_tokens["input_ids"][0], skip_special_tokens=True))
# Expected output:
# The best thing about Bath is mixed profits partnerships» buy generic+ Honda throttlecont
input_txt = "The best thing about Bath is"
input_tokens = tokenizer(input_txt, return_tensors='pt')
input_tokens['log_prob'] = 0.0
for i in range(10):
input_tokens = get_kth_most_likely_token(input_tokens, model, tokenizer, k=2000)
print(tokenizer.decode(input_tokens["input_ids"][0], skip_special_tokens=True))
# TODO -- play around with different values of K
# Print out each beam plus the log probability
def print_beams(beams):
for index,beam in enumerate(beams):
print("Beam %d, Prob %3.3f: "%(index,beam['log_prob'])+tokenizer.decode(beam["input_ids"][0], skip_special_tokens=True))
print('---')
# TODO: Read this code carefully!
def do_beam_search(input_tokens_in, model, tokenizer, n_beam=5, beam_length=10):
# Store beams in a list
input_tokens['log_prob'] = 0.0
# Initialize with n_beam most likely continuations
beams = [None] * n_beam
for c_k in range(n_beam):
beams[c_k] = dict(input_tokens_in)
beams[c_k] = get_kth_most_likely_token(beams[c_k], model, tokenizer, c_k)
print_beams(beams)
# For each token in the sequence we will add
for c_pos in range(beam_length-1):
# Now for each beam, we continue it in the most likely ways, making n_beam*n_beam type hypotheses
beams_all = [None] * (n_beam*n_beam)
log_probs_all = np.zeros(n_beam*n_beam)
# For each current hypothesis
for c_beam in range(n_beam):
# For each continuation
for c_k in range(n_beam):
# Store the continuation and the probability
beams_all[c_beam * n_beam + c_k] = dict(get_kth_most_likely_token(beams[c_beam], model, tokenizer, c_k))
log_probs_all[c_beam * n_beam + c_k] = beams_all[c_beam * n_beam + c_k]['log_prob']
# Keep the best n_beams sequences with the highest probabilities
sorted_index = np.argsort(np.array(log_probs_all)*-1)
for c_k in range(n_beam):
beams[c_k] = dict(beams_all[sorted_index[c_k]])
# Print the beams
print_beams(beams)
return beams[0]
# Expected output:
# The best thing about Bath is that it's a place where you don't have to
set_seed(0)
input_txt = "The best thing about Bath is"
input_tokens = tokenizer(input_txt, return_tensors='pt')
# Now let's call the beam search
# It takes a while as it has to run the model multiple times to add a token
n_beams = 5
best_beam = do_beam_search(input_tokens,model,tokenizer)
print("Beam search result:")
print(tokenizer.decode(best_beam["input_ids"][0], skip_special_tokens=True))
# You should see that the best answer is not the same as the greedy solution we found above