def collater_fn(batch: Dict[Text, tf.Tensor]) -> Dict[Text, tf.Tensor]:
      """Collater function for mention classification task. See BaseTask."""

      new_batch = {}

      # Sample mentions uniformly across batch
      mention_mask = tf.reshape(batch['mention_mask'], [n_candidate_mentions])
      sample_scores = tf.random.uniform(shape=[n_candidate_mentions]) * tf.cast(
          mention_mask, tf.float32)

      mention_target_indices = tf.reshape(batch['mention_target_indices'],
                                          [bsz])

      # We want to make sure that the target mentions always have a priority
      # when we sample `max_batch_mentions` out of all available mentions.
      # Additionally, we want these target mentions to be in the same order as
      # their samples. In other words, we want the first sampled mention to be
      # target mention from the first sample, the second sampled mention to be
      # tagret mention from the second sample, etc.

      # Positions of target mentions in the flat array
      mention_target_indices_flat = (
          tf.cast(
              tf.range(bsz) * max_mentions_per_sample,
              mention_target_indices.dtype) + mention_target_indices)
      # These extra score makes sure that target mentions have a priority and
      # will be sampled in the correct order.
      mention_target_extra_score_flat = tf.cast(
          tf.reverse(tf.range(bsz) + 1, axis=[0]), tf.float32)
      # The model assumes that there is only ONE target mention per sample.
      # Moreover,we want to select them according to the order of samples:
      # target mention from sample 0, target mention from sample 1, ..., etc.
      sample_scores = tf.tensor_scatter_nd_add(
          sample_scores, tf.expand_dims(mention_target_indices_flat, 1),
          mention_target_extra_score_flat)

      sampled_indices = tf.math.top_k(
          sample_scores, max_batch_mentions, sorted=True).indices

      # Double-check target mentions were selected correctly.
      assert_op = tf.assert_equal(
          sampled_indices[:bsz],
          tf.cast(mention_target_indices_flat, sampled_indices.dtype))

      with tf.control_dependencies([assert_op]):
        mention_mask = tf.gather(mention_mask, sampled_indices)
      dtype = batch['mention_start_positions'].dtype
      mention_start_positions = tf.gather(
          tf.reshape(batch['mention_start_positions'], [n_candidate_mentions]),
          sampled_indices)
      mention_end_positions = tf.gather(
          tf.reshape(batch['mention_end_positions'], [n_candidate_mentions]),
          sampled_indices)

      mention_batch_positions = tf.gather(
          tf.repeat(tf.range(bsz, dtype=dtype), max_mentions_per_sample),
          sampled_indices)

      new_batch['text_ids'] = batch['text_ids']
      new_batch['text_mask'] = batch['text_mask']
      new_batch['classifier_target'] = tf.reshape(
          batch['target'], [bsz, config.max_num_labels_per_sample])
      new_batch['classifier_target_mask'] = tf.reshape(
          batch['target_mask'], [bsz, config.max_num_labels_per_sample])

      new_batch['mention_mask'] = mention_mask
      new_batch['mention_start_positions'] = mention_start_positions
      new_batch['mention_end_positions'] = mention_end_positions
      new_batch['mention_batch_positions'] = mention_batch_positions
      new_batch['mention_target_indices'] = tf.range(bsz, dtype=dtype)

      if config.get('max_length_with_entity_tokens') is not None:
        batch_with_entity_tokens = mention_preprocess_utils.add_entity_tokens(
            text_ids=new_batch['text_ids'],
            text_mask=new_batch['text_mask'],
            mention_mask=new_batch['mention_mask'],
            mention_batch_positions=new_batch['mention_batch_positions'],
            mention_start_positions=new_batch['mention_start_positions'],
            mention_end_positions=new_batch['mention_end_positions'],
            new_length=config.max_length_with_entity_tokens,
        )
        # Update `text_ids`, `text_mask`, `mention_mask`, `mention_*_positions`
        new_batch.update(batch_with_entity_tokens)
        # Update `max_length`
        max_length = config.max_length_with_entity_tokens
      else:
        max_length = encoder_config.max_length

      new_batch['mention_target_batch_positions'] = tf.gather(
          new_batch['mention_batch_positions'],
          new_batch['mention_target_indices'])
      new_batch['mention_target_start_positions'] = tf.gather(
          new_batch['mention_start_positions'],
          new_batch['mention_target_indices'])
      new_batch['mention_target_end_positions'] = tf.gather(
          new_batch['mention_end_positions'],
          new_batch['mention_target_indices'])
      new_batch['mention_target_weights'] = tf.ones(bsz)

      # Fake IDs -- some encoders (ReadTwice) need them
      new_batch['mention_target_ids'] = tf.zeros(bsz)

      new_batch['segment_ids'] = tf.zeros_like(new_batch['text_ids'])

      position_ids = tf.expand_dims(tf.range(max_length, dtype=dtype), axis=0)
      new_batch['position_ids'] = tf.tile(position_ids, (bsz, 1))

      return new_batch
Collater function for mention classification task. See BaseTask.
language/mentionmemory/tasks/mention_classifier_task.py