Using token sequences as stop criteria does not work in mlx_lm about mlx-examples HOT 6 OPEN

Yeah, I think it was intended that it runs on every token generated, and we can throw away the buffer entirely. This would address most issues.

It still needs a special case for streaming, since we need to anticipate a stop word, or we might stream parts of it. I have a prototype for that check currently, but still doing some testing.

from mlx-examples.

I'll try to sum it up

max_stop_id_sequence_len is the length of the longest stop id sequence. Now let's assume the buffer is the same size as the longest stop sequence.

if len(stop_sequence_buffer) > max_stop_id_sequence_len:

Since this check is non-inclusive, it would loop one extra time before running what comes after the if statement. Meaning the length of the buffer would now be 1 larger than the longest stop sequence.

Now when we take into account, that the stop_criteria function only checks for perfect matches, where the tail of the "tokens" matches a stop sequence. It can no longer ever match, because the extra tokens was generated before calling stop_criteria.

from mlx-examples.

Additionally the condition for checking the stop criteria is:

if len(stop_sequence_buffer) > max_stop_id_sequence_len:

It is not inclusive, meaning there will always be an extra token appended. Meaning no stop sequences can ever be matched.

from mlx-examples.

Additionally the condition for checking the stop criteria is:

if len(stop_sequence_buffer) > max_stop_id_sequence_len:

It is not inclusive, meaning there will always be an extra token appended. Meaning no stop sequences can ever be matched.

Just curious, how did you find the issue? I ran a few tests before and didn't see any extra tokens added.

from mlx-examples.

I see, it would only happen when the model starts with the stop word. Maybe that's why it wasn't picked up by my testing. This edge case is a bit difficult to pick up with manual testing, but it would be more obvious when the model starts with a stop word.

from mlx-examples.

Yeah, you are right. The original implementation didn't have a buffer, so it ended up sending the stop word back to the client in the streaming. The buffer was introduced to solve that issue, but it seems like it wasn't well thought out in the implementation.

from mlx-examples.