Generative AI - Practical Coding: Implementing Transformers

Practical Coding: Implementing Transformers

Using TensorFlow and Keras, we will take a simple method to making a Transformer model from scratch. The Transformer model doesn't need recurring structures to handle the links between inputs; it instead uses attention processes. This makes it suitable for sequence-to-sequence tasks like translation and text generation.

1. Import Libraries

import tensorflow as tf
from tensorflow.keras import layers, models
import numpy as np

For making the model, we import TensorFlow and Keras. For math tasks, we import NumPy.

2. Define the Scaled Dot-Product Attention

class ScaledDotProductAttention(layers.Layer):
    def __init__(self):
        super(ScaledDotProductAttention, self).__init__()

    def call(self, query, key, value, mask=None):
        matmul_qk = tf.matmul(query, key, transpose_b=True)
        dk = tf.cast(tf.shape(key)[-1], tf.float32)
        scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)

        if mask is not None:
            scaled_attention_logits += (mask * -1e9)

        attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)
        output = tf.matmul(attention_weights, value)

        return output, attention_weights

This class describes how the scaled dot-product attention system working. After figuring out the attention scores and scaling them, it uses a softmax function to get the attention weights. This is followed by finding the result, which is the weighted sum of the value vectors.

3. Define the Multi-Head Attention

class MultiHeadAttention(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 = layers.Dense(d_model)
        self.wk = layers.Dense(d_model)
        self.wv = layers.Dense(d_model)
        self.dense = layers.Dense(d_model)

    def split_heads(self, x, batch_size):
        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)
        k = self.wk(k)
        v = self.wv(v)

        q = self.split_heads(q, batch_size)
        k = self.split_heads(k, batch_size)
        v = self.split_heads(v, batch_size)

        scaled_attention, attention_weights = ScaledDotProductAttention()(q, k, v, mask)

        scaled_attention = tf.transpose(scaled_attention, perm=[0, 2, 1, 3])
        concat_attention = tf.reshape(scaled_attention, (batch_size, -1, self.d_model))

        output = self.dense(concat_attention)

        return output, attention_weights

4. Define the Encoder Layer

class EncoderLayer(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 = tf.keras.Sequential([
            layers.Dense(dff, activation='relu'),
            layers.Dense(d_model)
        ])

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

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

    def call(self, x, training, mask):
        attn_output, _ = self.mha(x, x, x, mask)
        attn_output = self.dropout1(attn_output, training=training)
        out1 = self.layernorm1(x + attn_output)

        ffn_output = self.ffn(out1)
        ffn_output = self.dropout2(ffn_output, training=training)
        out2 = self.layernorm2(out1 + ffn_output)

        return out2

It describes a single decoder layer that has a feed-forward network and multi-head attention. For regularization and stability, layer normalization and dropout are used.

5. Define the Transformer Encoder

It includes embedding and positional encoding layers. This class describes the Transformer decoder. Along with its own attention and feed-forward networks, each decoder layer is stacked on top of the others.

6. Compile and Train the Model

# Parameters
num_layers = 4
d_model = 128
dff = 512
num_heads = 8
input_vocab_size = 8500
maximum_position_encoding = 10000

encoder = Encoder(num_layers, d_model, num_heads, dff, input_vocab_size, maximum_position_encoding)

# Dummy input for demonstration
sample_input = tf.random.uniform((64, 62), dtype=tf.int64, minval=0, maxval=200)

# Model Call
sample_output = encoder(sample_input, training=False, mask=None)
print(sample_output.shape)  # (batch_size, input_seq_len, d_model)

Explanation:

Set the maximal position encoding, model dimension, feed-forward dimension, number of heads, and input word size.

Encoding: Create an instance of the encoder with the given settings.

Create a fake input to show how the encoder works.

Model Call: Let's test the solution by running the sample input through the encoder and printing the shape that comes out.

Doing these steps will help you make a basic Transformer model from start while learning about the main parts and how they work together. You can better understand Transformers and their powerful attention systems by doing this useful code activity.