def make_prediction_fn(
      cls,
      config: ml_collections.ConfigDict) -> Callable[..., Dict[str, Array]]:
    """Creates task prediction function for inference."""

    def predict_fn(
        model_config: ml_collections.FrozenConfigDict,
        model_params: Dict[str, Any],
        model_vars: Dict[str, Any],
        batch: Dict[str, Any],
    ) -> Dict[str, Array]:
      """Model-specific prediction function.

      Args:
        model_config: contains model config hyperparameters.
        model_params: contains model parameters.
        model_vars: contains model variables (not optimized).
        batch: model input.

      Returns:
        Dict[str, Array]. predictions.
      """

      variable_dict = {'params': model_params}
      variable_dict.update(model_vars)

      loss_helpers, _ = cls.build_model(model_config).apply(
          variable_dict, batch, deterministic=True, rngs=None)

      return loss_helpers['memory_generation']

    return predict_fn
Creates task prediction function for inference.
language/mentionmemory/tasks/memory_generation_task.py