def mask_mentions_and_tokens_tf(
    text_ids: tf.Tensor,
    text_mask: tf.Tensor,
    dense_span_starts: tf.Tensor,
    dense_span_ends: tf.Tensor,
    non_mention_mask_rate: float,
    mention_mask_rate: float,
    max_mlm_targets: int,
    mask_token_id: int,
    vocab_size: int,
    random_replacement_prob: float = 0.1,
    identity_replacement_prob: float = 0.1,
) -> Dict[str, tf.Tensor]:
  """Randomly masks whole mentions and random tokens up to a maximum.

  First, mentions are masked according to mention mask rate. If a mention is
  masked, all tokens in the mention are replaced by the mask token. If the
  passage has at least one mention and the mention rask rate is greater than
  zero, we mask at least one mention.

  After masking mentions, if there are fewer masked tokens than maximum mlm
  targets, we additionally mask non-mention words. TODO: If a token in a word
  is masked, all tokens in the word are masked. Some proportion of targets are
  not masked to ameliorate pretrain-finetune mismatch. If there are insufficient
  masked tokens, the target array is padded up to max targets.

  Args:
    text_ids: [seq_length] tensor with token ids.
    text_mask: [seq_length] tensor with 1s for tokens and 0 for padding.
    dense_span_starts: [seq_length] tensor with 1s for mention start positions
      and 0 otherwise.
    dense_span_ends: [seq_length] tensor with 1s for mention end positions and 0
      otherwise.
    non_mention_mask_rate: percentage of non mention tokens to be masked.
    mention_mask_rate: percentage of mentions to be masked.
    max_mlm_targets: total number of mlm targets.
    mask_token_id: token id for mask token.
    vocab_size: vocabulary size.
    random_replacement_prob: probability that to-be-masked token will be
      replaced with a random token instead of [MASK].
    identity_replacement_prob: probability that to-be-masked token will be
      replaced with itself instead of [MASK].

  Returns:
    Dictionary with masked text, mask positions, target ids, target weights.
  """
  # Mask mentions
  mention_start_positions = non_zero_1d(dense_span_starts)
  mention_end_positions = non_zero_1d(dense_span_ends)
  mention_masked_positions = mask_tokens_by_spans(text_ids,
                                                  mention_start_positions,
                                                  mention_end_positions,
                                                  mention_mask_rate,
                                                  max_mlm_targets)

  dense_is_mention = get_dense_is_inside_for_dense_spans(
      dense_span_starts, dense_span_ends)
  dense_is_not_mention = 1 - dense_is_mention
  dense_is_not_mention = dense_is_not_mention * text_mask

  # Mask non-mentions
  non_mention_start_positions = non_zero_1d(dense_is_not_mention)
  # TODO(urikz): Implement whole-word masking
  non_mention_end_positions = non_mention_start_positions
  non_mention_masked_positions = mask_tokens_by_spans(
      text_ids, non_mention_start_positions, non_mention_end_positions,
      non_mention_mask_rate,
      max_mlm_targets - tf.shape(mention_masked_positions)[0])

  # Merge masked positions for mention and non-mention tokens
  mlm_target_positions = tf.concat(
      [mention_masked_positions, non_mention_masked_positions], -1)
  n_mlm_target_positions = tf.shape(mlm_target_positions)

  # Get target IDs, weights and other features
  mlm_target_ids = tf.gather(text_ids, mlm_target_positions)
  mlm_target_weights = tf.ones(n_mlm_target_positions, dtype=tf.int64)
  mlm_target_is_mention = tf.ones(
      tf.shape(mention_masked_positions), dtype=tf.int64)
  seq_length = tf.shape(text_ids)[0]
  dense_is_masked = sparse_to_dense_1d(mlm_target_positions, seq_length)

  # Replace masked tokens with [MASK], random or original tokens.
  replacement_scores = tf.random.uniform(n_mlm_target_positions)
  replacement_tokens = tf.where(
      replacement_scores > random_replacement_prob + identity_replacement_prob,
      # replace tokens with [MASK]
      tf.cast(
          tf.fill(n_mlm_target_positions, value=mask_token_id), dtype=tf.int64),
      tf.where(
          replacement_scores > random_replacement_prob,
          # keep original
          mlm_target_ids,
          # replace with random token
          tf.random.uniform(
              n_mlm_target_positions, maxval=vocab_size, dtype=tf.int64)))
  replacement_positions = tf.expand_dims(mlm_target_positions, 1)
  # Indicies should be tf.int32 only.
  replacement_positions = tf.cast(replacement_positions, tf.int32)
  replacement_tokens = tf.scatter_nd(replacement_positions, replacement_tokens,
                                     tf.shape(text_ids))
  masked_text_ids = (
      text_ids * (1 - dense_is_masked) + replacement_tokens * dense_is_masked)

  return {
      'masked_text_ids':
          masked_text_ids,
      'mlm_target_positions':
          dynamic_padding_1d(mlm_target_positions, max_mlm_targets),
      'mlm_target_ids':
          dynamic_padding_1d(mlm_target_ids, max_mlm_targets),
      'mlm_target_weights':
          dynamic_padding_1d(mlm_target_weights, max_mlm_targets),
      'mlm_target_is_mention':
          dynamic_padding_1d(mlm_target_is_mention, max_mlm_targets),
      'dense_is_masked':
          dense_is_masked,
  }
Randomly masks whole mentions and random tokens up to a maximum.

  First, mentions are masked according to mention mask rate. If a mention is
  masked, all tokens in the mention are replaced by the mask token. If the
  passage has at least one mention and the mention rask rate is greater than
  zero, we mask at least one mention.

  After masking mentions, if there are fewer masked tokens than maximum mlm
  targets, we additionally mask non-mention words. TODO: If a token in a word
  is masked, all tokens in the word are masked. Some proportion of targets are
  not masked to ameliorate pretrain-finetune mismatch. If there are insufficient
  masked tokens, the target array is padded up to max targets.

  Args:
    text_ids: [seq_length] tensor with token ids.
    text_mask: [seq_length] tensor with 1s for tokens and 0 for padding.
    dense_span_starts: [seq_length] tensor with 1s for mention start positions
      and 0 otherwise.
    dense_span_ends: [seq_length] tensor with 1s for mention end positions and 0
      otherwise.
    non_mention_mask_rate: percentage of non mention tokens to be masked.
    mention_mask_rate: percentage of mentions to be masked.
    max_mlm_targets: total number of mlm targets.
    mask_token_id: token id for mask token.
    vocab_size: vocabulary size.
    random_replacement_prob: probability that to-be-masked token will be
      replaced with a random token instead of [MASK].
    identity_replacement_prob: probability that to-be-masked token will be
      replaced with itself instead of [MASK].

  Returns:
    Dictionary with masked text, mask positions, target ids, target weights.
  
language/mentionmemory/utils/mention_preprocess_utils.py