docs: def model_fn(features, labels, mode, params):
    """Defines how to train, evaluate and predict from the transformer model."""
    with tf.variable_scope("model"):
      inputs, targets = features, labels
  
      # Create model and get output logits.
      model = transformer.Transformer(params, mode == tf.estimator.ModeKeys.TRAIN)
  
      output = model(inputs, targets)
  
      # When in prediction mode, the labels/targets is None. The model output
      # is the prediction
      if mode == tf.estimator.ModeKeys.PREDICT:
        return tf.estimator.EstimatorSpec(
            tf.estimator.ModeKeys.PREDICT,
            predictions=output)
  
      logits = output
  
      # Calculate model loss.
      xentropy, weights = metrics.padded_cross_entropy_loss(
          logits, targets, params.label_smoothing, params.vocab_size)
      loss = tf.reduce_sum(xentropy * weights) / tf.reduce_sum(weights)
  
      if mode == tf.estimator.ModeKeys.EVAL:
        return tf.estimator.EstimatorSpec(
            mode=mode, loss=loss, predictions={"predictions": logits},
            eval_metric_ops=metrics.get_eval_metrics(logits, labels, params))
      else:
        train_op = get_train_op(loss, params)
        return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)