def make_output_postprocess_fn(
      cls,
      config: ml_collections.ConfigDict  # pylint: disable=unused-argument
  ) -> Callable[[Dict[str, Any], Dict[str, Any]], Dict[str, Any]]:
    """Postprocess task samples (input and output). See BaseTask."""

    base_postprocess_fn = base_task.BaseTask.make_output_postprocess_fn(config)

    encoder_class = encoder_registry.get_registered_encoder(cls.encoder_name)
    encoder_postprocess_fn = encoder_class.make_output_postprocess_fn(config)

    def postprocess_fn(batch: Dict[str, Any],
                       auxiliary_output: Dict[str, Any]) -> Dict[str, Any]:
      """Function that prepares model's input and output for serialization."""

      new_auxiliary_output = {}
      new_auxiliary_output.update(auxiliary_output)
      encoder_specific_features = encoder_postprocess_fn(
          batch, new_auxiliary_output)
      new_auxiliary_output.update(encoder_specific_features)
      return base_postprocess_fn(batch, new_auxiliary_output)

    return postprocess_fn
Postprocess task samples (input and output). See BaseTask.
language/mentionmemory/tasks/mention_encoder_task.py