mask_messages

torchtune.data.mask_messages(messages: list[torchtune.data._messages.Message], masking_strategy: MaskingStrategy) → None

Set the masked attribute for each message in the list based on the specified masking strategy.

Parameters:

• messages (list[Message]) – a list of messages to mask.

• masking_strategy (MaskingStrategy) –

  masking strategy to use. Must be one of train_on_all, train_on_assistant, train_on_last.

  • train_on_all: both user and assistant messages are unmasked

  • train_on_assistant: user messages are masked, only assistant messages are unmasked

  • train_on_last: only the last assistant message is unmasked