title gpt text generation from scratch with kerasnlp jesse chanhttps github comjessechancy date created 20220725 last modified 20220725 description using kerasnlp to train a minigpt model for text generation accelerator gpu introduction in this example we will use kerasnlp to build a scaled down generative pretrained gpt model gpt is a transformerbased model that allows you to generate sophisticated text from a prompt we will train the model on the simplebooks92https arxiv orgabs1911 12391 corpus which is a dataset made from several novels it is a good dataset for this example since it has a small vocabulary and high word frequency which is beneficial when training a model with few parameters this example combines concepts from text generation with a miniature gpthttps keras ioexamplesgenerativetextgenerationwithminiaturegpt with kerasnlp abstractions we will demonstrate how kerasnlp tokenization layers and metrics simplify the training process and then show how to generate output text using the kerasnlp sampling utilities note if you are running this example on a colab make sure to enable gpu runtime for faster training this example requires kerasnlp you can install it via the following command pip install kerasnlp setup settings hyperparameters data model training inference load the data now let s download the dataset the simplebooks dataset consists of 1 573 gutenberg books and has one of the smallest vocabulary size to wordlevel tokens ratio it has a vocabulary size of 98k a third of wikitext103 s with around the same number of tokens 100m this makes it easy to fit a small model load simplebooks92 train set and filter out short lines load simplebooks92 validation set and filter out short lines train the tokenizer we train the tokenizer from the training dataset for a vocabulary size of vocabsize which is a tuned hyperparameter we want to limit the vocabulary as much as possible as we will see later on that it has a large effect on the number of model parameters we also don t want to include too few vocabulary terms or there would be too many outofvocabulary oov subwords in addition three tokens are reserved in the vocabulary pad for padding sequences to seqlen this token has index 0 in both reservedtokens and vocab since wordpiecetokenizer and other layers consider 0vocab0 as the default padding unk for oov subwords which should match the default oovtokenunk in wordpiecetokenizer bos stands for beginning of sentence but here technically it is a token representing the beginning of each line of training data train tokenizer vocabulary load tokenizer we use the vocabulary data to initialize kerasnlp tokenizers wordpiecetokenizer wordpiecetokenizer is an efficient implementation of the wordpiece algorithm used by bert and other models it will strip lowercase and do other irreversible preprocessing operations tokenize data we preprocess the dataset by tokenizing and splitting it into features and labels packer adds a start token tokenize and split into train and label sequences build the model we create our scaled down gpt model with the following layers one kerasnlp layers tokenandpositionembedding layer which combines the embedding for the token and its position multiple kerasnlp layers transformerdecoder layers with the default causal masking the layer has no crossattention when run with decoder sequence only one final dense linear layer embedding transformer decoders output let s take a look at our model summary a large majority of the parameters are in the tokenandpositionembedding and the output dense layer this means that the vocabulary size vocabsize has a large effect on the size of the model while the number of transformer decoder layers numlayers doesn t affect it as much training now that we have our model let s train it with the fit method inference with our trained model we can test it out to gauge its performance to do this we can seed our model with an input sequence starting with the bos token and progressively sample the model by making predictions for each subsequent token in a loop to start lets build a prompt with the same shape as our model inputs containing only the bos token the packer layers adds the bos token for us we will use the kerasnlp samplers module for inference which requires a callback function wrapping the model we just trained this wrapper calls the model and returns the logit predictions for the current token we are generating note there are two pieces of more advanced functionality available when defining your callback the first is the ability to take in a cache of states computed in previous generation steps which can be used to speed up generation the second is the ability to output the final dense hidden state of each generated token this is used by kerasnlp samplers contrastivesampler which avoids repetition by penalizing repeated hidden states both are optional and we will ignore them for now ignore hidden states for now only needed for contrastive search creating the wrapper function is the most complex part of using these functions now that it s done let s test out the different utilities starting with greedy search greedy search we greedily pick the most probable token at each timestep in other words we get the argmax of the model output as you can see greedy search starts out making some sense but quickly starts repeating itself this is a common problem with text generation that can be fixed by some of the probabilistic text generation utilities shown later on beam search at a highlevel beam search keeps track of the numbeams most probable sequences at each timestep and predicts the best next token from all sequences it is an improvement over greedy search since it stores more possibilities however it is less efficient than greedy search since it has to compute and store multiple potential sequences note beam search with numbeams1 is identical to greedy search similar to greedy search beam search quickly starts repeating itself since it is still a deterministic method random search random search is our first probabilistic method at each time step it samples the next token using the softmax probabilities provided by the model voil no repetitions however with random search we may see some nonsensical words appearing since any word in the vocabulary has a chance of appearing with this sampling method this is fixed by our next search utility topk search topk search similar to random search we sample the next token from the probability distribution provided by the model the only difference is that here we select out the top k most probable tokens and distribute the probability mass over them before sampling this way we won t be sampling from low probability tokens and hence we would have less nonsensical words topp search even with the topk search there is something to improve upon with topk search the number k is fixed which means it selects the same number of tokens for any probability distribution consider two scenarios one where the probability mass is concentrated over 2 words and another where the probability mass is evenly concentrated across 10 should we choose k2 or k10 there is no one size that fits all k here this is where topp search comes in instead of choosing a k we choose a probability p that we want the probabilities of the top tokens to sum up to this way we can dynamically adjust the k based on the probability distribution by setting p0 9 if 90 of the probability mass is concentrated on the top 2 tokens we can filter out the top 2 tokens to sample from if instead the 90 is distributed over 10 tokens it will similarly filter out the top 10 tokens to sample from using callbacks for text generation we can also wrap the utilities in a callback which allows you to print out a prediction sequence for every epoch of the model here is an example of a callback for topk search a callback to generate text from a trained model using topk def initself k self sampler kerasnlp samplers topksamplerk def onepochendself epoch logsnone outputtokens self sampler nextnext promptprompttokens index1 txt tokenizer detokenizeoutputtokens printftopk search generated text ntxtn textgenerationcallback topktextgeneratork10 dummy training loop to demonstrate callback model fit trainds take1 verbose2 epochs2 callbackstextgenerationcallback conclusion introduction in this example we will use kerasnlp to build a scaled down generative pre trained gpt model gpt is a transformer based model that allows you to generate sophisticated text from a prompt we will train the model on the simplebooks 92 https arxiv org abs 1911 12391 corpus which is a dataset made from several novels it is a good dataset for this example since it has a small vocabulary and high word frequency which is beneficial when training a model with few parameters this example combines concepts from text generation with a miniature gpt https keras io examples generative text_generation_with_miniature_gpt with kerasnlp abstractions we will demonstrate how kerasnlp tokenization layers and metrics simplify the training process and then show how to generate output text using the kerasnlp sampling utilities note if you are running this example on a colab make sure to enable gpu runtime for faster training this example requires kerasnlp you can install it via the following command pip install keras nlp setup settings hyperparameters data strings shorter than this will be discarded length of training sequences in tokens model limits parameters in model training inference load the data now let s download the dataset the simplebooks dataset consists of 1 573 gutenberg books and has one of the smallest vocabulary size to word level tokens ratio it has a vocabulary size of 98k a third of wikitext 103 s with around the same number of tokens 100m this makes it easy to fit a small model load simplebooks 92 train set and filter out short lines load simplebooks 92 validation set and filter out short lines train the tokenizer we train the tokenizer from the training dataset for a vocabulary size of vocab_size which is a tuned hyperparameter we want to limit the vocabulary as much as possible as we will see later on that it has a large effect on the number of model parameters we also don t want to include too few vocabulary terms or there would be too many out of vocabulary oov sub words in addition three tokens are reserved in the vocabulary pad for padding sequences to seq_len this token has index 0 in both reserved_tokens and vocab since wordpiecetokenizer and other layers consider 0 vocab 0 as the default padding unk for oov sub words which should match the default oov_token unk in wordpiecetokenizer bos stands for beginning of sentence but here technically it is a token representing the beginning of each line of training data train tokenizer vocabulary load tokenizer we use the vocabulary data to initialize keras_nlp tokenizers wordpiecetokenizer wordpiecetokenizer is an efficient implementation of the wordpiece algorithm used by bert and other models it will strip lower case and do other irreversible preprocessing operations tokenize data we preprocess the dataset by tokenizing and splitting it into features and labels packer adds a start token tokenize and split into train and label sequences build the model we create our scaled down gpt model with the following layers one keras_nlp layers tokenandpositionembedding layer which combines the embedding for the token and its position multiple keras_nlp layers transformerdecoder layers with the default causal masking the layer has no cross attention when run with decoder sequence only one final dense linear layer embedding transformer decoders giving one argument only skips cross attention output let s take a look at our model summary a large majority of the parameters are in the token_and_position_embedding and the output dense layer this means that the vocabulary size vocab_size has a large effect on the size of the model while the number of transformer decoder layers num_layers doesn t affect it as much training now that we have our model let s train it with the fit method inference with our trained model we can test it out to gauge its performance to do this we can seed our model with an input sequence starting with the bos token and progressively sample the model by making predictions for each subsequent token in a loop to start lets build a prompt with the same shape as our model inputs containing only the bos token the packer layers adds the bos token for us we will use the keras_nlp samplers module for inference which requires a callback function wrapping the model we just trained this wrapper calls the model and returns the logit predictions for the current token we are generating note there are two pieces of more advanced functionality available when defining your callback the first is the ability to take in a cache of states computed in previous generation steps which can be used to speed up generation the second is the ability to output the final dense hidden state of each generated token this is used by keras_nlp samplers contrastivesampler which avoids repetition by penalizing repeated hidden states both are optional and we will ignore them for now ignore hidden states for now only needed for contrastive search creating the wrapper function is the most complex part of using these functions now that it s done let s test out the different utilities starting with greedy search greedy search we greedily pick the most probable token at each timestep in other words we get the argmax of the model output start sampling immediately after the bos token as you can see greedy search starts out making some sense but quickly starts repeating itself this is a common problem with text generation that can be fixed by some of the probabilistic text generation utilities shown later on beam search at a high level beam search keeps track of the num_beams most probable sequences at each timestep and predicts the best next token from all sequences it is an improvement over greedy search since it stores more possibilities however it is less efficient than greedy search since it has to compute and store multiple potential sequences note beam search with num_beams 1 is identical to greedy search similar to greedy search beam search quickly starts repeating itself since it is still a deterministic method random search random search is our first probabilistic method at each time step it samples the next token using the softmax probabilities provided by the model voilà no repetitions however with random search we may see some nonsensical words appearing since any word in the vocabulary has a chance of appearing with this sampling method this is fixed by our next search utility top k search top k search similar to random search we sample the next token from the probability distribution provided by the model the only difference is that here we select out the top k most probable tokens and distribute the probability mass over them before sampling this way we won t be sampling from low probability tokens and hence we would have less nonsensical words top p search even with the top k search there is something to improve upon with top k search the number k is fixed which means it selects the same number of tokens for any probability distribution consider two scenarios one where the probability mass is concentrated over 2 words and another where the probability mass is evenly concentrated across 10 should we choose k 2 or k 10 there is no one size that fits all k here this is where top p search comes in instead of choosing a k we choose a probability p that we want the probabilities of the top tokens to sum up to this way we can dynamically adjust the k based on the probability distribution by setting p 0 9 if 90 of the probability mass is concentrated on the top 2 tokens we can filter out the top 2 tokens to sample from if instead the 90 is distributed over 10 tokens it will similarly filter out the top 10 tokens to sample from using callbacks for text generation we can also wrap the utilities in a callback which allows you to print out a prediction sequence for every epoch of the model here is an example of a callback for top k search a callback to generate text from a trained model using top k dummy training loop to demonstrate callback conclusion to recap in this example we use kerasnlp layers to train a sub word vocabulary tokenize training data create a miniature gpt model and perform inference with the text generation library if you would like to understand how transformers work or learn more about training the full gpt model here are some further readings attention is all you need vaswani et al 2017 https arxiv org abs 1706 03762 gpt 3 paper brown et al 2020 https arxiv org abs 2005 14165
import os
import keras_nlp
import keras

import tensorflow.data as tf_data
import tensorflow.strings as tf_strings




BATCH_SIZE = 64
MIN_STRING_LEN = 512  
SEQ_LEN = 128  


EMBED_DIM = 256
FEED_FORWARD_DIM = 128
NUM_HEADS = 3
NUM_LAYERS = 2
VOCAB_SIZE = 5000  


EPOCHS = 5


NUM_TOKENS_TO_GENERATE = 80



keras.utils.get_file(
    origin="https://dldata-public.s3.us-east-2.amazonaws.com/simplebooks.zip",
    extract=True,
)
dir = os.path.expanduser("~/.keras/datasets/simplebooks/")


raw_train_ds = (
    tf_data.TextLineDataset(dir + "simplebooks-92-raw/train.txt")
    .filter(lambda x: tf_strings.length(x) > MIN_STRING_LEN)
    .batch(BATCH_SIZE)
    .shuffle(buffer_size=256)
)


raw_val_ds = (
    tf_data.TextLineDataset(dir + "simplebooks-92-raw/valid.txt")
    .filter(lambda x: tf_strings.length(x) > MIN_STRING_LEN)
    .batch(BATCH_SIZE)
)




vocab = keras_nlp.tokenizers.compute_word_piece_vocabulary(
    raw_train_ds,
    vocabulary_size=VOCAB_SIZE,
    lowercase=True,
    reserved_tokens=["[PAD]", "[UNK]", "[BOS]"],
)



tokenizer = keras_nlp.tokenizers.WordPieceTokenizer(
    vocabulary=vocab,
    sequence_length=SEQ_LEN,
    lowercase=True,
)




start_packer = keras_nlp.layers.StartEndPacker(
    sequence_length=SEQ_LEN,
    start_value=tokenizer.token_to_id("[BOS]"),
)


def preprocess(inputs):
    outputs = tokenizer(inputs)
    features = start_packer(outputs)
    labels = outputs
    return features, labels



train_ds = raw_train_ds.map(
    preprocess, num_parallel_calls=tf_data.AUTOTUNE
).prefetch(tf_data.AUTOTUNE)
val_ds = raw_val_ds.map(
    preprocess, num_parallel_calls=tf_data.AUTOTUNE
).prefetch(tf_data.AUTOTUNE)



inputs = keras.layers.Input(shape=(None,), dtype="int32")

embedding_layer = keras_nlp.layers.TokenAndPositionEmbedding(
    vocabulary_size=VOCAB_SIZE,
    sequence_length=SEQ_LEN,
    embedding_dim=EMBED_DIM,
    mask_zero=True,
)
x = embedding_layer(inputs)

for _ in range(NUM_LAYERS):
    decoder_layer = keras_nlp.layers.TransformerDecoder(
        num_heads=NUM_HEADS,
        intermediate_dim=FEED_FORWARD_DIM,
    )
    x = decoder_layer(x)  

outputs = keras.layers.Dense(VOCAB_SIZE)(x)
model = keras.Model(inputs=inputs, outputs=outputs)
loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True)
perplexity = keras_nlp.metrics.Perplexity(from_logits=True, mask_token_id=0)
model.compile(optimizer="adam", loss=loss_fn, metrics=[perplexity])



model.summary()



model.fit(train_ds, validation_data=val_ds, epochs=EPOCHS)




prompt_tokens = start_packer(tokenizer([""]))
prompt_tokens




def next(prompt, cache, index):
    logits = model(prompt)[:, index - 1, :]
    
    hidden_states = None
    return logits, hidden_states, cache






sampler = keras_nlp.samplers.GreedySampler()
output_tokens = sampler(
    next=next,
    prompt=prompt_tokens,
    index=1,  
)
txt = tokenizer.detokenize(output_tokens)
print(f"Greedy search generated text: \n{txt}\n")





sampler = keras_nlp.samplers.BeamSampler(num_beams=10)
output_tokens = sampler(
    next=next,
    prompt=prompt_tokens,
    index=1,
)
txt = tokenizer.detokenize(output_tokens)
print(f"Beam search generated text: \n{txt}\n")





sampler = keras_nlp.samplers.RandomSampler()
output_tokens = sampler(
    next=next,
    prompt=prompt_tokens,
    index=1,
)
txt = tokenizer.detokenize(output_tokens)
print(f"Random search generated text: \n{txt}\n")





sampler = keras_nlp.samplers.TopKSampler(k=10)
output_tokens = sampler(
    next=next,
    prompt=prompt_tokens,
    index=1,
)
txt = tokenizer.detokenize(output_tokens)
print(f"Top-K search generated text: \n{txt}\n")



sampler = keras_nlp.samplers.TopPSampler(p=0.5)
output_tokens = sampler(
    next=next,
    prompt=prompt_tokens,
    index=1,
)
txt = tokenizer.detokenize(output_tokens)
print(f"Top-P search generated text: \n{txt}\n")




class TopKTextGenerator(keras.callbacks.Callback):
    

    def __init__(self, k):
        self.sampler = keras_nlp.samplers.TopKSampler(k)

    def on_epoch_end(self, epoch, logs=None):
        output_tokens = self.sampler(
            next=next,
            prompt=prompt_tokens,
            index=1,
        )
        txt = tokenizer.detokenize(output_tokens)
        print(f"Top-K search generated text: \n{txt}\n")


text_generation_callback = TopKTextGenerator(k=10)

model.fit(
    train_ds.take(1), verbose=2, epochs=2, callbacks=[text_generation_callback]
)