Add masking strategies to message transforms #2284
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/2284
Note: Links to docs will display an error until the docs builds have been completed. This comment was automatically generated by Dr. CI and updates every 15 minutes. |
facebook-github-bot
added
the
CLA Signed
supreethmanyam
force-pushed
the
add-chat-loss-masking-strategies
branch
from
db9a565 to
255bdcd
Compare
supreethmanyam
force-pushed
the
add-chat-loss-masking-strategies
branch
from
255bdcd to
ac6272f
Compare
|
@RdoubleA Added |
torchtune/data/_messages.py
Outdated
| @@ -176,12 +185,24 @@ class InputOutputToMessages(Transform): | |||
|
|
|||
| def __init__( | |||
| self, | |||
| train_on_input: bool = False, | |||
| masking_strategy: Optional[str] = "train_on_all", | |||
There was a problem hiding this comment.
if the default value for train_on_input was False, this would correspond to train_on_assistant. Let's keep that same default behavior.
There was a problem hiding this comment.
Updated default strategy
torchtune/data/_messages.py
Outdated
| @@ -305,11 +332,23 @@ class ChosenRejectedToMessages(Transform): | |||
|
|
|||
| def __init__( | |||
| self, | |||
| train_on_input: bool = False, | |||
| masking_strategy: Optional[str] = "train_on_all", | |||
There was a problem hiding this comment.
Updated default strategy
torchtune/data/_messages.py
Outdated
|
|
||
| Raises: | ||
| ValueError: If ``column_map`` is provided and ``conversations`` not in ``column_map``. | ||
| """ | ||
|
|
||
| def __init__( | ||
| self, | ||
| train_on_input: bool = False, | ||
| masking_strategy: Optional[str] = "train_on_all", |
There was a problem hiding this comment.
Updated default strategy
| """ | ||
|
|
||
| def __init__( | ||
| self, train_on_input: bool = True, column_map: Optional[Dict[str, str]] = None | ||
| self, | ||
| masking_strategy: Optional[str] = "train_on_all", |
There was a problem hiding this comment.
contrary to the other comment, this default value should remain train_on_all since train_on_input was defaulted to True. so let's keep this
torchtune/data/_messages.py
Outdated
| ValueError: If the masking strategy is not one of the supported strategies: | ||
| `train_on_all`, `train_on_assistant`, `train_on_last`. | ||
| """ | ||
| if masking_strategy not in ["train_on_all", "train_on_assistant", "train_on_last"]: |
There was a problem hiding this comment.
One suggestion that I didn't mention in the issue is that we could make masking_strategy an Enum so that it can be type checked more easily. users can still specify as a string in the config and in the message transforms above, but here in mask_messages you could convert to the Enum, and the ValueError will automatically be captured.
There was a problem hiding this comment.
Converted masking_strategy to Enum and raises ValueErroe automatically if doesn't match with mapping
|
In addition to converting masking_strategy to Enum and fixing default strategy in built-in transforms, I made following small changes:
|
Context
What is the purpose of this PR? Is it to
Please link to any issues this PR addresses.
#2261
Changelog
What are the changes made in this PR?
mask_messagesutility intorchtune/data/_messages.pyto mask each message given a list of messages based on user provided masking_strategyTest 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