StoppingCriterion Issue in Generation

[jcperdomo]

When generating data, the model produces a valid sequence of tokens for the label followed by the token. However, instead of returning immediately after the end completion token is produced, the model keeps generating until it hits the max_new_tokens upper bound.

For instance, I tried it on a problem with 1 token labels. If you set max number of tokens to 4 it would generate '0 ', if you set it to 10, it would do the same but have 9 end completion tokens.

Example of text:

Example of token sequence

5500, 26576, 24865, 12112, 18, 13, 578, 865, 18,
374, 482, 15, 13, 26083, 5500, 19416, 23642, 21310,
20, 13, 3639, 374, 279, 907, 315, 379, 30,
220, 8651, 18, 8651, 16, 8651, 15, 8651, 17,
8651, 128257, 16, 128256, 128256, 128256, 128256, 128256, 128256,
128256, 128256, 128256]]

128256 is the token

This was tested on v2_evals branch, the relevant code is in inference_utils. The way the stopping criterion is handled is the same as in the main branch, hence it is likely that the issue persists there too.