def mask_random_tokens(batch,
                       explicit_vocab_size=30522,
                       masking_prob=0.15,
                       cls_id=101,
                       sep_id=102,
                       mask_id=103,
                       vocab_start_id=999):
  """Prepares input for the masking task.

  Preparation consist in masking masking_prob percentage of non-special tokens
  at each input row; round(masking_prob * num_nonspecial_tokens) random tokens
  are selected out of which each token is either
  - replaced with [MASK] token with 80% probability,
  - replaced with random token with 10% probability,
  - or unchanged with 10%.
  The implentation is based on
  https://github.com/google-research/bert/blob/master/create_pretraining_data.py#L342

  Examples:
  - batch is a stream with each row having tuple (token_ids,). Function yields
  rows of form (modified_token_ids, original_tokens, token_weights), where
  modified_token_ids have [MASK] tokens or random tokens according to the
  procedure described above.
  - batch is a stream with each row having tuple (token_ids, segment_embeddings,
  nsp_label, nsp_weight).Function yields rows of form (modified_token_ids,
  segment_embeddings, nsp_label, nsp_weight, original_tokens, token_weights).

  Args:
    batch: stream of inputs. Each row in the stream is a tuple which first
      element is an array of tokens
    explicit_vocab_size: the total size of the vocabulary.
    masking_prob: Determines percent of non-special tokens to be selected for
      masking.
    cls_id: id of the special CLS token.
    sep_id: id of the special SEP token.
    mask_id: id of the special MASK token.
    vocab_start_id: id of first non-special token in the vocabulary.

  Yields:
    a stream with tokens masked for MLM training and 2 appended arrays:
      - original tokens: a copy of original tokens used as a label for mlm
      training
      - token_weights: weights distributed uniformly over selected tokens (sum
      is 1). Other tokens have 0 weight.
  """
  for token_ids, *row_rest in batch:
    original_tokens = token_ids.copy()

    # choose tokens for prediction. Chooses 0.15 of
    # all non-special tokens
    is_special_token = np.logical_or(token_ids == cls_id,
                                     token_ids == sep_id)  # CLS and SEP tokens
    is_special_token = np.logical_or(is_special_token,
                                     token_ids == 0)  # padding
    viable_ids = np.arange(token_ids.shape[0])[~is_special_token]
    num_to_sample = round(masking_prob * viable_ids.shape[0])
    if num_to_sample == 0:
      # sentence is too short to select given percentage of tokens to mask
      continue
    candidate_ids = np.random.choice(viable_ids, num_to_sample, replace=False)

    # create weights
    token_weights = np.zeros(token_ids.shape)
    token_weights[candidate_ids] = 1 / candidate_ids.shape[0]

    prob_scores = np.random.random(candidate_ids.shape)

    # change 80 % of tokens to [MASK]
    mask_token_ids = candidate_ids[prob_scores < 0.8]
    token_ids[mask_token_ids] = mask_id

    # change 10% of tokens to random token
    random_token_ids = candidate_ids[(0.8 <= prob_scores) & (prob_scores < 0.9)]
    token_ids[random_token_ids] = np.random.randint(vocab_start_id,
                                                    explicit_vocab_size,
                                                    random_token_ids.shape[0])

    # rest (10%) is left unchaged
    yield (token_ids, *row_rest, original_tokens, token_weights)
Prepares input for the masking task.

  Preparation consist in masking masking_prob percentage of non-special tokens
  at each input row; round(masking_prob * num_nonspecial_tokens) random tokens
  are selected out of which each token is either
  - replaced with [MASK] token with 80% probability,
  - replaced with random token with 10% probability,
  - or unchanged with 10%.
  The implentation is based on
  https://github.com/google-research/bert/blob/master/create_pretraining_data.py#L342

  Examples:
  - batch is a stream with each row having tuple (token_ids,). Function yields
  rows of form (modified_token_ids, original_tokens, token_weights), where
  modified_token_ids have [MASK] tokens or random tokens according to the
  procedure described above.
  - batch is a stream with each row having tuple (token_ids, segment_embeddings,
  nsp_label, nsp_weight).Function yields rows of form (modified_token_ids,
  segment_embeddings, nsp_label, nsp_weight, original_tokens, token_weights).

  Args:
    batch: stream of inputs. Each row in the stream is a tuple which first
      element is an array of tokens
    explicit_vocab_size: the total size of the vocabulary.
    masking_prob: Determines percent of non-special tokens to be selected for
      masking.
    cls_id: id of the special CLS token.
    sep_id: id of the special SEP token.
    mask_id: id of the special MASK token.
    vocab_start_id: id of first non-special token in the vocabulary.

  Yields:
    a stream with tokens masked for MLM training and 2 appended arrays:
      - original tokens: a copy of original tokens used as a label for mlm
      training
      - token_weights: weights distributed uniformly over selected tokens (sum
      is 1). Other tokens have 0 weight.
  
trax/data/tf_inputs.py