Chapter 13.3 - Collate Function

Overview

When the AI goes to school (training), it needs to bring its textbooks in perfectly sized backpacks. But sentences are all different lengths! The Collate Function is like a packing expert that makes sure all sentences fit into the exact same sized box.

---

🎯 Why we do it

Rationale

GPUs (the chips that train AI) love nice, neat boxes (matrices). If one sentence is 5 words long and another is 10 words long, the GPU gets confused. We use "Padding" to fill the short sentences with blank spaces so everything is exactly the same length!

🛠️ How we do it

Methodology

We find the longest sentence in our batch. Then, we add <PAD> tokens (which are just 0s) to the end of all the shorter sentences until they match the longest one.

# Two sentences of different lengths
sentence_1 = [45, 92, 12] # Length 3
sentence_2 = [88, 14, 99, 42, 11] # Length 5

# The Collate Function pads sentence_1 with zeros!
padded_sentence_1 = [45, 92, 12, 0, 0] # Now it's Length 5!

print("Perfectly matched backpacks!")
print(padded_sentence_1)
print(sentence_2)

---

💻 Code Implementation

Here is the exact PyTorch implementation for the concepts discussed above:

import torch
def collate_fn(
    batch,
    pad_token_id=50256,
    ignore_index=-100,
    allowed_max_length=None,
    device="cpu"
):
    # Find the longest sequence in the batch
    batch_max_length = max(len(item)+1 for item in batch)

    # Pad and prepare inputs and targets
    inputs_lst, targets_lst = [], []

    for item in batch:
        new_item = item.copy()
        # Add an <|endoftext|> token
        new_item += [pad_token_id]
        # Pad sequences to max_length
        padded = (
            new_item + [pad_token_id] *
            (batch_max_length - len(new_item))
        )
        inputs = torch.tensor(padded[:-1])  # Truncate the last token for inputs
        targets = torch.tensor(padded[1:])  # Shift +1 to the right for targets

        # New: Replace all but the first padding tokens in targets by ignore_index
        mask = targets == pad_token_id
        indices = torch.nonzero(mask).squeeze()
        if indices.numel() > 1:
            targets[indices[1:]] = ignore_index

        # New: Optionally truncate to maximum sequence length
        if allowed_max_length is not None:
            inputs = inputs[:allowed_max_length]
            targets = targets[:allowed_max_length]

        inputs_lst.append(inputs)
        targets_lst.append(targets)

    # Convert list of inputs and targets to tensors and transfer to target device
    inputs_tensor = torch.stack(inputs_lst).to(device)
    targets_tensor = torch.stack(targets_lst).to(device)

    return inputs_tensor, targets_tensor

#testing the function

# inputs_1 = [0, 1, 2, 3, 4]
# inputs_2 = [5, 6]
# inputs_3 = [7, 8, 9]

# batch = (
#     inputs_1,
#     inputs_2,
#     inputs_3
# )

# inputs, targets = collate_fn(batch)
# print(inputs)
# print(targets)