code: def mlm_mask_tokens(
      inputs: torch.LongTensor, tokenizer, mlm_probability
  ) -> Tuple[torch.LongTensor, torch.LongTensor]:
      """From HuggingFace"""
      device = inputs.device
      inputs = inputs.cpu().clone()
      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_probability)
      special_tokens_mask = [
          tokenizer.get_special_tokens_mask(
              val, already_has_special_tokens=True
          )
          for val in labels.tolist()
      ]
      probability_matrix.masked_fill_(
          torch.tensor(special_tokens_mask, dtype=torch.bool), value=0.0
      )
      # noinspection PyProtectedMember
      if tokenizer._pad_token is not None:
          padding_mask = labels.eq(tokenizer.pad_token_id)
          probability_matrix.masked_fill_(padding_mask, value=0.0)
      masked_indices = torch.bernoulli(probability_matrix).bool()
      labels[
          ~masked_indices
      ] = NON_MASKED_TOKEN_LABEL_ID  # We only compute loss on masked tokens
  
      # 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] = tokenizer.convert_tokens_to_ids(
          tokenizer.mask_token
      )
  
      # 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(
          len(tokenizer), labels.shape, dtype=torch.long
      )
      inputs[indices_random] = random_words[indices_random]
  
      # The rest of the time (10% of the time) we keep the masked input tokens unchanged
  
      # noinspection PyTypeChecker
      return inputs.to(device), labels.to(device)
  
imports: import numpy
  import torch
  from dataclasses import dataclass
  from typing import List, Tuple
  from jiant.utils.python.datastructures import ReusableGenerator
  from jiant.tasks.core import Task, TaskTypes, BaseExample, BaseTokenizedExample, BaseDataRow, BatchMixin
  from jiant.tasks.lib.templates.shared import construct_single_input_tokens_and_segment_ids, create_input_set_from_tokens_and_segments
label_desc: Use image node_id 1 for calling a global function with example usage: mlm_mask_tokens(inputs, tokenizer, mlm_probability) and returns: inputs,  labels