Simple Transformer for Text Classification

A step-by-step guide to using a Transformer for a basic text classification task.

Libraries Used

This pipeline uses a few key Python libraries. Here is a quick overview of their roles:

• tensorflow: The core deep learning framework.
• numpy: Used for numerical operations.
• tensorflow_text: Provides tools for text preprocessing.

1. Data Loading and Preprocessing

• What the code does: It uses the `IMDB` movie review dataset. The text is tokenized and padded to a fixed length, and an `TextVectorization` layer is created to handle this process.
• Behind the scenes: Transformers work with sequences of fixed length. The `TextVectorization` layer handles the conversion of raw text into integer sequences, which is a key part of the preprocessing pipeline.
• Abstraction: The `TextVectorization` layer and `tf.data.Dataset` API abstract the entire process of text preprocessing, tokenization, and batching, making it much easier to prepare text data for a model.

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

# Load the IMDB dataset
(train_data, train_labels), (test_data, test_labels) = tf.keras.datasets.imdb.load_data(num_words=10000)

# Create a text vectorization layer
max_features = 20000
sequence_length = 500

vectorize_layer = layers.TextVectorization(
    max_tokens=max_features,
    output_mode='int',
    output_sequence_length=sequence_length
)

# Convert integer data back to text for the vectorizer
word_index = tf.keras.datasets.imdb.get_word_index()
reverse_word_index = dict([(value, key) for (key, value) in word_index.items()])
def decode_review(text_list):
    return " ".join([reverse_word_index.get(i - 3, "?") for i in text_list])
train_texts = [decode_review(s) for s in train_data]

vectorize_layer.adapt(train_texts)

2. Building the Transformer Model

• What the code does: It defines a custom `TransformerBlock` and a final model. The transformer block contains multi-head self-attention and dense layers, which are the core components of the transformer architecture.
• Behind the scenes: Self-attention is a mechanism that allows the model to weigh the importance of different words in a sequence when processing a specific word. It learns contextual relationships without using recurrence.
• Abstraction: We create a custom `TransformerBlock` class to abstract the complex self-attention mechanism, including the multi-head attention and feed-forward networks, which are crucial for the model's performance.

def transformer_block(embed_dim, num_heads, ff_dim, rate=0.1):
    inputs = layers.Input(shape=(None, embed_dim))
    attn_output = layers.MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)(inputs, inputs)
    attn_output = layers.Dropout(rate)(attn_output)
    norm1 = layers.LayerNormalization(epsilon=1e-6)(inputs + attn_output)
    ffn_output = layers.Dense(ff_dim, activation="relu")(norm1)
    ffn_output = layers.Dense(embed_dim)(ffn_output)
    ffn_output = layers.Dropout(rate)(ffn_output)
    return layers.LayerNormalization(epsilon=1e-6)(norm1 + ffn_output)

def get_transformer_model():
    inputs = layers.Input(shape=(sequence_length,))
    x = vectorize_layer(inputs)
    x = layers.Embedding(max_features, 128)(x)
    x = transformer_block(128, 2, 32)
    x = layers.GlobalAveragePooling1D()(x)
    x = layers.Dropout(0.1)(x)
    x = layers.Dense(20, activation="relu")(x)
    x = layers.Dropout(0.1)(x)
    outputs = layers.Dense(1, activation="sigmoid")(x)
    return models.Model(inputs, outputs)

3. Training and Evaluation

• What the code does: It compiles the model with an optimizer and trains it on the prepared text data.
• Behind the scenes: The model is trained to minimize a loss function (`binary_crossentropy`) and maximize a metric (`accuracy`) on the sentiment classification task. The transformer's attention mechanism helps it weigh the most important words in a review to determine its sentiment.
• Abstraction: Keras's `model.compile` and `model.fit` methods manage the training loop, including backpropagation and parameter updates, making it possible to train a complex transformer model with a single command.

model = get_transformer_model()
model.compile("adam", "binary_crossentropy", metrics=["accuracy"])

# Convert test data to text for evaluation
test_texts = [decode_review(s) for s in test_data]
test_texts = np.array(test_texts)

history = model.fit(np.array(train_texts), train_labels, batch_size=32, epochs=2, validation_data=(test_texts, test_labels))

test_loss, test_acc = model.evaluate(test_texts, test_labels, verbose=2)
print(f"\nTest accuracy: {test_acc}")