def test_process_batchwise_mention_targets(
      self,
      batch_size: int,
      text_len: int,
      n_mentions: int,
      max_mentions: int,
      max_mention_targets: int,
  ):
    """Test batchwise mention processing for correctness and consistency."""
    max_mention_targets = min(max_mention_targets, max_mentions)
    dense_span_starts = np.zeros((batch_size, self.max_len), dtype=np.int32)
    dense_span_ends = np.zeros((batch_size, self.max_len), dtype=np.int32)
    dense_mention_ids_list = []
    dense_mention_mask_list = []
    dense_is_masked_list = []

    for index in range(batch_size):
      input_features = self._gen_input(
          text_len=text_len,
          n_mentions=n_mentions,
      )
      dense_span_starts[index, input_features['mention_start_positions']] = 1
      dense_span_ends[index, input_features['mention_end_positions']] = 1
      dense_mention_ids_list.append(input_features['dense_mention_ids'])
      dense_mention_mask_list.append(input_features['dense_mention_mask'])
      dense_is_masked_list.append(input_features['dense_is_masked'])

    def placeholder(array: np.ndarray) -> tf.Tensor:
      return tf.compat.v1.placeholder_with_default(
          array.astype(np.int32), shape=[None])

    dense_mention_ids = np.stack(dense_mention_ids_list)
    dense_mention_mask = np.stack(dense_mention_mask_list)
    dense_is_masked = np.stack(dense_is_masked_list)

    actual_dict_tf = mention_preprocess_utils.process_batchwise_mention_targets(
        dense_span_starts=placeholder(dense_span_starts),
        dense_span_ends=placeholder(dense_span_ends),
        dense_mention_ids=placeholder(dense_mention_ids),
        dense_linked_mention_mask=placeholder(dense_mention_mask),
        dense_is_masked=placeholder(dense_is_masked),
        max_mentions=max_mentions,
        max_mention_targets=max_mention_targets,
    )

    actual_dict_np = self.evaluate(actual_dict_tf)

    # Check if all arrays are the right length
    self.assertLen(actual_dict_np['mention_batch_positions'], max_mentions)
    self.assertLen(actual_dict_np['mention_start_positions'], max_mentions)
    self.assertLen(actual_dict_np['mention_end_positions'], max_mentions)
    self.assertLen(actual_dict_np['mention_mask'], max_mentions)
    self.assertLen(actual_dict_np['mention_target_weights'],
                   max_mention_targets)
    self.assertLen(actual_dict_np['mention_target_indices'],
                   max_mention_targets)
    self.assertLen(actual_dict_np['mention_target_ids'], max_mention_targets)
    self.assertLen(actual_dict_np['mention_target_is_masked'],
                   max_mention_targets)
    self.assertLen(actual_dict_np['mention_target_batch_positions'],
                   max_mention_targets)
    self.assertLen(actual_dict_np['mention_target_start_positions'],
                   max_mention_targets)
    self.assertLen(actual_dict_np['mention_target_end_positions'],
                   max_mention_targets)

    # Check every individual mention extracted
    num_sampled_mentions = 0
    for i in range(max_mentions):
      if actual_dict_np['mention_mask'][i] == 0:
        self.assertEqual(actual_dict_np['mention_batch_positions'][i], 0)
        self.assertEqual(actual_dict_np['mention_start_positions'][i], 0)
        self.assertEqual(actual_dict_np['mention_end_positions'][i], 0)
      else:
        num_sampled_mentions += 1
        batch_index = actual_dict_np['mention_batch_positions'][i]
        start_pos = actual_dict_np['mention_start_positions'][i]
        end_pos = actual_dict_np['mention_end_positions'][i]
        self.assertEqual(dense_span_starts[batch_index, start_pos], 1)
        self.assertEqual(dense_span_ends[batch_index, end_pos], 1)
        # Check start_pos and end_pos actually correspond to the same mention
        if start_pos < end_pos:
          self.assertEqual(
              dense_span_starts[batch_index, start_pos + 1:end_pos + 1].sum(),
              0)
          self.assertEqual(
              dense_span_ends[batch_index, start_pos + 1:end_pos].sum(), 0)

    total_mentions = dense_span_starts.sum()
    if max_mentions < total_mentions:
      self.assertEqual(num_sampled_mentions, max_mentions)
    else:
      self.assertEqual(num_sampled_mentions, total_mentions)

    # Check every individual linked mention extracted
    num_sampled_target_mentions = 0
    for j in range(max_mention_targets):
      if actual_dict_np['mention_target_weights'][j] == 0:
        self.assertEqual(actual_dict_np['mention_target_indices'][j], 0)
        self.assertEqual(actual_dict_np['mention_target_ids'][j], 0)
      else:
        num_sampled_target_mentions += 1
        i = actual_dict_np['mention_target_indices'][j]
        batch_index = actual_dict_np['mention_batch_positions'][i]
        start_pos = actual_dict_np['mention_start_positions'][i]
        end_pos = actual_dict_np['mention_end_positions'][i]
        self.assertEqual(dense_mention_mask[batch_index, start_pos], 1)
        self.assertEqual(dense_mention_ids[batch_index, start_pos],
                         actual_dict_np['mention_target_ids'][j])

    total_target_mentions = (dense_span_starts * dense_mention_mask).sum()
    if max_mention_targets < total_target_mentions:
      self.assertEqual(num_sampled_target_mentions, max_mention_targets)
    else:
      self.assertEqual(num_sampled_target_mentions, total_target_mentions)

    # Check mention masking computed correctly
    expected_value = dense_is_masked[
        actual_dict_np['mention_batch_positions'],
        actual_dict_np['mention_start_positions']][:max_mention_targets]
    self.assertTrue(
        np.all(actual_dict_np['mention_target_is_masked'] == expected_value))

    # Check mention_x_target_positions
    mention_target_indices = actual_dict_np['mention_target_indices']
    expected = actual_dict_np['mention_batch_positions'][mention_target_indices]
    self.assertTrue(
        np.all(actual_dict_np['mention_target_batch_positions'] == expected))
    expected = actual_dict_np['mention_start_positions'][mention_target_indices]
    self.assertTrue(
        np.all(actual_dict_np['mention_target_start_positions'] == expected))
    expected = actual_dict_np['mention_end_positions'][mention_target_indices]
    self.assertTrue(
        np.all(actual_dict_np['mention_target_end_positions'] == expected))
Test batchwise mention processing for correctness and consistency.
language/mentionmemory/utils/mention_preprocess_utils_test.py