BERT pretrains by hiding 15% of tokens and predicting them. Of the selected tokens, 80% become [MASK], 10% become a random token and 10% stay unchanged, so the model can't assume every prediction target is a [MASK].
Implement mask_tokens(ids, rng, mask_id, vocab_size, special_ids=(), p=0.15) returning (inputs, labels) as new lists:
rng.random(len(ids)) and select positions where the draw is below p, skipping special tokens (like [CLS] and [SEP]).r = rng.random(). If r < 0.8, use mask_id. If r < 0.9, use a random token from rng.integers(vocab_size) (one extra draw). Otherwise keep the token.labels hold the original token at selected positions and -100 everywhere else.