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

    See BaseTask.

    In the collater we subsample and flatten mentions across the batch, and
    sample mention targets for mention entity prediction or coreference tasks.

    Args:
      config: contains experiment hyperparameters.

    Returns:
      Function that preprocesses batches to be usable for the model
      (mod casting from tf to jnp dtype).
    """
    max_batch_mentions = config.max_mentions * config.per_device_batch_size
    max_batch_mention_targets = (
        config.max_mention_targets * config.per_device_batch_size)

    def collater_fn(batch: Dict[str, tf.Tensor]) -> Dict[str, tf.Tensor]:
      new_features = mention_preprocess_utils.process_batchwise_mention_targets(
          dense_span_starts=batch['dense_span_starts'],
          dense_span_ends=batch['dense_span_ends'],
          dense_mention_ids=batch['dense_mention_ids'],
          dense_linked_mention_mask=batch['dense_mention_mask'],
          dense_is_masked=batch.get('dense_is_masked'),
          max_mentions=max_batch_mentions,
          max_mention_targets=max_batch_mention_targets,
      )
      new_batch = {key: value for key, value in batch.items()}
      new_batch.update(new_features)

      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'],
            mlm_target_positions=new_batch.get('mlm_target_positions'),
            mlm_target_weights=new_batch.get('mlm_target_weights'),
            new_length=config.max_length_with_entity_tokens,
        )
        # Update `text_ids`, `text_mask`, `mention_mask`, `mention_*_positions`
        new_batch.update(batch_with_entity_tokens)

        # Need to update mention target features because some of the positions
        # and `mention_mask` have been updated.
        mention_target_features = mention_preprocess_utils.prepare_mention_target_features(
            new_batch['mention_batch_positions'],
            new_batch['mention_start_positions'],
            new_batch['mention_end_positions'], new_batch['mention_mask'],
            new_batch['mention_target_weights'],
            new_batch['mention_target_indices'])
        new_batch.update(mention_target_features)

        # Need to update `position_ids` and `segment_ids` because the max_length
        # has changed.
        new_batch['position_ids'] = tf.tile(
            tf.expand_dims(tf.range(config.max_length_with_entity_tokens), 0),
            [config.per_device_batch_size, 1])
        new_batch['position_ids'] = tf.cast(
            new_batch['position_ids'], dtype=new_batch['text_ids'].dtype)
        new_batch['segment_ids'] = tf.zeros(
            (config.per_device_batch_size,
             config.max_length_with_entity_tokens),
            dtype=new_batch['text_ids'].dtype)

      return new_batch

    return collater_fn
Produces function to preprocess batches.

    See BaseTask.

    In the collater we subsample and flatten mentions across the batch, and
    sample mention targets for mention entity prediction or coreference tasks.

    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/mention_encoder_task.py