def apply_masking(token_ids, target_token_ids, mask_indices, mask_token_id,
                  vocab_size):
  """Applies BERT masking.

  Args:
    token_ids (Tensor): 1-D Tensor of token IDs (ints)
    target_token_ids (Tensor): 1-D Tensor of token IDs (ints)
    mask_indices (Tensor): 1-D Tensor of indices (ints)
    mask_token_id (int): ID of [MASK] token.
    vocab_size (int): total size of vocabulary.

  Returns:
    token_ids_masked (Tensor): 1-D Tensor of token IDs, after target positions
      have been replaced with [MASK], a random token, or left alone.
    target_token_ids (Tensor): the original token IDs at the target positions.
  """
  num_to_mask = tf.size(mask_indices)

  mask_token_ids = tf.fill([num_to_mask], tf.cast(mask_token_id, tf.int64))
  random_token_ids = tf.random.uniform([num_to_mask],
                                       minval=0,
                                       maxval=vocab_size,
                                       dtype=tf.int64)

  # Uniform [0, 1) floats.
  randomness = tf.random.uniform([num_to_mask])

  # Replace target tokens with mask tokens.
  mask_values = tf.where(randomness < 0.8, mask_token_ids, target_token_ids)

  # Replace target tokens with random tokens.
  mask_values = tf.where(randomness > 0.9, random_token_ids, mask_values)

  # Mask out token_ids at mask_indices.
  token_ids_masked = tf.tensor_scatter_update(token_ids, mask_indices[:, None],
                                              mask_values)

  return token_ids_masked
Applies BERT masking.

  Args:
    token_ids (Tensor): 1-D Tensor of token IDs (ints)
    target_token_ids (Tensor): 1-D Tensor of token IDs (ints)
    mask_indices (Tensor): 1-D Tensor of indices (ints)
    mask_token_id (int): ID of [MASK] token.
    vocab_size (int): total size of vocabulary.

  Returns:
    token_ids_masked (Tensor): 1-D Tensor of token IDs, after target positions
      have been replaced with [MASK], a random token, or left alone.
    target_token_ids (Tensor): the original token IDs at the target positions.
  
language/conpono/reconstruct/preprocess.py