Chapter 14
Text generation
NotebookPython 386 cells
This is a companion notebook for the book Deep Learning with Python, Third Edition. For readability, it only contains runnable code blocks and section titles, and omits everything else in the book: text paragraphs, figures, and pseudocode.
If you want to be able to follow what's going on, I recommend reading the notebook side by side with your copy of the book.
The book's contents are available online at deeplearningwithpython.io.
In [0]python · cell 2
python
!pip install keras keras-hub --upgrade -qIn [0]python · cell 3
python
import os
os.environ["KERAS_BACKEND"] = "jax"In [0]python · cell 4
python
# @title
import os
from IPython.core.magic import register_cell_magic
@register_cell_magic
def backend(line, cell):
current, required = os.environ.get("KERAS_BACKEND", ""), line.split()[-1]
if current == required:
get_ipython().run_cell(cell)
else:
print(
f"This cell requires the {required} backend. To run it, change KERAS_BACKEND to "
f"\"{required}\" at the top of the notebook, restart the runtime, and rerun the notebook."
)Text generation
A brief history of sequence generation
Training a mini-GPT
In [0]python · cell 8
python
import os
# Free up more GPU memory on the Jax and TensorFlow backends.
os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = "1.00"In [0]python · cell 9
python
import keras
import pathlib
extract_dir = keras.utils.get_file(
fname="mini-c4",
origin=(
"https://hf.co/datasets/mattdangerw/mini-c4/resolve/main/mini-c4.zip"
),
extract=True,
)
extract_dir = pathlib.Path(extract_dir) / "mini-c4"In [0]python · cell 10
python
with open(extract_dir / "shard0.txt", "r") as f:
print(f.readline().replace("\\n", "\n")[:100])In [0]python · cell 11
python
import keras_hub
import numpy as np
vocabulary_file = keras.utils.get_file(
origin="https://hf.co/mattdangerw/spiece/resolve/main/vocabulary.proto",
)
tokenizer = keras_hub.tokenizers.SentencePieceTokenizer(vocabulary_file)In [0]python · cell 12
python
tokenizer.tokenize("The quick brown fox.")In [0]python · cell 13
python
tokenizer.detokenize([450, 4996, 17354, 1701, 29916, 29889])In [0]python · cell 14
python
import tensorflow as tf
batch_size = 64
sequence_length = 256
suffix = np.array([tokenizer.token_to_id("<|endoftext|>")])
def read_file(filename):
ds = tf.data.TextLineDataset(filename)
ds = ds.map(lambda x: tf.strings.regex_replace(x, r"\\n", "\n"))
ds = ds.map(tokenizer, num_parallel_calls=8)
return ds.map(lambda x: tf.concat([x, suffix], -1))
files = [str(file) for file in extract_dir.glob("*.txt")]
ds = tf.data.Dataset.from_tensor_slices(files)
ds = ds.interleave(read_file, cycle_length=32, num_parallel_calls=32)
ds = ds.rebatch(sequence_length + 1, drop_remainder=True)
ds = ds.map(lambda x: (x[:-1], x[1:]))
ds = ds.batch(batch_size).prefetch(8)In [0]python · cell 15
python
num_batches = 58746
num_val_batches = 500
num_train_batches = num_batches - num_val_batches
val_ds = ds.take(num_val_batches).repeat()
train_ds = ds.skip(num_val_batches).repeat()Building the model
In [0]python · cell 17
python
from keras import layers
class TransformerDecoder(keras.Layer):
def __init__(self, hidden_dim, intermediate_dim, num_heads):
super().__init__()
key_dim = hidden_dim // num_heads
self.self_attention = layers.MultiHeadAttention(
num_heads, key_dim, dropout=0.1
)
self.self_attention_layernorm = layers.LayerNormalization()
self.feed_forward_1 = layers.Dense(intermediate_dim, activation="relu")
self.feed_forward_2 = layers.Dense(hidden_dim)
self.feed_forward_layernorm = layers.LayerNormalization()
self.dropout = layers.Dropout(0.1)
def call(self, inputs):
residual = x = inputs
x = self.self_attention(query=x, key=x, value=x, use_causal_mask=True)
x = self.dropout(x)
x = x + residual
x = self.self_attention_layernorm(x)
residual = x
x = self.feed_forward_1(x)
x = self.feed_forward_2(x)
x = self.dropout(x)
x = x + residual
x = self.feed_forward_layernorm(x)
return xIn [0]python · cell 18
python
from keras import ops
class PositionalEmbedding(keras.Layer):
def __init__(self, sequence_length, input_dim, output_dim):
super().__init__()
self.token_embeddings = layers.Embedding(input_dim, output_dim)
self.position_embeddings = layers.Embedding(sequence_length, output_dim)
def call(self, inputs, reverse=False):
if reverse:
token_embeddings = self.token_embeddings.embeddings
return ops.matmul(inputs, ops.transpose(token_embeddings))
positions = ops.cumsum(ops.ones_like(inputs), axis=-1) - 1
embedded_tokens = self.token_embeddings(inputs)
embedded_positions = self.position_embeddings(positions)
return embedded_tokens + embedded_positionsIn [0]python · cell 19
python
keras.config.set_dtype_policy("mixed_float16")
vocab_size = tokenizer.vocabulary_size()
hidden_dim = 512
intermediate_dim = 2056
num_heads = 8
num_layers = 8
inputs = keras.Input(shape=(None,), dtype="int32", name="inputs")
embedding = PositionalEmbedding(sequence_length, vocab_size, hidden_dim)
x = embedding(inputs)
x = layers.LayerNormalization()(x)
for i in range(num_layers):
x = TransformerDecoder(hidden_dim, intermediate_dim, num_heads)(x)
outputs = embedding(x, reverse=True)
mini_gpt = keras.Model(inputs, outputs)Pretraining the model
In [0]python · cell 21
python
class WarmupSchedule(keras.optimizers.schedules.LearningRateSchedule):
def __init__(self):
self.rate = 2e-4
self.warmup_steps = 1_000.0
def __call__(self, step):
step = ops.cast(step, dtype="float32")
scale = ops.minimum(step / self.warmup_steps, 1.0)
return self.rate * scaleIn [0]python · cell 22
python
import matplotlib.pyplot as plt
schedule = WarmupSchedule()
x = range(0, 5_000, 100)
y = [ops.convert_to_numpy(schedule(step)) for step in x]
plt.plot(x, y)
plt.xlabel("Train Step")
plt.ylabel("Learning Rate")
plt.show()In [0]python · cell 23
python
# ⚠️NOTE⚠️: If you can run the following with a Colab Pro GPU, we suggest you
# do so. This fit() call will take many hours on free tier GPUs. You can also
# reduce steps_per_epoch to try the code with a less trained model.In [0]python · cell 24
python
num_epochs = 8
steps_per_epoch = num_train_batches // num_epochs
validation_steps = num_val_batches
mini_gpt.compile(
optimizer=keras.optimizers.Adam(schedule),
loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=["accuracy"],
)
mini_gpt.fit(
train_ds,
validation_data=val_ds,
epochs=num_epochs,
steps_per_epoch=steps_per_epoch,
validation_steps=validation_steps,
)Generative decoding
In [0]python · cell 26
python
def generate(prompt, max_length=64):
tokens = list(ops.convert_to_numpy(tokenizer(prompt)))
prompt_length = len(tokens)
for _ in range(max_length - prompt_length):
prediction = mini_gpt(ops.convert_to_numpy([tokens]))
prediction = ops.convert_to_numpy(prediction[0, -1])
tokens.append(np.argmax(prediction).item())
return tokenizer.detokenize(tokens)In [0]python · cell 27
python
prompt = "A piece of advice"
generate(prompt)In [0]python · cell 28
python
def compiled_generate(prompt, max_length=64):
tokens = list(ops.convert_to_numpy(tokenizer(prompt)))
prompt_length = len(tokens)
tokens = tokens + [0] * (max_length - prompt_length)
for i in range(prompt_length, max_length):
prediction = mini_gpt.predict(np.array([tokens]), verbose=0)
prediction = prediction[0, i - 1]
tokens[i] = np.argmax(prediction).item()
return tokenizer.detokenize(tokens)In [0]python · cell 29
python
import timeit
tries = 10
timeit.timeit(lambda: compiled_generate(prompt), number=tries) / triesSampling strategies
In [0]python · cell 31
python
def compiled_generate(prompt, sample_fn, max_length=64):
tokens = list(ops.convert_to_numpy(tokenizer(prompt)))
prompt_length = len(tokens)
tokens = tokens + [0] * (max_length - prompt_length)
for i in range(prompt_length, max_length):
prediction = mini_gpt.predict(np.array([tokens]), verbose=0)
prediction = prediction[0, i - 1]
next_token = ops.convert_to_numpy(sample_fn(prediction))
tokens[i] = np.array(next_token).item()
return tokenizer.detokenize(tokens)In [0]python · cell 32
python
def greedy_search(preds):
return ops.argmax(preds)
compiled_generate(prompt, greedy_search)In [0]python · cell 33
python
def random_sample(preds, temperature=1.0):
preds = preds / temperature
return keras.random.categorical(preds[None, :], num_samples=1)[0]In [0]python · cell 34
python
compiled_generate(prompt, random_sample)In [0]python · cell 35
python
from functools import partial
compiled_generate(prompt, partial(random_sample, temperature=2.0))In [0]python · cell 36
python
compiled_generate(prompt, partial(random_sample, temperature=0.8))In [0]python · cell 37
python
compiled_generate(prompt, partial(random_sample, temperature=0.2))In [0]python · cell 38
python
def top_k(preds, k=5, temperature=1.0):
preds = preds / temperature
top_preds, top_indices = ops.top_k(preds, k=k, sorted=False)
choice = keras.random.categorical(top_preds[None, :], num_samples=1)[0]
return ops.take_along_axis(top_indices, choice, axis=-1)In [0]python · cell 39
python
compiled_generate(prompt, partial(top_k, k=5))In [0]python · cell 40
python
compiled_generate(prompt, partial(top_k, k=20))In [0]python · cell 41
python
compiled_generate(prompt, partial(top_k, k=5, temperature=0.5))Using a pretrained LLM
Text generation with the Gemma model
In [0]python · cell 44
python
import kagglehub
kagglehub.login()In [0]python · cell 45
python
gemma_lm = keras_hub.models.CausalLM.from_preset(
"gemma3_1b",
dtype="float32",
)In [0]python · cell 46
python
gemma_lm.summary(line_length=80)In [0]python · cell 47
python
gemma_lm.compile(sampler="greedy")
gemma_lm.generate("A piece of advice", max_length=40)In [0]python · cell 48
python
gemma_lm.generate("How can I make brownies?", max_length=40)In [0]python · cell 49
python
gemma_lm.generate(
"The following brownie recipe is easy to make in just a few "
"steps.\n\nYou can start by",
max_length=40,
)In [0]python · cell 50
python
gemma_lm.generate(
"Tell me about the 542nd president of the United States.",
max_length=40,
)Instruction fine-tuning
In [0]python · cell 52
python
import json
PROMPT_TEMPLATE = """"[instruction]\n{}[end]\n[response]\n"""
RESPONSE_TEMPLATE = """{}[end]"""
dataset_path = keras.utils.get_file(
origin=(
"https://hf.co/datasets/databricks/databricks-dolly-15k/"
"resolve/main/databricks-dolly-15k.jsonl"
),
)
data = {"prompts": [], "responses": []}
with open(dataset_path) as file:
for line in file:
features = json.loads(line)
if features["context"]:
continue
data["prompts"].append(PROMPT_TEMPLATE.format(features["instruction"]))
data["responses"].append(RESPONSE_TEMPLATE.format(features["response"]))In [0]python · cell 53
python
data["prompts"][0]In [0]python · cell 54
python
data["responses"][0]In [0]python · cell 55
python
ds = tf.data.Dataset.from_tensor_slices(data).shuffle(2000).batch(2)
val_ds = ds.take(100)
train_ds = ds.skip(100)In [0]python · cell 56
python
preprocessor = gemma_lm.preprocessor
preprocessor.sequence_length = 512
batch = next(iter(train_ds))
x, y, sample_weight = preprocessor(batch)
x["token_ids"].shapeIn [0]python · cell 57
python
x["padding_mask"].shapeIn [0]python · cell 58
python
y.shapeIn [0]python · cell 59
python
sample_weight.shapeIn [0]python · cell 60
python
x["token_ids"][0, :5], y[0, :5]Low-Rank Adaptation (LoRA)
In [0]python · cell 62
python
gemma_lm.backbone.enable_lora(rank=8)In [0]python · cell 63
python
gemma_lm.summary(line_length=80)In [0]python · cell 64
python
gemma_lm.compile(
loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
optimizer=keras.optimizers.Adam(5e-5),
weighted_metrics=[keras.metrics.SparseCategoricalAccuracy()],
)
gemma_lm.fit(train_ds, validation_data=val_ds, epochs=1)In [0]python · cell 65
python
gemma_lm.generate(
"[instruction]\nHow can I make brownies?[end]\n"
"[response]\n",
max_length=512,
)In [0]python · cell 66
python
gemma_lm.generate(
"[instruction]\nWhat is a proper noun?[end]\n"
"[response]\n",
max_length=512,
)In [0]python · cell 67
python
gemma_lm.generate(
"[instruction]\nWho is the 542nd president of the United States?[end]\n"
"[response]\n",
max_length=512,
)Going further with LLMs
Reinforcement Learning with Human Feedback (RLHF)
Using a chatbot trained with RLHF
In [0]python · cell 71
python
# ⚠️NOTE⚠️: If you are running on the free tier Colab GPUs, you will need to
# restart your runtime and run the notebook from here to free up memory for
# this 4 billion parameter model.
import os
os.environ["KERAS_BACKEND"] = "jax"
# Free up more GPU memory on the Jax and TensorFlow backends.
os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = "1.00"
import keras
import keras_hub
import kagglehub
import numpy as np
kagglehub.login()In [0]python · cell 72
python
gemma_lm = keras_hub.models.CausalLM.from_preset(
"gemma3_instruct_4b",
dtype="bfloat16",
)In [0]python · cell 73
python
PROMPT_TEMPLATE = """<start_of_turn>user
{}<end_of_turn>
<start_of_turn>model
"""In [0]python · cell 74
python
prompt = "Why can't you assign values in Jax tensors? Be brief!"
gemma_lm.generate(PROMPT_TEMPLATE.format(prompt), max_length=512)In [0]python · cell 75
python
prompt = "Who is the 542nd president of the United States?"
gemma_lm.generate(PROMPT_TEMPLATE.format(prompt), max_length=512)Multimodal LLMs
In [0]python · cell 77
python
import matplotlib.pyplot as plt
image_url = (
"https://github.com/mattdangerw/keras-nlp-scripts/"
"blob/main/learned-python.png?raw=true"
)
image_path = keras.utils.get_file(origin=image_url)
image = np.array(keras.utils.load_img(image_path))
plt.axis("off")
plt.imshow(image)
plt.show()In [0]python · cell 78
python
gemma_lm.preprocessor.max_images_per_prompt = 1
gemma_lm.preprocessor.sequence_length = 512
prompt = "What is going on in this image? Be concise!<start_of_image>"
gemma_lm.generate({
"prompts": PROMPT_TEMPLATE.format(prompt),
"images": [image],
})In [0]python · cell 79
python
prompt = "What is the snake wearing?<start_of_image>"
gemma_lm.generate({
"prompts": PROMPT_TEMPLATE.format(prompt),
"images": [image],
})Foundation models
Retrieval Augmented Generation (RAG)
"Reasoning" models
In [0]python · cell 83
python
prompt = """Judy wrote a 2-page letter to 3 friends twice a week for 3 months.
How many letters did she write?
Be brief, and add "ANSWER:" before your final answer."""
gemma_lm.compile(sampler="random")In [0]python · cell 84
python
gemma_lm.generate(PROMPT_TEMPLATE.format(prompt))In [0]python · cell 85
python
gemma_lm.generate(PROMPT_TEMPLATE.format(prompt))