Skip entire header for llama3 decode #1656
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/1656
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 New FailureAs of commit e7bbac6 with merge base 57ab583 ( NEW FAILURE - The following job has failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
facebook-github-bot
added
the
CLA Signed
| elif token_ids[idx] == self.end_header_id: | ||
| # Mask out end header id and "\n\n" after it, then reset mask | ||
| not_header[idx] = False | ||
| mask = True |
There was a problem hiding this comment.
So after we see "end_header_id", mask is set to True, and will never be False again, is that right? If so, can we just exit the while loop? Or do we need to keep checking because of sample packing, for example?
| not_header = [True] * len(token_ids) | ||
| mask = True | ||
| idx = 0 | ||
| while idx < len(token_ids): |
There was a problem hiding this comment.
This feels like it would be more efficient with numpy arrays. What sucks is going list -> array -> list
But you could do something like:
indices_start = get_index_start_header(token_ids, start_header_id)
indices_end = get_index_start_header(token_ids, start_header_id)
for idx_start, index_end zip(indices_start, indices_end):
mask[idx_start:index_end+1] = False
that way the python for-loop is over number of header ids, and not over indices
ps: not saying it needs to be done. Just thinking outloud. I usually try to avoid python loops
| truncate_at_eos=truncate_at_eos, | ||
| skip_special_tokens=skip_special_tokens, | ||
| ) | ||
| if skip_special_tokens: |
There was a problem hiding this comment.
will this test pass?
text = decode(encode(text), skip_special_tokens=True)
Probably not well formulated, but i want to see if after every encode/decode we are adding \n\n
| idx += 1 | ||
|
|
||
| return self.tt_model.decode( | ||
| [token_ids[i] for i, m in enumerate(not_header) if m], |
There was a problem hiding this comment.
I suggested defining token_ids inside of the if condition. Then, you can have a single return and remove the if/else and duplicated code
eg:
instead of
if something:
return self.tt_model.decode(token_ids_a)
else:
return self.tt_model.decode(token_ids_b)
do:
if something:
token_ids = token_ids_a
return self.tt_model.decode(token_ids)
| truncate_at_eos=truncate_at_eos, | ||
| skip_special_tokens=skip_special_tokens, | ||
| ) | ||
| if skip_special_tokens: |
There was a problem hiding this comment.
I think that a string here is necessary to explain why you need this extra logic for "skip_special_tokens" if decode already has this flag. In other words: why do we skip special tokens tokens in two different places? Is there a more elegant way to solve this, like adding the special token to the tt_model directly?
|
@KaiserLeave has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
|
@KaiserLeave has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
| skip_special_tokens=False, | ||
| ) | ||
| if skip_special_tokens: |
There was a problem hiding this comment.
So I understand that the goal here is just to fix skip_special_tokens for Llama3 tokenizer decode, but it seems to me like we are doing something very unexpected here with skip_special_tokens. (a) we have it defined on the base tokenizer and it's basically a no-op now, and (b) we are now inconsistent on whether this needs to be defined on the ModelTokenizer and BaseTokenizer. If it is a function of tokenize_messages on the ModelTokenizer moreso than the BaseTokenizer, maybe we should update the Protocol along with other callsites?
| # We will remove special tokens manually via regex on the decoded string. | ||
| # This is because removing all special tokens does not remove the role and | ||
| # whitespace added from the special tokens, i.e., the "user" and "\n\n" in | ||
| # "<|start_header_id|>user<|end_header_id|>\n\n" |
There was a problem hiding this comment.
I would maybe move this comment up to where you define self._special_token_regex and self._special_token_header_regex
There was a problem hiding this comment.
this is where it actually happens, so makes more sense to keep it here? no strong opinions
There was a problem hiding this comment.
Yeah I feel the same, fine to keep it here then
|
@KaiserLeave has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
Context
What is the purpose of this PR? Is it to
With
skip_special_tokens=True, the llama3 tokenizer still included the role name because it is part of the header but is not a special token. This PR filters out all tokens between start_header_id and end_header_id. Bonus: also exposeskip_special_tokensin Phi3MiniTokenizer for consistency.Before:
After:
Test plan
Please make sure to do each of the following if applicable to your PR. If you're unsure about any one of these just ask and we will happily help. We also have a contributing page for some guidance on contributing.
pre-commit install)pytest testspytest tests -m integration_testUX
If your function changed a public API, please add a dummy example of what the user experience will look like when calling it.
Here is a docstring example
and a tutorial example