🎓 A1Learner
🔍

PyTorch Building a Transformer Classifier

A transformer classifier is a model that reads a whole review with an encoder and answers with a class. A pooled head turns the encoded positions into one vector and scores it.

Reviews become numbers

The vocabulary holds five positive words, five negative words, seven filler words and the word not. Each word becomes its number in the list.

import torch
from torch import nn
POSITIVE = ["great", "lovely", "bright", "fun", "warm"]
NEGATIVE = ["awful", "dull", "cold", "boring", "grim"]
FILLER = ["the", "a", "movie", "story", "scene", "really", "very"]
VOCAB = POSITIVE + NEGATIVE + FILLER + ["not"]
WORD_ID = {word: i for i, word in enumerate(VOCAB)}
review = ["not", "great", "the", "movie", "really"]
print("Vocabulary size:", len(VOCAB))
print("Encoded review:", [WORD_ID[w] for w in review])

Embeddings plus positions

The word vectors and the position vectors add together. The shape shows one review, five positions and width 16.

torch.manual_seed(0)
embedding = nn.Embedding(len(VOCAB), 16)
position = nn.Embedding(5, 16)
ids = torch.tensor([[WORD_ID[w] for w in review]])
h = embedding(ids) + position.weight
print("Embedded batch shape:", tuple(h.shape))

Encode, pool, score

The encoder refines all five positions. The mean over positions pools them into one vector, and the head scores it.

layer = nn.TransformerEncoderLayer(d_model=16, nhead=2, dim_feedforward=32, dropout=0.0, batch_first=True)
encoder = nn.TransformerEncoder(layer, num_layers=1)
encoded = encoder(h)
pooled = encoded.mean(dim=1)
head = nn.Linear(16, 2)
print("Encoded shape:", tuple(encoded.shape))
print("Pooled shape:", tuple(pooled.shape))
print("Scores shape:", tuple(head(pooled).shape))

Example

import random

import torch
from torch import nn

POSITIVE = ["great", "lovely", "bright", "fun", "warm"]
NEGATIVE = ["awful", "dull", "cold", "boring", "grim"]
FILLER = ["the", "a", "movie", "story", "scene", "really", "very"]
VOCAB = POSITIVE + NEGATIVE + FILLER + ["not"]
WORD_ID = {word: i for i, word in enumerate(VOCAB)}
LENGTH = 5


def make_data(rng, n_plain, n_negated):
    data = []
    for _ in range(n_plain):
        label = rng.randrange(2)
        side = POSITIVE if label == 1 else NEGATIVE
        words = [rng.choice(side) for _ in range(3)]
        words += [rng.choice(FILLER) for _ in range(LENGTH - 3)]
        rng.shuffle(words)
        data.append((words, label))
    for _ in range(n_negated):
        side_label = rng.randrange(2)
        side = POSITIVE if side_label == 1 else NEGATIVE
        rest = [rng.choice(side) for _ in range(2)]
        rest += [rng.choice(FILLER) for _ in range(LENGTH - 3)]
        rng.shuffle(rest)
        data.append((["not"] + rest, 1 - side_label))
    rng.shuffle(data)
    return data


class ReviewClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.embedding = nn.Embedding(len(VOCAB), 16)
        self.position = nn.Embedding(LENGTH, 16)
        layer = nn.TransformerEncoderLayer(
            d_model=16, nhead=2, dim_feedforward=32,
            dropout=0.0, batch_first=True)
        self.encoder = nn.TransformerEncoder(layer, num_layers=1)
        self.head = nn.Linear(16, 2)

    def forward(self, ids):
        h = self.encoder(self.embedding(ids) + self.position.weight)
        return self.head(h.mean(dim=1))


rng = random.Random(42)
train = make_data(rng, 96, 64)
x = torch.tensor([[WORD_ID[w] for w in words] for words, _ in train])
y = torch.tensor([label for _, label in train])

torch.manual_seed(0)
model = ReviewClassifier()
optimizer = torch.optim.Adam(model.parameters(), lr=0.02)
loss_fn = nn.CrossEntropyLoss()

for step in range(80):
    optimizer.zero_grad()
    loss = loss_fn(model(x), y)
    loss.backward()
    optimizer.step()

held = [
    (["not", "great", "the", "movie", "really"], 0),
    (["not", "lovely", "warm", "story", "a"], 0),
    (["not", "awful", "dull", "the", "scene"], 1),
    (["not", "cold", "grim", "movie", "very"], 1),
    (["great", "lovely", "the", "warm", "movie"], 1),
    (["awful", "cold", "a", "grim", "story"], 0),
]
names = {1: "positive", 0: "negative"}
model.eval()
hx = torch.tensor([[WORD_ID[w] for w in words] for words, _ in held])
with torch.no_grad():
    predicted = model(hx).argmax(dim=1).tolist()

print("Training reviews:", len(train))
print("Final loss:", f"{loss.item():.6f}")
right = 0
for (words, label), guess in zip(held, predicted):
    ok = guess == label
    right += ok
    print(" ".join(words), "| true:", names[label],
          "| predicted:", names[guess], "| right:", ok)
print("Held-out reviews right:", right, "/", len(held))

Key points

  • Pooling: averaging the encoded vectors over the positions to get one vector per review
  • Head: the last linear layer, which turns the pooled vector into class scores
  • Negation: a word like not, which flips the meaning of the words near it
  • Use a learned position table to give each position its own added vector
  • Attention lets the word not change the vectors of the words it flips

In this example, a seeded generator builds 160 training reviews. Plain reviews mix three sentiment words from one side with two filler words, then shuffle. Negated reviews start with not and take their label from the opposite side. A position table adds one learned vector per position, as the first lesson of this entry showed. One encoder layer refines the five positions of each review. The mean over positions pools the review into one vector, and a linear head scores it. Training runs 80 full passes with Adam at rate 0.02. The final loss prints as 0.000028. The six held-out reviews then print with their true and predicted labels. All six are right, and four of them are negation reviews that a bag of words cannot read.

Note

Every review in this lesson has five words, so no padding or mask is needed

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Why can this model classify the review 'not great' as negative?