text: import torch
  
  
  def mask_tokens(inputs, special_tokens_mask, pad_token_id, mask_token_id,
                  vocab_size, mlm_prob, rand_replace=True):
      labels = inputs.clone()
      # We sample a few tokens in each sequence for masked-LM training
      # (with probability args.mlm_probability defaults to 0.15 in Bert/RoBERTa)
      probability_matrix = torch.full(labels.shape, mlm_prob)
      probability_matrix.masked_fill_(special_tokens_mask, value=0.0)
      padding_mask = labels.eq(pad_token_id)
      probability_matrix.masked_fill_(padding_mask, value=0.0)
      masked_indices = torch.bernoulli(probability_matrix).bool()
      labels[~masked_indices] = -100  # We only compute loss on masked tokens
  
      if rand_replace:
          # 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK])
          indices_replaced = torch.bernoulli(torch.full(labels.shape,
                                                        0.8)).bool() & masked_indices
          inputs[indices_replaced] = mask_token_id
  
          # 10% of the time, we replace masked input tokens with random word
          indices_random = torch.bernoulli(torch.full(
              labels.shape, 0.5)).bool() & masked_indices & ~indices_replaced
          random_words = torch.randint(vocab_size, labels.shape, dtype=torch.long)
          inputs[indices_random] = random_words[indices_random]
      else:
          inputs[masked_indices] = mask_token_id
  
      # The rest of the time (10% of the time) we keep the masked input tokens unchanged
      return inputs, labels
  