Hardcoded GENERATION_REGEX logic (with return_assistant_tokens_mask) causes incorret masks production for chat examples in _chat_preprocess function

[DzmitryPihulski]

Describe the bug

The function _chat_preprocess is used to apply chat template to the hf chat examples.

The examples are in the format of messages. To correctly produce the mask and train only on assistant messages, the user/system/tool messages need to be masked (mask=0). The function logic is:

template_has_generation_kwd = GENERATION_REGEX.search(tokenizer.chat_template) is not None

tokenized_chat = tokenizer.apply_chat_template(
      chat,
      tools=tools,
      tokenize=True,
      return_dict=True,
      return_assistant_tokens_mask=template_has_generation_kwd,
  )

input_ids = tokenized_chat.get("input_ids")
if template_has_generation_kwd:
    mask = tokenized_chat["assistant_masks"]
else:
    mask = [1] * len(input_ids)

The issue here is that not all tokenizers has {% generation %} keyword, for example Nemotron 3 Nano model doesn't, which means we will get default attention_mask (all 1's). Such mask is not the best for SFT, cause it forces the model to be trained on every user/system/tool message.

As far as I understand this issue is only for hf chats, cause for not hf chats there is _preprocess function from the same file, that uses _mask_targets function to correctly mask messages.

Expected behavior

Maybe remove the return_assistant_tokens_mask parameter, it doesn't help with the modern chat templates and apply own logic of masking.

[DzmitryPihulski]

Also, if I understand correctly, since there is this logic in _chat_preprocess:

  if 0 in mask:
      # traverse the list backward for first occurrence of masked token
      context_end_idx = len(mask) - mask[::-1].index(0)
  else:
      context_end_idx = len(mask)

  context_ids = input_ids[:context_end_idx]
  answer_ids = input_ids[context_end_idx:]

  return dict(
      input_ids=torch.LongTensor(input_ids),
      loss_mask=torch.BoolTensor(mask),
      context_ids=torch.LongTensor(context_ids),
      answer_ids=torch.LongTensor(answer_ids),
  )

We do not have any answer_ids after the tokenization, because there are no zeros in mask. Which combined with answer_only_loss: bool = True in the create_sft_dataset function may suggest the training will not run at all.