"""Transformer from 'Attention is all you need' (Vaswani et al., 2017)"""
# Reference: https://www.tensorflow.org/text/tutorials/transformer
# Reference: https://keras.io/examples/nlp/text_classification_with_transformer/

import numpy as np
import tensorflow as tf


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

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

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

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

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

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

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

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

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

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

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


class Encoder(tf.keras.layers.Layer):
    """Transformer encoder from 'Attention is all you need' (Vaswani et al., 2017)

    Contains:
        1. Input Embedding
        2. Positional Encoding
        3. N encoder layers
    """

    def __init__(
        self,
        num_layers,
        d_model,
        num_heads,
        dff,
        input_vocab_size,
        maximum_position_encoding,
        rate=0.1,
    ):
        super(Encoder, self).__init__()

        self.d_model = d_model
        self.num_layers = num_layers

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

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

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

    def call(self, x, training, mask):
        seq_len = tf.shape(x)[1]

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

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

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

        return x  # (batch_size, input_seq_len, d_model)


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

        self.d_model = d_model
        self.num_layers = num_layers

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

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

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

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

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

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

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

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

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


class EncoderLayer(tf.keras.layers.Layer):
    """Transformer encoder layer from 'Attention is all you need' (Vaswani et al., 2017)

    One of the main difference between the transformer encoder from decoder is
    the self-attention. The reasons for it is detailed in the Section 4 and can
    be summarized as a way to reduce the path length between long-range depencies
    in the network.
    """

    def __init__(self, d_model=512, num_heads=8, dff=2048, rate=0.1):
        """Initializer a Transformer Encoder Layer

        Attributes
        ----------
        d_model : int
            Model dimension used on all sub-layers and embedding.
        num_heads : int
            Number of heads. Vaswani et al., 2017 describes as $h$
        dff : int
            FeedForward dimension.
        rate : float
            Dropout rate parameter applied after self-attention and
            FeedForward.
        """

        super(EncoderLayer, self).__init__()
        self.mha = MultiHeadAttention(num_heads=num_heads, d_model=d_model)
        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)
        return self.layernorm2(
            out1 + ffn_output
        )  # (batch_size, input_seq_len, d_model)


class DecoderLayer(tf.keras.layers.Layer):
    """Transformer decoder layer from 'Attention is all you need' (Vaswani et al., 2017)

    Decoder layer is similar to encoder but have a third sub-layer performing
    multi-head attention over the encoder stack. The self-attention sub-layer
    is modified preventing positions from attending to subsequent positions.
    Embeddings are also offset by one position, forcing predictions of
    position i to depend on the known outputs at positions less than i.
    """

    def __init__(self, d_model, num_heads, dff, rate=0.1):
        super(DecoderLayer, self).__init__()

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

        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


def _point_wise_feed_forward_network(d_model, dff):
    """Position-wise Feed-Forward Network

    It's a fully connnected feed-forward network applied to each position
    separately and identically represented by:
        ```
        FFN(x) = max(0, xW_1 + b_1)W_2 + b2$
        ```
    It contains two linear transformation with a ReLU activation in between.
    """
    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)
        ]
    )


def _create_padding_mask(seq):
    """Mask all the pad tokens in the batch of sequence"""

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

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


def _create_look_ahead_mask(size):
    """Mask the future tokens in a sequence"""
    mask = 1 - tf.linalg.band_part(tf.ones((size, size)), -1, 0)
    return mask  # (seq_len, seq_len)


def _positional_encoding(position, d_model):
    """Position Encoding (PE)

    Because the model contains no recurrence and convolution, positional
    encoding is inject to add information about absolute position of the
    tokens in the sequence. It can be fixed or learned, however, fixed
    has proven to be as efficient as learned.
    This is the fixed Positional Encoding and are derived from sine and
    cosine functions of different frequencies:
        $PE(pos, 2i) = sin(pos/10000^{2i/d_model})
        $PE(pos, 2i + 1) = cos(pos/10000^{2i/d_model})
    where pos is the absolute position of a token in the sequence and $i$
    is the dimension.
    """

    def get_angles(pos, i, d_model):
        angle_rates = 1 / np.power(10000, (2 * (i // 2)) / np.float32(d_model))
        return pos * angle_rates

    angle_rads = get_angles(
        np.arange(position)[:, np.newaxis], np.arange(d_model)[np.newaxis, :], d_model
    )

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

    pos_encoding = angle_rads[np.newaxis, ...]
    return tf.cast(pos_encoding, dtype=tf.float32)


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

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

    Returns:
      output, attention_weights
    """
    # (..., seq_len_q, seq_len_k)
    matmul_qk = tf.matmul(q, k, transpose_b=True)

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

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

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

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

    return output, attention_weights


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

        assert d_model % self.num_heads == 0

        self.depth = d_model // self.num_heads

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

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

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

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

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

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

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

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

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

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

        return output, attention_weights


def optimizer(d_model):
    """Adam optimizer as of Section 5.3"""

    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)
    return tf.keras.optimizers.Adam(
        learning_rate, beta_1=0.9, beta_2=0.98, epsilon=1e-9
    )

"""Transformer from 'Attention is all you need' (Vaswani et al., 2017)"""
# Reference: https://www.tensorflow.org/text/tutorials/transformer
# Reference: https://keras.io/examples/nlp/text_classification_with_transformer/

import numpy as np
import tensorflow as tf


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

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

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

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

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

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

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

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

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

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

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


class Encoder(tf.keras.layers.Layer):
    """Transformer encoder from 'Attention is all you need' (Vaswani et al., 2017)

    Contains:
        1. Input Embedding
        2. Positional Encoding
        3. N encoder layers
    """

    def __init__(
        self,
        num_layers,
        d_model,
        num_heads,
        dff,
        input_vocab_size,
        maximum_position_encoding,
        rate=0.1,
    ):
        super(Encoder, self).__init__()

        self.d_model = d_model
        self.num_layers = num_layers

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

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

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

    def call(self, x, training, mask):
        seq_len = tf.shape(x)[1]

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

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

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

        return x  # (batch_size, input_seq_len, d_model)


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

        self.d_model = d_model
        self.num_layers = num_layers

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

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

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

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

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

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

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

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

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


class EncoderLayer(tf.keras.layers.Layer):
    """Transformer encoder layer from 'Attention is all you need' (Vaswani et al., 2017)

    One of the main difference between the transformer encoder from decoder is
    the self-attention. The reasons for it is detailed in the Section 4 and can
    be summarized as a way to reduce the path length between long-range depencies
    in the network.
    """

    def __init__(self, d_model=512, num_heads=8, dff=2048, rate=0.1):
        """Initializer a Transformer Encoder Layer

        Attributes
        ----------
        d_model : int
            Model dimension used on all sub-layers and embedding.
        num_heads : int
            Number of heads. Vaswani et al., 2017 describes as $h$
        dff : int
            FeedForward dimension.
        rate : float
            Dropout rate parameter applied after self-attention and
            FeedForward.
        """

        super(EncoderLayer, self).__init__()
        self.mha = MultiHeadAttention(num_heads=num_heads, d_model=d_model)
        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)
        return self.layernorm2(
            out1 + ffn_output
        )  # (batch_size, input_seq_len, d_model)


class DecoderLayer(tf.keras.layers.Layer):
    """Transformer decoder layer from 'Attention is all you need' (Vaswani et al., 2017)

    Decoder layer is similar to encoder but have a third sub-layer performing
    multi-head attention over the encoder stack. The self-attention sub-layer
    is modified preventing positions from attending to subsequent positions.
    Embeddings are also offset by one position, forcing predictions of
    position i to depend on the known outputs at positions less than i.
    """

    def __init__(self, d_model, num_heads, dff, rate=0.1):
        super(DecoderLayer, self).__init__()

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

        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


def _point_wise_feed_forward_network(d_model, dff):
    """Position-wise Feed-Forward Network

    It's a fully connnected feed-forward network applied to each position
    separately and identically represented by:
        ```
        FFN(x) = max(0, xW_1 + b_1)W_2 + b2$
        ```
    It contains two linear transformation with a ReLU activation in between.
    """
    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)
        ]
    )


def _create_padding_mask(seq):
    """Mask all the pad tokens in the batch of sequence"""

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

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


def _create_look_ahead_mask(size):
    """Mask the future tokens in a sequence"""
    mask = 1 - tf.linalg.band_part(tf.ones((size, size)), -1, 0)
    return mask  # (seq_len, seq_len)


def _positional_encoding(position, d_model):
    """Position Encoding (PE)

    Because the model contains no recurrence and convolution, positional
    encoding is inject to add information about absolute position of the
    tokens in the sequence. It can be fixed or learned, however, fixed
    has proven to be as efficient as learned.
    This is the fixed Positional Encoding and are derived from sine and
    cosine functions of different frequencies:
        $PE(pos, 2i) = sin(pos/10000^{2i/d_model})
        $PE(pos, 2i + 1) = cos(pos/10000^{2i/d_model})
    where pos is the absolute position of a token in the sequence and $i$
    is the dimension.
    """

    def get_angles(pos, i, d_model):
        angle_rates = 1 / np.power(10000, (2 * (i // 2)) / np.float32(d_model))
        return pos * angle_rates

    angle_rads = get_angles(
        np.arange(position)[:, np.newaxis], np.arange(d_model)[np.newaxis, :], d_model
    )

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

    pos_encoding = angle_rads[np.newaxis, ...]
    return tf.cast(pos_encoding, dtype=tf.float32)


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

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

    Returns:
      output, attention_weights
    """
    # (..., seq_len_q, seq_len_k)
    matmul_qk = tf.matmul(q, k, transpose_b=True)

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

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

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

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

    return output, attention_weights


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

        assert d_model % self.num_heads == 0

        self.depth = d_model // self.num_heads

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

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

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

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

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

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

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

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

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

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

        return output, attention_weights


def optimizer(d_model):
    """Adam optimizer as of Section 5.3"""

    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)
    return tf.keras.optimizers.Adam(
        learning_rate, beta_1=0.9, beta_2=0.98, epsilon=1e-9
    )

Transformer from 'Attention is all you need' (Vaswani et al., 2017) # Reference: https://www.tensorflow.org/text/tutorials/transformer # Reference: https://keras.io/examples/nlp/text_classification_with_transformer/ # Keras models prefer if you pass all your inputs in the first argument # (batch_size, inp_seq_len, d_model) # dec_output.shape == (batch_size, tar_seq_len, d_model) # (batch_size, tar_seq_len, target_vocab_size) # Encoder padding mask # Used in the 2nd attention block in the decoder # This padding mask is used to mask the encoder outputs. # 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. Transformer encoder from 'Attention is all you need' (Vaswani et al., 2017)      Contains:         1. Input Embedding         2. Positional Encoding         3. N encoder layers # adding embedding and position encoding # (batch_size, input_seq_len, d_model) # (batch_size, input_seq_len, d_model) # (batch_size, target_seq_len, d_model) # x.shape == (batch_size, target_seq_len, d_model) Transformer encoder layer from 'Attention is all you need' (Vaswani et al., 2017)      One of the main difference between the transformer encoder from decoder is     the self-attention. The reasons for it is detailed in the Section 4 and can     be summarized as a way to reduce the path length between long-range depencies     in the network. Initializer a Transformer Encoder Layer          Attributes         ----------         d_model : int             Model dimension used on all sub-layers and embedding.         num_heads : int             Number of heads. Vaswani et al., 2017 describes as $h$         dff : int             FeedForward dimension.         rate : float             Dropout rate parameter applied after self-attention and             FeedForward. # (batch_size, input_seq_len, d_model) # (batch_size, input_seq_len, d_model) # (batch_size, input_seq_len, d_model) # (batch_size, input_seq_len, d_model) Transformer decoder layer from 'Attention is all you need' (Vaswani et al., 2017)      Decoder layer is similar to encoder but have a third sub-layer performing     multi-head attention over the encoder stack. The self-attention sub-layer     is modified preventing positions from attending to subsequent positions.     Embeddings are also offset by one position, forcing predictions of     position i to depend on the known outputs at positions less than i. # enc_output.shape == (batch_size, input_seq_len, d_model) # (batch_size, target_seq_len, d_model) # (batch_size, target_seq_len, d_model) # (batch_size, target_seq_len, d_model) # (batch_size, target_seq_len, d_model) # (batch_size, target_seq_len, d_model) Position-wise Feed-Forward Network      It's a fully connnected feed-forward network applied to each position     separately and identically represented by:         ```         FFN(x) = max(0, xW_1 + b_1)W_2 + b2$         ```     It contains two linear transformation with a ReLU activation in between. # (batch_size, seq_len, dff) # (batch_size, seq_len, d_model) Mask all the pad tokens in the batch of sequence # add extra dimensions to add the padding # to the attention logits. # (batch_size, 1, 1, seq_len) Mask the future tokens in a sequence # (seq_len, seq_len) Position Encoding (PE)      Because the model contains no recurrence and convolution, positional     encoding is inject to add information about absolute position of the     tokens in the sequence. It can be fixed or learned, however, fixed     has proven to be as efficient as learned.     This is the fixed Positional Encoding and are derived from sine and     cosine functions of different frequencies:         $PE(pos, 2i) = sin(pos/10000^{2i/d_model})         $PE(pos, 2i + 1) = cos(pos/10000^{2i/d_model})     where pos is the absolute position of a token in the sequence and $i$     is the dimension. # apply sin to even indices in the array; 2i # apply cos to odd indices in the array; 2i+1 Calculate the attention weights.     q, k, v must have matching leading dimensions.     k, v must have matching penultimate dimension, i.e.: seq_len_k = seq_len_v.     The mask has different shapes depending on its type(padding or look ahead)     but it must be broadcastable for addition.      Args:       q: query shape == (..., seq_len_q, depth)       k: key shape == (..., seq_len_k, depth)       v: value shape == (..., seq_len_v, depth_v)       mask: Float tensor with shape broadcastable             to (..., seq_len_q, seq_len_k). Defaults to None.      Returns:       output, attention_weights # (..., seq_len_q, seq_len_k) # scale matmul_qk # add the mask to the scaled tensor. # softmax is normalized on the last axis (seq_len_k) so that the scores # add up to 1. # (..., seq_len_q, seq_len_k) # (..., seq_len_q, depth_v) Split the last dimension into (num_heads, depth).         Transpose the result such that the shape is (batch_size, num_heads, seq_len, depth) # (batch_size, seq_len, d_model) # (batch_size, seq_len, d_model) # (batch_size, seq_len, d_model) # (batch_size, num_heads, seq_len_q, depth) # (batch_size, num_heads, seq_len_k, depth) # (batch_size, num_heads, seq_len_v, depth) # scaled_attention.shape == (batch_size, num_heads, seq_len_q, depth) # attention_weights.shape == (batch_size, num_heads, seq_len_q, seq_len_k) # (batch_size, seq_len_q, num_heads, depth) # (batch_size, seq_len_q, d_model) # (batch_size, seq_len_q, d_model) Adam optimizer as of Section 5.3