🎓 A1Learner
🔍

PyTorch Training a Text Classifier

A text classifier is a model that reads words and picks one label. This one averages word embeddings into a single vector and classifies that.

Padding makes rows equal

Reviews differ in length. Padding fills the short row with zeros up to the shared width.

pair = [(["great", "warm"], 1), (["the", "awful", "cold", "movie"], 0)]
x, y = encode(pair, 4)
print(x.tolist())
print(y.tolist())

The bag of embeddings

The forward pass averages the embeddings of the real words. Padding adds nothing and counts for nothing.

def forward(self, ids):
    mask = (ids != 0).float().unsqueeze(-1)
    summed = (self.embedding(ids) * mask).sum(dim=1)
    counts = mask.sum(dim=1).clamp(min=1)
    return self.linear(summed / counts)

Two scores per review

One forward pass gives one score pair per review. The larger score is the model's pick.

torch.manual_seed(0)
model = BagClassifier(len(WORDS), 8)
logits = model(x)
print(tuple(logits.shape))

Example

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"]
WORDS = ["<pad>"] + POSITIVE + NEGATIVE + FILLER
WORD_ID = {word: i for i, word in enumerate(WORDS)}


def make_review(gen):
    label = int(torch.randint(0, 2, (1,), generator=gen))
    side = POSITIVE if label == 1 else NEGATIVE
    n_sent = int(torch.randint(2, 4, (1,), generator=gen))
    n_fill = int(torch.randint(1, 4, (1,), generator=gen))
    picks = torch.randperm(len(side), generator=gen)[:n_sent]
    words = [side[i] for i in picks.tolist()]
    picks = torch.randint(0, len(FILLER), (n_fill,), generator=gen)
    words += [FILLER[i] for i in picks.tolist()]
    order = torch.randperm(len(words), generator=gen)
    return [words[i] for i in order.tolist()], label


def encode(reviews, width):
    rows = []
    for words, label in reviews:
        ids = [WORD_ID[w] for w in words]
        ids += [0] * (width - len(ids))
        rows.append((ids, label))
    x = torch.tensor([r[0] for r in rows])
    y = torch.tensor([r[1] for r in rows])
    return x, y


class BagClassifier(nn.Module):
    def __init__(self, vocab_size, dim):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, dim, padding_idx=0)
        self.linear = nn.Linear(dim, 2)

    def forward(self, ids):
        mask = (ids != 0).float().unsqueeze(-1)
        summed = (self.embedding(ids) * mask).sum(dim=1)
        counts = mask.sum(dim=1).clamp(min=1)
        return self.linear(summed / counts)


gen = torch.Generator().manual_seed(7)
reviews = [make_review(gen) for _ in range(200)]
width = max(len(words) for words, _ in reviews)
x_train, y_train = encode(reviews[:160], width)
x_val, y_val = encode(reviews[160:], width)

torch.manual_seed(0)
model = BagClassifier(len(WORDS), 8)
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.05)

for epoch in range(60):
    logits = model(x_train)
    loss = loss_fn(logits, y_train)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

with torch.no_grad():
    pred = model(x_val).argmax(dim=1)
    correct = (pred == y_val).sum().item()

print("Train reviews:", len(y_train))
print("Held-out reviews:", len(y_val))
print("Final train loss:", round(loss.item(), 4))
print("Validation accuracy:", correct, "/", len(y_val), "=", round(correct / len(y_val), 3))

Key points

  • Embedding: a learned vector that stands for one word
  • Bag of embeddings: the mean of the word vectors in a review
  • Padding fills short reviews with zeros up to one shared width
  • The mask keeps padding out of the mean
  • Use nn.CrossEntropyLoss to train the two-score output
  • Validation accuracy: the share of held-out reviews labelled correctly
  • The bag of embeddings is the whole model in this lesson, an embedding layer, a mean, and one linear layer

In this example, the generator from the last lesson rebuilds the same 200 reviews. The encode function turns each review into word numbers and pads every row to the widest review. The BagClassifier holds an embedding layer with one 8-number vector per word. Its forward pass masks out the padding and averages the real word vectors. A linear layer turns that mean into two scores. Training runs 60 full passes with Adam at rate 0.05. The final train loss prints as 0.0002. The model then labels the 40 held-out reviews. It gets all 40 right, so the validation accuracy prints as 1.0. The task is separable because each review uses words from one side only. The next lesson feeds the model reviews that break that rule.

Note

Word order never reaches this model, so two orders of one review always score the same

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does the mask do in the bag model's forward pass?