TensorFlow Transformers and KerasHub

Real-World Scenario

A support team wants to route incoming tickets by reading them, and a Transformer is the default architecture for that today. Using one well means knowing what happens inside: what attention computes, why padding needs a mask, why the model needs to be told about word order, and how the pretrained-model API fits together. This lesson builds each piece from Keras layers, trains a small encoder on movie reviews, then walks through the KerasHub BERT API.

All examples ran with TensorFlow 2.21, Keras 3.15, and KerasHub 0.32 on a CPU-only machine. TensorFlow’s startup log lines on stderr (no CUDA driver found, oneDNN enabled) are left out of the outputs below.

About pretrained weights: this machine could not reach Kaggle or the Hugging Face Hub, where KerasHub downloads its presets, so no pretrained transformer weights were used. The KerasHub section builds a small BERT with randomly initialized weights and a hand-written vocabulary to show the API; from_preset would replace the hand-built tokenizer and backbone with pretrained ones, and the compile/fit/predict code stays the same. Also note that Hugging Face Transformers v5 no longer includes TensorFlow model classes — the TFAutoModel code in older tutorials needs a 4.x release.

Scaled Dot-Product Attention by Hand

Each query scores every key with a dot product, softmax turns the scores into weights, and the output is the weighted average of the values.

import numpy as np
import keras
from keras import ops

rng = np.random.default_rng(0)
q = rng.normal(size=(1, 3, 4)).astype("float32")      # (batch, queries, dim)
k = rng.normal(size=(1, 5, 4)).astype("float32")      # (batch, keys, dim)
v = rng.normal(size=(1, 5, 4)).astype("float32")

scores = ops.matmul(q, ops.transpose(k, (0, 2, 1))) / np.sqrt(4)
weights = ops.softmax(scores, axis=-1)
by_hand = ops.matmul(weights, v)

print(weights.shape)
print(ops.sum(weights, axis=-1).numpy().round(4))
print(by_hand.shape)
# keras.ops expects (batch, seq, heads, dim): add a heads axis of 1
fused = ops.dot_product_attention(q[:, :, None], k[:, :, None], v[:, :, None])
print(np.allclose(by_hand, fused[:, :, 0], atol=1e-6))

$ python attention.py
(1, 3, 5)
[[1. 1. 1.]]
(1, 3, 4)
True

Three queries against five keys give a 3×5 weight matrix, and each row sums to 1 — every output is a weighted average of the five value vectors. Dividing by √dim keeps the dot products from growing with the vector size, which would push softmax towards putting all weight on one key. keras.ops.dot_product_attention computes the same thing in one fused call; it expects a heads axis, which is why each input gets a None axis added.

The MultiHeadAttention Layer

keras.layers.MultiHeadAttention adds learned projections around that computation and runs several heads side by side.

import numpy as np
import keras

x = np.random.default_rng(0).normal(size=(2, 6, 16)).astype("float32")
mha = keras.layers.MultiHeadAttention(num_heads=4, key_dim=4)

out, scores = mha(x, x, return_attention_scores=True)     # self-attention
print("output:", out.shape)
print("scores:", scores.shape)
print("params:", mha.count_params())
for w in mha.weights:
    print(f"  {w.path:45s} {w.shape}")

$ python mha.py
output: (2, 6, 16)
scores: (2, 4, 6, 6)
params: 1088
  multi_head_attention/query/kernel             (16, 4, 4)
  multi_head_attention/query/bias               (4, 4)
  multi_head_attention/key/kernel               (16, 4, 4)
  multi_head_attention/key/bias                 (4, 4)
  multi_head_attention/value/kernel             (16, 4, 4)
  multi_head_attention/value/bias               (4, 4)
  multi_head_attention/attention_output/kernel  (4, 4, 16)
  multi_head_attention/attention_output/bias    (16,)

Passing x as both query and value makes this self-attention: each of the 6 positions attends to all 6. The scores are (batch, heads, queries, keys) — one 6×6 attention matrix per head. The query, key, and value kernels are (input_dim, heads, key_dim): each projects the 16-wide input into 4 heads of width 4, and the output kernel maps the heads back to 16. That is 3 × (16×16 + 16) + (16×16 + 16) = 1,088 parameters. Unlike PyTorch’s layer, key_dim is set per head and need not divide the input width.

Padding Masks

Batches are padded to one length. Without a mask, attention happily puts weight on the padding.

import numpy as np
import keras

keras.utils.set_random_seed(0)
ids = np.array([[4, 7, 2, 0, 0]])                        # last two are padding
x = np.random.default_rng(0).normal(size=(1, 5, 8)).astype("float32") * 3
mha = keras.layers.MultiHeadAttention(num_heads=1, key_dim=8)

_, plain = mha(x, x, return_attention_scores=True)
_, masked = mha(x, x, return_attention_scores=True,
                attention_mask=(ids != 0)[:, None, :])   # (batch, queries, keys)
print(plain[0, 0, 0].numpy().round(3))
print(masked[0, 0, 0].numpy().round(3))

embed = keras.layers.Embedding(10, 8, mask_zero=True)
e = embed(ids)
print(embed.compute_mask(ids).numpy())
_, auto = mha(e, e, return_attention_scores=True)        # mask travels with e
print(auto[0, 0, 0].numpy().round(3))

$ python padding_mask.py
[0.    0.003 0.    0.996 0.   ]
[0.026 0.96  0.014 0.    0.   ]
[[ True  True  True False False]]
[0.333 0.333 0.334 0.    0.   ]

Each line is the first token’s attention over the five positions. Unmasked, it puts 99.6% of its weight on position 3 — padding. With attention_mask, a boolean (batch, queries, keys) array where True means “may attend”, the two padded keys get exactly zero and the weight is redistributed over the real tokens. An Embedding with mask_zero=True attaches the same mask to its output, and MultiHeadAttention picks it up without being told. Those last weights are nearly uniform because freshly initialized embeddings are tiny, so all the scores are close to zero.

Attention Ignores Word Order

Shuffle the tokens and compare the outputs — first without, then with position information.

import numpy as np
import keras

x = np.random.default_rng(0).normal(size=(1, 4, 8)).astype("float32")
perm = [2, 0, 3, 1]
mha = keras.layers.MultiHeadAttention(num_heads=2, key_dim=4)

out = mha(x, x).numpy()
out_shuffled = mha(x[:, perm], x[:, perm]).numpy()
print(np.allclose(out[:, perm], out_shuffled, atol=1e-6))

pos = keras.layers.Embedding(4, 8)(np.arange(4))          # one vector per position
with_pos = mha(x + pos, x + pos).numpy()
with_pos_shuffled = mha(x[:, perm] + pos, x[:, perm] + pos).numpy()
print(np.allclose(with_pos[:, perm], with_pos_shuffled, atol=1e-6))

$ python order_blind.py
True
False

Shuffling the input tokens only shuffles the outputs the same way: self-attention sees a set, not a sequence. Adding a vector that depends on the position — here a learned Embedding indexed by 0, 1, 2, 3 — breaks that symmetry, because the same token now looks different at different positions. Every Transformer adds some form of positional information before its first attention layer.

A Transformer Encoder in Keras

Two custom layers make up the model. The first adds a learned position embedding to each token embedding and reports id 0 as padding through compute_mask. Saved as embedding.py:

import numpy as np
import keras
from keras import ops

@keras.saving.register_keras_serializable()
class TokenAndPosition(keras.layers.Layer):
    def __init__(self, maxlen, vocab, dim, **kwargs):
        super().__init__(**kwargs)
        self.config = dict(maxlen=maxlen, vocab=vocab, dim=dim)
        self.tokens = keras.layers.Embedding(vocab, dim)
        self.positions = keras.layers.Embedding(maxlen, dim)

    def call(self, ids):
        return self.tokens(ids) + self.positions(ops.arange(ops.shape(ids)[-1]))

    def compute_mask(self, ids, mask=None):
        return ops.not_equal(ids, 0)                     # 0 is padding

    def get_config(self):
        return {**super().get_config(), **self.config}

if __name__ == "__main__":
    layer = TokenAndPosition(maxlen=6, vocab=10, dim=4)
    ids = np.array([[4, 7, 2, 0, 0, 0]])
    print(layer(ids).shape, layer.compute_mask(ids).numpy())

$ python embedding.py
(1, 6, 4) [[ True  True  True False False False]]

The second is one encoder block: self-attention, then a two-layer feed-forward network, each wrapped in a residual connection and layer normalization. Saved as encoder.py:

import numpy as np
import keras

@keras.saving.register_keras_serializable()
class EncoderBlock(keras.layers.Layer):
    def __init__(self, dim, heads, ff_dim, **kwargs):
        super().__init__(**kwargs)
        self.config = dict(dim=dim, heads=heads, ff_dim=ff_dim)
        self.attn = keras.layers.MultiHeadAttention(heads, key_dim=dim // heads)
        self.ff = keras.Sequential([keras.layers.Dense(ff_dim, activation="relu"),
                                    keras.layers.Dense(dim)])
        self.norm1, self.norm2 = keras.layers.LayerNormalization(), keras.layers.LayerNormalization()
        self.drop = keras.layers.Dropout(0.1)
        self.supports_masking = True                     # pass the mask on

    def call(self, x, mask=None, training=False):
        attn_mask = None if mask is None else mask[:, None, :]  # hide padded keys
        x = self.norm1(x + self.drop(self.attn(x, x, attention_mask=attn_mask), training=training))
        return self.norm2(x + self.drop(self.ff(x), training=training))
    def get_config(self):
        return {**super().get_config(), **self.config}

if __name__ == "__main__":
    block = EncoderBlock(dim=32, heads=2, ff_dim=64)
    print(block(np.zeros((2, 200, 32), dtype="float32")).shape, block.count_params())

$ python encoder.py
(2, 200, 32) 8544

The block returns the same shape it receives, so blocks can be stacked. Its 8,544 parameters are 4,224 in attention, 4,192 in the feed-forward network, and 128 in the two layer norms. Keras hands the incoming mask to call as mask, shaped (batch, keys); indexing with [:, None, :] turns it into the (batch, queries, keys) form attention_mask expects. supports_masking = True passes the mask on to the next layer — without it Keras warns that the layer “will therefore destroy the mask information”. register_keras_serializable and get_config are what let load_model rebuild these layers from a .keras file.

Training the Encoder on IMDB

The same 25,000 IMDB training reviews as the LSTM in TensorFlow RNNs and LSTMs, cut to 200 word ids each.

import time
import keras
from embedding import TokenAndPosition
from encoder import EncoderBlock

VOCAB, MAXLEN = 10_000, 200
(x_train, y_train), (x_test, y_test) = keras.datasets.imdb.load_data(num_words=VOCAB)
x_train, x_test = (keras.utils.pad_sequences(d, maxlen=MAXLEN) for d in (x_train, x_test))

keras.utils.set_random_seed(0)
inputs = keras.Input(shape=(MAXLEN,))
x = TokenAndPosition(MAXLEN, VOCAB, 32)(inputs)
x = EncoderBlock(32, heads=2, ff_dim=64)(x)
x = keras.layers.GlobalAveragePooling1D()(x)
outputs = keras.layers.Dense(1, activation="sigmoid")(x)
model = keras.Model(inputs, outputs)
model.compile(optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"])
print(f"parameters: {model.count_params():,}")
stop = keras.callbacks.EarlyStopping(patience=1, restore_best_weights=True)
start = time.perf_counter()
model.fit(x_train, y_train, epochs=5, batch_size=64, validation_split=0.2,
          callbacks=[stop], verbose=2)
print(f"trained in {time.perf_counter() - start:.0f}s")
print(f"test accuracy: {model.evaluate(x_test, y_test, verbose=0)[1]:.4f}")
model.save("imdb_transformer.keras")

$ python imdb_transformer.py
parameters: 334,977
Epoch 1/5
313/313 - 15s - 48ms/step - accuracy: 0.8189 - loss: 0.3879 - val_accuracy: 0.8302 - val_loss: 0.3887
Epoch 2/5
313/313 - 13s - 41ms/step - accuracy: 0.9097 - loss: 0.2251 - val_accuracy: 0.8284 - val_loss: 0.4552
trained in 28s
test accuracy: 0.8140

GlobalAveragePooling1D averages the 200 output vectors into one, skipping padded positions because the mask reaches it. Validation loss rose in epoch 2 while training accuracy jumped to 91%, so EarlyStopping(patience=1) stopped training and restore_best_weights=True rolled back to the epoch-1 weights, which score 81.4% on the test set. That is below the bidirectional LSTM’s 86.8% on the same data. Of the 334,977 parameters, 320,000 are the token embedding table, learned from scratch on 20,000 reviews; this is the gap pretraining closes — a pretrained model arrives already knowing what words mean.

The KerasHub BERT API

KerasHub packages pretrained models as a tokenizer, a preprocessor, a backbone, and a task model. With network access, from_preset("bert_tiny_en_uncased") loads all of them with pretrained weights. Here the tokenizer gets a hand-written 16-entry WordPiece vocabulary instead. Saved as hub_tokenizer.py:

import keras
import keras_hub

vocab = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]",
         "the", "movie", "was", "great", "bad", "plot", "dull", "fun", "##s", "act", "##ing"]
tokenizer = keras_hub.tokenizers.BertTokenizer(vocabulary=vocab, lowercase=True)
preprocessor = keras_hub.models.BertTextClassifierPreprocessor(tokenizer, sequence_length=8)

if __name__ == "__main__":
    print(tokenizer("The movie was GREAT"))
    print(tokenizer("the acting plots"), tokenizer.detokenize(tokenizer("the acting plots")))
    print(tokenizer("the popcorn"))
    for name, value in preprocessor(["the movie was fun", "bad plot"]).items():
        print(f"{name:12s}", keras.ops.convert_to_numpy(value).astype(int).tolist())

$ python hub_tokenizer.py
[5 6 7 8]
[ 5 14 15 10 13] the acting plots
[5 1]
token_ids    [[2, 5, 6, 7, 12, 3, 0, 0], [2, 9, 10, 3, 0, 0, 0, 0]]
padding_mask [[1, 1, 1, 1, 1, 1, 0, 0], [1, 1, 1, 1, 0, 0, 0, 0]]
segment_ids  [[0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0]]

WordPiece splits words it does not know whole into known pieces: acting becomes act + ##ing and plots becomes plot + ##s, where ## marks a continuation. A word that cannot be built entirely from vocabulary pieces, like popcorn, becomes a single [UNK] (id 1). The preprocessor adds BERT’s [CLS] (2) and [SEP] (3) markers, pads to sequence_length, and returns the three inputs a BERT backbone takes: the ids, a padding mask, and segment ids that distinguish the two sentences of a sentence-pair task.

import keras
import keras_hub
from hub_tokenizer import tokenizer, preprocessor

keras.utils.set_random_seed(0)
backbone = keras_hub.models.BertBackbone(          # random weights, tiny config
    vocabulary_size=tokenizer.vocabulary_size(), num_layers=2, num_heads=2,
    hidden_dim=32, intermediate_dim=64, max_sequence_length=16)
# with network access: BertTextClassifier.from_preset("bert_tiny_en_uncased", num_classes=2)
classifier = keras_hub.models.BertTextClassifier(backbone, num_classes=2,
                                                 preprocessor=preprocessor)
print(f"parameters: {classifier.count_params():,}")

texts = ["the movie was great", "great fun", "the acting was great",
         "the movie was bad", "dull plot", "the acting was bad"] * 8
labels = [1, 1, 1, 0, 0, 0] * 8
classifier.compile(optimizer=keras.optimizers.Adam(1e-3), metrics=["accuracy"])
history = classifier.fit(x=texts, y=labels, batch_size=8, epochs=10, verbose=0)
print("final accuracy:", history.history["accuracy"][-1])

probs = keras.ops.softmax(classifier.predict(["great plot", "dull movie"], verbose=0))
print(probs.numpy().round(3))
classifier.save("tiny_bert.keras")
restored = keras.models.load_model("tiny_bert.keras")
print(restored.predict(["great plot"], verbose=0).round(3))

$ python hub_classifier.py
parameters: 19,362
final accuracy: 1.0
[[0.06 0.94]
 [0.94 0.06]]
[[-1.379  1.369]]

Because the preprocessor is part of the classifier, fit and predict take raw strings — tokenization happens inside the model, and the saved .keras file carries the vocabulary with it. The classifier outputs logits, so softmax turns them into probabilities: “great plot” is 94% positive and “dull movie” 94% negative, phrases it never saw whole in training. The reloaded model’s logits [-1.379, 1.369] are the same prediction before softmax. With a pretrained preset the backbone would be far larger and compile would use a learning rate around 2e-5 to avoid overwriting what pretraining learned; this randomly initialized model needs 1e-3 to learn anything in ten epochs.

Practice

import numpy as np
import keras

mha = keras.layers.MultiHeadAttention(num_heads=8, key_dim=16)
x = np.zeros((1, 10, 64), dtype="float32")
mha(x, x)                                               # builds the weights

qkv = 3 * (64 * 8 * 16 + 8 * 16)
out = 8 * 16 * 64 + 64
print(mha.count_params(), qkv + out)

$ python p1_params.py
33216 33216

import numpy as np
import keras

keras.utils.set_random_seed(0)
x = np.random.default_rng(0).normal(size=(1, 4, 8)).astype("float32")
mha = keras.layers.MultiHeadAttention(num_heads=1, key_dim=8)

_, scores = mha(x, x, use_causal_mask=True, return_attention_scores=True)
print(scores[0, 0].numpy().round(2))

$ python p2_causal.py
[[1.   0.   0.   0.  ]
 [0.41 0.59 0.   0.  ]
 [0.29 0.42 0.29 0.  ]
 [0.38 0.19 0.23 0.21]]

import numpy as np
import keras

rng = np.random.default_rng(0)
decoder = rng.normal(size=(2, 3, 16)).astype("float32")    # 3 target positions
encoder = rng.normal(size=(2, 7, 16)).astype("float32")    # 7 source positions
mha = keras.layers.MultiHeadAttention(num_heads=2, key_dim=8)

out, scores = mha(query=decoder, value=encoder, return_attention_scores=True)
print(out.shape, scores.shape)

$ python p3_cross.py
(2, 3, 16) (2, 2, 3, 7)

import time
import keras
from encoder import EncoderBlock

VOCAB, MAXLEN = 10_000, 200
(x_train, y_train), (x_test, y_test) = keras.datasets.imdb.load_data(num_words=VOCAB)
x_train = keras.utils.pad_sequences(x_train, maxlen=MAXLEN)
x_test = keras.utils.pad_sequences(x_test, maxlen=MAXLEN)

keras.utils.set_random_seed(0)
inputs = keras.Input(shape=(MAXLEN,))
x = keras.layers.Embedding(VOCAB, 32, mask_zero=True)(inputs)    # no positions
x = EncoderBlock(32, heads=2, ff_dim=64)(x)
x = keras.layers.GlobalAveragePooling1D()(x)
outputs = keras.layers.Dense(1, activation="sigmoid")(x)
model = keras.Model(inputs, outputs)
model.compile(optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"])
print(f"parameters: {model.count_params():,}")

start = time.perf_counter()
stop = keras.callbacks.EarlyStopping(patience=1, restore_best_weights=True)
model.fit(x_train, y_train, epochs=5, batch_size=64, validation_split=0.2,
          callbacks=[stop], verbose=2)
print(f"trained in {time.perf_counter() - start:.0f}s")
print(f"test accuracy: {model.evaluate(x_test, y_test, verbose=0)[1]:.4f}")

$ python p4_no_positions.py
parameters: 328,577
Epoch 1/5
313/313 - 16s - 51ms/step - accuracy: 0.8301 - loss: 0.3714 - val_accuracy: 0.8322 - val_loss: 0.3866
Epoch 2/5
313/313 - 13s - 42ms/step - accuracy: 0.9104 - loss: 0.2205 - val_accuracy: 0.8160 - val_loss: 0.4638
trained in 29s
test accuracy: 0.8144