def _gen_batch_for_test_add_entity_tokens(
      self, batch_size: int, min_length: int, max_length: int, vocab_size: int,
      max_mentions_per_sample: int, n_mentions: int, p_mention_to_pad: float,
      min_mlm_targets: int, max_mlm_targets: int,
      masked_token_id: int) -> Dict[str, np.ndarray]:
    """Generate batch with batchwise mentions to test `add_entity_tokens`."""

    features = {}

    # Generate text and mentions per sample
    features['text_ids'] = np.zeros((batch_size, max_length), dtype=np.int64)
    features['text_mask'] = np.zeros((batch_size, max_length), dtype=np.int64)
    features['mlm_target_positions'] = np.zeros((batch_size, max_mlm_targets),
                                                dtype=np.int64)
    features['mlm_target_weights'] = np.zeros((batch_size, max_mlm_targets),
                                              dtype=np.int64)
    all_mentions_in_batch = []

    for batch_index in range(batch_size):
      text_length = np.random.randint(low=min_length, high=max_length)
      text_ids = np.random.randint(
          low=1, high=vocab_size, size=text_length, dtype=np.int64)
      features['text_ids'][batch_index, :text_length] = text_ids
      features['text_mask'][batch_index, :text_length] = np.ones_like(text_ids)

      n_mlm_targets = np.random.randint(
          low=min_mlm_targets, high=max_mlm_targets)
      n_mlm_targets = min(n_mlm_targets, text_length)

      mlm_target_positions = np.random.choice(
          text_length, size=(n_mlm_targets), replace=False)
      features['mlm_target_positions'][
          batch_index, :n_mlm_targets] = mlm_target_positions
      features['mlm_target_weights'][batch_index, :n_mlm_targets] = 1

      features['text_ids'][batch_index, mlm_target_positions] = masked_token_id

      mention_positions = np.random.choice(
          text_length, size=(2 * max_mentions_per_sample), replace=False)
      mention_positions.sort()
      for index in range(max_mentions_per_sample):
        all_mentions_in_batch.append((batch_index, mention_positions[2 * index],
                                      mention_positions[2 * index + 1]))

    # Subsample at most `n_mentions` for the entire batch.
    np.random.shuffle(all_mentions_in_batch)

    def get_array_from_tuple(array_of_tuples, index):
      return np.array([x[index] for x in array_of_tuples], dtype=np.int64)

    all_mentions_in_batch = all_mentions_in_batch[:n_mentions]
    features['mention_batch_positions'] = get_array_from_tuple(
        all_mentions_in_batch, 0)
    self.assertLen(features['mention_batch_positions'], n_mentions)
    features['mention_start_positions'] = get_array_from_tuple(
        all_mentions_in_batch, 1)
    self.assertLen(features['mention_start_positions'], n_mentions)
    features['mention_end_positions'] = get_array_from_tuple(
        all_mentions_in_batch, 2)
    self.assertLen(features['mention_end_positions'], n_mentions)
    features['mention_mask'] = np.ones_like(features['mention_batch_positions'])

    # Mask fraction of `p_mention_to_pad` mentions.
    n_mentions_to_mask = int(p_mention_to_pad * n_mentions)
    if n_mentions_to_mask > 0:
      mentions_to_mask = np.random.choice(
          n_mentions, size=n_mentions_to_mask, replace=False)
      features['mention_mask'][mentions_to_mask] = 0
    self.assertEqual(features['mention_mask'].sum(),
                     n_mentions - n_mentions_to_mask)

    return features
Generate batch with batchwise mentions to test `add_entity_tokens`.
language/mentionmemory/utils/mention_preprocess_utils_test.py