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. ❌ 5 New Failures, 4 Cancelled JobsAs of commit 4e97a57 with merge base 779569e ( NEW FAILURES - The following jobs have failed:
CANCELLED JOBS - The following jobs were cancelled. Please retry:
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 |
| @@ -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.
| @@ -305,11 +332,23 @@ class ChosenRejectedToMessages(Transform): | |||
|
|
|||
| def __init__( | |||
| self, | |||
| train_on_input: bool = False, | |||
| masking_strategy: Optional[str] = "train_on_all", | |||
|
|
||
| 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", |
| return {"messages": updated_messages} | ||
|
|
||
|
|
||
| def validate_messages( |
There was a problem hiding this comment.
not sure if this is still used but let's keep this here for now
| """ | ||
|
|
||
| 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
| 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.
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