Class prediction in chapter 6 depends on the number of paddings added

[itdxer]

Hello,

I was trying to play around with a classifier trained in chapter 6 and found out that it was always predicting "spam" for any message I give it. When I looked into it I found that predicted class actually depends on the number of padding tokens added into the sequence. If such a model is going to be used in production the same padding size must always be applied to the model. Two graphs below show probability of spam prediction next to a token up to which a sequence was considered

You can see that final token in the message has probability close to 1, but it changes with each extra padding token added to the sequence. In the code, all sequences have the same length so in order to get desired probability we will need to have a correct padding size, otherwise predicted probability can give a very poor accuracy.

I'm not sure if this classifies as a bug, since, in a way, the model is doing what it was asked to do, but at the same time model's behavior is unintuitive. I made small adjustments to the model in order to get the desired behavior which shows how probability evolves with every token added in a single inference step and doesn't depend on number of padding tokens

Main changes:

# create one label per each token in order to make partial predictions
self.labels = [
    [label] * len(encoded_text) + [ignore_index] * (self.max_length - len(encoded_text))
    for label, encoded_text in zip(self.labels, self.encoded_texts)
]
... 
logits = model(input_batch)
loss = torch.nn.functional.cross_entropy(logits.flatten(0, 1), target_batch.flatten())

Modified code (not everything works from the old code, but it runs fine): https://gist.github.com/itdxer/b30ceedb4ac0f3fd2b3e37fb54f71398
Original code: https://github.com/rasbt/LLMs-from-scratch/blob/main/ch06/01_main-chapter-code/gpt_class_finetune.py

Is there a better way to address this issue?