def mask_token(self, idx, tokens, np_rng):
    '\n        Helper function to mask `idx` token from `tokens` according to\n        section 3.3.1 of https://arxiv.org/pdf/1810.04805.pdf\n        '
    label = tokens[idx]
    if (np_rng.random() < 0.8):
        new_label = self.mask_id
    elif (np_rng.random() < 0.5):
        new_label = label
    else:
        new_label = np_rng.choice(self.vocab_id_list)
    tokens[idx] = new_label
    return label
Helper function to mask `idx` token from `tokens` according to
section 3.3.1 of https://arxiv.org/pdf/1810.04805.pdf
def mask_token(self, idx, tokens, np_rng):
    '\n        Helper function to mask `idx` token from `tokens` according to\n        section 3.3.1 of https://arxiv.org/pdf/1810.04805.pdf\n        '
    label = tokens[idx]
    if (np_rng.random() < 0.8):
        new_label = self.mask_id
    elif (np_rng.random() < 0.5):
        new_label = label
    else:
        new_label = np_rng.choice(self.vocab_id_list)
    tokens[idx] = new_label
    return label