Data Collator with Padding

A collator combines individual examples into a batch. This helper takes integer sequences of different lengths and pads them to a common length, so the result can later be converted to a rectangular array.

Your task

Implement pad_batch(sequences, pad_id=0). sequences is a Python list of Python lists of integers. pad_id is the integer to append as padding; it may be zero, positive or negative.

Find the longest sequence length in this batch. Return a new list of lists where every sequence has that length, adding pad_id only at the right end of shorter sequences. Preserve row order and the order and values of every original token. Do not truncate, sort or modify the input lists.

Build the result explicitly with Python list operations. Do not call torch.nn.utils.rnn.pad_sequence. The return value is a list of lists, not a PyTorch tensor or JAX array.

Example

For sequences = [[7, 8], [9], []] and pad_id = -1, the longest length is 2. Return [[7, 8], [9, -1], [-1, -1]].

If a token already equals pad_id, retain it. For [[0, 1], [0]] with pad_id = 0, return [[0, 1], [0, 0]]. Padding depends on the original lengths, not on whether a token has the padding value.

Empty inputs and boundaries

An empty batch [] returns []; there is no longest row to find.
A batch of empty rows, such as [[], []], returns the same number of empty rows.
One sequence or equal-length sequences need no added padding, but the input lists must still remain unchanged.

For a nonempty batch with N rows and longest length T, the output has N rows of length T. An empty batch has no rows; do not invent a row or a nonzero padding length.

What padding does not do

This function adds values only. It does not return an attention mask, sequence lengths or targets, and padding alone does not tell a model which positions to ignore. If a later step needs that information, retain the original lengths or construct a separate mask. Comparing tokens with pad_id is insufficient when that value can also occur in real data.

Right-padding is the rule for this exercise. The appropriate padding side and mask convention in a full training or generation pipeline depend on the model and how its outputs are consumed.

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Convert text into token IDs

A tokenizer decides which units of text the model sees. Whitespace splitting, greedy subword matching, and byte-pair encoding (BPE) use different rules. They can agree on one string and produce different results on another.

A greedy longest-prefix tokenizer chooses the longest vocabulary entry matching the current position, emits its ID, then advances. With vocabulary a, ab, and b, the text ab becomes one token. When no entry matches, handle the unknown input and advance according to the task's rule. Otherwise the scan can loop forever. A simplified exercise may omit WordPiece features such as continuation markers and whole-word unknown handling.

Separate BPE training from encoding

BPE training learns merges by repeatedly combining frequent neighboring pairs. Encoding uses the learned merge order: apply eligible pairs according to their rank. It does not count frequencies afresh to learn a new tokenizer for each input. Follow the specified handling of overlapping pairs within a merge pass.

Pad a batch without treating padding as content

Sequences with lengths two and four can share a length-four tensor by adding two padding entries to the shorter sequence. An attention mask identifies which positions contain real tokens. Left and right padding place those tokens at different indices, which can affect position IDs and generation logic.

Keep the truncation length, unknown-token ID, and padding ID explicit. If the contract allows the padding ID to appear as real content, token values alone cannot identify padding; retain the lengths or mask. Return the list or tensor representation requested by the exercise.

Detecting runtime…