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

    See BaseTask.

    Batches text identifiers after standard mention preprocessing. Also masks
    out mentions that are too close to a passage boundary, and for which we may
    not have enough context to generate a meaningful encoding.

    Args:
      config: contains experiment hyperparameters.

    Returns:
      Function that preprocesses batches to be usable for the model
      (mod casting from tf to jnp dtype).
    """
    mention_collater_fn = mention_encoder_task.MentionEncoderTask.make_collater_fn(config)  # pylint: disable=line-too-long
    min_distance_from_passage_boundary = config.min_distance_from_passage_boundary
    bsz = config.per_device_batch_size
    max_mentions_per_sample = config.max_mentions_per_sample

    def collater_fn(batch: Dict[str, tf.Tensor]) -> Dict[str, tf.Tensor]:
      new_batch = mention_collater_fn(batch)
      # Only generate text identifiers and mention hashes for
      # the target (linked) mentions.
      new_batch['target_text_identifiers'] = tf.gather(
          new_batch['text_identifiers'],
          new_batch['mention_target_batch_positions'])
      new_batch[
          'target_mention_hashes'] = mention_preprocess_utils.modified_cantor_pairing(
              new_batch['mention_target_start_positions'],
              new_batch['target_text_identifiers'])

      seq_len = tf.shape(batch['text_ids'])[1]
      starts_far_from_passage_boundary = tf.greater_equal(
          new_batch['mention_target_start_positions'],
          min_distance_from_passage_boundary)
      ends_far_from_passage_boundary = tf.less(
          new_batch['mention_target_end_positions'],
          tf.cast(seq_len, new_batch['mention_target_end_positions'].dtype) -
          min_distance_from_passage_boundary)
      far_from_passage_boundary = tf.logical_and(
          starts_far_from_passage_boundary, ends_far_from_passage_boundary)
      far_from_passage_boundary = tf.cast(
          far_from_passage_boundary,
          dtype=new_batch['mention_target_weights'].dtype)
      new_batch['mention_target_weights'] = (
          new_batch['mention_target_weights'] * far_from_passage_boundary)

      # Collect unique mention IDs per sample in the batch
      unique_mention_ids = []
      # Mask-out not linked entities.
      dense_mention_ids = batch['dense_mention_ids'] * batch[
          'dense_mention_mask']
      for i in range(bsz):
        unique_mention_ids_per_i = tf.unique(dense_mention_ids[i]).y
        unique_mention_ids_per_i = tf.cast(unique_mention_ids_per_i, tf.int32)
        unique_mention_ids_per_i = mention_preprocess_utils.dynamic_padding_1d(
            unique_mention_ids_per_i, max_mentions_per_sample)
        unique_mention_ids.append(unique_mention_ids_per_i)
      new_batch['unique_mention_ids'] = tf.stack(unique_mention_ids)
      return new_batch

    return collater_fn
Produces function to preprocess batches.

    See BaseTask.

    Batches text identifiers after standard mention preprocessing. Also masks
    out mentions that are too close to a passage boundary, and for which we may
    not have enough context to generate a meaningful encoding.

    Args:
      config: contains experiment hyperparameters.

    Returns:
      Function that preprocesses batches to be usable for the model
      (mod casting from tf to jnp dtype).
    
language/mentionmemory/tasks/memory_generation_task.py