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

    See BaseTask.

    Our encoders take a flat tensor of mentions as input for the whole batch of
    passages to save memory. Here we aggregate mentions from different examples,
    and subsample if there are too many total mentions. In particular, we sample
    a random score for each mention, zeroing out padded mentions, then select
    the k highest-scoring. Similarly, we flatten mention features (passage start
    and end positions and padding mask), keeping track of original passage id
    through batch position feature.

    The encoders process the mention features in flattened form. Mention batch
    positions and start/end positions are then used to incorporate the results
    into the unflattened token representation through a scatter-type operation,
    taking the batch positions as the index into the first axis and the
    start/end positions as index into the second axis.

    Args:
      config: contains experiment hyperparameters.

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

    def collater_fn(batch: Dict[Text, tf.Tensor]) -> Dict[Text, tf.Tensor]:

      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)
      sampled_indices = tf.math.top_k(
          sample_scores, max_batch_mentions, sorted=False).indices

      dtype = batch['mention_start_positions'].dtype
      mention_mask = tf.gather(mention_mask, sampled_indices)
      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_sample_mentions),
          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])

      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

      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['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.

    See BaseTask.

    Our encoders take a flat tensor of mentions as input for the whole batch of
    passages to save memory. Here we aggregate mentions from different examples,
    and subsample if there are too many total mentions. In particular, we sample
    a random score for each mention, zeroing out padded mentions, then select
    the k highest-scoring. Similarly, we flatten mention features (passage start
    and end positions and padding mask), keeping track of original passage id
    through batch position feature.

    The encoders process the mention features in flattened form. Mention batch
    positions and start/end positions are then used to incorporate the results
    into the unflattened token representation through a scatter-type operation,
    taking the batch positions as the index into the first axis and the
    start/end positions as index into the second axis.

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