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.

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.

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.

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.

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.

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.

Probably, I have kind of the same problem. The 'generate' function outputs a single key per token, here is some pseudocode for the problem:

from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B")

temperature = 0.2
prompt = "Extract in JSON format: "

for token in generate(model, tokenizer, prompt, temp, max_tokens, verbose=True):

    generated_tokens.append(tokenizer.encode(token))

text = "Here is the extracted text"
generated_tokens = [[220], [39], [68], [81], [68], [220], [72], [82], [220], [83], [71], [68], [220], [68], [87], [83], [81], [64], [66], [83], [68], [67], [220], [83], [68], [87], [83]]

how_it_should_be = tokenizer.encode(text)

output = [8586, 374, 279, 28532, 1495]

Is there an easier way /example to retrieve the produced tokens per generated token such as 8586 or 374, and use those for stop criteria? Ideally multiple tokens as a stop criteria would be best because of stopping synthetic JSON generation just as it finishes.