Potential masking issue in `get_masked_and_orig_text` function

[aldopareja]

File: src/instructlab/training/data_process.py

In the get_masked_and_orig_text function, there's a potential issue with the masking mechanism:

def get_masked_and_orig_text(sample):
    labels = sample["labels"]
    input_ids = sample["input_ids"]
    label = [pad_tk[0] if tk == -100 else tk for tk in labels]
    text = tokenizer.decode(label).replace(pad_str, "<mask>")
    orig_text = tokenizer.decode(input_ids)
    return text, orig_text

Issue Description

The function uses pad_str to identify and replace masked tokens with "". However, if pad_str happens to be the same as another special token in the vocabulary, it could lead to unintended replacements, making it appear that some tokens are masked when they actually aren't.

Recommendation

To avoid this potential confusion, consider one of the following approaches:

1. Use a unique string for masking that's guaranteed not to appear in the tokenizer's vocabulary.
2. Utilize a dedicated special token for masking (e.g., "") and add it to the tokenizer's special tokens.
3. Implement the masking logic directly on the token IDs before decoding, ensuring only the intended tokens are masked.

this could fix it, but we should instead add it when adding <|pretrain|> and the other one.

def get_masked_and_orig_text(sample):
    labels = sample["labels"]
    input_ids = sample["input_ids"]
    tokenizer.add_special_tokens({"additional_special_tokens": ["<MASK>"]})
    mask_token_id = tokenizer.convert_tokens_to_ids("<MASK>")
    label = [mask_token_id if tk == -100 else tk for tk in labels]
    text = tokenizer.decode(label)
    orig_text = tokenizer.decode(input_ids)
    return text, orig_text

[github-actions]

This issue has been automatically marked as stale because it has not had activity within 90 days. It will be automatically closed if no further activity occurs within 30 days.

[github-actions]

This issue has been automatically closed due to inactivity. Please feel free to reopen if you feel it is still relevant!