def sample_mask_indices(tokens, mask_rate, mask_blacklist, max_num_to_mask):
  """Samples indices to mask.

  Args:
    tokens (Tensor): 1-D string Tensor.
    mask_rate (float): percentage of tokens to mask.
    mask_blacklist (Tensor): 1-D string Tensor of tokens to NEVER mask.
    max_num_to_mask (int): max # of masks.

  Returns:
    mask_indices (Tensor): 1-D int32 Tensor of indices to mask.
  """
  if mask_rate < 0 or mask_rate > 1:
    raise ValueError("mask_rate must be within [0, 1].")

  # Compute how many tokens to mask.
  num_tokens = tf.size(tokens)
  num_to_mask = tf.to_int32(tf.ceil(mask_rate * tf.to_float(num_tokens)))

  if mask_rate > 0:
    # If masking is enabled, then mask at least one, no matter what.
    # Original BERT code does this too.
    num_to_mask = tf.maximum(num_to_mask, 1)

  num_to_mask = tf.minimum(num_to_mask, max_num_to_mask)

  # If there are any [CLS] or [SEP], we count these as part of num_tokens.
  # Note that the original implementation of BERT does this as well.

  all_indices = tf.range(num_tokens)

  # Filter out indices containing CLS and SEP.
  allow_masking = tf.reduce_all(
      tf.not_equal(tokens, mask_blacklist[:, None]), axis=0)

  filtered_indices = tf.boolean_mask(all_indices, allow_masking)

  # Randomly select indices without replacement.
  shuffled_indices = tf.random.shuffle(filtered_indices)
  mask_indices = shuffled_indices[:num_to_mask]

  return mask_indices
Samples indices to mask.

  Args:
    tokens (Tensor): 1-D string Tensor.
    mask_rate (float): percentage of tokens to mask.
    mask_blacklist (Tensor): 1-D string Tensor of tokens to NEVER mask.
    max_num_to_mask (int): max # of masks.

  Returns:
    mask_indices (Tensor): 1-D int32 Tensor of indices to mask.
  
language/conpono/reconstruct/preprocess.py