def test_mask_mentions_and_tokens_numpy(
      self,
      mask_rate: float,
      mention_mask_rate: float,
      text_len: int,
      n_mentions: int,
  ):
    """Test masking for correctness and consistency."""

    input_features = self._gen_input(
        text_len=text_len,
        n_mentions=n_mentions,
    )

    masked_dict = mention_preprocess_utils.mask_mentions_and_tokens(
        text_ids=input_features['text_ids'],
        text_mask=input_features['text_mask'],
        mention_start_positions=input_features['mention_start_positions'],
        mention_end_positions=input_features['mention_end_positions'],
        mask_rate=mask_rate,
        mention_mask_rate=mention_mask_rate,
        max_mlm_targets=self.max_mlm_targets,
        mask_token_id=self.mask_token_id,
    )

    # Check if arrays have correct length
    self.assertLen(masked_dict['masked_text_ids'], self.max_len)
    self.assertLen(masked_dict['mlm_target_positions'], self.max_mlm_targets)
    self.assertLen(masked_dict['mlm_target_ids'], self.max_mlm_targets)
    self.assertLen(masked_dict['mlm_target_weights'], self.max_mlm_targets)
    self.assertLen(masked_dict['mlm_target_is_mention'], self.max_mlm_targets)

    # Check if number of masked tokens corresponds to nonzero weights
    self.assertEqual(
        (masked_dict['masked_text_ids'] == self.mask_token_id).sum(),
        int(masked_dict['mlm_target_weights'].sum()))

    # Check if target ids consistent with token ids at target positions
    self.assertTrue(
        np.all((masked_dict['mlm_target_ids'] == input_features['text_ids'][
            masked_dict['mlm_target_positions']]
               )[:int(masked_dict['mlm_target_weights'].sum())]))
Test masking for correctness and consistency.
language/mentionmemory/utils/mention_preprocess_utils_test.py