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

    See BaseTask.

    During preprocessing we mask entity mentions as well as non-mention tokens.

    Args:
      config: contains experiment hyperparameters.

    Returns:
      Function that preprocesses samples to be usable for the model
      (mod casting from tf to jnp dtype).
    """
    max_length = config.model_config.encoder_config.max_length
    mask_token_id = getattr(config, 'mask_token_id', default_values.MASK_TOKEN)
    apply_answer_mask = config.apply_answer_mask

    def preprocess_fn(example: Dict[str, tf.Tensor]) -> Dict[str, tf.Tensor]:
      """Performs preprocessing for individual sample."""

      int_type = example['text_ids'].dtype
      dense_is_masked = tf.equal(example['text_ids'], mask_token_id)
      dense_is_masked = tf.cast(dense_is_masked, dtype=int_type)
      example['dense_is_masked'] = dense_is_masked
      example['position_ids'] = tf.range(max_length, dtype=int_type)
      example['segment_ids'] = tf.zeros(max_length, dtype=int_type)

      if apply_answer_mask:
        example['dense_mention_ids'] = (
            example['dense_mention_ids'] * example['dense_answer_mask'])
      return example

    return preprocess_fn
Produces function to preprocess samples.

    See BaseTask.

    During preprocessing we mask entity mentions as well as non-mention tokens.

    Args:
      config: contains experiment hyperparameters.

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