Transformer Models From Scratch in Python

Transformer neural networks (Vaswani et al. 2017) are a type of neural network specialized for modern natural language processing (NLP) tasks, such as translation, text generation, and chatbots.

They consist of an encoder and a decoder that work as an all-around NLP solution. In terms of the architecture, the main innovation behind the transformer is the self-attention mechanism—originally created for machine translation with RNNs.

The attention mechanism allows transformers to weigh the importance of different tokens (words or symbols) in an input sequence when making predictions.

In this issue, we’ll look at creating a transformer encoder from scratch using Python and Pytorch. We’ll dive into the full code and conceptual base behind the transformer, the architecture of choice for the most popular large language models (such as GPT or BERT).

The first thing we’ll do to build this architecture is to inspect how it works—how the data flows through the structure. We notice that inputs are embedded and positionally encoded before being passed into the encoder.

The encoder is composed of two main components: self-attention (orange) and a feedforward layer (blue).

Let’s say our input sentence is a simple one, such as “Deep learning is simple.”

The first thing we need to do is to start with all of our imports. We’ll import all libraries at once to have them ready for later—we don’t need to import them one by one:

from transformers import AutoTokenizer
import torch
from torch import nn
import torch.nn.functional as F
import math

We need to process the data first. Let’s tokenize the sentence, then embed the tokens:

tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')
text = 'Deep learning is simple.'
print(tokenizer(text, add_special_tokens=False, return_tensors='pt'))
input = tokenizer(text, add_special_tokens=False, return_tensors='pt')