More straightforward custom_collate_function (chapter 7)

[behzadshomali]

Hello!

I personally found the introduced custom_collate_funciton in the book a bit confusing. To be more specific, the way we added and ignored the pad_token_id was not really straightforward for me (the idea behind it was clear, but in order to understand the code I had to read it multiple times), therefore I modified it slightly and here I'm going to share it (I validated it by comparing its output against the output of the original collate function introduced in the book for several inpouts and got the same output):

def custom_collate_fn(
    batch,
    pad_token_id=50256,
    ignore_index=-100,
    allowed_max_length=None,
    device="cpu"
):
    batch_max_length = max(len(item) for item in batch)
    
    inputs_lst, targets_lst = [], []
    for item in batch:
        padded_item = (
            item.copy() + 
            [pad_token_id] * (batch_max_length - len(item))
        )
        inputs = torch.tensor(padded_item)
        targets = torch.tensor(padded_item[1:] + [pad_token_id]) # add an extra pad_token_id (the main one)
        
        mask = targets == pad_token_id
        indices = torch.nonzero(mask).squeeze()
        if indices.numel() > 1:
            targets[indices[1:]] = ignore_index # keep the first one

        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)

    return torch.stack(inputs_lst).to(device), torch.stack(targets_lst).to(device)

Was there any specific reason that the padding was done in that way in the book?