def prepare_mention_target_features(
    mention_batch_positions: tf.Tensor,
    mention_start_positions: tf.Tensor,
    mention_end_positions: tf.Tensor,
    mention_mask: tf.Tensor,
    mention_target_weights: tf.Tensor,
    mention_target_indices: tf.Tensor,
) -> Dict[str, tf.Tensor]:
  """Produce mention target features based on batchwise mention features."""
  mention_target_weights = mention_target_weights * tf.gather(
      mention_mask, mention_target_indices)
  mention_target_batch_positions = tf.gather(mention_batch_positions,
                                             mention_target_indices)
  mention_target_start_positions = tf.gather(mention_start_positions,
                                             mention_target_indices)
  mention_target_end_positions = tf.gather(mention_end_positions,
                                           mention_target_indices)
  return {
      'mention_target_batch_positions': mention_target_batch_positions,
      'mention_target_start_positions': mention_target_start_positions,
      'mention_target_end_positions': mention_target_end_positions,
      'mention_target_weights': mention_target_weights,
  }
Produce mention target features based on batchwise mention features.
language/mentionmemory/utils/mention_preprocess_utils.py