def make_collater_fn(
      config: ml_collections.ConfigDict
  ) -> Callable[[Dict[Text, tf.Tensor]], Dict[Text, tf.Tensor]]:
    """Produces function to preprocess batches for relation classification task.

    This function samples and flattens mentions from input data.

    Args:
      config: task configuration.

    Returns:
      collater function.
    """
    encoder_config = config.model_config.encoder_config
    bsz = config.per_device_batch_size
    max_batch_mentions = config.max_mentions * bsz
    n_candidate_mentions = config.max_mentions_per_sample * bsz

    if config.max_mentions < 2:
      raise ValueError('Need at least two mentions per sample in order to '
                       'include object and subject mentions.')

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

      def flatten_bsz(tensor):
        return tf.reshape(tensor, [bsz])

      new_batch = {
          'text_ids': batch['text_ids'],
          'text_mask': batch['text_mask'],
          'classifier_target': flatten_bsz(batch['target']),
      }

      # Sample mentions across batch

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

      subj_index = flatten_bsz(batch['subject_mention_indices'])
      obj_index = flatten_bsz(batch['object_mention_indices'])

      # Adjust subject / object mention positions in individual samples to their
      # positions in flattened mentions.
      shift = tf.range(
          bsz, dtype=obj_index.dtype) * config.max_mentions_per_sample
      mention_target_indices = tf.reshape(
          tf.stack([subj_index + shift, obj_index + shift], axis=1), [-1])

      # Sample the rest of the mentions uniformly across batch
      scores = tf.random.uniform(shape=tf.shape(batch['mention_mask']))
      scores = scores * tf.cast(batch['mention_mask'], tf.float32)

      # We want to adjust scores for target mentions so they don't get sampled
      # for the second time. We achive this by making their scores negative.
      def set_negative_scores(scores, indices):
        indices_2d = tf.stack([tf.range(bsz, dtype=indices.dtype), indices],
                              axis=1)
        return tf.tensor_scatter_nd_update(scores, indices_2d,
                                           tf.fill(tf.shape(indices), -1.0))

      # Note that since we're using 2D scores (not yet flattened for simplicity)
      # we use unadjusted `subj_index` and `obj_index`.
      scores = set_negative_scores(scores, subj_index)
      scores = set_negative_scores(scores, obj_index)

      # There are `2 * bsz` target mentions which were already chosen
      num_to_sample = tf.maximum(max_batch_mentions - 2 * bsz, 0)
      sampled_scores, sampled_indices = tf.math.top_k(
          tf.reshape(scores, [-1]), num_to_sample, sorted=True)

      # Note that negative scores indicate that we have double-sampled some of
      # the target mentions (we set their scores to negative right above).
      # In this case, we remove them.
      num_not_double_sampled = tf.reduce_sum(
          tf.cast(tf.not_equal(sampled_scores, -1), tf.int32))
      sampled_indices = sampled_indices[:num_not_double_sampled]

      # Combine target mentions (subject / object) with sampled mentions
      mention_target_indices = tf.cast(mention_target_indices,
                                       sampled_indices.dtype)
      sampled_indices = tf.concat([mention_target_indices, sampled_indices],
                                  axis=0)

      sampled_indices = mention_preprocess_utils.dynamic_padding_1d(
          sampled_indices, max_batch_mentions)

      dtype = batch['mention_start_positions'].dtype
      mention_mask = tf.reshape(batch['mention_mask'], [n_candidate_mentions])
      new_batch['mention_mask'] = tf.gather(mention_mask, sampled_indices)
      new_batch['mention_start_positions'] = tf.gather(
          tf.reshape(batch['mention_start_positions'], [n_candidate_mentions]),
          sampled_indices)
      new_batch['mention_end_positions'] = tf.gather(
          tf.reshape(batch['mention_end_positions'], [n_candidate_mentions]),
          sampled_indices)
      new_batch['mention_batch_positions'] = tf.gather(
          tf.repeat(tf.range(bsz, dtype=dtype), config.max_mentions_per_sample),
          sampled_indices)

      new_batch['mention_target_indices'] = tf.range(2 * bsz, dtype=dtype)
      new_batch['mention_subject_indices'] = tf.range(bsz, dtype=dtype) * 2
      new_batch['mention_object_indices'] = tf.range(bsz, dtype=dtype) * 2 + 1

      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(2 * bsz)

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

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

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

      return new_batch

    return collater_fn
Produces function to preprocess batches for relation classification task.

    This function samples and flattens mentions from input data.

    Args:
      config: task configuration.

    Returns:
      collater function.
    
language/mentionmemory/tasks/relation_classifier_task.py