Transformer model for language understanding

View on TensorFlow.org Run in Google Colab View source on GitHub Download notebook

This tutorial trains a Transformer model to translate a Portuguese to English dataset. This is an advanced example that assumes knowledge of text generation and attention.

The core idea behind the Transformer model is self-attention—the ability to attend to different positions of the input sequence to compute a representation of that sequence. Transformer creates stacks of self-attention layers and is explained below in the sections Scaled dot product attention and Multi-head attention.

A transformer model handles variable-sized input using stacks of self-attention layers instead of RNNs or CNNs. This general architecture has a number of advantages:

  • It makes no assumptions about the temporal/spatial relationships across the data. This is ideal for processing a set of objects (for example, StarCraft units).
  • Layer outputs can be calculated in parallel, instead of a series like an RNN.
  • Distant items can affect each other's output without passing through many RNN-steps, or convolution layers (see Scene Memory Transformer for example).
  • It can learn long-range dependencies. This is a challenge in many sequence tasks.

The downsides of this architecture are:

  • For a time-series, the output for a time-step is calculated from the entire history instead of only the inputs and current hidden-state. This may be less efficient.
  • If the input does have a temporal/spatial relationship, like text, some positional encoding must be added or the model will effectively see a bag of words.

After training the model in this notebook, you will be able to input a Portuguese sentence and return the English translation.

Attention heatmap

Setup

pip install tensorflow_datasets
pip install -U tensorflow-text
import collections
import logging
import os
import pathlib
import re
import string
import sys
import time

import numpy as np
import matplotlib.pyplot as plt

import tensorflow_datasets as tfds
import tensorflow_text as text
import tensorflow as tf
logging.getLogger('tensorflow').setLevel(logging.ERROR)  # suppress warnings

Download the Dataset

Use TensorFlow datasets to load the Portuguese-English translation dataset from the TED Talks Open Translation Project.

This dataset contains approximately 50000 training examples, 1100 validation examples, and 2000 test examples.

examples, metadata = tfds.load('ted_hrlr_translate/pt_to_en', with_info=True,
                               as_supervised=True)
train_examples, val_examples = examples['train'], examples['validation']

The tf.data.Dataset object returned by TensorFlow datasets yields pairs of text examples:

for pt_examples, en_examples in train_examples.batch(3).take(1):
  for pt in pt_examples.numpy():
    print(pt.decode('utf-8'))

  print()

  for en in en_examples.numpy():
    print(en.decode('utf-8'))
e quando melhoramos a procura , tiramos a única vantagem da impressão , que é a serendipidade .
mas e se estes fatores fossem ativos ?
mas eles não tinham a curiosidade de me testar .

and when you improve searchability , you actually take away the one advantage of print , which is serendipity .
but what if it were active ?
but they did n't test for curiosity .

Text tokenization & detokenization

You can't train a model directly on text. The text needs to be converted to some numeric representation first. Typically, you convert the text to sequences of token IDs, which are used as indices into an embedding.

One popular implementation is demonstrated in the Subword tokenizer tutorial builds subword tokenizers (text.BertTokenizer) optimized for this dataset and exports them in a saved_model.

Download and unzip and import the saved_model:

model_name = "ted_hrlr_translate_pt_en_converter"
tf.keras.utils.get_file(
    f"{model_name}.zip",
    f"https://storage.googleapis.com/download.tensorflow.org/models/{model_name}.zip",
    cache_dir='.', cache_subdir='', extract=True
)
Downloading data from https://storage.googleapis.com/download.tensorflow.org/models/ted_hrlr_translate_pt_en_converter.zip
188416/184801 [==============================] - 0s 0us/step
196608/184801 [===============================] - 0s 0us/step
'./ted_hrlr_translate_pt_en_converter.zip'
tokenizers = tf.saved_model.load(model_name)

The tf.saved_model contains two text tokenizers, one for English and one for Portuguese. Both have the same methods:

[item for item in dir(tokenizers.en) if not item.startswith('_')]
['detokenize',
 'get_reserved_tokens',
 'get_vocab_path',
 'get_vocab_size',
 'lookup',
 'tokenize',
 'tokenizer',
 'vocab']

The tokenize method converts a batch of strings to a padded-batch of token IDs. This method splits punctuation, lowercases and unicode-normalizes the input before tokenizing. That standardization is not visible here because the input data is already standardized.

for en in en_examples.numpy():
  print(en.decode('utf-8'))
and when you improve searchability , you actually take away the one advantage of print , which is serendipity .
but what if it were active ?
but they did n't test for curiosity .
encoded = tokenizers.en.tokenize(en_examples)

for row in encoded.to_list():
  print(row)
[2, 72, 117, 79, 1259, 1491, 2362, 13, 79, 150, 184, 311, 71, 103, 2308, 74, 2679, 13, 148, 80, 55, 4840, 1434, 2423, 540, 15, 3]
[2, 87, 90, 107, 76, 129, 1852, 30, 3]
[2, 87, 83, 149, 50, 9, 56, 664, 85, 2512, 15, 3]

The detokenize method attempts to convert these token IDs back to human readable text:

round_trip = tokenizers.en.detokenize(encoded)
for line in round_trip.numpy():
  print(line.decode('utf-8'))
and when you improve searchability , you actually take away the one advantage of print , which is serendipity .
but what if it were active ?
but they did n ' t test for curiosity .

The lower level lookup method converts from token-IDs to token text:

tokens = tokenizers.en.lookup(encoded)
tokens
<tf.RaggedTensor [[b'[START]', b'and', b'when', b'you', b'improve', b'search', b'##ability', b',', b'you', b'actually', b'take', b'away', b'the', b'one', b'advantage', b'of', b'print', b',', b'which', b'is', b's', b'##ere', b'##nd', b'##ip', b'##ity', b'.', b'[END]'], [b'[START]', b'but', b'what', b'if', b'it', b'were', b'active', b'?', b'[END]'], [b'[START]', b'but', b'they', b'did', b'n', b"'", b't', b'test', b'for', b'curiosity', b'.', b'[END]']]>

Here you can see the "subword" aspect of the tokenizers. The word "searchability" is decomposed into "search ##ability" and the word "serendipity" into "s ##ere ##nd ##ip ##ity"

Setup input pipeline

To build an input pipeline suitable for training you'll apply some transformations to the dataset.

This function will be used to encode the batches of raw text:

def tokenize_pairs(pt, en):
    pt = tokenizers.pt.tokenize(pt)
    # Convert from ragged to dense, padding with zeros.
    pt = pt.to_tensor()

    en = tokenizers.en.tokenize(en)
    # Convert from ragged to dense, padding with zeros.
    en = en.to_tensor()
    return pt, en

Here's a simple input pipeline that processes, shuffles and batches the data:

BUFFER_SIZE = 20000
BATCH_SIZE = 64
def make_batches(ds):
  return (
      ds
      .cache()
      .shuffle(BUFFER_SIZE)
      .batch(BATCH_SIZE)
      .map(tokenize_pairs, num_parallel_calls=tf.data.AUTOTUNE)
      .prefetch(tf.data.AUTOTUNE))


train_batches = make_batches(train_examples)
val_batches = make_batches(val_examples)

Positional encoding

Attention layers see their input as a set of vectors, with no sequential order. This model also doesn't contain any recurrent or convolutional layers. Because of this a "positional encoding" is added to give the model some information about the relative position of the tokens in the sentence.

The positional encoding vector is added to the embedding vector. Embeddings represent a token in a d-dimensional space where tokens with similar meaning will be closer to each other. But the embeddings do not encode the relative position of tokens in a sentence. So after adding the positional encoding, tokens will be closer to each other based on the similarity of their meaning and their position in the sentence, in the d-dimensional space.

The formula for calculating the positional encoding is as follows:

$$\Large{PE_{(pos, 2i)} = \sin(pos / 10000^{2i / d_{model} })} $$
$$\Large{PE_{(pos, 2i+1)} = \cos(pos / 10000^{2i / d_{model} })} $$
def get_angles(pos, i, d_model):
  angle_rates = 1 / np.power(10000, (2 * (i//2)) / np.float32(d_model))
  return pos * angle_rates
def positional_encoding(position, d_model):
  angle_rads = get_angles(np.arange(position)[:, np.newaxis],
                          np.arange(d_model)[np.newaxis, :],
                          d_model)

  # apply sin to even indices in the array; 2i
  angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])

  # apply cos to odd indices in the array; 2i+1
  angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])

  pos_encoding = angle_rads[np.newaxis, ...]

  return tf.cast(pos_encoding, dtype=tf.float32)
n, d = 2048, 512
pos_encoding = positional_encoding(n, d)
print(pos_encoding.shape)
pos_encoding = pos_encoding[0]

# Juggle the dimensions for the plot
pos_encoding = tf.reshape(pos_encoding, (n, d//2, 2))
pos_encoding = tf.transpose(pos_encoding, (2, 1, 0))
pos_encoding = tf.reshape(pos_encoding, (d, n))

plt.pcolormesh(pos_encoding, cmap='RdBu')
plt.ylabel('Depth')
plt.xlabel('Position')
plt.colorbar()
plt.show()
(1, 2048, 512)

png

Masking

Mask all the pad tokens in the batch of sequence. It ensures that the model does not treat padding as the input. The mask indicates where pad value 0 is present: it outputs a 1 at those locations, and a 0 otherwise.

def create_padding_mask(seq):
  seq = tf.cast(tf.math.equal(seq, 0), tf.float32)

  # add extra dimensions to add the padding
  # to the attention logits.
  return seq[:, tf.newaxis, tf.newaxis, :]  # (batch_size, 1, 1, seq_len)
x = tf.constant([[7, 6, 0, 0, 1], [1, 2, 3, 0, 0], [0, 0, 0, 4, 5]])
create_padding_mask(x)
<tf.Tensor: shape=(3, 1, 1, 5), dtype=float32, numpy=
array([[[[0., 0., 1., 1., 0.]]],


       [[[0., 0., 0., 1., 1.]]],


       [[[1., 1., 1., 0., 0.]]]], dtype=float32)>

The look-ahead mask is used to mask the future tokens in a sequence. In other words, the mask indicates which entries should not be used.

This means that to predict the third token, only the first and second token will be used. Similarly to predict the fourth token, only the first, second and the third tokens will be used and so on.

def create_look_ahead_mask(size):
  mask = 1 - tf.linalg.band_part(tf.ones((size, size)), -1, 0)
  return mask  # (seq_len, seq_len)
x = tf.random.uniform((1, 3))
temp = create_look_ahead_mask(x.shape[1])
temp
<tf.Tensor: shape=(3, 3), dtype=float32, numpy=
array([[0., 1., 1.],
       [0., 0., 1.],
       [0., 0., 0.]], dtype=float32)>

Scaled dot product attention

scaled_dot_product_attention

The attention function used by the transformer takes three inputs: Q (query), K (key), V (value). The equation used to calculate the attention weights is:

$$\Large{Attention(Q, K, V) = softmax_k\left(\frac{QK^T}{\sqrt{d_k} }\right) V} $$

The dot-product attention is scaled by a factor of square root of the depth. This is done because for large values of depth, the dot product grows large in magnitude pushing the softmax function where it has small gradients resulting in a very hard softmax.

For example, consider that Q and K have a mean of 0 and variance of 1. Their matrix multiplication will have a mean of 0 and variance of dk. So the square root of dk is used for scaling, so you get a consistent variance regardless of the value of dk. If the variance is too low the output may be too flat to optimize effectively. If the variance is too high the softmax may saturate at initialization making it difficult to learn.

The mask is multiplied with -1e9 (close to negative infinity). This is done because the mask is summed with the scaled matrix multiplication of Q and K and is applied immediately before a softmax. The goal is to zero out these cells, and large negative inputs to softmax are near zero in the output.

def scaled_dot_product_attention(q, k, v, mask):
  """Calculate the attention weights.
  q, k, v must have matching leading dimensions.
  k, v must have matching penultimate dimension, i.e.: seq_len_k = seq_len_v.
  The mask has different shapes depending on its type(padding or look ahead)
  but it must be broadcastable for addition.

  Args:
    q: query shape == (..., seq_len_q, depth)
    k: key shape == (..., seq_len_k, depth)
    v: value shape == (..., seq_len_v, depth_v)
    mask: Float tensor with shape broadcastable
          to (..., seq_len_q, seq_len_k). Defaults to None.

  Returns:
    output, attention_weights
  """

  matmul_qk = tf.matmul(q, k, transpose_b=True)  # (..., seq_len_q, seq_len_k)

  # scale matmul_qk
  dk = tf.cast(tf.shape(k)[-1], tf.float32)
  scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)

  # add the mask to the scaled tensor.
  if mask is not None:
    scaled_attention_logits += (mask * -1e9)

  # softmax is normalized on the last axis (seq_len_k) so that the scores
  # add up to 1.
  attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)  # (..., seq_len_q, seq_len_k)

  output = tf.matmul(attention_weights, v)  # (..., seq_len_q, depth_v)

  return output, attention_weights

As the softmax normalization is done on K, its values decide the amount of importance given to Q.

The output represents the multiplication of the attention weights and the V (value) vector. This ensures that the tokens you want to focus on are kept as-is and the irrelevant tokens are flushed out.

def print_out(q, k, v):
  temp_out, temp_attn = scaled_dot_product_attention(
      q, k, v, None)
  print('Attention weights are:')
  print(temp_attn)
  print('Output is:')
  print(temp_out)
np.set_printoptions(suppress=True)

temp_k = tf.constant([[10, 0, 0],
                      [0, 10, 0],
                      [0, 0, 10],
                      [0, 0, 10]], dtype=tf.float32)  # (4, 3)

temp_v = tf.constant([[1, 0],
                      [10, 0],
                      [100, 5],
                      [1000, 6]], dtype=tf.float32)  # (4, 2)

# This `query` aligns with the second `key`,
# so the second `value` is returned.
temp_q = tf.constant([[0, 10, 0]], dtype=tf.float32)  # (1, 3)
print_out(temp_q, temp_k, temp_v)
Attention weights are:
tf.Tensor([[0. 1. 0. 0.]], shape=(1, 4), dtype=float32)
Output is:
tf.Tensor([[10.  0.]], shape=(1, 2), dtype=float32)
# This query aligns with a repeated key (third and fourth),
# so all associated values get averaged.
temp_q = tf.constant([[0, 0, 10]], dtype=tf.float32)  # (1, 3)
print_out(temp_q, temp_k, temp_v)
Attention weights are:
tf.Tensor([[0.  0.  0.5 0.5]], shape=(1, 4), dtype=float32)
Output is:
tf.Tensor([[550.    5.5]], shape=(1, 2), dtype=float32)
# This query aligns equally with the first and second key,
# so their values get averaged.
temp_q = tf.constant([[10, 10, 0]], dtype=tf.float32)  # (1, 3)
print_out(temp_q, temp_k, temp_v)
Attention weights are:
tf.Tensor([[0.5 0.5 0.  0. ]], shape=(1, 4), dtype=float32)
Output is:
tf.Tensor([[5.5 0. ]], shape=(1, 2), dtype=float32)

Pass all the queries together.

temp_q = tf.constant([[0, 0, 10],
                      [0, 10, 0],
                      [10, 10, 0]], dtype=tf.float32)  # (3, 3)
print_out(temp_q, temp_k, temp_v)
Attention weights are:
tf.Tensor(
[[0.  0.  0.5 0.5]
 [0.  1.  0.  0. ]
 [0.5 0.5 0.  0. ]], shape=(3, 4), dtype=float32)
Output is:
tf.Tensor(
[[550.    5.5]
 [ 10.    0. ]
 [  5.5   0. ]], shape=(3, 2), dtype=float32)

Multi-head attention

multi-head attention

Multi-head attention consists of four parts:

  • Linear layers.
  • Scaled dot-product attention.
  • Final linear layer.

Each multi-head attention block gets three inputs; Q (query), K (key), V (value). These are put through linear (Dense) layers before the multi-head attention function.

In the diagram above (K,Q,V) are passed through sepearte linear (Dense) layers for each attention head. For simplicity/efficiency the code below implements this using a single dense layer with num_heads times as many outputs. The output is rearranged to a shape of (batch, num_heads, ...) before applying the attention function.

The scaled_dot_product_attention function defined above is applied in a single call, broadcasted for efficiency. An appropriate mask must be used in the attention step. The attention output for each head is then concatenated (using tf.transpose, and tf.reshape) and put through a final Dense layer.

Instead of one single attention head, Q, K, and V are split into multiple heads because it allows the model to jointly attend to information from different representation subspaces at different positions. After the split each head has a reduced dimensionality, so the total computation cost is the same as a single head attention with full dimensionality.

class MultiHeadAttention(tf.keras.layers.Layer):
  def __init__(self, d_model, num_heads):
    super(MultiHeadAttention, self).__init__()
    self.num_heads = num_heads
    self.d_model = d_model

    assert d_model % self.num_heads == 0

    self.depth = d_model // self.num_heads

    self.wq = tf.keras.layers.Dense(d_model)
    self.wk = tf.keras.layers.Dense(d_model)
    self.wv = tf.keras.layers.Dense(d_model)

    self.dense = tf.keras.layers.Dense(d_model)

  def split_heads(self, x, batch_size):
    """Split the last dimension into (num_heads, depth).
    Transpose the result such that the shape is (batch_size, num_heads, seq_len, depth)
    """
    x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))
    return tf.transpose(x, perm=[0, 2, 1, 3])

  def call(self, v, k, q, mask):
    batch_size = tf.shape(q)[0]

    q = self.wq(q)  # (batch_size, seq_len, d_model)
    k = self.wk(k)  # (batch_size, seq_len, d_model)
    v = self.wv(v)  # (batch_size, seq_len, d_model)

    q = self.split_heads(q, batch_size)  # (batch_size, num_heads, seq_len_q, depth)
    k = self.split_heads(k, batch_size)  # (batch_size, num_heads, seq_len_k, depth)
    v = self.split_heads(v, batch_size)  # (batch_size, num_heads, seq_len_v, depth)

    # scaled_attention.shape == (batch_size, num_heads, seq_len_q, depth)
    # attention_weights.shape == (batch_size, num_heads, seq_len_q, seq_len_k)
    scaled_attention, attention_weights = scaled_dot_product_attention(
        q, k, v, mask)

    scaled_attention = tf.transpose(scaled_attention, perm=[0, 2, 1, 3])  # (batch_size, seq_len_q, num_heads, depth)

    concat_attention = tf.reshape(scaled_attention,
                                  (batch_size, -1, self.d_model))  # (batch_size, seq_len_q, d_model)

    output = self.dense(concat_attention)  # (batch_size, seq_len_q, d_model)

    return output, attention_weights

Create a MultiHeadAttention layer to try out. At each location in the sequence, y, the MultiHeadAttention runs all 8 attention heads across all other locations in the sequence, returning a new vector of the same length at each location.

temp_mha = MultiHeadAttention(d_model=512, num_heads=8)
y = tf.random.uniform((1, 60, 512))  # (batch_size, encoder_sequence, d_model)
out, attn = temp_mha(y, k=y, q=y, mask=None)
out.shape, attn.shape
(TensorShape([1, 60, 512]), TensorShape([1, 8, 60, 60]))

Point wise feed forward network

Point wise feed forward network consists of two fully-connected layers with a ReLU activation in between.

def point_wise_feed_forward_network(d_model, dff):
  return tf.keras.Sequential([
      tf.keras.layers.Dense(dff, activation='relu'),  # (batch_size, seq_len, dff)
      tf.keras.layers.Dense(d_model)  # (batch_size, seq_len, d_model)
  ])
sample_ffn = point_wise_feed_forward_network(512, 2048)
sample_ffn(tf.random.uniform((64, 50, 512))).shape
TensorShape([64, 50, 512])

Encoder and decoder

transformer

The transformer model follows the same general pattern as a standard sequence to sequence with attention model.

  • The input sentence is passed through N encoder layers that generates an output for each token in the sequence.
  • The decoder attends to the encoder's output and its own input (self-attention) to predict the next word.

Encoder layer

Each encoder layer consists of sublayers:

  1. Multi-head attention (with padding mask)
  2. Point wise feed forward networks.

Each of these sublayers has a residual connection around it followed by a layer normalization. Residual connections help in avoiding the vanishing gradient problem in deep networks.

The output of each sublayer is LayerNorm(x + Sublayer(x)). The normalization is done on the d_model (last) axis. There are N encoder layers in the transformer.

class EncoderLayer(tf.keras.layers.Layer):
  def __init__(self, d_model, num_heads, dff, rate=0.1):
    super(EncoderLayer, self).__init__()

    self.mha = MultiHeadAttention(d_model, num_heads)
    self.ffn = point_wise_feed_forward_network(d_model, dff)

    self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
    self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)

    self.dropout1 = tf.keras.layers.Dropout(rate)
    self.dropout2 = tf.keras.layers.Dropout(rate)

  def call(self, x, training, mask):

    attn_output, _ = self.mha(x, x, x, mask)  # (batch_size, input_seq_len, d_model)
    attn_output = self.dropout1(attn_output, training=training)
    out1 = self.layernorm1(x + attn_output)  # (batch_size, input_seq_len, d_model)

    ffn_output = self.ffn(out1)  # (batch_size, input_seq_len, d_model)
    ffn_output = self.dropout2(ffn_output, training=training)
    out2 = self.layernorm2(out1 + ffn_output)  # (batch_size, input_seq_len, d_model)

    return out2
sample_encoder_layer = EncoderLayer(512, 8, 2048)

sample_encoder_layer_output = sample_encoder_layer(
    tf.random.uniform((64, 43, 512)), False, None)

sample_encoder_layer_output.shape  # (batch_size, input_seq_len, d_model)
TensorShape([64, 43, 512])

Decoder layer

Each decoder layer consists of sublayers:

  1. Masked multi-head attention (with look ahead mask and padding mask)
  2. Multi-head attention (with padding mask). V (value) and K (key) receive the encoder output as inputs. Q (query) receives the output from the masked multi-head attention sublayer.
  3. Point wise feed forward networks

Each of these sublayers has a residual connection around it followed by a layer normalization. The output of each sublayer is LayerNorm(x + Sublayer(x)). The normalization is done on the d_model (last) axis.

There are N decoder layers in the transformer.

As Q receives the output from decoder's first attention block, and K receives the encoder output, the attention weights represent the importance given to the decoder's input based on the encoder's output. In other words, the decoder predicts the next token by looking at the encoder output and self-attending to its own output. See the demonstration above in the scaled dot product attention section.

class DecoderLayer(tf.keras.layers.Layer):
  def __init__(self, d_model, num_heads, dff, rate=0.1):
    super(DecoderLayer, self).__init__()

    self.mha1 = MultiHeadAttention(d_model, num_heads)
    self.mha2 = MultiHeadAttention(d_model, num_heads)

    self.ffn = point_wise_feed_forward_network(d_model, dff)

    self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
    self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
    self.layernorm3 = tf.keras.layers.LayerNormalization(epsilon=1e-6)

    self.dropout1 = tf.keras.layers.Dropout(rate)
    self.dropout2 = tf.keras.layers.Dropout(rate)
    self.dropout3 = tf.keras.layers.Dropout(rate)

  def call(self, x, enc_output, training,
           look_ahead_mask, padding_mask):
    # enc_output.shape == (batch_size, input_seq_len, d_model)

    attn1, attn_weights_block1 = self.mha1(x, x, x, look_ahead_mask)  # (batch_size, target_seq_len, d_model)
    attn1 = self.dropout1(attn1, training=training)
    out1 = self.layernorm1(attn1 + x)

    attn2, attn_weights_block2 = self.mha2(
        enc_output, enc_output, out1, padding_mask)  # (batch_size, target_seq_len, d_model)
    attn2 = self.dropout2(attn2, training=training)
    out2 = self.layernorm2(attn2 + out1)  # (batch_size, target_seq_len, d_model)

    ffn_output = self.ffn(out2)  # (batch_size, target_seq_len, d_model)
    ffn_output = self.dropout3(ffn_output, training=training)
    out3 = self.layernorm3(ffn_output + out2)  # (batch_size, target_seq_len, d_model)

    return out3, attn_weights_block1, attn_weights_block2
sample_decoder_layer = DecoderLayer(512, 8, 2048)

sample_decoder_layer_output, _, _ = sample_decoder_layer(
    tf.random.uniform((64, 50, 512)), sample_encoder_layer_output,
    False, None, None)

sample_decoder_layer_output.shape  # (batch_size, target_seq_len, d_model)
TensorShape([64, 50, 512])

Encoder

The Encoder consists of:

  1. Input Embedding
  2. Positional Encoding
  3. N encoder layers

The input is put through an embedding which is summed with the positional encoding. The output of this summation is the input to the encoder layers. The output of the encoder is the input to the decoder.

class Encoder(tf.keras.layers.Layer):
  def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size,
               maximum_position_encoding, rate=0.1):
    super(Encoder, self).__init__()

    self.d_model = d_model
    self.num_layers = num_layers

    self.embedding = tf.keras.layers.Embedding(input_vocab_size, d_model)
    self.pos_encoding = positional_encoding(maximum_position_encoding,
                                            self.d_model)

    self.enc_layers = [EncoderLayer(d_model, num_heads, dff, rate)
                       for _ in range(num_layers)]

    self.dropout = tf.keras.layers.Dropout(rate)

  def call(self, x, training, mask):

    seq_len = tf.shape(x)[1]

    # adding embedding and position encoding.
    x = self.embedding(x)  # (batch_size, input_seq_len, d_model)
    x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
    x += self.pos_encoding[:, :seq_len, :]

    x = self.dropout(x, training=training)

    for i in range(self.num_layers):
      x = self.enc_layers[i](x, training, mask)

    return x  # (batch_size, input_seq_len, d_model)
sample_encoder = Encoder(num_layers=2, d_model=512, num_heads=8,
                         dff=2048, input_vocab_size=8500,
                         maximum_position_encoding=10000)
temp_input = tf.random.uniform((64, 62), dtype=tf.int64, minval=0, maxval=200)

sample_encoder_output = sample_encoder(temp_input, training=False, mask=None)

print(sample_encoder_output.shape)  # (batch_size, input_seq_len, d_model)
(64, 62, 512)

Decoder

The Decoder consists of:

  1. Output Embedding
  2. Positional Encoding
  3. N decoder layers

The target is put through an embedding which is summed with the positional encoding. The output of this summation is the input to the decoder layers. The output of the decoder is the input to the final linear layer.

class Decoder(tf.keras.layers.Layer):
  def __init__(self, num_layers, d_model, num_heads, dff, target_vocab_size,
               maximum_position_encoding, rate=0.1):
    super(Decoder, self).__init__()

    self.d_model = d_model
    self.num_layers = num_layers

    self.embedding = tf.keras.layers.Embedding(target_vocab_size, d_model)
    self.pos_encoding = positional_encoding(maximum_position_encoding, d_model)

    self.dec_layers = [DecoderLayer(d_model, num_heads, dff, rate)
                       for _ in range(num_layers)]
    self.dropout = tf.keras.layers.Dropout(rate)

  def call(self, x, enc_output, training,
           look_ahead_mask, padding_mask):

    seq_len = tf.shape(x)[1]
    attention_weights = {}

    x = self.embedding(x)  # (batch_size, target_seq_len, d_model)
    x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
    x += self.pos_encoding[:, :seq_len, :]

    x = self.dropout(x, training=training)

    for i in range(self.num_layers):
      x, block1, block2 = self.dec_layers[i](x, enc_output, training,
                                             look_ahead_mask, padding_mask)

      attention_weights[f'decoder_layer{i+1}_block1'] = block1
      attention_weights[f'decoder_layer{i+1}_block2'] = block2

    # x.shape == (batch_size, target_seq_len, d_model)
    return x, attention_weights
sample_decoder = Decoder(num_layers=2, d_model=512, num_heads=8,
                         dff=2048, target_vocab_size=8000,
                         maximum_position_encoding=5000)
temp_input = tf.random.uniform((64, 26), dtype=tf.int64, minval=0, maxval=200)

output, attn = sample_decoder(temp_input,
                              enc_output=sample_encoder_output,
                              training=False,
                              look_ahead_mask=None,
                              padding_mask=None)

output.shape, attn['decoder_layer2_block2'].shape
(TensorShape([64, 26, 512]), TensorShape([64, 8, 26, 62]))

Create the Transformer

Transformer consists of the encoder, decoder and a final linear layer. The output of the decoder is the input to the linear layer and its output is returned.

class Transformer(tf.keras.Model):
  def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size,
               target_vocab_size, pe_input, pe_target, rate=0.1):
    super().__init__()
    self.encoder = Encoder(num_layers, d_model, num_heads, dff,
                             input_vocab_size, pe_input, rate)

    self.decoder = Decoder(num_layers, d_model, num_heads, dff,
                           target_vocab_size, pe_target, rate)

    self.final_layer = tf.keras.layers.Dense(target_vocab_size)

  def call(self, inputs, training):
    # Keras models prefer if you pass all your inputs in the first argument
    inp, tar = inputs

    enc_padding_mask, look_ahead_mask, dec_padding_mask = self.create_masks(inp, tar)

    enc_output = self.encoder(inp, training, enc_padding_mask)  # (batch_size, inp_seq_len, d_model)

    # dec_output.shape == (batch_size, tar_seq_len, d_model)
    dec_output, attention_weights = self.decoder(
        tar, enc_output, training, look_ahead_mask, dec_padding_mask)

    final_output = self.final_layer(dec_output)  # (batch_size, tar_seq_len, target_vocab_size)

    return final_output, attention_weights

  def create_masks(self, inp, tar):
    # Encoder padding mask
    enc_padding_mask = create_padding_mask(inp)

    # Used in the 2nd attention block in the decoder.
    # This padding mask is used to mask the encoder outputs.
    dec_padding_mask = create_padding_mask(inp)

    # Used in the 1st attention block in the decoder.
    # It is used to pad and mask future tokens in the input received by
    # the decoder.
    look_ahead_mask = create_look_ahead_mask(tf.shape(tar)[1])
    dec_target_padding_mask = create_padding_mask(tar)
    look_ahead_mask = tf.maximum(dec_target_padding_mask, look_ahead_mask)

    return enc_padding_mask, look_ahead_mask, dec_padding_mask
sample_transformer = Transformer(
    num_layers=2, d_model=512, num_heads=8, dff=2048,
    input_vocab_size=8500, target_vocab_size=8000,
    pe_input=10000, pe_target=6000)

temp_input = tf.random.uniform((64, 38), dtype=tf.int64, minval=0, maxval=200)
temp_target = tf.random.uniform((64, 36), dtype=tf.int64, minval=0, maxval=200)

fn_out, _ = sample_transformer([temp_input, temp_target], training=False)

fn_out.shape  # (batch_size, tar_seq_len, target_vocab_size)
TensorShape([64, 36, 8000])

Set hyperparameters

To keep this example small and relatively fast, the values for num_layers, d_model, dff have been reduced.

The base model described in the paper used: num_layers=6, d_model=512, dff=2048.

num_layers = 4
d_model = 128
dff = 512
num_heads = 8
dropout_rate = 0.1

Optimizer

Use the Adam optimizer with a custom learning rate scheduler according to the formula in the paper.

$$\Large{lrate = d_{model}^{-0.5} * \min(step{\_}num^{-0.5}, step{\_}num \cdot warmup{\_}steps^{-1.5})}$$
class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):
  def __init__(self, d_model, warmup_steps=4000):
    super(CustomSchedule, self).__init__()

    self.d_model = d_model
    self.d_model = tf.cast(self.d_model, tf.float32)

    self.warmup_steps = warmup_steps

  def __call__(self, step):
    arg1 = tf.math.rsqrt(step)
    arg2 = step * (self.warmup_steps ** -1.5)

    return tf.math.rsqrt(self.d_model) * tf.math.minimum(arg1, arg2)
learning_rate = CustomSchedule(d_model)

optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=0.9, beta_2=0.98,
                                     epsilon=1e-9)
temp_learning_rate_schedule = CustomSchedule(d_model)

plt.plot(temp_learning_rate_schedule(tf.range(40000, dtype=tf.float32)))
plt.ylabel("Learning Rate")
plt.xlabel("Train Step")
Text(0.5, 0, 'Train Step')

png

Loss and metrics

Since the target sequences are padded, it is important to apply a padding mask when calculating the loss.

loss_object = tf.keras.losses.SparseCategoricalCrossentropy(
    from_logits=True, reduction='none')
def loss_function(real, pred):
  mask = tf.math.logical_not(tf.math.equal(real, 0))
  loss_ = loss_object(real, pred)

  mask = tf.cast(mask, dtype=loss_.dtype)
  loss_ *= mask

  return tf.reduce_sum(loss_)/tf.reduce_sum(mask)


def accuracy_function(real, pred):
  accuracies = tf.equal(real, tf.argmax(pred, axis=2))

  mask = tf.math.logical_not(tf.math.equal(real, 0))
  accuracies = tf.math.logical_and(mask, accuracies)

  accuracies = tf.cast(accuracies, dtype=tf.float32)
  mask = tf.cast(mask, dtype=tf.float32)
  return tf.reduce_sum(accuracies)/tf.reduce_sum(mask)
train_loss = tf.keras.metrics.Mean(name='train_loss')
train_accuracy = tf.keras.metrics.Mean(name='train_accuracy')

Training and checkpointing

transformer = Transformer(
    num_layers=num_layers,
    d_model=d_model,
    num_heads=num_heads,
    dff=dff,
    input_vocab_size=tokenizers.pt.get_vocab_size().numpy(),
    target_vocab_size=tokenizers.en.get_vocab_size().numpy(),
    pe_input=1000,
    pe_target=1000,
    rate=dropout_rate)

Create the checkpoint path and the checkpoint manager. This will be used to save checkpoints every n epochs.

checkpoint_path = "./checkpoints/train"

ckpt = tf.train.Checkpoint(transformer=transformer,
                           optimizer=optimizer)

ckpt_manager = tf.train.CheckpointManager(ckpt, checkpoint_path, max_to_keep=5)

# if a checkpoint exists, restore the latest checkpoint.
if ckpt_manager.latest_checkpoint:
  ckpt.restore(ckpt_manager.latest_checkpoint)
  print('Latest checkpoint restored!!')

The target is divided into tar_inp and tar_real. tar_inp is passed as an input to the decoder. tar_real is that same input shifted by 1: At each location in tar_input, tar_real contains the next token that should be predicted.

For example, sentence = "SOS A lion in the jungle is sleeping EOS"

tar_inp = "SOS A lion in the jungle is sleeping"

tar_real = "A lion in the jungle is sleeping EOS"

The transformer is an auto-regressive model: it makes predictions one part at a time, and uses its output so far to decide what to do next.

During training this example uses teacher-forcing (like in the text generation tutorial). Teacher forcing is passing the true output to the next time step regardless of what the model predicts at the current time step.

As the transformer predicts each token, self-attention allows it to look at the previous tokens in the input sequence to better predict the next token.

To prevent the model from peeking at the expected output the model uses a look-ahead mask.

EPOCHS = 20
# The @tf.function trace-compiles train_step into a TF graph for faster
# execution. The function specializes to the precise shape of the argument
# tensors. To avoid re-tracing due to the variable sequence lengths or variable
# batch sizes (the last batch is smaller), use input_signature to specify
# more generic shapes.

train_step_signature = [
    tf.TensorSpec(shape=(None, None), dtype=tf.int64),
    tf.TensorSpec(shape=(None, None), dtype=tf.int64),
]


@tf.function(input_signature=train_step_signature)
def train_step(inp, tar):
  tar_inp = tar[:, :-1]
  tar_real = tar[:, 1:]

  with tf.GradientTape() as tape:
    predictions, _ = transformer([inp, tar_inp],
                                 training = True)
    loss = loss_function(tar_real, predictions)

  gradients = tape.gradient(loss, transformer.trainable_variables)
  optimizer.apply_gradients(zip(gradients, transformer.trainable_variables))

  train_loss(loss)
  train_accuracy(accuracy_function(tar_real, predictions))

Portuguese is used as the input language and English is the target language.

for epoch in range(EPOCHS):
  start = time.time()

  train_loss.reset_states()
  train_accuracy.reset_states()

  # inp -> portuguese, tar -> english
  for (batch, (inp, tar)) in enumerate(train_batches):
    train_step(inp, tar)

    if batch % 50 == 0:
      print(f'Epoch {epoch + 1} Batch {batch} Loss {train_loss.result():.4f} Accuracy {train_accuracy.result():.4f}')

  if (epoch + 1) % 5 == 0:
    ckpt_save_path = ckpt_manager.save()
    print(f'Saving checkpoint for epoch {epoch+1} at {ckpt_save_path}')

  print(f'Epoch {epoch + 1} Loss {train_loss.result():.4f} Accuracy {train_accuracy.result():.4f}')

  print(f'Time taken for 1 epoch: {time.time() - start:.2f} secs\n')
Epoch 1 Batch 0 Loss 8.8605 Accuracy 0.0007
Epoch 1 Batch 50 Loss 8.7995 Accuracy 0.0115
Epoch 1 Batch 100 Loss 8.6947 Accuracy 0.0286
Epoch 1 Batch 150 Loss 8.5783 Accuracy 0.0345
Epoch 1 Batch 200 Loss 8.4360 Accuracy 0.0377
Epoch 1 Batch 250 Loss 8.2655 Accuracy 0.0401
Epoch 1 Batch 300 Loss 8.0774 Accuracy 0.0454
Epoch 1 Batch 350 Loss 7.8808 Accuracy 0.0523
Epoch 1 Batch 400 Loss 7.6919 Accuracy 0.0598
Epoch 1 Batch 450 Loss 7.5265 Accuracy 0.0672
Epoch 1 Batch 500 Loss 7.3797 Accuracy 0.0743
Epoch 1 Batch 550 Loss 7.2445 Accuracy 0.0818
Epoch 1 Batch 600 Loss 7.1217 Accuracy 0.0895
Epoch 1 Batch 650 Loss 7.0085 Accuracy 0.0966
Epoch 1 Batch 700 Loss 6.9018 Accuracy 0.1034
Epoch 1 Batch 750 Loss 6.8018 Accuracy 0.1096
Epoch 1 Batch 800 Loss 6.7103 Accuracy 0.1153
Epoch 1 Loss 6.6942 Accuracy 0.1164
Time taken for 1 epoch: 55.30 secs

Epoch 2 Batch 0 Loss 5.3761 Accuracy 0.2022
Epoch 2 Batch 50 Loss 5.2412 Accuracy 0.2111
Epoch 2 Batch 100 Loss 5.2087 Accuracy 0.2148
Epoch 2 Batch 150 Loss 5.1818 Accuracy 0.2177
Epoch 2 Batch 200 Loss 5.1579 Accuracy 0.2198
Epoch 2 Batch 250 Loss 5.1306 Accuracy 0.2229
Epoch 2 Batch 300 Loss 5.1053 Accuracy 0.2255
Epoch 2 Batch 350 Loss 5.0841 Accuracy 0.2276
Epoch 2 Batch 400 Loss 5.0599 Accuracy 0.2298
Epoch 2 Batch 450 Loss 5.0396 Accuracy 0.2319
Epoch 2 Batch 500 Loss 5.0174 Accuracy 0.2340
Epoch 2 Batch 550 Loss 4.9966 Accuracy 0.2359
Epoch 2 Batch 600 Loss 4.9747 Accuracy 0.2379
Epoch 2 Batch 650 Loss 4.9591 Accuracy 0.2394
Epoch 2 Batch 700 Loss 4.9417 Accuracy 0.2410
Epoch 2 Batch 750 Loss 4.9244 Accuracy 0.2426
Epoch 2 Batch 800 Loss 4.9069 Accuracy 0.2441
Epoch 2 Loss 4.9038 Accuracy 0.2444
Time taken for 1 epoch: 44.39 secs

Epoch 3 Batch 0 Loss 4.7541 Accuracy 0.2551
Epoch 3 Batch 50 Loss 4.5856 Accuracy 0.2717
Epoch 3 Batch 100 Loss 4.5931 Accuracy 0.2712
Epoch 3 Batch 150 Loss 4.5763 Accuracy 0.2730
Epoch 3 Batch 200 Loss 4.5587 Accuracy 0.2744
Epoch 3 Batch 250 Loss 4.5486 Accuracy 0.2749
Epoch 3 Batch 300 Loss 4.5353 Accuracy 0.2766
Epoch 3 Batch 350 Loss 4.5217 Accuracy 0.2781
Epoch 3 Batch 400 Loss 4.5050 Accuracy 0.2802
Epoch 3 Batch 450 Loss 4.4907 Accuracy 0.2821
Epoch 3 Batch 500 Loss 4.4749 Accuracy 0.2841
Epoch 3 Batch 550 Loss 4.4612 Accuracy 0.2856
Epoch 3 Batch 600 Loss 4.4462 Accuracy 0.2874
Epoch 3 Batch 650 Loss 4.4284 Accuracy 0.2894
Epoch 3 Batch 700 Loss 4.4148 Accuracy 0.2912
Epoch 3 Batch 750 Loss 4.3985 Accuracy 0.2932
Epoch 3 Batch 800 Loss 4.3822 Accuracy 0.2951
Epoch 3 Loss 4.3786 Accuracy 0.2956
Time taken for 1 epoch: 43.81 secs

Epoch 4 Batch 0 Loss 4.1349 Accuracy 0.3099
Epoch 4 Batch 50 Loss 4.0319 Accuracy 0.3351
Epoch 4 Batch 100 Loss 3.9986 Accuracy 0.3385
Epoch 4 Batch 150 Loss 3.9924 Accuracy 0.3394
Epoch 4 Batch 200 Loss 3.9861 Accuracy 0.3402
Epoch 4 Batch 250 Loss 3.9726 Accuracy 0.3422
Epoch 4 Batch 300 Loss 3.9573 Accuracy 0.3443
Epoch 4 Batch 350 Loss 3.9411 Accuracy 0.3465
Epoch 4 Batch 400 Loss 3.9292 Accuracy 0.3484
Epoch 4 Batch 450 Loss 3.9128 Accuracy 0.3506
Epoch 4 Batch 500 Loss 3.9003 Accuracy 0.3522
Epoch 4 Batch 550 Loss 3.8855 Accuracy 0.3542
Epoch 4 Batch 600 Loss 3.8702 Accuracy 0.3563
Epoch 4 Batch 650 Loss 3.8566 Accuracy 0.3582
Epoch 4 Batch 700 Loss 3.8431 Accuracy 0.3600
Epoch 4 Batch 750 Loss 3.8278 Accuracy 0.3620
Epoch 4 Batch 800 Loss 3.8119 Accuracy 0.3641
Epoch 4 Loss 3.8098 Accuracy 0.3644
Time taken for 1 epoch: 43.84 secs

Epoch 5 Batch 0 Loss 3.4650 Accuracy 0.4146
Epoch 5 Batch 50 Loss 3.4995 Accuracy 0.4013
Epoch 5 Batch 100 Loss 3.4873 Accuracy 0.4025
Epoch 5 Batch 150 Loss 3.4809 Accuracy 0.4033
Epoch 5 Batch 200 Loss 3.4767 Accuracy 0.4036
Epoch 5 Batch 250 Loss 3.4669 Accuracy 0.4048
Epoch 5 Batch 300 Loss 3.4555 Accuracy 0.4068
Epoch 5 Batch 350 Loss 3.4474 Accuracy 0.4080
Epoch 5 Batch 400 Loss 3.4375 Accuracy 0.4096
Epoch 5 Batch 450 Loss 3.4291 Accuracy 0.4106
Epoch 5 Batch 500 Loss 3.4226 Accuracy 0.4114
Epoch 5 Batch 550 Loss 3.4110 Accuracy 0.4128
Epoch 5 Batch 600 Loss 3.4017 Accuracy 0.4137
Epoch 5 Batch 650 Loss 3.3891 Accuracy 0.4155
Epoch 5 Batch 700 Loss 3.3788 Accuracy 0.4169
Epoch 5 Batch 750 Loss 3.3691 Accuracy 0.4183
Epoch 5 Batch 800 Loss 3.3623 Accuracy 0.4192
Saving checkpoint for epoch 5 at ./checkpoints/train/ckpt-1
Epoch 5 Loss 3.3605 Accuracy 0.4194
Time taken for 1 epoch: 44.11 secs

Epoch 6 Batch 0 Loss 2.8712 Accuracy 0.4718
Epoch 6 Batch 50 Loss 3.1137 Accuracy 0.4453
Epoch 6 Batch 100 Loss 3.0904 Accuracy 0.4499
Epoch 6 Batch 150 Loss 3.0816 Accuracy 0.4517
Epoch 6 Batch 200 Loss 3.0681 Accuracy 0.4533
Epoch 6 Batch 250 Loss 3.0602 Accuracy 0.4542
Epoch 6 Batch 300 Loss 3.0545 Accuracy 0.4555
Epoch 6 Batch 350 Loss 3.0500 Accuracy 0.4564
Epoch 6 Batch 400 Loss 3.0435 Accuracy 0.4573
Epoch 6 Batch 450 Loss 3.0323 Accuracy 0.4590
Epoch 6 Batch 500 Loss 3.0229 Accuracy 0.4602
Epoch 6 Batch 550 Loss 3.0127 Accuracy 0.4617
Epoch 6 Batch 600 Loss 3.0027 Accuracy 0.4629
Epoch 6 Batch 650 Loss 2.9953 Accuracy 0.4637
Epoch 6 Batch 700 Loss 2.9859 Accuracy 0.4652
Epoch 6 Batch 750 Loss 2.9791 Accuracy 0.4663
Epoch 6 Batch 800 Loss 2.9712 Accuracy 0.4674
Epoch 6 Loss 2.9698 Accuracy 0.4675
Time taken for 1 epoch: 43.91 secs

Epoch 7 Batch 0 Loss 2.7064 Accuracy 0.4926
Epoch 7 Batch 50 Loss 2.7271 Accuracy 0.4952
Epoch 7 Batch 100 Loss 2.7132 Accuracy 0.4978
Epoch 7 Batch 150 Loss 2.6918 Accuracy 0.5018
Epoch 7 Batch 200 Loss 2.6918 Accuracy 0.5026
Epoch 7 Batch 250 Loss 2.6846 Accuracy 0.5038
Epoch 7 Batch 300 Loss 2.6766 Accuracy 0.5051
Epoch 7 Batch 350 Loss 2.6739 Accuracy 0.5053
Epoch 7 Batch 400 Loss 2.6714 Accuracy 0.5055
Epoch 7 Batch 450 Loss 2.6704 Accuracy 0.5056
Epoch 7 Batch 500 Loss 2.6673 Accuracy 0.5062
Epoch 7 Batch 550 Loss 2.6606 Accuracy 0.5071
Epoch 7 Batch 600 Loss 2.6571 Accuracy 0.5078
Epoch 7 Batch 650 Loss 2.6527 Accuracy 0.5085
Epoch 7 Batch 700 Loss 2.6496 Accuracy 0.5090
Epoch 7 Batch 750 Loss 2.6464 Accuracy 0.5094
Epoch 7 Batch 800 Loss 2.6430 Accuracy 0.5101
Epoch 7 Loss 2.6419 Accuracy 0.5102
Time taken for 1 epoch: 43.94 secs

Epoch 8 Batch 0 Loss 2.3879 Accuracy 0.5353
Epoch 8 Batch 50 Loss 2.4168 Accuracy 0.5395
Epoch 8 Batch 100 Loss 2.4016 Accuracy 0.5418
Epoch 8 Batch 150 Loss 2.4192 Accuracy 0.5387
Epoch 8 Batch 200 Loss 2.4204 Accuracy 0.5385
Epoch 8 Batch 250 Loss 2.4224 Accuracy 0.5381
Epoch 8 Batch 300 Loss 2.4164 Accuracy 0.5392
Epoch 8 Batch 350 Loss 2.4207 Accuracy 0.5386
Epoch 8 Batch 400 Loss 2.4216 Accuracy 0.5387
Epoch 8 Batch 450 Loss 2.4182 Accuracy 0.5390
Epoch 8 Batch 500 Loss 2.4157 Accuracy 0.5392
Epoch 8 Batch 550 Loss 2.4129 Accuracy 0.5398
Epoch 8 Batch 600 Loss 2.4128 Accuracy 0.5399
Epoch 8 Batch 650 Loss 2.4100 Accuracy 0.5404
Epoch 8 Batch 700 Loss 2.4082 Accuracy 0.5408
Epoch 8 Batch 750 Loss 2.4084 Accuracy 0.5407
Epoch 8 Batch 800 Loss 2.4066 Accuracy 0.5412
Epoch 8 Loss 2.4059 Accuracy 0.5413
Time taken for 1 epoch: 45.21 secs

Epoch 9 Batch 0 Loss 2.3713 Accuracy 0.5347
Epoch 9 Batch 50 Loss 2.2134 Accuracy 0.5653
Epoch 9 Batch 100 Loss 2.2173 Accuracy 0.5657
Epoch 9 Batch 150 Loss 2.2200 Accuracy 0.5654
Epoch 9 Batch 200 Loss 2.2266 Accuracy 0.5645
Epoch 9 Batch 250 Loss 2.2307 Accuracy 0.5643
Epoch 9 Batch 300 Loss 2.2351 Accuracy 0.5635
Epoch 9 Batch 350 Loss 2.2342 Accuracy 0.5637
Epoch 9 Batch 400 Loss 2.2330 Accuracy 0.5637
Epoch 9 Batch 450 Loss 2.2329 Accuracy 0.5638
Epoch 9 Batch 500 Loss 2.2330 Accuracy 0.5640
Epoch 9 Batch 550 Loss 2.2332 Accuracy 0.5640
Epoch 9 Batch 600 Loss 2.2311 Accuracy 0.5646
Epoch 9 Batch 650 Loss 2.2303 Accuracy 0.5649
Epoch 9 Batch 700 Loss 2.2313 Accuracy 0.5650
Epoch 9 Batch 750 Loss 2.2326 Accuracy 0.5649
Epoch 9 Batch 800 Loss 2.2337 Accuracy 0.5648
Epoch 9 Loss 2.2331 Accuracy 0.5649
Time taken for 1 epoch: 46.60 secs

Epoch 10 Batch 0 Loss 2.0186 Accuracy 0.5855
Epoch 10 Batch 50 Loss 2.0522 Accuracy 0.5890
Epoch 10 Batch 100 Loss 2.0522 Accuracy 0.5899
Epoch 10 Batch 150 Loss 2.0653 Accuracy 0.5877
Epoch 10 Batch 200 Loss 2.0744 Accuracy 0.5863
Epoch 10 Batch 250 Loss 2.0765 Accuracy 0.5861
Epoch 10 Batch 300 Loss 2.0789 Accuracy 0.5860
Epoch 10 Batch 350 Loss 2.0828 Accuracy 0.5855
Epoch 10 Batch 400 Loss 2.0808 Accuracy 0.5858
Epoch 10 Batch 450 Loss 2.0823 Accuracy 0.5853
Epoch 10 Batch 500 Loss 2.0821 Accuracy 0.5855
Epoch 10 Batch 550 Loss 2.0850 Accuracy 0.5853
Epoch 10 Batch 600 Loss 2.0884 Accuracy 0.5849
Epoch 10 Batch 650 Loss 2.0908 Accuracy 0.5845
Epoch 10 Batch 700 Loss 2.0902 Accuracy 0.5846
Epoch 10 Batch 750 Loss 2.0913 Accuracy 0.5845
Epoch 10 Batch 800 Loss 2.0916 Accuracy 0.5846
Saving checkpoint for epoch 10 at ./checkpoints/train/ckpt-2
Epoch 10 Loss 2.0924 Accuracy 0.5845
Time taken for 1 epoch: 45.44 secs

Epoch 11 Batch 0 Loss 2.3480 Accuracy 0.5353
Epoch 11 Batch 50 Loss 1.9758 Accuracy 0.6008
Epoch 11 Batch 100 Loss 1.9599 Accuracy 0.6030
Epoch 11 Batch 150 Loss 1.9656 Accuracy 0.6020
Epoch 11 Batch 200 Loss 1.9595 Accuracy 0.6030
Epoch 11 Batch 250 Loss 1.9623 Accuracy 0.6024
Epoch 11 Batch 300 Loss 1.9632 Accuracy 0.6020
Epoch 11 Batch 350 Loss 1.9654 Accuracy 0.6017
Epoch 11 Batch 400 Loss 1.9647 Accuracy 0.6017
Epoch 11 Batch 450 Loss 1.9634 Accuracy 0.6021
Epoch 11 Batch 500 Loss 1.9652 Accuracy 0.6019
Epoch 11 Batch 550 Loss 1.9688 Accuracy 0.6013
Epoch 11 Batch 600 Loss 1.9715 Accuracy 0.6011
Epoch 11 Batch 650 Loss 1.9724 Accuracy 0.6012
Epoch 11 Batch 700 Loss 1.9759 Accuracy 0.6009
Epoch 11 Batch 750 Loss 1.9791 Accuracy 0.6004
Epoch 11 Batch 800 Loss 1.9799 Accuracy 0.6005
Epoch 11 Loss 1.9802 Accuracy 0.6005
Time taken for 1 epoch: 45.27 secs

Epoch 12 Batch 0 Loss 2.0515 Accuracy 0.5702
Epoch 12 Batch 50 Loss 1.8557 Accuracy 0.6184
Epoch 12 Batch 100 Loss 1.8429 Accuracy 0.6198
Epoch 12 Batch 150 Loss 1.8551 Accuracy 0.6181
Epoch 12 Batch 200 Loss 1.8670 Accuracy 0.6161
Epoch 12 Batch 250 Loss 1.8688 Accuracy 0.6160
Epoch 12 Batch 300 Loss 1.8681 Accuracy 0.6162
Epoch 12 Batch 350 Loss 1.8681 Accuracy 0.6166
Epoch 12 Batch 400 Loss 1.8695 Accuracy 0.6163
Epoch 12 Batch 450 Loss 1.8692 Accuracy 0.6164
Epoch 12 Batch 500 Loss 1.8753 Accuracy 0.6152
Epoch 12 Batch 550 Loss 1.8773 Accuracy 0.6149
Epoch 12 Batch 600 Loss 1.8772 Accuracy 0.6149
Epoch 12 Batch 650 Loss 1.8806 Accuracy 0.6144
Epoch 12 Batch 700 Loss 1.8801 Accuracy 0.6146
Epoch 12 Batch 750 Loss 1.8835 Accuracy 0.6143
Epoch 12 Batch 800 Loss 1.8864 Accuracy 0.6140
Epoch 12 Loss 1.8852 Accuracy 0.6142
Time taken for 1 epoch: 45.93 secs

Epoch 13 Batch 0 Loss 1.7169 Accuracy 0.6391
Epoch 13 Batch 50 Loss 1.7965 Accuracy 0.6262
Epoch 13 Batch 100 Loss 1.7888 Accuracy 0.6273
Epoch 13 Batch 150 Loss 1.7768 Accuracy 0.6293
Epoch 13 Batch 200 Loss 1.7762 Accuracy 0.6298
Epoch 13 Batch 250 Loss 1.7812 Accuracy 0.6293
Epoch 13 Batch 300 Loss 1.7808 Accuracy 0.6293
Epoch 13 Batch 350 Loss 1.7831 Accuracy 0.6289
Epoch 13 Batch 400 Loss 1.7852 Accuracy 0.6287
Epoch 13 Batch 450 Loss 1.7881 Accuracy 0.6282
Epoch 13 Batch 500 Loss 1.7878 Accuracy 0.6282
Epoch 13 Batch 550 Loss 1.7927 Accuracy 0.6274
Epoch 13 Batch 600 Loss 1.7941 Accuracy 0.6272
Epoch 13 Batch 650 Loss 1.7962 Accuracy 0.6268
Epoch 13 Batch 700 Loss 1.7989 Accuracy 0.6263
Epoch 13 Batch 750 Loss 1.8018 Accuracy 0.6260
Epoch 13 Batch 800 Loss 1.8065 Accuracy 0.6254
Epoch 13 Loss 1.8061 Accuracy 0.6254
Time taken for 1 epoch: 45.97 secs

Epoch 14 Batch 0 Loss 1.6530 Accuracy 0.6440
Epoch 14 Batch 50 Loss 1.6887 Accuracy 0.6422
Epoch 14 Batch 100 Loss 1.6872 Accuracy 0.6432
Epoch 14 Batch 150 Loss 1.7022 Accuracy 0.6410
Epoch 14 Batch 200 Loss 1.7091 Accuracy 0.6401
Epoch 14 Batch 250 Loss 1.7053 Accuracy 0.6406
Epoch 14 Batch 300 Loss 1.7081 Accuracy 0.6403
Epoch 14 Batch 350 Loss 1.7122 Accuracy 0.6398
Epoch 14 Batch 400 Loss 1.7154 Accuracy 0.6393
Epoch 14 Batch 450 Loss 1.7225 Accuracy 0.6378
Epoch 14 Batch 500 Loss 1.7243 Accuracy 0.6376
Epoch 14 Batch 550 Loss 1.7243 Accuracy 0.6377
Epoch 14 Batch 600 Loss 1.7270 Accuracy 0.6374
Epoch 14 Batch 650 Loss 1.7305 Accuracy 0.6370
Epoch 14 Batch 700 Loss 1.7323 Accuracy 0.6368
Epoch 14 Batch 750 Loss 1.7336 Accuracy 0.6367
Epoch 14 Batch 800 Loss 1.7355 Accuracy 0.6364
Epoch 14 Loss 1.7357 Accuracy 0.6364
Time taken for 1 epoch: 45.54 secs

Epoch 15 Batch 0 Loss 1.7611 Accuracy 0.6201
Epoch 15 Batch 50 Loss 1.6252 Accuracy 0.6524
Epoch 15 Batch 100 Loss 1.6315 Accuracy 0.6511
Epoch 15 Batch 150 Loss 1.6338 Accuracy 0.6516
Epoch 15 Batch 200 Loss 1.6360 Accuracy 0.6509
Epoch 15 Batch 250 Loss 1.6393 Accuracy 0.6505
Epoch 15 Batch 300 Loss 1.6450 Accuracy 0.6496
Epoch 15 Batch 350 Loss 1.6459 Accuracy 0.6497
Epoch 15 Batch 400 Loss 1.6494 Accuracy 0.6492
Epoch 15 Batch 450 Loss 1.6542 Accuracy 0.6484
Epoch 15 Batch 500 Loss 1.6550 Accuracy 0.6482
Epoch 15 Batch 550 Loss 1.6563 Accuracy 0.6480
Epoch 15 Batch 600 Loss 1.6598 Accuracy 0.6475
Epoch 15 Batch 650 Loss 1.6633 Accuracy 0.6470
Epoch 15 Batch 700 Loss 1.6659 Accuracy 0.6466
Epoch 15 Batch 750 Loss 1.6692 Accuracy 0.6462
Epoch 15 Batch 800 Loss 1.6728 Accuracy 0.6456
Saving checkpoint for epoch 15 at ./checkpoints/train/ckpt-3
Epoch 15 Loss 1.6739 Accuracy 0.6454
Time taken for 1 epoch: 45.77 secs

Epoch 16 Batch 0 Loss 1.5074 Accuracy 0.6766
Epoch 16 Batch 50 Loss 1.5829 Accuracy 0.6601
Epoch 16 Batch 100 Loss 1.5843 Accuracy 0.6604
Epoch 16 Batch 150 Loss 1.5878 Accuracy 0.6603
Epoch 16 Batch 200 Loss 1.5837 Accuracy 0.6609
Epoch 16 Batch 250 Loss 1.5878 Accuracy 0.6598
Epoch 16 Batch 300 Loss 1.5942 Accuracy 0.6586
Epoch 16 Batch 350 Loss 1.5981 Accuracy 0.6579
Epoch 16 Batch 400 Loss 1.5998 Accuracy 0.6574
Epoch 16 Batch 450 Loss 1.6022 Accuracy 0.6571
Epoch 16 Batch 500 Loss 1.6045 Accuracy 0.6566
Epoch 16 Batch 550 Loss 1.6054 Accuracy 0.6563
Epoch 16 Batch 600 Loss 1.6077 Accuracy 0.6561
Epoch 16 Batch 650 Loss 1.6110 Accuracy 0.6556
Epoch 16 Batch 700 Loss 1.6147 Accuracy 0.6549
Epoch 16 Batch 750 Loss 1.6164 Accuracy 0.6548
Epoch 16 Batch 800 Loss 1.6197 Accuracy 0.6544
Epoch 16 Loss 1.6204 Accuracy 0.6542
Time taken for 1 epoch: 45.91 secs

Epoch 17 Batch 0 Loss 1.5031 Accuracy 0.6719
Epoch 17 Batch 50 Loss 1.4991 Accuracy 0.6722
Epoch 17 Batch 100 Loss 1.5160 Accuracy 0.6691
Epoch 17 Batch 150 Loss 1.5271 Accuracy 0.6671
Epoch 17 Batch 200 Loss 1.5291 Accuracy 0.6668
Epoch 17 Batch 250 Loss 1.5347 Accuracy 0.6661
Epoch 17 Batch 300 Loss 1.5399 Accuracy 0.6653
Epoch 17 Batch 350 Loss 1.5423 Accuracy 0.6649
Epoch 17 Batch 400 Loss 1.5464 Accuracy 0.6642
Epoch 17 Batch 450 Loss 1.5504 Accuracy 0.6636
Epoch 17 Batch 500 Loss 1.5509 Accuracy 0.6636
Epoch 17 Batch 550 Loss 1.5514 Accuracy 0.6637
Epoch 17 Batch 600 Loss 1.5539 Accuracy 0.6633
Epoch 17 Batch 650 Loss 1.5572 Accuracy 0.6629
Epoch 17 Batch 700 Loss 1.5614 Accuracy 0.6624
Epoch 17 Batch 750 Loss 1.5659 Accuracy 0.6618
Epoch 17 Batch 800 Loss 1.5691 Accuracy 0.6614
Epoch 17 Loss 1.5696 Accuracy 0.6613
Time taken for 1 epoch: 45.74 secs

Epoch 18 Batch 0 Loss 1.5180 Accuracy 0.6631
Epoch 18 Batch 50 Loss 1.4671 Accuracy 0.6788
Epoch 18 Batch 100 Loss 1.4752 Accuracy 0.6771
Epoch 18 Batch 150 Loss 1.4848 Accuracy 0.6756
Epoch 18 Batch 200 Loss 1.4952 Accuracy 0.6740
Epoch 18 Batch 250 Loss 1.4974 Accuracy 0.6734
Epoch 18 Batch 300 Loss 1.5011 Accuracy 0.6727
Epoch 18 Batch 350 Loss 1.5044 Accuracy 0.6719
Epoch 18 Batch 400 Loss 1.5045 Accuracy 0.6718
Epoch 18 Batch 450 Loss 1.5090 Accuracy 0.6710
Epoch 18 Batch 500 Loss 1.5095 Accuracy 0.6709
Epoch 18 Batch 550 Loss 1.5117 Accuracy 0.6706
Epoch 18 Batch 600 Loss 1.5171 Accuracy 0.6698
Epoch 18 Batch 650 Loss 1.5183 Accuracy 0.6697
Epoch 18 Batch 700 Loss 1.5210 Accuracy 0.6694
Epoch 18 Batch 750 Loss 1.5237 Accuracy 0.6689
Epoch 18 Batch 800 Loss 1.5261 Accuracy 0.6686
Epoch 18 Loss 1.5262 Accuracy 0.6686
Time taken for 1 epoch: 45.53 secs

Epoch 19 Batch 0 Loss 1.4657 Accuracy 0.6641
Epoch 19 Batch 50 Loss 1.4264 Accuracy 0.6826
Epoch 19 Batch 100 Loss 1.4386 Accuracy 0.6817
Epoch 19 Batch 150 Loss 1.4445 Accuracy 0.6809
Epoch 19 Batch 200 Loss 1.4401 Accuracy 0.6823
Epoch 19 Batch 250 Loss 1.4458 Accuracy 0.6813
Epoch 19 Batch 300 Loss 1.4516 Accuracy 0.6803
Epoch 19 Batch 350 Loss 1.4565 Accuracy 0.6792
Epoch 19 Batch 400 Loss 1.4593 Accuracy 0.6786
Epoch 19 Batch 450 Loss 1.4636 Accuracy 0.6780
Epoch 19 Batch 500 Loss 1.4656 Accuracy 0.6778
Epoch 19 Batch 550 Loss 1.4682 Accuracy 0.6774
Epoch 19 Batch 600 Loss 1.4698 Accuracy 0.6771
Epoch 19 Batch 650 Loss 1.4753 Accuracy 0.6763
Epoch 19 Batch 700 Loss 1.4810 Accuracy 0.6755
Epoch 19 Batch 750 Loss 1.4831 Accuracy 0.6752
Epoch 19 Batch 800 Loss 1.4865 Accuracy 0.6745
Epoch 19 Loss 1.4879 Accuracy 0.6743
Time taken for 1 epoch: 46.16 secs

Epoch 20 Batch 0 Loss 1.3513 Accuracy 0.6976
Epoch 20 Batch 50 Loss 1.3936 Accuracy 0.6905
Epoch 20 Batch 100 Loss 1.3964 Accuracy 0.6894
Epoch 20 Batch 150 Loss 1.4042 Accuracy 0.6885
Epoch 20 Batch 200 Loss 1.4112 Accuracy 0.6871
Epoch 20 Batch 250 Loss 1.4174 Accuracy 0.6856
Epoch 20 Batch 300 Loss 1.4184 Accuracy 0.6855
Epoch 20 Batch 350 Loss 1.4205 Accuracy 0.6853
Epoch 20 Batch 400 Loss 1.4240 Accuracy 0.6846
Epoch 20 Batch 450 Loss 1.4307 Accuracy 0.6834
Epoch 20 Batch 500 Loss 1.4343 Accuracy 0.6828
Epoch 20 Batch 550 Loss 1.4353 Accuracy 0.6826
Epoch 20 Batch 600 Loss 1.4384 Accuracy 0.6820
Epoch 20 Batch 650 Loss 1.4407 Accuracy 0.6818
Epoch 20 Batch 700 Loss 1.4431 Accuracy 0.6814
Epoch 20 Batch 750 Loss 1.4459 Accuracy 0.6811
Epoch 20 Batch 800 Loss 1.4484 Accuracy 0.6807
Saving checkpoint for epoch 20 at ./checkpoints/train/ckpt-4
Epoch 20 Loss 1.4495 Accuracy 0.6806
Time taken for 1 epoch: 47.55 secs

Run inference

The following steps are used for inference:

  • Encode the input sentence using the Portuguese tokenizer (tokenizers.pt). This is the encoder input.
  • The decoder input is initialized to the [START] token.
  • Calculate the padding masks and the look ahead masks.
  • The decoder then outputs the predictions by looking at the encoder output and its own output (self-attention).
  • Concatenate the predicted token to the decoder input and pass it to the decoder.
  • In this approach, the decoder predicts the next token based on the previous tokens it predicted.
class Translator(tf.Module):
  def __init__(self, tokenizers, transformer):
    self.tokenizers = tokenizers
    self.transformer = transformer

  def __call__(self, sentence, max_length=20):
    # input sentence is portuguese, hence adding the start and end token
    assert isinstance(sentence, tf.Tensor)
    if len(sentence.shape) == 0:
      sentence = sentence[tf.newaxis]

    sentence = self.tokenizers.pt.tokenize(sentence).to_tensor()

    encoder_input = sentence

    # as the target is english, the first token to the transformer should be the
    # english start token.
    start_end = self.tokenizers.en.tokenize([''])[0]
    start = start_end[0][tf.newaxis]
    end = start_end[1][tf.newaxis]

    # `tf.TensorArray` is required here (instead of a python list) so that the
    # dynamic-loop can be traced by `tf.function`.
    output_array = tf.TensorArray(dtype=tf.int64, size=0, dynamic_size=True)
    output_array = output_array.write(0, start)

    for i in tf.range(max_length):
      output = tf.transpose(output_array.stack())
      predictions, _ = self.transformer([encoder_input, output], training=False)

      # select the last token from the seq_len dimension
      predictions = predictions[:, -1:, :]  # (batch_size, 1, vocab_size)

      predicted_id = tf.argmax(predictions, axis=-1)

      # concatentate the predicted_id to the output which is given to the decoder
      # as its input.
      output_array = output_array.write(i+1, predicted_id[0])

      if predicted_id == end:
        break

    output = tf.transpose(output_array.stack())
    # output.shape (1, tokens)
    text = tokenizers.en.detokenize(output)[0]  # shape: ()

    tokens = tokenizers.en.lookup(output)[0]

    # `tf.function` prevents us from using the attention_weights that were
    # calculated on the last iteration of the loop. So recalculate them outside
    # the loop.
    _, attention_weights = self.transformer([encoder_input, output[:,:-1]], training=False)

    return text, tokens, attention_weights

Create an instance of this Translator class, and try it out a few times:

translator = Translator(tokenizers, transformer)
def print_translation(sentence, tokens, ground_truth):
  print(f'{"Input:":15s}: {sentence}')
  print(f'{"Prediction":15s}: {tokens.numpy().decode("utf-8")}')
  print(f'{"Ground truth":15s}: {ground_truth}')
sentence = "este é um problema que temos que resolver."
ground_truth = "this is a problem we have to solve ."

translated_text, translated_tokens, attention_weights = translator(
    tf.constant(sentence))
print_translation(sentence, translated_text, ground_truth)
Input:         : este é um problema que temos que resolver.
Prediction     : this is a problem that we have to solve .
Ground truth   : this is a problem we have to solve .
sentence = "os meus vizinhos ouviram sobre esta ideia."
ground_truth = "and my neighboring homes heard about this idea ."

translated_text, translated_tokens, attention_weights = translator(
    tf.constant(sentence))
print_translation(sentence, translated_text, ground_truth)
Input:         : os meus vizinhos ouviram sobre esta ideia.
Prediction     : my neighbors heard about this idea .
Ground truth   : and my neighboring homes heard about this idea .
sentence = "vou então muito rapidamente partilhar convosco algumas histórias de algumas coisas mágicas que aconteceram."
ground_truth = "so i \'ll just share with you some stories very quickly of some magical things that have happened ."

translated_text, translated_tokens, attention_weights = translator(
    tf.constant(sentence))
print_translation(sentence, translated_text, ground_truth)
Input:         : vou então muito rapidamente partilhar convosco algumas histórias de algumas coisas mágicas que aconteceram.
Prediction     : so i ' m going to share with you a few very quickly stories that happened .
Ground truth   : so i 'll just share with you some stories very quickly of some magical things that have happened .

Attention plots

The Translator class returns a dictionary of attention maps you can use to visualize the internal working of the model:

sentence = "este é o primeiro livro que eu fiz."
ground_truth = "this is the first book i've ever done."

translated_text, translated_tokens, attention_weights = translator(
    tf.constant(sentence))
print_translation(sentence, translated_text, ground_truth)
Input:         : este é o primeiro livro que eu fiz.
Prediction     : this is the first book i did .
Ground truth   : this is the first book i've ever done.
def plot_attention_head(in_tokens, translated_tokens, attention):
  # The plot is of the attention when a token was generated.
  # The model didn't generate `<START>` in the output. Skip it.
  translated_tokens = translated_tokens[1:]

  ax = plt.gca()
  ax.matshow(attention)
  ax.set_xticks(range(len(in_tokens)))
  ax.set_yticks(range(len(translated_tokens)))

  labels = [label.decode('utf-8') for label in in_tokens.numpy()]
  ax.set_xticklabels(
      labels, rotation=90)

  labels = [label.decode('utf-8') for label in translated_tokens.numpy()]
  ax.set_yticklabels(labels)
head = 0
# shape: (batch=1, num_heads, seq_len_q, seq_len_k)
attention_heads = tf.squeeze(
  attention_weights['decoder_layer4_block2'], 0)
attention = attention_heads[head]
attention.shape
TensorShape([9, 11])
in_tokens = tf.convert_to_tensor([sentence])
in_tokens = tokenizers.pt.tokenize(in_tokens).to_tensor()
in_tokens = tokenizers.pt.lookup(in_tokens)[0]
in_tokens
<tf.Tensor: shape=(11,), dtype=string, numpy=
array([b'[START]', b'este', b'e', b'o', b'primeiro', b'livro', b'que',
       b'eu', b'fiz', b'.', b'[END]'], dtype=object)>
translated_tokens
<tf.Tensor: shape=(10,), dtype=string, numpy=
array([b'[START]', b'this', b'is', b'the', b'first', b'book', b'i',
       b'did', b'.', b'[END]'], dtype=object)>
plot_attention_head(in_tokens, translated_tokens, attention)

png

def plot_attention_weights(sentence, translated_tokens, attention_heads):
  in_tokens = tf.convert_to_tensor([sentence])
  in_tokens = tokenizers.pt.tokenize(in_tokens).to_tensor()
  in_tokens = tokenizers.pt.lookup(in_tokens)[0]
  in_tokens

  fig = plt.figure(figsize=(16, 8))

  for h, head in enumerate(attention_heads):
    ax = fig.add_subplot(2, 4, h+1)

    plot_attention_head(in_tokens, translated_tokens, head)

    ax.set_xlabel(f'Head {h+1}')

  plt.tight_layout()
  plt.show()
plot_attention_weights(sentence, translated_tokens,
                       attention_weights['decoder_layer4_block2'][0])

png

The model does okay on unfamiliar words. Neither "triceratops" or "encyclopedia" are in the input dataset and the model almost learns to transliterate them, even without a shared vocabulary:

sentence = "Eu li sobre triceratops na enciclopédia."
ground_truth = "I read about triceratops in the encyclopedia."

translated_text, translated_tokens, attention_weights = translator(
    tf.constant(sentence))
print_translation(sentence, translated_text, ground_truth)

plot_attention_weights(sentence, translated_tokens,
                       attention_weights['decoder_layer4_block2'][0])
Input:         : Eu li sobre triceratops na enciclopédia.
Prediction     : i read about trijacopters in the encyclopedia .
Ground truth   : I read about triceratops in the encyclopedia.

png

Export

That inference model is working, so next you'll export it as a tf.saved_model.

To do that, wrap it in yet another tf.Module sub-class, this time with a tf.function on the __call__ method:

class ExportTranslator(tf.Module):
  def __init__(self, translator):
    self.translator = translator

  @tf.function(input_signature=[tf.TensorSpec(shape=[], dtype=tf.string)])
  def __call__(self, sentence):
    (result, 
     tokens,
     attention_weights) = self.translator(sentence, max_length=100)

    return result

In the above tf.function only the output sentence is returned. Thanks to the non-strict execution in tf.function any unnecessary values are never computed.

translator = ExportTranslator(translator)

Since the model is decoding the predictions using tf.argmax the predictions are deterministic. The original model and one reloaded from its SavedModel should give identical predictions:

translator("este é o primeiro livro que eu fiz.").numpy()
b'this is the first book i did .'
tf.saved_model.save(translator, export_dir='translator')
2021-08-25 11:27:05.726890: W tensorflow/python/util/util.cc:348] Sets are not currently considered sequences, but this may change in the future, so consider avoiding using them.
WARNING:absl:Found untraced functions such as embedding_4_layer_call_fn, embedding_4_layer_call_and_return_conditional_losses, dropout_37_layer_call_fn, dropout_37_layer_call_and_return_conditional_losses, embedding_5_layer_call_fn while saving (showing 5 of 560). These functions will not be directly callable after loading.
reloaded = tf.saved_model.load('translator')
reloaded("este é o primeiro livro que eu fiz.").numpy()
b'this is the first book i did .'

Summary

In this tutorial, you learned about positional encoding, multi-head attention, the importance of masking and how to create a transformer.

Try using a different dataset to train the transformer. You can also create the base transformer or transformer XL by changing the hyperparameters above. You can also use the layers defined here to create BERT and train state of the art models. Furthermore, you can implement beam search to get better predictions.