Creating Padded PyTorch Batches and Valid-Token Masks

Hello! In the last lesson, you assigned integer IDs to characters and reserved separate IDs for [PAD] and [MASK]. Now you’ll put several encoded texts into one PyTorch tensor. The key is to keep track of which positions contain real characters, so the next lesson can corrupt text without touching padding.

This should take about 15–20 minutes in Colab. Keep the vocabulary, encode, decode, PAD_ID, and MASK_ID cells from the previous lesson available.

Why a batch needs padding

Texts rarely have equal lengths, but a batch tensor needs a rectangular shape. We’ll give every row 11 positions: short texts will have [PAD] IDs on the right, while their original characters stay unchanged.

Watch the opening of Hugging Face’s What is dynamic padding? for the visual distinction between padding every batch to a chosen length and padding each batch only to its longest member.

Padding makes the tensor rectangular, but it does not turn the extra positions into characters. We need a second tensor to record that distinction.

Make a fixed-length batch in PyTorch

Run this cell after the previous lesson’s vocabulary code. torch.full starts with an all-padding tensor; each encoded text then replaces the beginning of one row. If a text is too long, the function raises an error rather than silently dropping characters.

import torch

def make_batch(texts, max_length):
    if not texts:
        raise ValueError("The batch must contain at least one text.")

    encoded = [encode(text) for text in texts]
    if any(len(row) > max_length for row in encoded):
        raise ValueError("A text exceeds max_length; choose a larger length.")

    batch = torch.full(
        (len(encoded), max_length),
        PAD_ID,
        dtype=torch.long,
    )

    for row_index, token_ids in enumerate(encoded):
        batch[row_index, :len(token_ids)] = torch.tensor(
            token_ids, dtype=torch.long
        )

    valid_mask = batch != PAD_ID
    return batch, valid_mask

texts = ["cat", "dog naps.", "cat naps."]
batch, valid_mask = make_batch(texts, max_length=11)

print(batch)
print("Token shape:", batch.shape, "dtype:", batch.dtype)
print("Mask shape:", valid_mask.shape, "dtype:", valid_mask.dtype)

You should get three rows and eleven columns in both tensors. batch has integer dtype torch.int64 (also called torch.long), suitable for token IDs. valid_mask has Boolean dtype torch.bool.

The first row begins with the three IDs for cat and ends with eight PAD_IDs. The other two texts each occupy nine positions and have two PAD_IDs. Choosing 11 is a small, explicit batch-format decision—not a property of the vocabulary.

Read the Boolean mask

At each position, valid_mask is True if the token ID is not PAD_ID. A space or period is a real character, so its position is True too.

print(valid_mask)
print("Characters per row:", valid_mask.sum(dim=1).tolist())

assert batch.shape == valid_mask.shape == (3, 11)
assert valid_mask.dtype == torch.bool
assert valid_mask.sum(dim=1).tolist() == [3, 9, 9]
assert torch.all(batch[~valid_mask] == PAD_ID)

The counts [3, 9, 9] are the original text lengths. The last assertion checks every position marked invalid: all of them must contain padding.

Be precise about what this mask means. valid_mask distinguishes text positions from padding positions. It is not the corruption mask that will identify hidden characters in the next lesson. In particular, if a valid character is later replaced by MASK_ID, its position still belongs to the text.

Let a DataLoader form successive batches

For these three examples, one call to make_batch is enough. When a corpus grows, PyTorch’s DataLoader can collect examples in groups and call our function for each group.

Here, collate_texts is that custom function. Every batch still has 11 columns, even when the final batch contains only one example.

from torch.utils.data import DataLoader

def collate_texts(items):
    return make_batch(items, max_length=11)

loader = DataLoader(
    texts,
    batch_size=2,
    shuffle=False,
    collate_fn=collate_texts,
)

for token_ids, is_valid in loader:
    print(token_ids.shape, is_valid.sum(dim=1).tolist())

The output should be:

torch.Size([2, 11]) [3, 9]
torch.Size([1, 11]) [9]

The number of rows changes in the last batch; the fixed number of columns does not.

Takeaway

You can now turn unequal-length texts into a fixed-length integer tensor and a same-shaped Boolean mask. True marks a real text position; False marks padding. Keep both tensors: next lesson, you’ll sample corruption only at valid positions and return a separate mask showing which characters were actually hidden.