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).
    """
    model_config = config.model_config
    encoder_config = model_config.encoder_config
    max_mlm_targets = config.max_mlm_targets
    max_length = encoder_config.max_length
    vocab_size = encoder_config.vocab_size

    mask_rate = config.mask_rate
    mention_mask_rate = config.mention_mask_rate
    mask_token_id = getattr(config, 'mask_token_id', 103)

    def preprocess_fn(example: Dict[str, tf.Tensor]) -> Dict[str, tf.Tensor]:
      """Performs preprocessing for individual sample."""
      new_example = {}
      new_example['text_mask'] = example['text_mask']
      new_example['dense_span_starts'] = example['dense_span_starts']
      new_example['dense_span_ends'] = example['dense_span_ends']
      new_example['dense_mention_ids'] = example['dense_mention_ids']
      new_example['dense_mention_mask'] = example['dense_mention_mask']

      if max_mlm_targets > 0:
        # Perform masking
        masked_dict = mention_preprocess_utils.mask_mentions_and_tokens_tf(
            text_ids=example['text_ids'],
            text_mask=example['text_mask'],
            dense_span_starts=example['dense_span_starts'],
            dense_span_ends=example['dense_span_ends'],
            non_mention_mask_rate=mask_rate,
            mention_mask_rate=mention_mask_rate,
            max_mlm_targets=max_mlm_targets,
            mask_token_id=mask_token_id,
            vocab_size=vocab_size,
            random_replacement_prob=0.1,
            identity_replacement_prob=0.1,
        )

        new_example['text_ids'] = masked_dict['masked_text_ids']
        new_example['mlm_target_positions'] = masked_dict[
            'mlm_target_positions']
        new_example['mlm_target_ids'] = masked_dict['mlm_target_ids']
        new_example['mlm_target_weights'] = masked_dict['mlm_target_weights']
        new_example['mlm_target_is_mention'] = masked_dict[
            'mlm_target_is_mention']
        new_example['dense_is_masked'] = masked_dict['dense_is_masked']
      else:
        # Don't perform masking if 0 `max_mlm_targets` were requested.
        new_example['text_ids'] = example['text_ids']
        new_example['dense_is_masked'] = np.zeros(max_length, dtype=np.int32)

      new_example['position_ids'] = tf.range(max_length)
      new_example['segment_ids'] = tf.zeros(
          max_length, dtype=example['text_ids'].dtype)
      return new_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/mention_encoder_task.py