Sequences in a batch have different lengths, but tensors must be rectangular. A collate function pads them and builds the attention mask so the model ignores padding.
Implement collate(seqs, pad_id=0, max_len=None, side="right"):
max_len (keep the first tokens). Without max_len, pad to the longest sequence.side="left" (decoder models generating text need the real tokens to end at the same position).input_ids, attention_mask (1 for real tokens) and position_ids, which count real tokens from 0 and are 0 at padding positions.Then implement padding_waste(seqs): the fraction of the padded batch made of padding, which is what sorting by length (bucketing) reduces.