def mask_mentions_and_tokens(
    text_ids: np.ndarray,
    text_mask: np.ndarray,
    mention_start_positions: np.ndarray,
    mention_end_positions: np.ndarray,
    mask_rate: float,
    mention_mask_rate: float,
    max_mlm_targets: int,
    mask_token_id: int,
) -> Dict[str, np.ndarray]:
  """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: array of token ids.
    text_mask: array with 1s for tokens and 0 for padding.
    mention_start_positions: array of sparse mention start positions.
    mention_end_positions: array of sparse mention end positions.
    mask_rate: percentage of normal 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.

  Returns:
    Masked text, mask positions, target ids, target weights
  """

  n_tokens = text_mask.sum(dtype=np.int64)

  n_mentions = len(mention_start_positions)
  n_sample_mentions = int(mention_mask_rate * n_mentions)
  # Mask at least one mention
  if n_sample_mentions == 0 and n_mentions > 0 and mention_mask_rate > 0.0:
    n_sample_mentions = 1

  # Mask mentions
  if n_sample_mentions > 0:
    sample_mention_mask_indices = np.random.choice(
        n_mentions, n_sample_mentions, replace=False)
    sample_mention_mask_indices.sort()
    mention_token_mask_positions = np.concatenate([
        np.arange(mention_start_positions[idx], mention_end_positions[idx] + 1)
        for idx in sample_mention_mask_indices
    ])
  else:
    mention_token_mask_positions = np.zeros(0, dtype=np.int64)

  if n_mentions == 0:
    non_mention_positions = np.arange(n_tokens)
  else:
    non_mention_spans = [
        np.arange(mention_end_positions[idx] + 1,
                  mention_start_positions[idx + 1])
        for idx in range(n_mentions - 1)
    ]

    non_mention_spans = [np.arange(mention_start_positions[0])
                        ] + non_mention_spans + [
                            np.arange(mention_end_positions[-1] + 1, n_tokens)
                        ]
    non_mention_positions = np.concatenate(non_mention_spans)

  # Mask non-mention words
  n_remaining_masks = max_mlm_targets - len(mention_token_mask_positions)
  if n_remaining_masks < 0:
    mention_token_mask_positions = mention_token_mask_positions[:
                                                                max_mlm_targets]
    all_mask_positions = mention_token_mask_positions
  else:
    n_non_mention_masks = min(n_remaining_masks,
                              int(mask_rate * len(non_mention_positions)))
    non_mention_mask_positions = np.random.choice(
        non_mention_positions, n_non_mention_masks, replace=False)
    all_mask_positions = np.concatenate(
        (mention_token_mask_positions, non_mention_mask_positions))

  masked_text_ids = np.copy(text_ids)
  masked_text_ids[all_mask_positions] = mask_token_id

  mlm_pad_shape = (0, max_mlm_targets - len(all_mask_positions))
  mlm_target_positions = np.pad(
      all_mask_positions, mlm_pad_shape, mode='constant')
  mlm_target_ids = text_ids[all_mask_positions]
  mlm_target_ids = np.pad(mlm_target_ids, mlm_pad_shape, mode='constant')
  mlm_target_weights = np.ones(len(all_mask_positions), dtype=np.float32)
  mlm_target_weights = np.pad(
      mlm_target_weights, mlm_pad_shape, mode='constant')

  mlm_target_is_mention = np.zeros_like(mlm_target_ids)
  mlm_target_is_mention[:len(mention_token_mask_positions)] = 1

  dense_is_masked = np.zeros_like(text_ids)
  dense_is_masked[all_mask_positions] = 1

  return {
      'masked_text_ids': masked_text_ids,
      'mlm_target_positions': mlm_target_positions,
      'mlm_target_ids': mlm_target_ids,
      'mlm_target_weights': mlm_target_weights,
      'mlm_target_is_mention': mlm_target_is_mention,
      '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: array of token ids.
    text_mask: array with 1s for tokens and 0 for padding.
    mention_start_positions: array of sparse mention start positions.
    mention_end_positions: array of sparse mention end positions.
    mask_rate: percentage of normal 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.

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