Warning message in DataCollatorForCompletionOnlyLM is misleading when only response_template is missing in the batch

[Aratako]

In the current implementation of DataCollatorForCompletionOnlyLM, there is an issue with the warning messages that are displayed when response_template is missing from the batch labels.

Currently, when the response_template is not found, the code sets the entire batch["labels"][i, :] to self.ignore_index and then proceeds to check for the instruction_template. This results in a misleading warning message saying the instruction_template is missing, even if it is actually present in the data.

Here's the relevant code snippet:

if len(response_token_ids_idxs) == 0:
    warnings.warn(...)
    batch["labels"][i, :] = self.ignore_index

human_token_ids = self.instruction_token_ids
for human_idx in np.where(batch["labels"][i] == human_token_ids[0])[0]:
    if human_token_ids == batch["labels"][i][human_idx : human_idx + len(human_token_ids)].tolist():
        human_token_ids_idxs.append(human_idx)

if len(human_token_ids_idxs) == 0:
    warnings.warn(...)
    batch["labels"][i, :] = self.ignore_index

To improve this, I propose a simple change. After setting the batch["labels"][i, :] to self.ignore_index when either template is missing, we can immediately continue to the next iteration of the loop. This will prevent the unnecessary check for the other template and the associated warning.

Here's the modified code:

if len(response_token_ids_idxs) == 0:
    warnings.warn(...)
    batch["labels"][i, :] = self.ignore_index
    continue

human_token_ids = self.instruction_token_ids
for human_idx in np.where(batch["labels"][i] == human_token_ids[0])[0]:
    if human_token_ids == batch["labels"][i][human_idx : human_idx + len(human_token_ids)].tolist():
        human_token_ids_idxs.append(human_idx)

if len(human_token_ids_idxs) == 0:
    warnings.warn(...)
    batch["labels"][i, :] = self.ignore_index
    continue

This change ensures that the warning messages accurately reflect which template is missing and avoids any confusion. The logic of ignoring instances with missing templates in the loss calculation remains intact.
Please let me know your thoughts on this proposed change. I'm happy to submit a pull request with the updated code if you agree it would be an improvement.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.