<p dir="auto">transformer:<a href="https://www.tensorflow.org/alpha/tutorials/text/transformer#top_of_page" rel="nofollow">https://www.tensorflow.org/alpha/tutorials/text/transformer#top_of_page</a></p>
<p dir="auto">Has anyone run this experiment, and the results of my run have not reached the official results. I posted my code and helped me find the reason.</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import tensorflow_datasets as tfds
import tensorflow as tf

import time
import numpy as np
import matplotlib.pyplot as plt

from tqdm.auto import tqdm
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

print(tf.__version__)

if tf.test.is_gpu_available():
    device = &quot;/gpu:0&quot;
else:
    device = &quot;/cpu:0&quot;


print(&quot;(1):Reading dataset and token......&quot;)
examples, metadata = tfds.load('ted_hrlr_translate/pt_to_en', with_info=True,as_supervised=True)
train_examples, val_examples = examples['train'], examples['validation']

tokenizer_en = tfds.features.text.SubwordTextEncoder.build_from_corpus(
    (en.numpy() for pt, en in train_examples), target_vocab_size=2 ** 13)

tokenizer_pt = tfds.features.text.SubwordTextEncoder.build_from_corpus(
    (pt.numpy() for pt, en in train_examples), target_vocab_size=2 ** 13)


BUFFER_SIZE = 20000
BATCH_SIZE = 64

&quot;&quot;&quot;Add a start and end token to the input and target.&quot;&quot;&quot;


def encode(lang1, lang2):
    lang1 = [tokenizer_pt.vocab_size] + tokenizer_pt.encode(
        lang1.numpy()) + [tokenizer_pt.vocab_size + 1]

    lang2 = [tokenizer_en.vocab_size] + tokenizer_en.encode(
        lang2.numpy()) + [tokenizer_en.vocab_size + 1]

    return lang1, lang2


&quot;&quot;&quot;Note: To keep this example small and relatively fast, drop examples with a length of over 40 tokens.&quot;&quot;&quot;

MAX_LENGTH = 40


def filter_max_length(x, y, max_length=MAX_LENGTH):
    return tf.logical_and(tf.size(x) &lt;= max_length,
                          tf.size(y) &lt;= max_length)


&quot;&quot;&quot;Operations inside `.map()` run in graph mode and receive a graph tensor that do not have a numpy attribute. The `tokenizer` expects a string or Unicode symbol to encode it into integers. Hence, you need to run the encoding inside a `tf.py_function`, which receives an eager tensor having a numpy attribute that contains the string value.&quot;&quot;&quot;


def tf_encode(pt, en):
    return tf.py_function(encode, [pt, en], [tf.int64, tf.int64])

print(&quot;(2):Encode and padded batch......&quot;)
train_dataset = train_examples.map(tf_encode)
train_dataset = train_dataset.filter(filter_max_length)
# cache the dataset to memory to get a speedup while reading from it.
train_dataset = train_dataset.cache()
train_dataset = train_dataset.shuffle(BUFFER_SIZE).padded_batch(
    BATCH_SIZE, padded_shapes=([-1], [-1]))
train_dataset = train_dataset.prefetch(tf.data.experimental.AUTOTUNE)

val_dataset = val_examples.map(tf_encode)
val_dataset = val_dataset.filter(filter_max_length).padded_batch(
    BATCH_SIZE, padded_shapes=([-1], [-1]))



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
    sines = np.sin(angle_rads[:, 0::2])

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

    pos_encoding = np.concatenate([sines, cosines], axis=-1)

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

    return tf.cast(pos_encoding, dtype=tf.float32)


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

    # add extra dimensions so that we can add the padding
    # to the attention logits.
    return seq[:, tf.newaxis, tf.newaxis, :]  # (batch_size, 1, 1, seq_len)



def create_look_ahead_mask(size):
    mask = 1 - tf.linalg.band_part(tf.ones((size, size)), -1, 0)
    return mask  # (seq_len, seq_len)



def scaled_dot_product_attention(q, k, v, mask):
    &quot;&quot;&quot;Calculate the attention weights.
    q, k, v must have matching leading dimensions.
    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)
      mask: Float tensor with shape broadcastable
            to (..., seq_len_q, seq_len_k). Defaults to None.

    Returns:
      output, attention_weights
    &quot;&quot;&quot;

    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_v, depth)

    return output, attention_weights

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):
        &quot;&quot;&quot;Split the last dimension into (num_heads, depth).
        Transpose the result such that the shape is (batch_size, num_heads, seq_len, depth)
        &quot;&quot;&quot;
        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_v, 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_v, num_heads, depth)

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

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

        return output, attention_weights


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)
    ])

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


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


class Encoder(tf.keras.layers.Layer):
    def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size,
                 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(input_vocab_size, 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)


class Decoder(tf.keras.layers.Layer):
    def __init__(self, num_layers, d_model, num_heads, dff, target_vocab_size,
                 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(target_vocab_size, self.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['decoder_layer{}_block1'.format(i + 1)] = block1
            attention_weights['decoder_layer{}_block2'.format(i + 1)] = block2

        # x.shape == (batch_size, target_seq_len, d_model)
        return x, attention_weights


&quot;&quot;&quot;## 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.
&quot;&quot;&quot;


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

        self.encoder = Encoder(num_layers, d_model, num_heads, dff,
                               input_vocab_size, rate)

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

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

    def call(self, inp, tar, training, enc_padding_mask,
             look_ahead_mask, dec_padding_mask):
        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


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

input_vocab_size = tokenizer_pt.vocab_size + 2
target_vocab_size = tokenizer_en.vocab_size + 2
dropout_rate = 0.1


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(&quot;Learning Rate&quot;)
plt.xlabel(&quot;Train Step&quot;)

&quot;&quot;&quot;## Loss and metrics

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

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_mean(loss_)


train_loss = tf.keras.metrics.Mean(name='train_loss')
train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')
val_loss = tf.keras.metrics.Mean(name='val_loss')
val_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='val_accuracy')

&quot;&quot;&quot;## Training and checkpointing&quot;&quot;&quot;

transformer = Transformer(num_layers, d_model, num_heads, dff,
                          input_vocab_size, target_vocab_size, dropout_rate)


def create_masks(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)
    combined_mask = tf.maximum(dec_target_padding_mask, look_ahead_mask)

    return enc_padding_mask, combined_mask, dec_padding_mask


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

checkpoint_path = &quot;./checkpoints/train&quot;

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!!')


EPOCHS = 200


train_num = len([1 for _, _ in train_dataset])

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

    enc_padding_mask, combined_mask, dec_padding_mask = create_masks(inp, tar_inp)

    with tf.GradientTape() as tape:
        predictions, _ = transformer(inp, tar_inp,
                                     True,
                                     enc_padding_mask,
                                     combined_mask,
                                     dec_padding_mask)
        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(tar_real, predictions)

@tf.function
def val_step(inp, tar):
    tar_inp = tar[:, :-1]
    tar_real = tar[:, 1:]
    enc_padding_mask, combined_mask, dec_padding_mask = create_masks(inp, tar_inp)
    predictions, _ = transformer(inp,tar_inp,
                                enc_padding_mask=enc_padding_mask,
                                look_ahead_mask=combined_mask,
                                dec_padding_mask=dec_padding_mask,
                                training=False)
    loss = loss_function(tar_real, predictions)
    ppl = tf.exp(loss)
    val_loss(ppl)
    val_accuracy(tar_real, predictions)

print(&quot;(3):Traning model......&quot;)
&quot;&quot;&quot;Portuguese is used as the input language and English is the target language.&quot;&quot;&quot;

for epoch in range(EPOCHS):
    train_loss.reset_states()
    train_accuracy.reset_states()
    val_loss.reset_states()
    val_accuracy.reset_states()

    print('Epoch {}'.format(epoch + 1))
    start = time.time()
    # inp -&gt; portuguese, tar -&gt; english
    with tqdm(total=train_num * BATCH_SIZE) as pbar:
        for inp, tar in train_dataset:
            train_step(inp, tar)
            pbar.update(BATCH_SIZE)

    for inp, tar in val_dataset:
        val_step(inp, tar)

    end = time.time()
    print('train_loss {:.4f}\ttrain_acc {:.2f}\t'
          'val_loss {:.4f}\tval_acc {:.2f}\t'
          'time {:.2f}s'.format(train_loss.result(),
                                train_accuracy.result() * 100,
                                val_loss.result(),
                                val_accuracy.result() * 100,
                                end - start,
                                ))


def evaluate(inp_sentence):
    start_token = [tokenizer_pt.vocab_size]
    end_token = [tokenizer_pt.vocab_size + 1]

    # inp sentence is portuguese, hence adding the start and end token
    inp_sentence = start_token + tokenizer_pt.encode(inp_sentence) + end_token
    encoder_input = tf.expand_dims(inp_sentence, 0)

    # as the target is english, the first word to the transformer should be the
    # english start token.
    decoder_input = [tokenizer_en.vocab_size]
    output = tf.expand_dims(decoder_input, 0)

    for i in range(MAX_LENGTH):
        enc_padding_mask, combined_mask, dec_padding_mask = create_masks(
            encoder_input, output)

        # predictions.shape == (batch_size, seq_len, vocab_size)
        predictions, attention_weights = transformer(encoder_input,
                                                     output,
                                                     False,
                                                     enc_padding_mask,
                                                     combined_mask,
                                                     dec_padding_mask)

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

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

        # return the result if the predicted_id is equal to the end token
        if tf.equal(predicted_id, tokenizer_en.vocab_size + 1):
            return tf.squeeze(output, axis=0), attention_weights

        # concatentate the predicted_id to the output which is given to the decoder
        # as its input.
        output = tf.concat([output, predicted_id], axis=-1)

    return tf.squeeze(output, axis=0), attention_weights


def plot_attention_weights(attention, sentence, result, layer):
    fig = plt.figure(figsize=(16, 8))

    sentence = tokenizer_pt.encode(sentence)

    attention = tf.squeeze(attention[layer], axis=0)

    for head in range(attention.shape[0]):
        ax = fig.add_subplot(2, 4, head + 1)

        # plot the attention weights
        ax.matshow(attention[head][:-1, :], cmap='viridis')

        fontdict = {'fontsize': 10}

        ax.set_xticks(range(len(sentence) + 2))
        ax.set_yticks(range(len(result)))

        ax.set_ylim(len(result) - 1.5, -0.5)

        ax.set_xticklabels(
            ['&lt;start&gt;'] + [tokenizer_pt.decode([i]) for i in sentence] + ['&lt;end&gt;'],
            fontdict=fontdict, rotation=90)

        ax.set_yticklabels([tokenizer_en.decode([i]) for i in result
                            if i &lt; tokenizer_en.vocab_size],
                           fontdict=fontdict)

        ax.set_xlabel('Head {}'.format(head + 1))

    plt.tight_layout()
    plt.show()


def translate(sentence, plot=''):
    result, attention_weights = evaluate(sentence)

    predicted_sentence = tokenizer_en.decode([i for i in result
                                              if i &lt; tokenizer_en.vocab_size])

    print('Input: {}'.format(sentence))
    print('Predicted translation: {}'.format(predicted_sentence))

    if plot:
        plot_attention_weights(attention_weights, sentence, result, plot)

print(&quot;(4):Evaluate model......&quot;)
translate(&quot;este é um problema que temos que resolver.&quot;)
print(&quot;Real translation: this is a problem we have to solve .&quot;)

translate(&quot;os meus vizinhos ouviram sobre esta ideia.&quot;)
print(&quot;Real translation: and my neighboring homes heard about this idea .&quot;)

translate(&quot;vou então muito rapidamente partilhar convosco algumas histórias de algumas coisas mágicas que aconteceram.&quot;)
print(
    &quot;Real translation: so i 'll just share with you some stories very quickly of some magical things that have happened .&quot;)

&quot;&quot;&quot;You can pass different layers and attention blocks of the decoder to the `plot` parameter.&quot;&quot;&quot;

translate(&quot;este é o primeiro livro que eu fiz.&quot;, plot='decoder_layer4_block2')
print(&quot;Real translation: this is the first book i've ever done.&quot;)
"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">tensorflow_datasets</span> <span class="pl-k">as</span> <span class="pl-s1">tfds</span>
<span class="pl-k">import</span> <span class="pl-s1">tensorflow</span> <span class="pl-k">as</span> <span class="pl-s1">tf</span>

<span class="pl-k">import</span> <span class="pl-s1">time</span>
<span class="pl-k">import</span> <span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">np</span>
<span class="pl-k">import</span> <span class="pl-s1">matplotlib</span>.<span class="pl-s1">pyplot</span> <span class="pl-k">as</span> <span class="pl-s1">plt</span>

<span class="pl-k">from</span> <span class="pl-s1">tqdm</span>.<span class="pl-s1">auto</span> <span class="pl-k">import</span> <span class="pl-s1">tqdm</span>
<span class="pl-k">import</span> <span class="pl-s1">os</span>
<span class="pl-s1">os</span>.<span class="pl-s1">environ</span>[<span class="pl-s">'TF_CPP_MIN_LOG_LEVEL'</span>] <span class="pl-c1">=</span> <span class="pl-s">'3'</span>

<span class="pl-en">print</span>(<span class="pl-s1">tf</span>.<span class="pl-s1">__version__</span>)

<span class="pl-k">if</span> <span class="pl-s1">tf</span>.<span class="pl-s1">test</span>.<span class="pl-en">is_gpu_available</span>():
    <span class="pl-s1">device</span> <span class="pl-c1">=</span> <span class="pl-s">"/gpu:0"</span>
<span class="pl-k">else</span>:
    <span class="pl-s1">device</span> <span class="pl-c1">=</span> <span class="pl-s">"/cpu:0"</span>


<span class="pl-en">print</span>(<span class="pl-s">"(1):Reading dataset and token......"</span>)
<span class="pl-s1">examples</span>, <span class="pl-s1">metadata</span> <span class="pl-c1">=</span> <span class="pl-s1">tfds</span>.<span class="pl-en">load</span>(<span class="pl-s">'ted_hrlr_translate/pt_to_en'</span>, <span class="pl-s1">with_info</span><span class="pl-c1">=</span><span class="pl-c1">True</span>,<span class="pl-s1">as_supervised</span><span class="pl-c1">=</span><span class="pl-c1">True</span>)
<span class="pl-s1">train_examples</span>, <span class="pl-s1">val_examples</span> <span class="pl-c1">=</span> <span class="pl-s1">examples</span>[<span class="pl-s">'train'</span>], <span class="pl-s1">examples</span>[<span class="pl-s">'validation'</span>]

<span class="pl-s1">tokenizer_en</span> <span class="pl-c1">=</span> <span class="pl-s1">tfds</span>.<span class="pl-s1">features</span>.<span class="pl-s1">text</span>.<span class="pl-v">SubwordTextEncoder</span>.<span class="pl-en">build_from_corpus</span>(
    (<span class="pl-s1">en</span>.<span class="pl-en">numpy</span>() <span class="pl-k">for</span> <span class="pl-s1">pt</span>, <span class="pl-s1">en</span> <span class="pl-c1">in</span> <span class="pl-s1">train_examples</span>), <span class="pl-s1">target_vocab_size</span><span class="pl-c1">=</span><span class="pl-c1">2</span> <span class="pl-c1">**</span> <span class="pl-c1">13</span>)

<span class="pl-s1">tokenizer_pt</span> <span class="pl-c1">=</span> <span class="pl-s1">tfds</span>.<span class="pl-s1">features</span>.<span class="pl-s1">text</span>.<span class="pl-v">SubwordTextEncoder</span>.<span class="pl-en">build_from_corpus</span>(
    (<span class="pl-s1">pt</span>.<span class="pl-en">numpy</span>() <span class="pl-k">for</span> <span class="pl-s1">pt</span>, <span class="pl-s1">en</span> <span class="pl-c1">in</span> <span class="pl-s1">train_examples</span>), <span class="pl-s1">target_vocab_size</span><span class="pl-c1">=</span><span class="pl-c1">2</span> <span class="pl-c1">**</span> <span class="pl-c1">13</span>)


<span class="pl-v">BUFFER_SIZE</span> <span class="pl-c1">=</span> <span class="pl-c1">20000</span>
<span class="pl-v">BATCH_SIZE</span> <span class="pl-c1">=</span> <span class="pl-c1">64</span>

<span class="pl-s">"""Add a start and end token to the input and target."""</span>


<span class="pl-k">def</span> <span class="pl-en">encode</span>(<span class="pl-s1">lang1</span>, <span class="pl-s1">lang2</span>):
    <span class="pl-s1">lang1</span> <span class="pl-c1">=</span> [<span class="pl-s1">tokenizer_pt</span>.<span class="pl-s1">vocab_size</span>] <span class="pl-c1">+</span> <span class="pl-s1">tokenizer_pt</span>.<span class="pl-en">encode</span>(
        <span class="pl-s1">lang1</span>.<span class="pl-en">numpy</span>()) <span class="pl-c1">+</span> [<span class="pl-s1">tokenizer_pt</span>.<span class="pl-s1">vocab_size</span> <span class="pl-c1">+</span> <span class="pl-c1">1</span>]

    <span class="pl-s1">lang2</span> <span class="pl-c1">=</span> [<span class="pl-s1">tokenizer_en</span>.<span class="pl-s1">vocab_size</span>] <span class="pl-c1">+</span> <span class="pl-s1">tokenizer_en</span>.<span class="pl-en">encode</span>(
        <span class="pl-s1">lang2</span>.<span class="pl-en">numpy</span>()) <span class="pl-c1">+</span> [<span class="pl-s1">tokenizer_en</span>.<span class="pl-s1">vocab_size</span> <span class="pl-c1">+</span> <span class="pl-c1">1</span>]

    <span class="pl-k">return</span> <span class="pl-s1">lang1</span>, <span class="pl-s1">lang2</span>


<span class="pl-s">"""Note: To keep this example small and relatively fast, drop examples with a length of over 40 tokens."""</span>

<span class="pl-v">MAX_LENGTH</span> <span class="pl-c1">=</span> <span class="pl-c1">40</span>


<span class="pl-k">def</span> <span class="pl-en">filter_max_length</span>(<span class="pl-s1">x</span>, <span class="pl-s1">y</span>, <span class="pl-s1">max_length</span><span class="pl-c1">=</span><span class="pl-v">MAX_LENGTH</span>):
    <span class="pl-k">return</span> <span class="pl-s1">tf</span>.<span class="pl-en">logical_and</span>(<span class="pl-s1">tf</span>.<span class="pl-en">size</span>(<span class="pl-s1">x</span>) <span class="pl-c1">&lt;=</span> <span class="pl-s1">max_length</span>,
                          <span class="pl-s1">tf</span>.<span class="pl-en">size</span>(<span class="pl-s1">y</span>) <span class="pl-c1">&lt;=</span> <span class="pl-s1">max_length</span>)


<span class="pl-s">"""Operations inside `.map()` run in graph mode and receive a graph tensor that do not have a numpy attribute. The `tokenizer` expects a string or Unicode symbol to encode it into integers. Hence, you need to run the encoding inside a `tf.py_function`, which receives an eager tensor having a numpy attribute that contains the string value."""</span>


<span class="pl-k">def</span> <span class="pl-en">tf_encode</span>(<span class="pl-s1">pt</span>, <span class="pl-s1">en</span>):
    <span class="pl-k">return</span> <span class="pl-s1">tf</span>.<span class="pl-en">py_function</span>(<span class="pl-s1">encode</span>, [<span class="pl-s1">pt</span>, <span class="pl-s1">en</span>], [<span class="pl-s1">tf</span>.<span class="pl-s1">int64</span>, <span class="pl-s1">tf</span>.<span class="pl-s1">int64</span>])

<span class="pl-en">print</span>(<span class="pl-s">"(2):Encode and padded batch......"</span>)
<span class="pl-s1">train_dataset</span> <span class="pl-c1">=</span> <span class="pl-s1">train_examples</span>.<span class="pl-en">map</span>(<span class="pl-s1">tf_encode</span>)
<span class="pl-s1">train_dataset</span> <span class="pl-c1">=</span> <span class="pl-s1">train_dataset</span>.<span class="pl-en">filter</span>(<span class="pl-s1">filter_max_length</span>)
<span class="pl-c"># cache the dataset to memory to get a speedup while reading from it.</span>
<span class="pl-s1">train_dataset</span> <span class="pl-c1">=</span> <span class="pl-s1">train_dataset</span>.<span class="pl-en">cache</span>()
<span class="pl-s1">train_dataset</span> <span class="pl-c1">=</span> <span class="pl-s1">train_dataset</span>.<span class="pl-en">shuffle</span>(<span class="pl-v">BUFFER_SIZE</span>).<span class="pl-en">padded_batch</span>(
    <span class="pl-v">BATCH_SIZE</span>, <span class="pl-s1">padded_shapes</span><span class="pl-c1">=</span>([<span class="pl-c1">-</span><span class="pl-c1">1</span>], [<span class="pl-c1">-</span><span class="pl-c1">1</span>]))
<span class="pl-s1">train_dataset</span> <span class="pl-c1">=</span> <span class="pl-s1">train_dataset</span>.<span class="pl-en">prefetch</span>(<span class="pl-s1">tf</span>.<span class="pl-s1">data</span>.<span class="pl-s1">experimental</span>.<span class="pl-v">AUTOTUNE</span>)

<span class="pl-s1">val_dataset</span> <span class="pl-c1">=</span> <span class="pl-s1">val_examples</span>.<span class="pl-en">map</span>(<span class="pl-s1">tf_encode</span>)
<span class="pl-s1">val_dataset</span> <span class="pl-c1">=</span> <span class="pl-s1">val_dataset</span>.<span class="pl-en">filter</span>(<span class="pl-s1">filter_max_length</span>).<span class="pl-en">padded_batch</span>(
    <span class="pl-v">BATCH_SIZE</span>, <span class="pl-s1">padded_shapes</span><span class="pl-c1">=</span>([<span class="pl-c1">-</span><span class="pl-c1">1</span>], [<span class="pl-c1">-</span><span class="pl-c1">1</span>]))



<span class="pl-k">def</span> <span class="pl-en">get_angles</span>(<span class="pl-s1">pos</span>, <span class="pl-s1">i</span>, <span class="pl-s1">d_model</span>):
    <span class="pl-s1">angle_rates</span> <span class="pl-c1">=</span> <span class="pl-c1">1</span> <span class="pl-c1">/</span> <span class="pl-s1">np</span>.<span class="pl-en">power</span>(<span class="pl-c1">10000</span>, (<span class="pl-c1">2</span> <span class="pl-c1">*</span> (<span class="pl-s1">i</span> <span class="pl-c1">//</span> <span class="pl-c1">2</span>)) <span class="pl-c1">/</span> <span class="pl-s1">np</span>.<span class="pl-en">float32</span>(<span class="pl-s1">d_model</span>))
    <span class="pl-k">return</span> <span class="pl-s1">pos</span> <span class="pl-c1">*</span> <span class="pl-s1">angle_rates</span>


<span class="pl-k">def</span> <span class="pl-en">positional_encoding</span>(<span class="pl-s1">position</span>, <span class="pl-s1">d_model</span>):
    <span class="pl-s1">angle_rads</span> <span class="pl-c1">=</span> <span class="pl-en">get_angles</span>(<span class="pl-s1">np</span>.<span class="pl-en">arange</span>(<span class="pl-s1">position</span>)[:, <span class="pl-s1">np</span>.<span class="pl-s1">newaxis</span>],
                            <span class="pl-s1">np</span>.<span class="pl-en">arange</span>(<span class="pl-s1">d_model</span>)[<span class="pl-s1">np</span>.<span class="pl-s1">newaxis</span>, :],
                            <span class="pl-s1">d_model</span>)

    <span class="pl-c"># apply sin to even indices in the array; 2i</span>
    <span class="pl-s1">sines</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">sin</span>(<span class="pl-s1">angle_rads</span>[:, <span class="pl-c1">0</span>::<span class="pl-c1">2</span>])

    <span class="pl-c"># apply cos to odd indices in the array; 2i+1</span>
    <span class="pl-s1">cosines</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">cos</span>(<span class="pl-s1">angle_rads</span>[:, <span class="pl-c1">1</span>::<span class="pl-c1">2</span>])

    <span class="pl-s1">pos_encoding</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">concatenate</span>([<span class="pl-s1">sines</span>, <span class="pl-s1">cosines</span>], <span class="pl-s1">axis</span><span class="pl-c1">=</span><span class="pl-c1">-</span><span class="pl-c1">1</span>)

    <span class="pl-s1">pos_encoding</span> <span class="pl-c1">=</span> <span class="pl-s1">pos_encoding</span>[<span class="pl-s1">np</span>.<span class="pl-s1">newaxis</span>, ...]

    <span class="pl-k">return</span> <span class="pl-s1">tf</span>.<span class="pl-en">cast</span>(<span class="pl-s1">pos_encoding</span>, <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">tf</span>.<span class="pl-s1">float32</span>)


<span class="pl-k">def</span> <span class="pl-en">create_padding_mask</span>(<span class="pl-s1">seq</span>):
    <span class="pl-s1">seq</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-en">cast</span>(<span class="pl-s1">tf</span>.<span class="pl-s1">math</span>.<span class="pl-en">equal</span>(<span class="pl-s1">seq</span>, <span class="pl-c1">0</span>), <span class="pl-s1">tf</span>.<span class="pl-s1">float32</span>)

    <span class="pl-c"># add extra dimensions so that we can add the padding</span>
    <span class="pl-c"># to the attention logits.</span>
    <span class="pl-k">return</span> <span class="pl-s1">seq</span>[:, <span class="pl-s1">tf</span>.<span class="pl-s1">newaxis</span>, <span class="pl-s1">tf</span>.<span class="pl-s1">newaxis</span>, :]  <span class="pl-c"># (batch_size, 1, 1, seq_len)</span>



<span class="pl-k">def</span> <span class="pl-en">create_look_ahead_mask</span>(<span class="pl-s1">size</span>):
    <span class="pl-s1">mask</span> <span class="pl-c1">=</span> <span class="pl-c1">1</span> <span class="pl-c1">-</span> <span class="pl-s1">tf</span>.<span class="pl-s1">linalg</span>.<span class="pl-en">band_part</span>(<span class="pl-s1">tf</span>.<span class="pl-en">ones</span>((<span class="pl-s1">size</span>, <span class="pl-s1">size</span>)), <span class="pl-c1">-</span><span class="pl-c1">1</span>, <span class="pl-c1">0</span>)
    <span class="pl-k">return</span> <span class="pl-s1">mask</span>  <span class="pl-c"># (seq_len, seq_len)</span>



<span class="pl-k">def</span> <span class="pl-en">scaled_dot_product_attention</span>(<span class="pl-s1">q</span>, <span class="pl-s1">k</span>, <span class="pl-s1">v</span>, <span class="pl-s1">mask</span>):
    <span class="pl-s">"""Calculate the attention weights.</span>
<span class="pl-s">    q, k, v must have matching leading dimensions.</span>
<span class="pl-s">    The mask has different shapes depending on its type(padding or look ahead)</span>
<span class="pl-s">    but it must be broadcastable for addition.</span>
<span class="pl-s"></span>
<span class="pl-s">    Args:</span>
<span class="pl-s">      q: query shape == (..., seq_len_q, depth)</span>
<span class="pl-s">      k: key shape == (..., seq_len_k, depth)</span>
<span class="pl-s">      v: value shape == (..., seq_len_v, depth)</span>
<span class="pl-s">      mask: Float tensor with shape broadcastable</span>
<span class="pl-s">            to (..., seq_len_q, seq_len_k). Defaults to None.</span>
<span class="pl-s"></span>
<span class="pl-s">    Returns:</span>
<span class="pl-s">      output, attention_weights</span>
<span class="pl-s">    """</span>

    <span class="pl-s1">matmul_qk</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-en">matmul</span>(<span class="pl-s1">q</span>, <span class="pl-s1">k</span>, <span class="pl-s1">transpose_b</span><span class="pl-c1">=</span><span class="pl-c1">True</span>)  <span class="pl-c"># (..., seq_len_q, seq_len_k)</span>

    <span class="pl-c"># scale matmul_qk</span>
    <span class="pl-s1">dk</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-en">cast</span>(<span class="pl-s1">tf</span>.<span class="pl-en">shape</span>(<span class="pl-s1">k</span>)[<span class="pl-c1">-</span><span class="pl-c1">1</span>], <span class="pl-s1">tf</span>.<span class="pl-s1">float32</span>)
    <span class="pl-s1">scaled_attention_logits</span> <span class="pl-c1">=</span> <span class="pl-s1">matmul_qk</span> <span class="pl-c1">/</span> <span class="pl-s1">tf</span>.<span class="pl-s1">math</span>.<span class="pl-en">sqrt</span>(<span class="pl-s1">dk</span>)

    <span class="pl-c"># add the mask to the scaled tensor.</span>
    <span class="pl-k">if</span> <span class="pl-s1">mask</span> <span class="pl-c1">is</span> <span class="pl-c1">not</span> <span class="pl-c1">None</span>:
        <span class="pl-s1">scaled_attention_logits</span> <span class="pl-c1">+=</span> (<span class="pl-s1">mask</span> <span class="pl-c1">*</span> <span class="pl-c1">-</span><span class="pl-c1">1e9</span>)

    <span class="pl-c"># softmax is normalized on the last axis (seq_len_k) so that the scores</span>
    <span class="pl-c"># add up to 1.</span>
    <span class="pl-s1">attention_weights</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">nn</span>.<span class="pl-en">softmax</span>(<span class="pl-s1">scaled_attention_logits</span>, <span class="pl-s1">axis</span><span class="pl-c1">=</span><span class="pl-c1">-</span><span class="pl-c1">1</span>)  <span class="pl-c"># (..., seq_len_q, seq_len_k)</span>

    <span class="pl-s1">output</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-en">matmul</span>(<span class="pl-s1">attention_weights</span>, <span class="pl-s1">v</span>)  <span class="pl-c"># (..., seq_len_v, depth)</span>

    <span class="pl-k">return</span> <span class="pl-s1">output</span>, <span class="pl-s1">attention_weights</span>

<span class="pl-k">class</span> <span class="pl-v">MultiHeadAttention</span>(<span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">layers</span>.<span class="pl-v">Layer</span>):
    <span class="pl-k">def</span> <span class="pl-en">__init__</span>(<span class="pl-s1">self</span>, <span class="pl-s1">d_model</span>, <span class="pl-s1">num_heads</span>):
        <span class="pl-en">super</span>(<span class="pl-v">MultiHeadAttention</span>, <span class="pl-s1">self</span>).<span class="pl-en">__init__</span>()
        <span class="pl-s1">self</span>.<span class="pl-s1">num_heads</span> <span class="pl-c1">=</span> <span class="pl-s1">num_heads</span>
        <span class="pl-s1">self</span>.<span class="pl-s1">d_model</span> <span class="pl-c1">=</span> <span class="pl-s1">d_model</span>

        <span class="pl-k">assert</span> <span class="pl-s1">d_model</span> <span class="pl-c1">%</span> <span class="pl-s1">self</span>.<span class="pl-s1">num_heads</span> <span class="pl-c1">==</span> <span class="pl-c1">0</span>

        <span class="pl-s1">self</span>.<span class="pl-s1">depth</span> <span class="pl-c1">=</span> <span class="pl-s1">d_model</span> <span class="pl-c1">//</span> <span class="pl-s1">self</span>.<span class="pl-s1">num_heads</span>

        <span class="pl-s1">self</span>.<span class="pl-s1">wq</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">layers</span>.<span class="pl-v">Dense</span>(<span class="pl-s1">d_model</span>)
        <span class="pl-s1">self</span>.<span class="pl-s1">wk</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">layers</span>.<span class="pl-v">Dense</span>(<span class="pl-s1">d_model</span>)
        <span class="pl-s1">self</span>.<span class="pl-s1">wv</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">layers</span>.<span class="pl-v">Dense</span>(<span class="pl-s1">d_model</span>)

        <span class="pl-s1">self</span>.<span class="pl-s1">dense</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">layers</span>.<span class="pl-v">Dense</span>(<span class="pl-s1">d_model</span>)

    <span class="pl-k">def</span> <span class="pl-en">split_heads</span>(<span class="pl-s1">self</span>, <span class="pl-s1">x</span>, <span class="pl-s1">batch_size</span>):
        <span class="pl-s">"""Split the last dimension into (num_heads, depth).</span>
<span class="pl-s">        Transpose the result such that the shape is (batch_size, num_heads, seq_len, depth)</span>
<span class="pl-s">        """</span>
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-en">reshape</span>(<span class="pl-s1">x</span>, (<span class="pl-s1">batch_size</span>, <span class="pl-c1">-</span><span class="pl-c1">1</span>, <span class="pl-s1">self</span>.<span class="pl-s1">num_heads</span>, <span class="pl-s1">self</span>.<span class="pl-s1">depth</span>))
        <span class="pl-k">return</span> <span class="pl-s1">tf</span>.<span class="pl-en">transpose</span>(<span class="pl-s1">x</span>, <span class="pl-s1">perm</span><span class="pl-c1">=</span>[<span class="pl-c1">0</span>, <span class="pl-c1">2</span>, <span class="pl-c1">1</span>, <span class="pl-c1">3</span>])

    <span class="pl-k">def</span> <span class="pl-en">call</span>(<span class="pl-s1">self</span>, <span class="pl-s1">v</span>, <span class="pl-s1">k</span>, <span class="pl-s1">q</span>, <span class="pl-s1">mask</span>):
        <span class="pl-s1">batch_size</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-en">shape</span>(<span class="pl-s1">q</span>)[<span class="pl-c1">0</span>]

        <span class="pl-s1">q</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">wq</span>(<span class="pl-s1">q</span>)  <span class="pl-c"># (batch_size, seq_len, d_model)</span>
        <span class="pl-s1">k</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">wk</span>(<span class="pl-s1">k</span>)  <span class="pl-c"># (batch_size, seq_len, d_model)</span>
        <span class="pl-s1">v</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">wv</span>(<span class="pl-s1">v</span>)  <span class="pl-c"># (batch_size, seq_len, d_model)</span>

        <span class="pl-s1">q</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">split_heads</span>(<span class="pl-s1">q</span>, <span class="pl-s1">batch_size</span>)  <span class="pl-c"># (batch_size, num_heads, seq_len_q, depth)</span>
        <span class="pl-s1">k</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">split_heads</span>(<span class="pl-s1">k</span>, <span class="pl-s1">batch_size</span>)  <span class="pl-c"># (batch_size, num_heads, seq_len_k, depth)</span>
        <span class="pl-s1">v</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">split_heads</span>(<span class="pl-s1">v</span>, <span class="pl-s1">batch_size</span>)  <span class="pl-c"># (batch_size, num_heads, seq_len_v, depth)</span>

        <span class="pl-c"># scaled_attention.shape == (batch_size, num_heads, seq_len_v, depth)</span>
        <span class="pl-c"># attention_weights.shape == (batch_size, num_heads, seq_len_q, seq_len_k)</span>
        <span class="pl-s1">scaled_attention</span>, <span class="pl-s1">attention_weights</span> <span class="pl-c1">=</span> <span class="pl-en">scaled_dot_product_attention</span>(
            <span class="pl-s1">q</span>, <span class="pl-s1">k</span>, <span class="pl-s1">v</span>, <span class="pl-s1">mask</span>)

        <span class="pl-s1">scaled_attention</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-en">transpose</span>(<span class="pl-s1">scaled_attention</span>,
                                        <span class="pl-s1">perm</span><span class="pl-c1">=</span>[<span class="pl-c1">0</span>, <span class="pl-c1">2</span>, <span class="pl-c1">1</span>, <span class="pl-c1">3</span>])  <span class="pl-c"># (batch_size, seq_len_v, num_heads, depth)</span>

        <span class="pl-s1">concat_attention</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-en">reshape</span>(<span class="pl-s1">scaled_attention</span>,
                                      (<span class="pl-s1">batch_size</span>, <span class="pl-c1">-</span><span class="pl-c1">1</span>, <span class="pl-s1">self</span>.<span class="pl-s1">d_model</span>))  <span class="pl-c"># (batch_size, seq_len_v, d_model)</span>

        <span class="pl-s1">output</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">dense</span>(<span class="pl-s1">concat_attention</span>)  <span class="pl-c"># (batch_size, seq_len_v, d_model)</span>

        <span class="pl-k">return</span> <span class="pl-s1">output</span>, <span class="pl-s1">attention_weights</span>


<span class="pl-k">def</span> <span class="pl-en">point_wise_feed_forward_network</span>(<span class="pl-s1">d_model</span>, <span class="pl-s1">dff</span>):
    <span class="pl-k">return</span> <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-v">Sequential</span>([
        <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">layers</span>.<span class="pl-v">Dense</span>(<span class="pl-s1">dff</span>, <span class="pl-s1">activation</span><span class="pl-c1">=</span><span class="pl-s">'relu'</span>),  <span class="pl-c"># (batch_size, seq_len, dff)</span>
        <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">layers</span>.<span class="pl-v">Dense</span>(<span class="pl-s1">d_model</span>)  <span class="pl-c"># (batch_size, seq_len, d_model)</span>
    ])

<span class="pl-k">class</span> <span class="pl-v">EncoderLayer</span>(<span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">layers</span>.<span class="pl-v">Layer</span>):
    <span class="pl-k">def</span> <span class="pl-en">__init__</span>(<span class="pl-s1">self</span>, <span class="pl-s1">d_model</span>, <span class="pl-s1">num_heads</span>, <span class="pl-s1">dff</span>, <span class="pl-s1">rate</span><span class="pl-c1">=</span><span class="pl-c1">0.1</span>):
        <span class="pl-en">super</span>(<span class="pl-v">EncoderLayer</span>, <span class="pl-s1">self</span>).<span class="pl-en">__init__</span>()

        <span class="pl-s1">self</span>.<span class="pl-s1">mha</span> <span class="pl-c1">=</span> <span class="pl-v">MultiHeadAttention</span>(<span class="pl-s1">d_model</span>, <span class="pl-s1">num_heads</span>)
        <span class="pl-s1">self</span>.<span class="pl-s1">ffn</span> <span class="pl-c1">=</span> <span class="pl-en">point_wise_feed_forward_network</span>(<span class="pl-s1">d_model</span>, <span class="pl-s1">dff</span>)

        <span class="pl-s1">self</span>.<span class="pl-s1">layernorm1</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">layers</span>.<span class="pl-v">LayerNormalization</span>(<span class="pl-s1">epsilon</span><span class="pl-c1">=</span><span class="pl-c1">1e-6</span>)
        <span class="pl-s1">self</span>.<span class="pl-s1">layernorm2</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">layers</span>.<span class="pl-v">LayerNormalization</span>(<span class="pl-s1">epsilon</span><span class="pl-c1">=</span><span class="pl-c1">1e-6</span>)

        <span class="pl-s1">self</span>.<span class="pl-s1">dropout1</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">layers</span>.<span class="pl-v">Dropout</span>(<span class="pl-s1">rate</span>)
        <span class="pl-s1">self</span>.<span class="pl-s1">dropout2</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">layers</span>.<span class="pl-v">Dropout</span>(<span class="pl-s1">rate</span>)

    <span class="pl-k">def</span> <span class="pl-en">call</span>(<span class="pl-s1">self</span>, <span class="pl-s1">x</span>, <span class="pl-s1">training</span>, <span class="pl-s1">mask</span>):
        <span class="pl-s1">attn_output</span>, <span class="pl-s1">_</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">mha</span>(<span class="pl-s1">x</span>, <span class="pl-s1">x</span>, <span class="pl-s1">x</span>, <span class="pl-s1">mask</span>)  <span class="pl-c"># (batch_size, input_seq_len, d_model)</span>
        <span class="pl-s1">attn_output</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">dropout1</span>(<span class="pl-s1">attn_output</span>, <span class="pl-s1">training</span><span class="pl-c1">=</span><span class="pl-s1">training</span>)
        <span class="pl-s1">out1</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">layernorm1</span>(<span class="pl-s1">x</span> <span class="pl-c1">+</span> <span class="pl-s1">attn_output</span>)  <span class="pl-c"># (batch_size, input_seq_len, d_model)</span>

        <span class="pl-s1">ffn_output</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">ffn</span>(<span class="pl-s1">out1</span>)  <span class="pl-c"># (batch_size, input_seq_len, d_model)</span>
        <span class="pl-s1">ffn_output</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">dropout2</span>(<span class="pl-s1">ffn_output</span>, <span class="pl-s1">training</span><span class="pl-c1">=</span><span class="pl-s1">training</span>)
        <span class="pl-s1">out2</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">layernorm2</span>(<span class="pl-s1">out1</span> <span class="pl-c1">+</span> <span class="pl-s1">ffn_output</span>)  <span class="pl-c"># (batch_size, input_seq_len, d_model)</span>

        <span class="pl-k">return</span> <span class="pl-s1">out2</span>


<span class="pl-k">class</span> <span class="pl-v">DecoderLayer</span>(<span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">layers</span>.<span class="pl-v">Layer</span>):
    <span class="pl-k">def</span> <span class="pl-en">__init__</span>(<span class="pl-s1">self</span>, <span class="pl-s1">d_model</span>, <span class="pl-s1">num_heads</span>, <span class="pl-s1">dff</span>, <span class="pl-s1">rate</span><span class="pl-c1">=</span><span class="pl-c1">0.1</span>):
        <span class="pl-en">super</span>(<span class="pl-v">DecoderLayer</span>, <span class="pl-s1">self</span>).<span class="pl-en">__init__</span>()

        <span class="pl-s1">self</span>.<span class="pl-s1">mha1</span> <span class="pl-c1">=</span> <span class="pl-v">MultiHeadAttention</span>(<span class="pl-s1">d_model</span>, <span class="pl-s1">num_heads</span>)
        <span class="pl-s1">self</span>.<span class="pl-s1">mha2</span> <span class="pl-c1">=</span> <span class="pl-v">MultiHeadAttention</span>(<span class="pl-s1">d_model</span>, <span class="pl-s1">num_heads</span>)

        <span class="pl-s1">self</span>.<span class="pl-s1">ffn</span> <span class="pl-c1">=</span> <span class="pl-en">point_wise_feed_forward_network</span>(<span class="pl-s1">d_model</span>, <span class="pl-s1">dff</span>)

        <span class="pl-s1">self</span>.<span class="pl-s1">layernorm1</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">layers</span>.<span class="pl-v">LayerNormalization</span>(<span class="pl-s1">epsilon</span><span class="pl-c1">=</span><span class="pl-c1">1e-6</span>)
        <span class="pl-s1">self</span>.<span class="pl-s1">layernorm2</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">layers</span>.<span class="pl-v">LayerNormalization</span>(<span class="pl-s1">epsilon</span><span class="pl-c1">=</span><span class="pl-c1">1e-6</span>)
        <span class="pl-s1">self</span>.<span class="pl-s1">layernorm3</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">layers</span>.<span class="pl-v">LayerNormalization</span>(<span class="pl-s1">epsilon</span><span class="pl-c1">=</span><span class="pl-c1">1e-6</span>)

        <span class="pl-s1">self</span>.<span class="pl-s1">dropout1</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">layers</span>.<span class="pl-v">Dropout</span>(<span class="pl-s1">rate</span>)
        <span class="pl-s1">self</span>.<span class="pl-s1">dropout2</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">layers</span>.<span class="pl-v">Dropout</span>(<span class="pl-s1">rate</span>)
        <span class="pl-s1">self</span>.<span class="pl-s1">dropout3</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">layers</span>.<span class="pl-v">Dropout</span>(<span class="pl-s1">rate</span>)

    <span class="pl-k">def</span> <span class="pl-en">call</span>(<span class="pl-s1">self</span>, <span class="pl-s1">x</span>, <span class="pl-s1">enc_output</span>, <span class="pl-s1">training</span>,
             <span class="pl-s1">look_ahead_mask</span>, <span class="pl-s1">padding_mask</span>):
        <span class="pl-c"># enc_output.shape == (batch_size, input_seq_len, d_model)</span>

        <span class="pl-s1">attn1</span>, <span class="pl-s1">attn_weights_block1</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">mha1</span>(<span class="pl-s1">x</span>, <span class="pl-s1">x</span>, <span class="pl-s1">x</span>, <span class="pl-s1">look_ahead_mask</span>)  <span class="pl-c"># (batch_size, target_seq_len, d_model)</span>
        <span class="pl-s1">attn1</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">dropout1</span>(<span class="pl-s1">attn1</span>, <span class="pl-s1">training</span><span class="pl-c1">=</span><span class="pl-s1">training</span>)
        <span class="pl-s1">out1</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">layernorm1</span>(<span class="pl-s1">attn1</span> <span class="pl-c1">+</span> <span class="pl-s1">x</span>)

        <span class="pl-s1">attn2</span>, <span class="pl-s1">attn_weights_block2</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">mha2</span>(
            <span class="pl-s1">enc_output</span>, <span class="pl-s1">enc_output</span>, <span class="pl-s1">out1</span>, <span class="pl-s1">padding_mask</span>)  <span class="pl-c"># (batch_size, target_seq_len, d_model)</span>
        <span class="pl-s1">attn2</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">dropout2</span>(<span class="pl-s1">attn2</span>, <span class="pl-s1">training</span><span class="pl-c1">=</span><span class="pl-s1">training</span>)
        <span class="pl-s1">out2</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">layernorm2</span>(<span class="pl-s1">attn2</span> <span class="pl-c1">+</span> <span class="pl-s1">out1</span>)  <span class="pl-c"># (batch_size, target_seq_len, d_model)</span>

        <span class="pl-s1">ffn_output</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">ffn</span>(<span class="pl-s1">out2</span>)  <span class="pl-c"># (batch_size, target_seq_len, d_model)</span>
        <span class="pl-s1">ffn_output</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">dropout3</span>(<span class="pl-s1">ffn_output</span>, <span class="pl-s1">training</span><span class="pl-c1">=</span><span class="pl-s1">training</span>)
        <span class="pl-s1">out3</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">layernorm3</span>(<span class="pl-s1">ffn_output</span> <span class="pl-c1">+</span> <span class="pl-s1">out2</span>)  <span class="pl-c"># (batch_size, target_seq_len, d_model)</span>

        <span class="pl-k">return</span> <span class="pl-s1">out3</span>, <span class="pl-s1">attn_weights_block1</span>, <span class="pl-s1">attn_weights_block2</span>


<span class="pl-k">class</span> <span class="pl-v">Encoder</span>(<span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">layers</span>.<span class="pl-v">Layer</span>):
    <span class="pl-k">def</span> <span class="pl-en">__init__</span>(<span class="pl-s1">self</span>, <span class="pl-s1">num_layers</span>, <span class="pl-s1">d_model</span>, <span class="pl-s1">num_heads</span>, <span class="pl-s1">dff</span>, <span class="pl-s1">input_vocab_size</span>,
                 <span class="pl-s1">rate</span><span class="pl-c1">=</span><span class="pl-c1">0.1</span>):
        <span class="pl-en">super</span>(<span class="pl-v">Encoder</span>, <span class="pl-s1">self</span>).<span class="pl-en">__init__</span>()

        <span class="pl-s1">self</span>.<span class="pl-s1">d_model</span> <span class="pl-c1">=</span> <span class="pl-s1">d_model</span>
        <span class="pl-s1">self</span>.<span class="pl-s1">num_layers</span> <span class="pl-c1">=</span> <span class="pl-s1">num_layers</span>

        <span class="pl-s1">self</span>.<span class="pl-s1">embedding</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">layers</span>.<span class="pl-v">Embedding</span>(<span class="pl-s1">input_vocab_size</span>, <span class="pl-s1">d_model</span>)
        <span class="pl-s1">self</span>.<span class="pl-s1">pos_encoding</span> <span class="pl-c1">=</span> <span class="pl-en">positional_encoding</span>(<span class="pl-s1">input_vocab_size</span>, <span class="pl-s1">self</span>.<span class="pl-s1">d_model</span>)

        <span class="pl-s1">self</span>.<span class="pl-s1">enc_layers</span> <span class="pl-c1">=</span> [<span class="pl-v">EncoderLayer</span>(<span class="pl-s1">d_model</span>, <span class="pl-s1">num_heads</span>, <span class="pl-s1">dff</span>, <span class="pl-s1">rate</span>)
                           <span class="pl-k">for</span> <span class="pl-s1">_</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-s1">num_layers</span>)]

        <span class="pl-s1">self</span>.<span class="pl-s1">dropout</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">layers</span>.<span class="pl-v">Dropout</span>(<span class="pl-s1">rate</span>)

    <span class="pl-k">def</span> <span class="pl-en">call</span>(<span class="pl-s1">self</span>, <span class="pl-s1">x</span>, <span class="pl-s1">training</span>, <span class="pl-s1">mask</span>):
        <span class="pl-s1">seq_len</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-en">shape</span>(<span class="pl-s1">x</span>)[<span class="pl-c1">1</span>]

        <span class="pl-c"># adding embedding and position encoding.</span>
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">embedding</span>(<span class="pl-s1">x</span>)  <span class="pl-c"># (batch_size, input_seq_len, d_model)</span>
        <span class="pl-s1">x</span> <span class="pl-c1">*=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">math</span>.<span class="pl-en">sqrt</span>(<span class="pl-s1">tf</span>.<span class="pl-en">cast</span>(<span class="pl-s1">self</span>.<span class="pl-s1">d_model</span>, <span class="pl-s1">tf</span>.<span class="pl-s1">float32</span>))
        <span class="pl-s1">x</span> <span class="pl-c1">+=</span> <span class="pl-s1">self</span>.<span class="pl-s1">pos_encoding</span>[:, :<span class="pl-s1">seq_len</span>, :]

        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">dropout</span>(<span class="pl-s1">x</span>, <span class="pl-s1">training</span><span class="pl-c1">=</span><span class="pl-s1">training</span>)

        <span class="pl-k">for</span> <span class="pl-s1">i</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-s1">self</span>.<span class="pl-s1">num_layers</span>):
            <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-s1">enc_layers</span>[<span class="pl-s1">i</span>](<span class="pl-s1">x</span>, <span class="pl-s1">training</span>, <span class="pl-s1">mask</span>)

        <span class="pl-k">return</span> <span class="pl-s1">x</span>  <span class="pl-c"># (batch_size, input_seq_len, d_model)</span>


<span class="pl-k">class</span> <span class="pl-v">Decoder</span>(<span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">layers</span>.<span class="pl-v">Layer</span>):
    <span class="pl-k">def</span> <span class="pl-en">__init__</span>(<span class="pl-s1">self</span>, <span class="pl-s1">num_layers</span>, <span class="pl-s1">d_model</span>, <span class="pl-s1">num_heads</span>, <span class="pl-s1">dff</span>, <span class="pl-s1">target_vocab_size</span>,
                 <span class="pl-s1">rate</span><span class="pl-c1">=</span><span class="pl-c1">0.1</span>):
        <span class="pl-en">super</span>(<span class="pl-v">Decoder</span>, <span class="pl-s1">self</span>).<span class="pl-en">__init__</span>()

        <span class="pl-s1">self</span>.<span class="pl-s1">d_model</span> <span class="pl-c1">=</span> <span class="pl-s1">d_model</span>
        <span class="pl-s1">self</span>.<span class="pl-s1">num_layers</span> <span class="pl-c1">=</span> <span class="pl-s1">num_layers</span>

        <span class="pl-s1">self</span>.<span class="pl-s1">embedding</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">layers</span>.<span class="pl-v">Embedding</span>(<span class="pl-s1">target_vocab_size</span>, <span class="pl-s1">d_model</span>)
        <span class="pl-s1">self</span>.<span class="pl-s1">pos_encoding</span> <span class="pl-c1">=</span> <span class="pl-en">positional_encoding</span>(<span class="pl-s1">target_vocab_size</span>, <span class="pl-s1">self</span>.<span class="pl-s1">d_model</span>)

        <span class="pl-s1">self</span>.<span class="pl-s1">dec_layers</span> <span class="pl-c1">=</span> [<span class="pl-v">DecoderLayer</span>(<span class="pl-s1">d_model</span>, <span class="pl-s1">num_heads</span>, <span class="pl-s1">dff</span>, <span class="pl-s1">rate</span>)
                           <span class="pl-k">for</span> <span class="pl-s1">_</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-s1">num_layers</span>)]
        <span class="pl-s1">self</span>.<span class="pl-s1">dropout</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">layers</span>.<span class="pl-v">Dropout</span>(<span class="pl-s1">rate</span>)

    <span class="pl-k">def</span> <span class="pl-en">call</span>(<span class="pl-s1">self</span>, <span class="pl-s1">x</span>, <span class="pl-s1">enc_output</span>, <span class="pl-s1">training</span>,
             <span class="pl-s1">look_ahead_mask</span>, <span class="pl-s1">padding_mask</span>):
        <span class="pl-s1">seq_len</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-en">shape</span>(<span class="pl-s1">x</span>)[<span class="pl-c1">1</span>]
        <span class="pl-s1">attention_weights</span> <span class="pl-c1">=</span> {}

        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">embedding</span>(<span class="pl-s1">x</span>)  <span class="pl-c"># (batch_size, target_seq_len, d_model)</span>
        <span class="pl-s1">x</span> <span class="pl-c1">*=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">math</span>.<span class="pl-en">sqrt</span>(<span class="pl-s1">tf</span>.<span class="pl-en">cast</span>(<span class="pl-s1">self</span>.<span class="pl-s1">d_model</span>, <span class="pl-s1">tf</span>.<span class="pl-s1">float32</span>))
        <span class="pl-s1">x</span> <span class="pl-c1">+=</span> <span class="pl-s1">self</span>.<span class="pl-s1">pos_encoding</span>[:, :<span class="pl-s1">seq_len</span>, :]

        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">dropout</span>(<span class="pl-s1">x</span>, <span class="pl-s1">training</span><span class="pl-c1">=</span><span class="pl-s1">training</span>)

        <span class="pl-k">for</span> <span class="pl-s1">i</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-s1">self</span>.<span class="pl-s1">num_layers</span>):
            <span class="pl-s1">x</span>, <span class="pl-s1">block1</span>, <span class="pl-s1">block2</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-s1">dec_layers</span>[<span class="pl-s1">i</span>](<span class="pl-s1">x</span>, <span class="pl-s1">enc_output</span>, <span class="pl-s1">training</span>,
                                                   <span class="pl-s1">look_ahead_mask</span>, <span class="pl-s1">padding_mask</span>)

            <span class="pl-s1">attention_weights</span>[<span class="pl-s">'decoder_layer{}_block1'</span>.<span class="pl-en">format</span>(<span class="pl-s1">i</span> <span class="pl-c1">+</span> <span class="pl-c1">1</span>)] <span class="pl-c1">=</span> <span class="pl-s1">block1</span>
            <span class="pl-s1">attention_weights</span>[<span class="pl-s">'decoder_layer{}_block2'</span>.<span class="pl-en">format</span>(<span class="pl-s1">i</span> <span class="pl-c1">+</span> <span class="pl-c1">1</span>)] <span class="pl-c1">=</span> <span class="pl-s1">block2</span>

        <span class="pl-c"># x.shape == (batch_size, target_seq_len, d_model)</span>
        <span class="pl-k">return</span> <span class="pl-s1">x</span>, <span class="pl-s1">attention_weights</span>


<span class="pl-s">"""## Create the Transformer</span>
<span class="pl-s"></span>
<span class="pl-s">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.</span>
<span class="pl-s">"""</span>


<span class="pl-k">class</span> <span class="pl-v">Transformer</span>(<span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-v">Model</span>):
    <span class="pl-k">def</span> <span class="pl-en">__init__</span>(<span class="pl-s1">self</span>, <span class="pl-s1">num_layers</span>, <span class="pl-s1">d_model</span>, <span class="pl-s1">num_heads</span>, <span class="pl-s1">dff</span>, <span class="pl-s1">input_vocab_size</span>,
                 <span class="pl-s1">target_vocab_size</span>, <span class="pl-s1">rate</span><span class="pl-c1">=</span><span class="pl-c1">0.1</span>):
        <span class="pl-en">super</span>(<span class="pl-v">Transformer</span>, <span class="pl-s1">self</span>).<span class="pl-en">__init__</span>()

        <span class="pl-s1">self</span>.<span class="pl-s1">encoder</span> <span class="pl-c1">=</span> <span class="pl-v">Encoder</span>(<span class="pl-s1">num_layers</span>, <span class="pl-s1">d_model</span>, <span class="pl-s1">num_heads</span>, <span class="pl-s1">dff</span>,
                               <span class="pl-s1">input_vocab_size</span>, <span class="pl-s1">rate</span>)

        <span class="pl-s1">self</span>.<span class="pl-s1">decoder</span> <span class="pl-c1">=</span> <span class="pl-v">Decoder</span>(<span class="pl-s1">num_layers</span>, <span class="pl-s1">d_model</span>, <span class="pl-s1">num_heads</span>, <span class="pl-s1">dff</span>,
                               <span class="pl-s1">target_vocab_size</span>, <span class="pl-s1">rate</span>)

        <span class="pl-s1">self</span>.<span class="pl-s1">final_layer</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">layers</span>.<span class="pl-v">Dense</span>(<span class="pl-s1">target_vocab_size</span>)

    <span class="pl-k">def</span> <span class="pl-en">call</span>(<span class="pl-s1">self</span>, <span class="pl-s1">inp</span>, <span class="pl-s1">tar</span>, <span class="pl-s1">training</span>, <span class="pl-s1">enc_padding_mask</span>,
             <span class="pl-s1">look_ahead_mask</span>, <span class="pl-s1">dec_padding_mask</span>):
        <span class="pl-s1">enc_output</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">encoder</span>(<span class="pl-s1">inp</span>, <span class="pl-s1">training</span>, <span class="pl-s1">enc_padding_mask</span>)  <span class="pl-c"># (batch_size, inp_seq_len, d_model)</span>

        <span class="pl-c"># dec_output.shape == (batch_size, tar_seq_len, d_model)</span>
        <span class="pl-s1">dec_output</span>, <span class="pl-s1">attention_weights</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">decoder</span>(
            <span class="pl-s1">tar</span>, <span class="pl-s1">enc_output</span>, <span class="pl-s1">training</span>, <span class="pl-s1">look_ahead_mask</span>, <span class="pl-s1">dec_padding_mask</span>)

        <span class="pl-s1">final_output</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">final_layer</span>(<span class="pl-s1">dec_output</span>)  <span class="pl-c"># (batch_size, tar_seq_len, target_vocab_size)</span>

        <span class="pl-k">return</span> <span class="pl-s1">final_output</span>, <span class="pl-s1">attention_weights</span>


<span class="pl-s1">num_layers</span> <span class="pl-c1">=</span> <span class="pl-c1">4</span>
<span class="pl-s1">d_model</span> <span class="pl-c1">=</span> <span class="pl-c1">128</span>
<span class="pl-s1">dff</span> <span class="pl-c1">=</span> <span class="pl-c1">512</span>
<span class="pl-s1">num_heads</span> <span class="pl-c1">=</span> <span class="pl-c1">8</span>

<span class="pl-s1">input_vocab_size</span> <span class="pl-c1">=</span> <span class="pl-s1">tokenizer_pt</span>.<span class="pl-s1">vocab_size</span> <span class="pl-c1">+</span> <span class="pl-c1">2</span>
<span class="pl-s1">target_vocab_size</span> <span class="pl-c1">=</span> <span class="pl-s1">tokenizer_en</span>.<span class="pl-s1">vocab_size</span> <span class="pl-c1">+</span> <span class="pl-c1">2</span>
<span class="pl-s1">dropout_rate</span> <span class="pl-c1">=</span> <span class="pl-c1">0.1</span>


<span class="pl-k">class</span> <span class="pl-v">CustomSchedule</span>(<span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">optimizers</span>.<span class="pl-s1">schedules</span>.<span class="pl-v">LearningRateSchedule</span>):
    <span class="pl-k">def</span> <span class="pl-en">__init__</span>(<span class="pl-s1">self</span>, <span class="pl-s1">d_model</span>, <span class="pl-s1">warmup_steps</span><span class="pl-c1">=</span><span class="pl-c1">4000</span>):
        <span class="pl-en">super</span>(<span class="pl-v">CustomSchedule</span>, <span class="pl-s1">self</span>).<span class="pl-en">__init__</span>()

        <span class="pl-s1">self</span>.<span class="pl-s1">d_model</span> <span class="pl-c1">=</span> <span class="pl-s1">d_model</span>
        <span class="pl-s1">self</span>.<span class="pl-s1">d_model</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-en">cast</span>(<span class="pl-s1">self</span>.<span class="pl-s1">d_model</span>, <span class="pl-s1">tf</span>.<span class="pl-s1">float32</span>)

        <span class="pl-s1">self</span>.<span class="pl-s1">warmup_steps</span> <span class="pl-c1">=</span> <span class="pl-s1">warmup_steps</span>

    <span class="pl-k">def</span> <span class="pl-en">__call__</span>(<span class="pl-s1">self</span>, <span class="pl-s1">step</span>):
        <span class="pl-s1">arg1</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">math</span>.<span class="pl-en">rsqrt</span>(<span class="pl-s1">step</span>)
        <span class="pl-s1">arg2</span> <span class="pl-c1">=</span> <span class="pl-s1">step</span> <span class="pl-c1">*</span> (<span class="pl-s1">self</span>.<span class="pl-s1">warmup_steps</span> <span class="pl-c1">**</span> <span class="pl-c1">-</span><span class="pl-c1">1.5</span>)

        <span class="pl-k">return</span> <span class="pl-s1">tf</span>.<span class="pl-s1">math</span>.<span class="pl-en">rsqrt</span>(<span class="pl-s1">self</span>.<span class="pl-s1">d_model</span>) <span class="pl-c1">*</span> <span class="pl-s1">tf</span>.<span class="pl-s1">math</span>.<span class="pl-en">minimum</span>(<span class="pl-s1">arg1</span>, <span class="pl-s1">arg2</span>)


<span class="pl-s1">learning_rate</span> <span class="pl-c1">=</span> <span class="pl-v">CustomSchedule</span>(<span class="pl-s1">d_model</span>)

<span class="pl-s1">optimizer</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">optimizers</span>.<span class="pl-v">Adam</span>(<span class="pl-s1">learning_rate</span>, <span class="pl-s1">beta_1</span><span class="pl-c1">=</span><span class="pl-c1">0.9</span>, <span class="pl-s1">beta_2</span><span class="pl-c1">=</span><span class="pl-c1">0.98</span>,
                                     <span class="pl-s1">epsilon</span><span class="pl-c1">=</span><span class="pl-c1">1e-9</span>)

<span class="pl-s1">temp_learning_rate_schedule</span> <span class="pl-c1">=</span> <span class="pl-v">CustomSchedule</span>(<span class="pl-s1">d_model</span>)

<span class="pl-s1">plt</span>.<span class="pl-en">plot</span>(<span class="pl-en">temp_learning_rate_schedule</span>(<span class="pl-s1">tf</span>.<span class="pl-en">range</span>(<span class="pl-c1">40000</span>, <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">tf</span>.<span class="pl-s1">float32</span>)))
<span class="pl-s1">plt</span>.<span class="pl-en">ylabel</span>(<span class="pl-s">"Learning Rate"</span>)
<span class="pl-s1">plt</span>.<span class="pl-en">xlabel</span>(<span class="pl-s">"Train Step"</span>)

<span class="pl-s">"""## Loss and metrics</span>
<span class="pl-s"></span>
<span class="pl-s">Since the target sequences are padded, it is important to apply a padding mask when calculating the loss.</span>
<span class="pl-s">"""</span>

<span class="pl-s1">loss_object</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">losses</span>.<span class="pl-v">SparseCategoricalCrossentropy</span>(
    <span class="pl-s1">from_logits</span><span class="pl-c1">=</span><span class="pl-c1">True</span>, <span class="pl-s1">reduction</span><span class="pl-c1">=</span><span class="pl-s">'none'</span>)


<span class="pl-k">def</span> <span class="pl-en">loss_function</span>(<span class="pl-s1">real</span>, <span class="pl-s1">pred</span>):
    <span class="pl-s1">mask</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">math</span>.<span class="pl-en">logical_not</span>(<span class="pl-s1">tf</span>.<span class="pl-s1">math</span>.<span class="pl-en">equal</span>(<span class="pl-s1">real</span>, <span class="pl-c1">0</span>))
    <span class="pl-s1">loss_</span> <span class="pl-c1">=</span> <span class="pl-en">loss_object</span>(<span class="pl-s1">real</span>, <span class="pl-s1">pred</span>)

    <span class="pl-s1">mask</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-en">cast</span>(<span class="pl-s1">mask</span>, <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">loss_</span>.<span class="pl-s1">dtype</span>)
    <span class="pl-s1">loss_</span> <span class="pl-c1">*=</span> <span class="pl-s1">mask</span>

    <span class="pl-k">return</span> <span class="pl-s1">tf</span>.<span class="pl-en">reduce_mean</span>(<span class="pl-s1">loss_</span>)


<span class="pl-s1">train_loss</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">metrics</span>.<span class="pl-v">Mean</span>(<span class="pl-s1">name</span><span class="pl-c1">=</span><span class="pl-s">'train_loss'</span>)
<span class="pl-s1">train_accuracy</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">metrics</span>.<span class="pl-v">SparseCategoricalAccuracy</span>(<span class="pl-s1">name</span><span class="pl-c1">=</span><span class="pl-s">'train_accuracy'</span>)
<span class="pl-s1">val_loss</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">metrics</span>.<span class="pl-v">Mean</span>(<span class="pl-s1">name</span><span class="pl-c1">=</span><span class="pl-s">'val_loss'</span>)
<span class="pl-s1">val_accuracy</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">metrics</span>.<span class="pl-v">SparseCategoricalAccuracy</span>(<span class="pl-s1">name</span><span class="pl-c1">=</span><span class="pl-s">'val_accuracy'</span>)

<span class="pl-s">"""## Training and checkpointing"""</span>

<span class="pl-s1">transformer</span> <span class="pl-c1">=</span> <span class="pl-v">Transformer</span>(<span class="pl-s1">num_layers</span>, <span class="pl-s1">d_model</span>, <span class="pl-s1">num_heads</span>, <span class="pl-s1">dff</span>,
                          <span class="pl-s1">input_vocab_size</span>, <span class="pl-s1">target_vocab_size</span>, <span class="pl-s1">dropout_rate</span>)


<span class="pl-k">def</span> <span class="pl-en">create_masks</span>(<span class="pl-s1">inp</span>, <span class="pl-s1">tar</span>):
    <span class="pl-c"># Encoder padding mask</span>
    <span class="pl-s1">enc_padding_mask</span> <span class="pl-c1">=</span> <span class="pl-en">create_padding_mask</span>(<span class="pl-s1">inp</span>)

    <span class="pl-c"># Used in the 2nd attention block in the decoder.</span>
    <span class="pl-c"># This padding mask is used to mask the encoder outputs.</span>
    <span class="pl-s1">dec_padding_mask</span> <span class="pl-c1">=</span> <span class="pl-en">create_padding_mask</span>(<span class="pl-s1">inp</span>)

    <span class="pl-c"># Used in the 1st attention block in the decoder.</span>
    <span class="pl-c"># It is used to pad and mask future tokens in the input received by</span>
    <span class="pl-c"># the decoder.</span>
    <span class="pl-s1">look_ahead_mask</span> <span class="pl-c1">=</span> <span class="pl-en">create_look_ahead_mask</span>(<span class="pl-s1">tf</span>.<span class="pl-en">shape</span>(<span class="pl-s1">tar</span>)[<span class="pl-c1">1</span>])
    <span class="pl-s1">dec_target_padding_mask</span> <span class="pl-c1">=</span> <span class="pl-en">create_padding_mask</span>(<span class="pl-s1">tar</span>)
    <span class="pl-s1">combined_mask</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-en">maximum</span>(<span class="pl-s1">dec_target_padding_mask</span>, <span class="pl-s1">look_ahead_mask</span>)

    <span class="pl-k">return</span> <span class="pl-s1">enc_padding_mask</span>, <span class="pl-s1">combined_mask</span>, <span class="pl-s1">dec_padding_mask</span>


<span class="pl-s">"""Create the checkpoint path and the checkpoint manager. This will be used to save checkpoints every `n` epochs."""</span>

<span class="pl-s1">checkpoint_path</span> <span class="pl-c1">=</span> <span class="pl-s">"./checkpoints/train"</span>

<span class="pl-s1">ckpt</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">train</span>.<span class="pl-v">Checkpoint</span>(<span class="pl-s1">transformer</span><span class="pl-c1">=</span><span class="pl-s1">transformer</span>,
                           <span class="pl-s1">optimizer</span><span class="pl-c1">=</span><span class="pl-s1">optimizer</span>)

<span class="pl-s1">ckpt_manager</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">train</span>.<span class="pl-v">CheckpointManager</span>(<span class="pl-s1">ckpt</span>, <span class="pl-s1">checkpoint_path</span>, <span class="pl-s1">max_to_keep</span><span class="pl-c1">=</span><span class="pl-c1">5</span>)

<span class="pl-c"># if a checkpoint exists, restore the latest checkpoint.</span>
<span class="pl-k">if</span> <span class="pl-s1">ckpt_manager</span>.<span class="pl-s1">latest_checkpoint</span>:
    <span class="pl-s1">ckpt</span>.<span class="pl-en">restore</span>(<span class="pl-s1">ckpt_manager</span>.<span class="pl-s1">latest_checkpoint</span>)
    <span class="pl-en">print</span>(<span class="pl-s">'Latest checkpoint restored!!'</span>)


<span class="pl-v">EPOCHS</span> <span class="pl-c1">=</span> <span class="pl-c1">200</span>


<span class="pl-s1">train_num</span> <span class="pl-c1">=</span> <span class="pl-en">len</span>([<span class="pl-c1">1</span> <span class="pl-k">for</span> <span class="pl-s1">_</span>, <span class="pl-s1">_</span> <span class="pl-c1">in</span> <span class="pl-s1">train_dataset</span>])

<span class="pl-en">@<span class="pl-s1">tf</span>.<span class="pl-s1">function</span></span>
<span class="pl-k">def</span> <span class="pl-en">train_step</span>(<span class="pl-s1">inp</span>, <span class="pl-s1">tar</span>):
    <span class="pl-s1">tar_inp</span> <span class="pl-c1">=</span> <span class="pl-s1">tar</span>[:, :<span class="pl-c1">-</span><span class="pl-c1">1</span>]
    <span class="pl-s1">tar_real</span> <span class="pl-c1">=</span> <span class="pl-s1">tar</span>[:, <span class="pl-c1">1</span>:]

    <span class="pl-s1">enc_padding_mask</span>, <span class="pl-s1">combined_mask</span>, <span class="pl-s1">dec_padding_mask</span> <span class="pl-c1">=</span> <span class="pl-en">create_masks</span>(<span class="pl-s1">inp</span>, <span class="pl-s1">tar_inp</span>)

    <span class="pl-k">with</span> <span class="pl-s1">tf</span>.<span class="pl-v">GradientTape</span>() <span class="pl-k">as</span> <span class="pl-s1">tape</span>:
        <span class="pl-s1">predictions</span>, <span class="pl-s1">_</span> <span class="pl-c1">=</span> <span class="pl-en">transformer</span>(<span class="pl-s1">inp</span>, <span class="pl-s1">tar_inp</span>,
                                     <span class="pl-c1">True</span>,
                                     <span class="pl-s1">enc_padding_mask</span>,
                                     <span class="pl-s1">combined_mask</span>,
                                     <span class="pl-s1">dec_padding_mask</span>)
        <span class="pl-s1">loss</span> <span class="pl-c1">=</span> <span class="pl-en">loss_function</span>(<span class="pl-s1">tar_real</span>, <span class="pl-s1">predictions</span>)

    <span class="pl-s1">gradients</span> <span class="pl-c1">=</span> <span class="pl-s1">tape</span>.<span class="pl-en">gradient</span>(<span class="pl-s1">loss</span>, <span class="pl-s1">transformer</span>.<span class="pl-s1">trainable_variables</span>)
    <span class="pl-s1">optimizer</span>.<span class="pl-en">apply_gradients</span>(<span class="pl-en">zip</span>(<span class="pl-s1">gradients</span>, <span class="pl-s1">transformer</span>.<span class="pl-s1">trainable_variables</span>))

    <span class="pl-en">train_loss</span>(<span class="pl-s1">loss</span>)
    <span class="pl-en">train_accuracy</span>(<span class="pl-s1">tar_real</span>, <span class="pl-s1">predictions</span>)

<span class="pl-en">@<span class="pl-s1">tf</span>.<span class="pl-s1">function</span></span>
<span class="pl-k">def</span> <span class="pl-en">val_step</span>(<span class="pl-s1">inp</span>, <span class="pl-s1">tar</span>):
    <span class="pl-s1">tar_inp</span> <span class="pl-c1">=</span> <span class="pl-s1">tar</span>[:, :<span class="pl-c1">-</span><span class="pl-c1">1</span>]
    <span class="pl-s1">tar_real</span> <span class="pl-c1">=</span> <span class="pl-s1">tar</span>[:, <span class="pl-c1">1</span>:]
    <span class="pl-s1">enc_padding_mask</span>, <span class="pl-s1">combined_mask</span>, <span class="pl-s1">dec_padding_mask</span> <span class="pl-c1">=</span> <span class="pl-en">create_masks</span>(<span class="pl-s1">inp</span>, <span class="pl-s1">tar_inp</span>)
    <span class="pl-s1">predictions</span>, <span class="pl-s1">_</span> <span class="pl-c1">=</span> <span class="pl-en">transformer</span>(<span class="pl-s1">inp</span>,<span class="pl-s1">tar_inp</span>,
                                <span class="pl-s1">enc_padding_mask</span><span class="pl-c1">=</span><span class="pl-s1">enc_padding_mask</span>,
                                <span class="pl-s1">look_ahead_mask</span><span class="pl-c1">=</span><span class="pl-s1">combined_mask</span>,
                                <span class="pl-s1">dec_padding_mask</span><span class="pl-c1">=</span><span class="pl-s1">dec_padding_mask</span>,
                                <span class="pl-s1">training</span><span class="pl-c1">=</span><span class="pl-c1">False</span>)
    <span class="pl-s1">loss</span> <span class="pl-c1">=</span> <span class="pl-en">loss_function</span>(<span class="pl-s1">tar_real</span>, <span class="pl-s1">predictions</span>)
    <span class="pl-s1">ppl</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-en">exp</span>(<span class="pl-s1">loss</span>)
    <span class="pl-en">val_loss</span>(<span class="pl-s1">ppl</span>)
    <span class="pl-en">val_accuracy</span>(<span class="pl-s1">tar_real</span>, <span class="pl-s1">predictions</span>)

<span class="pl-en">print</span>(<span class="pl-s">"(3):Traning model......"</span>)
<span class="pl-s">"""Portuguese is used as the input language and English is the target language."""</span>

<span class="pl-k">for</span> <span class="pl-s1">epoch</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-v">EPOCHS</span>):
    <span class="pl-s1">train_loss</span>.<span class="pl-en">reset_states</span>()
    <span class="pl-s1">train_accuracy</span>.<span class="pl-en">reset_states</span>()
    <span class="pl-s1">val_loss</span>.<span class="pl-en">reset_states</span>()
    <span class="pl-s1">val_accuracy</span>.<span class="pl-en">reset_states</span>()

    <span class="pl-en">print</span>(<span class="pl-s">'Epoch {}'</span>.<span class="pl-en">format</span>(<span class="pl-s1">epoch</span> <span class="pl-c1">+</span> <span class="pl-c1">1</span>))
    <span class="pl-s1">start</span> <span class="pl-c1">=</span> <span class="pl-s1">time</span>.<span class="pl-en">time</span>()
    <span class="pl-c"># inp -&gt; portuguese, tar -&gt; english</span>
    <span class="pl-k">with</span> <span class="pl-en">tqdm</span>(<span class="pl-s1">total</span><span class="pl-c1">=</span><span class="pl-s1">train_num</span> <span class="pl-c1">*</span> <span class="pl-v">BATCH_SIZE</span>) <span class="pl-k">as</span> <span class="pl-s1">pbar</span>:
        <span class="pl-k">for</span> <span class="pl-s1">inp</span>, <span class="pl-s1">tar</span> <span class="pl-c1">in</span> <span class="pl-s1">train_dataset</span>:
            <span class="pl-en">train_step</span>(<span class="pl-s1">inp</span>, <span class="pl-s1">tar</span>)
            <span class="pl-s1">pbar</span>.<span class="pl-en">update</span>(<span class="pl-v">BATCH_SIZE</span>)

    <span class="pl-k">for</span> <span class="pl-s1">inp</span>, <span class="pl-s1">tar</span> <span class="pl-c1">in</span> <span class="pl-s1">val_dataset</span>:
        <span class="pl-en">val_step</span>(<span class="pl-s1">inp</span>, <span class="pl-s1">tar</span>)

    <span class="pl-s1">end</span> <span class="pl-c1">=</span> <span class="pl-s1">time</span>.<span class="pl-en">time</span>()
    <span class="pl-en">print</span>(<span class="pl-s">'train_loss {:.4f}<span class="pl-cce">\t</span>train_acc {:.2f}<span class="pl-cce">\t</span>'</span>
          <span class="pl-s">'val_loss {:.4f}<span class="pl-cce">\t</span>val_acc {:.2f}<span class="pl-cce">\t</span>'</span>
          <span class="pl-s">'time {:.2f}s'</span>.<span class="pl-en">format</span>(<span class="pl-s1">train_loss</span>.<span class="pl-en">result</span>(),
                                <span class="pl-s1">train_accuracy</span>.<span class="pl-en">result</span>() <span class="pl-c1">*</span> <span class="pl-c1">100</span>,
                                <span class="pl-s1">val_loss</span>.<span class="pl-en">result</span>(),
                                <span class="pl-s1">val_accuracy</span>.<span class="pl-en">result</span>() <span class="pl-c1">*</span> <span class="pl-c1">100</span>,
                                <span class="pl-s1">end</span> <span class="pl-c1">-</span> <span class="pl-s1">start</span>,
                                ))


<span class="pl-k">def</span> <span class="pl-en">evaluate</span>(<span class="pl-s1">inp_sentence</span>):
    <span class="pl-s1">start_token</span> <span class="pl-c1">=</span> [<span class="pl-s1">tokenizer_pt</span>.<span class="pl-s1">vocab_size</span>]
    <span class="pl-s1">end_token</span> <span class="pl-c1">=</span> [<span class="pl-s1">tokenizer_pt</span>.<span class="pl-s1">vocab_size</span> <span class="pl-c1">+</span> <span class="pl-c1">1</span>]

    <span class="pl-c"># inp sentence is portuguese, hence adding the start and end token</span>
    <span class="pl-s1">inp_sentence</span> <span class="pl-c1">=</span> <span class="pl-s1">start_token</span> <span class="pl-c1">+</span> <span class="pl-s1">tokenizer_pt</span>.<span class="pl-en">encode</span>(<span class="pl-s1">inp_sentence</span>) <span class="pl-c1">+</span> <span class="pl-s1">end_token</span>
    <span class="pl-s1">encoder_input</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-en">expand_dims</span>(<span class="pl-s1">inp_sentence</span>, <span class="pl-c1">0</span>)

    <span class="pl-c"># as the target is english, the first word to the transformer should be the</span>
    <span class="pl-c"># english start token.</span>
    <span class="pl-s1">decoder_input</span> <span class="pl-c1">=</span> [<span class="pl-s1">tokenizer_en</span>.<span class="pl-s1">vocab_size</span>]
    <span class="pl-s1">output</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-en">expand_dims</span>(<span class="pl-s1">decoder_input</span>, <span class="pl-c1">0</span>)

    <span class="pl-k">for</span> <span class="pl-s1">i</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-v">MAX_LENGTH</span>):
        <span class="pl-s1">enc_padding_mask</span>, <span class="pl-s1">combined_mask</span>, <span class="pl-s1">dec_padding_mask</span> <span class="pl-c1">=</span> <span class="pl-en">create_masks</span>(
            <span class="pl-s1">encoder_input</span>, <span class="pl-s1">output</span>)

        <span class="pl-c"># predictions.shape == (batch_size, seq_len, vocab_size)</span>
        <span class="pl-s1">predictions</span>, <span class="pl-s1">attention_weights</span> <span class="pl-c1">=</span> <span class="pl-en">transformer</span>(<span class="pl-s1">encoder_input</span>,
                                                     <span class="pl-s1">output</span>,
                                                     <span class="pl-c1">False</span>,
                                                     <span class="pl-s1">enc_padding_mask</span>,
                                                     <span class="pl-s1">combined_mask</span>,
                                                     <span class="pl-s1">dec_padding_mask</span>)

        <span class="pl-c"># select the last word from the seq_len dimension</span>
        <span class="pl-s1">predictions</span> <span class="pl-c1">=</span> <span class="pl-s1">predictions</span>[:, <span class="pl-c1">-</span><span class="pl-c1">1</span>:, :]  <span class="pl-c"># (batch_size, 1, vocab_size)</span>

        <span class="pl-s1">predicted_id</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-en">cast</span>(<span class="pl-s1">tf</span>.<span class="pl-en">argmax</span>(<span class="pl-s1">predictions</span>, <span class="pl-s1">axis</span><span class="pl-c1">=</span><span class="pl-c1">-</span><span class="pl-c1">1</span>), <span class="pl-s1">tf</span>.<span class="pl-s1">int32</span>)

        <span class="pl-c"># return the result if the predicted_id is equal to the end token</span>
        <span class="pl-k">if</span> <span class="pl-s1">tf</span>.<span class="pl-en">equal</span>(<span class="pl-s1">predicted_id</span>, <span class="pl-s1">tokenizer_en</span>.<span class="pl-s1">vocab_size</span> <span class="pl-c1">+</span> <span class="pl-c1">1</span>):
            <span class="pl-k">return</span> <span class="pl-s1">tf</span>.<span class="pl-en">squeeze</span>(<span class="pl-s1">output</span>, <span class="pl-s1">axis</span><span class="pl-c1">=</span><span class="pl-c1">0</span>), <span class="pl-s1">attention_weights</span>

        <span class="pl-c"># concatentate the predicted_id to the output which is given to the decoder</span>
        <span class="pl-c"># as its input.</span>
        <span class="pl-s1">output</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-en">concat</span>([<span class="pl-s1">output</span>, <span class="pl-s1">predicted_id</span>], <span class="pl-s1">axis</span><span class="pl-c1">=</span><span class="pl-c1">-</span><span class="pl-c1">1</span>)

    <span class="pl-k">return</span> <span class="pl-s1">tf</span>.<span class="pl-en">squeeze</span>(<span class="pl-s1">output</span>, <span class="pl-s1">axis</span><span class="pl-c1">=</span><span class="pl-c1">0</span>), <span class="pl-s1">attention_weights</span>


<span class="pl-k">def</span> <span class="pl-en">plot_attention_weights</span>(<span class="pl-s1">attention</span>, <span class="pl-s1">sentence</span>, <span class="pl-s1">result</span>, <span class="pl-s1">layer</span>):
    <span class="pl-s1">fig</span> <span class="pl-c1">=</span> <span class="pl-s1">plt</span>.<span class="pl-en">figure</span>(<span class="pl-s1">figsize</span><span class="pl-c1">=</span>(<span class="pl-c1">16</span>, <span class="pl-c1">8</span>))

    <span class="pl-s1">sentence</span> <span class="pl-c1">=</span> <span class="pl-s1">tokenizer_pt</span>.<span class="pl-en">encode</span>(<span class="pl-s1">sentence</span>)

    <span class="pl-s1">attention</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-en">squeeze</span>(<span class="pl-s1">attention</span>[<span class="pl-s1">layer</span>], <span class="pl-s1">axis</span><span class="pl-c1">=</span><span class="pl-c1">0</span>)

    <span class="pl-k">for</span> <span class="pl-s1">head</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-s1">attention</span>.<span class="pl-s1">shape</span>[<span class="pl-c1">0</span>]):
        <span class="pl-s1">ax</span> <span class="pl-c1">=</span> <span class="pl-s1">fig</span>.<span class="pl-en">add_subplot</span>(<span class="pl-c1">2</span>, <span class="pl-c1">4</span>, <span class="pl-s1">head</span> <span class="pl-c1">+</span> <span class="pl-c1">1</span>)

        <span class="pl-c"># plot the attention weights</span>
        <span class="pl-s1">ax</span>.<span class="pl-en">matshow</span>(<span class="pl-s1">attention</span>[<span class="pl-s1">head</span>][:<span class="pl-c1">-</span><span class="pl-c1">1</span>, :], <span class="pl-s1">cmap</span><span class="pl-c1">=</span><span class="pl-s">'viridis'</span>)

        <span class="pl-s1">fontdict</span> <span class="pl-c1">=</span> {<span class="pl-s">'fontsize'</span>: <span class="pl-c1">10</span>}

        <span class="pl-s1">ax</span>.<span class="pl-en">set_xticks</span>(<span class="pl-en">range</span>(<span class="pl-en">len</span>(<span class="pl-s1">sentence</span>) <span class="pl-c1">+</span> <span class="pl-c1">2</span>))
        <span class="pl-s1">ax</span>.<span class="pl-en">set_yticks</span>(<span class="pl-en">range</span>(<span class="pl-en">len</span>(<span class="pl-s1">result</span>)))

        <span class="pl-s1">ax</span>.<span class="pl-en">set_ylim</span>(<span class="pl-en">len</span>(<span class="pl-s1">result</span>) <span class="pl-c1">-</span> <span class="pl-c1">1.5</span>, <span class="pl-c1">-</span><span class="pl-c1">0.5</span>)

        <span class="pl-s1">ax</span>.<span class="pl-en">set_xticklabels</span>(
            [<span class="pl-s">'&lt;start&gt;'</span>] <span class="pl-c1">+</span> [<span class="pl-s1">tokenizer_pt</span>.<span class="pl-en">decode</span>([<span class="pl-s1">i</span>]) <span class="pl-k">for</span> <span class="pl-s1">i</span> <span class="pl-c1">in</span> <span class="pl-s1">sentence</span>] <span class="pl-c1">+</span> [<span class="pl-s">'&lt;end&gt;'</span>],
            <span class="pl-s1">fontdict</span><span class="pl-c1">=</span><span class="pl-s1">fontdict</span>, <span class="pl-s1">rotation</span><span class="pl-c1">=</span><span class="pl-c1">90</span>)

        <span class="pl-s1">ax</span>.<span class="pl-en">set_yticklabels</span>([<span class="pl-s1">tokenizer_en</span>.<span class="pl-en">decode</span>([<span class="pl-s1">i</span>]) <span class="pl-k">for</span> <span class="pl-s1">i</span> <span class="pl-c1">in</span> <span class="pl-s1">result</span>
                            <span class="pl-k">if</span> <span class="pl-s1">i</span> <span class="pl-c1">&lt;</span> <span class="pl-s1">tokenizer_en</span>.<span class="pl-s1">vocab_size</span>],
                           <span class="pl-s1">fontdict</span><span class="pl-c1">=</span><span class="pl-s1">fontdict</span>)

        <span class="pl-s1">ax</span>.<span class="pl-en">set_xlabel</span>(<span class="pl-s">'Head {}'</span>.<span class="pl-en">format</span>(<span class="pl-s1">head</span> <span class="pl-c1">+</span> <span class="pl-c1">1</span>))

    <span class="pl-s1">plt</span>.<span class="pl-en">tight_layout</span>()
    <span class="pl-s1">plt</span>.<span class="pl-en">show</span>()


<span class="pl-k">def</span> <span class="pl-en">translate</span>(<span class="pl-s1">sentence</span>, <span class="pl-s1">plot</span><span class="pl-c1">=</span><span class="pl-s">''</span>):
    <span class="pl-s1">result</span>, <span class="pl-s1">attention_weights</span> <span class="pl-c1">=</span> <span class="pl-en">evaluate</span>(<span class="pl-s1">sentence</span>)

    <span class="pl-s1">predicted_sentence</span> <span class="pl-c1">=</span> <span class="pl-s1">tokenizer_en</span>.<span class="pl-en">decode</span>([<span class="pl-s1">i</span> <span class="pl-k">for</span> <span class="pl-s1">i</span> <span class="pl-c1">in</span> <span class="pl-s1">result</span>
                                              <span class="pl-k">if</span> <span class="pl-s1">i</span> <span class="pl-c1">&lt;</span> <span class="pl-s1">tokenizer_en</span>.<span class="pl-s1">vocab_size</span>])

    <span class="pl-en">print</span>(<span class="pl-s">'Input: {}'</span>.<span class="pl-en">format</span>(<span class="pl-s1">sentence</span>))
    <span class="pl-en">print</span>(<span class="pl-s">'Predicted translation: {}'</span>.<span class="pl-en">format</span>(<span class="pl-s1">predicted_sentence</span>))

    <span class="pl-k">if</span> <span class="pl-s1">plot</span>:
        <span class="pl-en">plot_attention_weights</span>(<span class="pl-s1">attention_weights</span>, <span class="pl-s1">sentence</span>, <span class="pl-s1">result</span>, <span class="pl-s1">plot</span>)

<span class="pl-en">print</span>(<span class="pl-s">"(4):Evaluate model......"</span>)
<span class="pl-en">translate</span>(<span class="pl-s">"este é um problema que temos que resolver."</span>)
<span class="pl-en">print</span>(<span class="pl-s">"Real translation: this is a problem we have to solve ."</span>)

<span class="pl-en">translate</span>(<span class="pl-s">"os meus vizinhos ouviram sobre esta ideia."</span>)
<span class="pl-en">print</span>(<span class="pl-s">"Real translation: and my neighboring homes heard about this idea ."</span>)

<span class="pl-en">translate</span>(<span class="pl-s">"vou então muito rapidamente partilhar convosco algumas histórias de algumas coisas mágicas que aconteceram."</span>)
<span class="pl-en">print</span>(
    <span class="pl-s">"Real translation: so i 'll just share with you some stories very quickly of some magical things that have happened ."</span>)

<span class="pl-s">"""You can pass different layers and attention blocks of the decoder to the `plot` parameter."""</span>

<span class="pl-en">translate</span>(<span class="pl-s">"este é o primeiro livro que eu fiz."</span>, <span class="pl-s1">plot</span><span class="pl-c1">=</span><span class="pl-s">'decoder_layer4_block2'</span>)
<span class="pl-en">print</span>(<span class="pl-s">"Real translation: this is the first book i've ever done."</span>)</pre></div>
<p dir="auto"><strong>System information</strong></p>
<ul dir="auto">
<li>OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Windows 10 64bit</li>
<li>TensorFlow installed from (source or binary): pip</li>
<li>TensorFlow version: 1.12.0</li>
<li>Python version: 3.6.7 64bit</li>
<li>Installed using virtualenv? pip? conda?: No virtualenv. Direct install with pip. No conda. Virtualenv tried, but result is the same.</li>
<li>CUDA/cuDNN version: cuda_9.0.176_win10, cudnn-9.0-windows10-x64-v7.5.0.56</li>
<li>GPU model and memory: RTX2070 8G</li>
<li>CPU model: intel i5-9400f (CPU without integrated GPU)</li>
</ul>
<p dir="auto"><strong>Describe the problem</strong></p>
<p dir="auto">To be clear, this is NOT a issue about the dll files that can not be found. Files are there, but something is wrong with the previlige.</p>
<p dir="auto">I recently built a new windows 10 PC and installed tensorflow on it. The installation was successful, but I had to run the code with Administrator. Otherwise a simple import code like "import tensorflow as tf" would result in "ImportError: DLL load failed: Access is denied" if I try to run it without "Run as Administrator". The code did work flawlessly if I choose to run as admin though. But on my old notebook I could execute without admin previlige.</p>
<p dir="auto">The installations on my new and old computer were the same, so the software version and system environment variables should not be the problem. I am wondering why this is happening and is there any way I can get rid of the "run as admin" process which is rather danger and annoying for every time I want to run the code.</p>
<p dir="auto"><strong>Any other info / logs</strong></p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="3.6.7 (v3.6.7:6ec5cf24b7, Oct 20 2018, 13:35:33) [MSC v.1900 64 bit (AMD64)]
Traceback (most recent call last):
  File &quot;C:\Python36\lib\site-packages\tensorflow\python\pywrap_tensorflow.py&quot;, line 58, in &lt;module&gt;
    from tensorflow.python.pywrap_tensorflow_internal import *
  File &quot;C:\Python36\lib\site-packages\tensorflow\python\pywrap_tensorflow_internal.py&quot;, line 28, in &lt;module&gt;
    _pywrap_tensorflow_internal = swig_import_helper()
  File &quot;C:\Python36\lib\site-packages\tensorflow\python\pywrap_tensorflow_internal.py&quot;, line 24, in swig_import_helper
    _mod = imp.load_module('_pywrap_tensorflow_internal', fp, pathname, description)
  File &quot;C:\Python36\lib\imp.py&quot;, line 243, in load_module
    return load_dynamic(name, filename, file)
  File &quot;C:\Python36\lib\imp.py&quot;, line 343, in load_dynamic
    return _load(spec)
ImportError: DLL load failed: Access is denied.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File &quot;hello tensorflow.py&quot;, line 12, in &lt;module&gt;
    import tensorflow as tf
  File &quot;C:\Python36\lib\site-packages\tensorflow\__init__.py&quot;, line 24, in &lt;module&gt;
    from tensorflow.python import pywrap_tensorflow  # pylint: disable=unused-import
  File &quot;C:\Python36\lib\site-packages\tensorflow\python\__init__.py&quot;, line 49, in &lt;module&gt;
    from tensorflow.python import pywrap_tensorflow
  File &quot;C:\Python36\lib\site-packages\tensorflow\python\pywrap_tensorflow.py&quot;, line 74, in &lt;module&gt;
    raise ImportError(msg)
ImportError: Traceback (most recent call last):
  File &quot;C:\Python36\lib\site-packages\tensorflow\python\pywrap_tensorflow.py&quot;, line 58, in &lt;module&gt;
    from tensorflow.python.pywrap_tensorflow_internal import *
  File &quot;C:\Python36\lib\site-packages\tensorflow\python\pywrap_tensorflow_internal.py&quot;, line 28, in &lt;module&gt;
    _pywrap_tensorflow_internal = swig_import_helper()
  File &quot;C:\Python36\lib\site-packages\tensorflow\python\pywrap_tensorflow_internal.py&quot;, line 24, in swig_import_helper
    _mod = imp.load_module('_pywrap_tensorflow_internal', fp, pathname, description)
  File &quot;C:\Python36\lib\imp.py&quot;, line 243, in load_module
    return load_dynamic(name, filename, file)
  File &quot;C:\Python36\lib\imp.py&quot;, line 343, in load_dynamic
    return _load(spec)
ImportError: DLL load failed: Access is denied.


Failed to load the native TensorFlow runtime.

See https://www.tensorflow.org/install/errors

for some common reasons and solutions.  Include the entire stack trace
above this error message when asking for help."><pre class="notranslate"><code class="notranslate">3.6.7 (v3.6.7:6ec5cf24b7, Oct 20 2018, 13:35:33) [MSC v.1900 64 bit (AMD64)]
Traceback (most recent call last):
  File "C:\Python36\lib\site-packages\tensorflow\python\pywrap_tensorflow.py", line 58, in &lt;module&gt;
    from tensorflow.python.pywrap_tensorflow_internal import *
  File "C:\Python36\lib\site-packages\tensorflow\python\pywrap_tensorflow_internal.py", line 28, in &lt;module&gt;
    _pywrap_tensorflow_internal = swig_import_helper()
  File "C:\Python36\lib\site-packages\tensorflow\python\pywrap_tensorflow_internal.py", line 24, in swig_import_helper
    _mod = imp.load_module('_pywrap_tensorflow_internal', fp, pathname, description)
  File "C:\Python36\lib\imp.py", line 243, in load_module
    return load_dynamic(name, filename, file)
  File "C:\Python36\lib\imp.py", line 343, in load_dynamic
    return _load(spec)
ImportError: DLL load failed: Access is denied.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "hello tensorflow.py", line 12, in &lt;module&gt;
    import tensorflow as tf
  File "C:\Python36\lib\site-packages\tensorflow\__init__.py", line 24, in &lt;module&gt;
    from tensorflow.python import pywrap_tensorflow  # pylint: disable=unused-import
  File "C:\Python36\lib\site-packages\tensorflow\python\__init__.py", line 49, in &lt;module&gt;
    from tensorflow.python import pywrap_tensorflow
  File "C:\Python36\lib\site-packages\tensorflow\python\pywrap_tensorflow.py", line 74, in &lt;module&gt;
    raise ImportError(msg)
ImportError: Traceback (most recent call last):
  File "C:\Python36\lib\site-packages\tensorflow\python\pywrap_tensorflow.py", line 58, in &lt;module&gt;
    from tensorflow.python.pywrap_tensorflow_internal import *
  File "C:\Python36\lib\site-packages\tensorflow\python\pywrap_tensorflow_internal.py", line 28, in &lt;module&gt;
    _pywrap_tensorflow_internal = swig_import_helper()
  File "C:\Python36\lib\site-packages\tensorflow\python\pywrap_tensorflow_internal.py", line 24, in swig_import_helper
    _mod = imp.load_module('_pywrap_tensorflow_internal', fp, pathname, description)
  File "C:\Python36\lib\imp.py", line 243, in load_module
    return load_dynamic(name, filename, file)
  File "C:\Python36\lib\imp.py", line 343, in load_dynamic
    return _load(spec)
ImportError: DLL load failed: Access is denied.


Failed to load the native TensorFlow runtime.

See https://www.tensorflow.org/install/errors

for some common reasons and solutions.  Include the entire stack trace
above this error message when asking for help.
</code></pre></div>