title gpt2 text generation with kerasnlp chen qian date created 04172023 last modified 04172023 description use kerasnlp gpt2 model and samplers to do text generation accelerator gpu in this tutorial you will learn to use kerasnlphttps keras iokerasnlp to load a pretrained large language model llm gpt2 modelhttps openai comresearchbetterlanguagemodels originally invented by openai finetune it to a specific text style and generate text based on users input also known as prompt you will also learn how gpt2 adapts quickly to nonenglish languages such as chinese before we begin colab offers different kinds of runtimes make sure to go to runtime change runtime type and choose the gpu hardware accelerator runtime which should have 12g host ram and 15g gpu ram since you will finetune the gpt2 model running this tutorial on cpu runtime will take hours install kerasnlp choose backend and import dependencies this examples uses keras 3 0https keras iokerascore to work in any of tensorflow jax or torch support for keras core is baked into kerasnlp simply change the kerasbackend environment variable to select the backend of your choice we select the jax backend below shell pip install githttps github comkerasteamkerasnlp git q introduction to generative large language models llms large language models llms are a type of machine learning models that are trained on a large corpus of text data to generate outputs for various natural language processing nlp tasks such as text generation question answering and machine translation generative llms are typically based on deep learning neural networks such as the transformer architecturehttps arxiv orgabs1706 03762 invented by google researchers in 2017 and are trained on massive amounts of text data often involving billions of words these models such as google lamdahttps blog googletechnologyailamda and palmhttps ai googleblog com202204pathwayslanguagemodelpalmscalingto html are trained with a large dataset from various data sources which allows them to generate output for many tasks the core of generative llms is predicting the next word in a sentence often referred as causal lm pretraining in this way llms can generate coherent text based on user prompts for a more pedagogical discussion on language models you can refer to the stanford cs324 llm classhttps stanfordcs324 github iowinter2022lecturesintroduction introduction to kerasnlp large language models are complex to build and expensive to train from scratch luckily there are pretrained llms available for use right away kerasnlphttps keras iokerasnlp provides a large number of pretrained checkpoints that allow you to experiment with sota models without needing to train them yourself kerasnlp is a natural language processing library that supports users through their entire development cycle kerasnlp offers both pretrained models and modularized building blocks so developers could easily reuse pretrained models or stack their own llm in a nutshell for generative llm kerasnlp offers pretrained models with generate method e g kerasnlp models gpt2causallm and kerasnlp models optcausallm sampler class that implements generation algorithms such as topk beam and contrastive search these samplers can be used to generate text with custom models load a pretrained gpt2 model and generate some text kerasnlp provides a number of pretrained models such as google berthttps ai googleblog com201811opensourcingbertstateofartpre html and gpt2https openai comresearchbetterlanguagemodels you can see the list of models available in the kerasnlp repositoryhttps github comkerasteamkerasnlptreemasterkerasnlpmodels it s very easy to load the gpt2 model as you can see below to speed up training and generation we use preprocessor of length 128 instead of full length 1024 once the model is loaded you can use it to generate some text right away run the cells below to give it a try it s as simple as calling a single function generate try another one notice how much faster the second call is this is because the computational graph is xla compiledhttps www tensorflow orgxla in the 1st run and reused in the 2nd behind the scenes the quality of the generated text looks ok but we can improve it via finetuning more on the gpt2 model from kerasnlp next up we will actually finetune the model to update its parameters but before we do let s take a look at the full set of tools we have to for working with for gpt2 the code of gpt2 can be found herehttps github comkerasteamkerasnlpblobmasterkerasnlpmodelsgpt2 conceptually the gpt2causallm can be hierarchically broken down into several modules in kerasnlp all of which have a frompreset function that loads a pretrained model kerasnlp models gpt2tokenizer the tokenizer used by gpt2 model which is a bytepair encoderhttps huggingface cocoursechapter65 fwpt kerasnlp models gpt2causallmpreprocessor the preprocessor used by gpt2 causal lm training it does the tokenization along with other preprocessing works such as creating the label and appending the end token kerasnlp models gpt2backbone the gpt2 model which is a stack of kerasnlp layers transformerdecoder this is usually just referred as gpt2 kerasnlp models gpt2causallm wraps gpt2backbone it multiplies the output of gpt2backbone by embedding matrix to generate logits over vocab tokens finetune on reddit dataset now you have the knowledge of the gpt2 model from kerasnlp you can take one step further to finetune the model so that it generates text in a specific style short or long strict or casual in this tutorial we will use reddit dataset for example let s take a look inside sample data from the reddit tensorflow dataset there are two features document text of the post title the title in our case we are performing next word prediction in a language model so we only need the document feature now you can finetune the model using the familiar fit function note that preprocessor will be automatically called inside fit method since gpt2causallm is a kerasnlp models task instance this step takes quite a bit of gpu memory and a long time if we were to train it all the way to a fully trained state here we just use part of the dataset for demo purposes linearly decaying learning rate after finetuning is finished you can again generate text using the same generate function this time the text will be closer to reddit writing style and the generated length will be close to our preset length in the training set into the sampling method in kerasnlp we offer a few sampling methods e g contrastive search topk and beam sampling by default our gpt2causallm uses topk search but you can choose your own sampling method much like optimizer and activations there are two ways to specify your custom sampler use a string identifier such as greedy you are using the default configuration via this way pass a kerasnlp samplers sampler instance you can use custom configuration via this way use a string identifier use a sampler instance greedysampler tends to repeat itself for more details on kerasnlp sampler class you can check the code herehttps github comkerasteamkerasnlptreemasterkerasnlpsamplers finetune on chinese poem dataset we can also finetune gpt2 on nonenglish datasets for readers knowing chinese this part illustrates how to finetune gpt2 on chinese poem dataset to teach our model to become a poet because gpt2 uses bytepair encoder and the original pretraining dataset contains some chinese characters we can use the original vocab to finetune on chinese dataset shell load chinese poetry dataset git clone https github comchinesepoetrychinesepoetry git load text from the json file we only usefor demo purposes let s take a look at sample data similar as reddit example we convert to tf dataset and only use partial data to train running through the whole dataset takes long only take 500 and run 1 epochs for demo purposes let s check the result not bad in this tutorial you will learn to use kerasnlp https keras io keras_nlp to load a pre trained large language model llm gpt 2 model https openai com research better language models originally invented by openai finetune it to a specific text style and generate text based on users input also known as prompt you will also learn how gpt2 adapts quickly to non english languages such as chinese before we begin colab offers different kinds of runtimes make sure to go to runtime change runtime type and choose the gpu hardware accelerator runtime which should have 12g host ram and 15g gpu ram since you will finetune the gpt 2 model running this tutorial on cpu runtime will take hours install kerasnlp choose backend and import dependencies this examples uses keras 3 0 https keras io keras_core to work in any of tensorflow jax or torch support for keras core is baked into kerasnlp simply change the keras_backend environment variable to select the backend of your choice we select the jax backend below shell pip install git https github com keras team keras nlp git q introduction to generative large language models llms large language models llms are a type of machine learning models that are trained on a large corpus of text data to generate outputs for various natural language processing nlp tasks such as text generation question answering and machine translation generative llms are typically based on deep learning neural networks such as the transformer architecture https arxiv org abs 1706 03762 invented by google researchers in 2017 and are trained on massive amounts of text data often involving billions of words these models such as google lamda https blog google technology ai lamda and palm https ai googleblog com 2022 04 pathways language model palm scaling to html are trained with a large dataset from various data sources which allows them to generate output for many tasks the core of generative llms is predicting the next word in a sentence often referred as causal lm pretraining in this way llms can generate coherent text based on user prompts for a more pedagogical discussion on language models you can refer to the stanford cs324 llm class https stanford cs324 github io winter2022 lectures introduction introduction to kerasnlp large language models are complex to build and expensive to train from scratch luckily there are pretrained llms available for use right away kerasnlp https keras io keras_nlp provides a large number of pre trained checkpoints that allow you to experiment with sota models without needing to train them yourself kerasnlp is a natural language processing library that supports users through their entire development cycle kerasnlp offers both pretrained models and modularized building blocks so developers could easily reuse pretrained models or stack their own llm in a nutshell for generative llm kerasnlp offers pretrained models with generate method e g keras_nlp models gpt2causallm and keras_nlp models optcausallm sampler class that implements generation algorithms such as top k beam and contrastive search these samplers can be used to generate text with custom models load a pre trained gpt 2 model and generate some text kerasnlp provides a number of pre trained models such as google bert https ai googleblog com 2018 11 open sourcing bert state of art pre html and gpt 2 https openai com research better language models you can see the list of models available in the kerasnlp repository https github com keras team keras nlp tree master keras_nlp models it s very easy to load the gpt 2 model as you can see below to speed up training and generation we use preprocessor of length 128 instead of full length 1024 once the model is loaded you can use it to generate some text right away run the cells below to give it a try it s as simple as calling a single function generate try another one notice how much faster the second call is this is because the computational graph is xla compiled https www tensorflow org xla in the 1st run and re used in the 2nd behind the scenes the quality of the generated text looks ok but we can improve it via fine tuning more on the gpt 2 model from kerasnlp next up we will actually fine tune the model to update its parameters but before we do let s take a look at the full set of tools we have to for working with for gpt2 the code of gpt2 can be found here https github com keras team keras nlp blob master keras_nlp models gpt2 conceptually the gpt2causallm can be hierarchically broken down into several modules in kerasnlp all of which have a from_preset function that loads a pretrained model keras_nlp models gpt2tokenizer the tokenizer used by gpt2 model which is a byte pair encoder https huggingface co course chapter6 5 fw pt keras_nlp models gpt2causallmpreprocessor the preprocessor used by gpt2 causal lm training it does the tokenization along with other preprocessing works such as creating the label and appending the end token keras_nlp models gpt2backbone the gpt2 model which is a stack of keras_nlp layers transformerdecoder this is usually just referred as gpt2 keras_nlp models gpt2causallm wraps gpt2backbone it multiplies the output of gpt2backbone by embedding matrix to generate logits over vocab tokens finetune on reddit dataset now you have the knowledge of the gpt 2 model from kerasnlp you can take one step further to finetune the model so that it generates text in a specific style short or long strict or casual in this tutorial we will use reddit dataset for example let s take a look inside sample data from the reddit tensorflow dataset there are two features __document__ text of the post __title__ the title in our case we are performing next word prediction in a language model so we only need the document feature now you can finetune the model using the familiar fit function note that preprocessor will be automatically called inside fit method since gpt2causallm is a keras_nlp models task instance this step takes quite a bit of gpu memory and a long time if we were to train it all the way to a fully trained state here we just use part of the dataset for demo purposes linearly decaying learning rate after fine tuning is finished you can again generate text using the same generate function this time the text will be closer to reddit writing style and the generated length will be close to our preset length in the training set into the sampling method in kerasnlp we offer a few sampling methods e g contrastive search top k and beam sampling by default our gpt2causallm uses top k search but you can choose your own sampling method much like optimizer and activations there are two ways to specify your custom sampler use a string identifier such as greedy you are using the default configuration via this way pass a keras_nlp samplers sampler instance you can use custom configuration via this way use a string identifier use a sampler instance greedysampler tends to repeat itself for more details on kerasnlp sampler class you can check the code here https github com keras team keras nlp tree master keras_nlp samplers finetune on chinese poem dataset we can also finetune gpt2 on non english datasets for readers knowing chinese this part illustrates how to fine tune gpt2 on chinese poem dataset to teach our model to become a poet because gpt2 uses byte pair encoder and the original pretraining dataset contains some chinese characters we can use the original vocab to finetune on chinese dataset shell load chinese poetry dataset git clone https github com chinese poetry chinese poetry git load text from the json file we only use 全唐诗 for demo purposes let s take a look at sample data similar as reddit example we convert to tf dataset and only use partial data to train running through the whole dataset takes long only take 500 and run 1 epochs for demo purposes let s check the result not bad
import keras_nlp
import tensorflow as tf
import json
import keras
import os
import tensorflow_datasets as tfds
import time










preprocessor = keras_nlp.models.GPT2CausalLMPreprocessor.from_preset(
    "gpt2_base_en",
    sequence_length=128,
)
gpt2_lm = keras_nlp.models.GPT2CausalLM.from_preset(
    "gpt2_base_en", preprocessor=preprocessor
)



start = time.time()

output = gpt2_lm.generate("My trip to Yosemite was", max_length=200)
print("\nGPT-2 output:")
print(output)

end = time.time()
print(f"TOTAL TIME ELAPSED: {end - start:.2f}s")



start = time.time()

output = gpt2_lm.generate("That Italian restaurant is", max_length=200)
print("\nGPT-2 output:")
print(output)

end = time.time()
print(f"TOTAL TIME ELAPSED: {end - start:.2f}s")







reddit_ds = tfds.load("reddit_tifu", split="train", as_supervised=True)



for document, title in reddit_ds:
    print(document.numpy())
    print(title.numpy())
    break



train_ds = (
    reddit_ds.map(lambda document, _: document)
    .batch(32)
    .cache()
    .prefetch(tf.data.AUTOTUNE)
)



train_ds = train_ds.take(500)
num_epochs = 1


learning_rate = keras.optimizers.schedules.PolynomialDecay(
    5e-5,
    decay_steps=train_ds.cardinality() * num_epochs,
    end_learning_rate=0.0,
)
loss = keras.losses.SparseCategoricalCrossentropy(from_logits=True)
gpt2_lm.compile(
    optimizer=keras.optimizers.Adam(learning_rate),
    loss=loss,
    weighted_metrics=["accuracy"],
)

gpt2_lm.fit(train_ds, epochs=num_epochs)



start = time.time()

output = gpt2_lm.generate("I like basketball", max_length=200)
print("\nGPT-2 output:")
print(output)

end = time.time()
print(f"TOTAL TIME ELAPSED: {end - start:.2f}s")




gpt2_lm.compile(sampler="top_k")
output = gpt2_lm.generate("I like basketball", max_length=200)
print("\nGPT-2 output:")
print(output)


greedy_sampler = keras_nlp.samplers.GreedySampler()
gpt2_lm.compile(sampler=greedy_sampler)

output = gpt2_lm.generate("I like basketball", max_length=200)
print("\nGPT-2 output:")
print(output)









poem_collection = []
for file in os.listdir("chinese-poetry/全唐诗"):
    if ".json" not in file or "poet" not in file:
        continue
    full_filename = f"{'chinese-poetry/全唐诗'}/{file}"
    with open(full_filename, "r") as f:
        content = json.load(f)
        poem_collection.extend(content)

paragraphs = ["".join(data["paragraphs"]) for data in poem_collection]



print(paragraphs[0])



train_ds = (
    tf.data.Dataset.from_tensor_slices(paragraphs)
    .batch(16)
    .cache()
    .prefetch(tf.data.AUTOTUNE)
)



train_ds = train_ds.take(500)
num_epochs = 1

learning_rate = keras.optimizers.schedules.PolynomialDecay(
    5e-4,
    decay_steps=train_ds.cardinality() * num_epochs,
    end_learning_rate=0.0,
)
loss = keras.losses.SparseCategoricalCrossentropy(from_logits=True)
gpt2_lm.compile(
    optimizer=keras.optimizers.Adam(learning_rate),
    loss=loss,
    weighted_metrics=["accuracy"],
)

gpt2_lm.fit(train_ds, epochs=num_epochs)



output = gpt2_lm.generate("昨夜雨疏风骤", max_length=200)
print(output)