def sari_score(predictions, labels, features, **unused_kwargs):
  """Computes the SARI scores from the given source, prediction and targets.

  An approximate SARI scoring method since we do not glue word pieces or
  decode the ids and tokenize the output. By default, we use ngram order of 4.
  Also, this does not have beam search.

  Args:
    predictions: tensor, model predictions.
    labels: tensor, gold output.
    features: dict, containing inputs.

  Returns:
    sari: int, approx sari score
  """
  if "inputs" not in features:
    raise ValueError("sari_score requires inputs feature")

  # Convert the inputs and outputs to a [batch_size, sequence_length] tensor.
  inputs = tf.squeeze(features["inputs"], axis=[-1, -2])
  outputs = tf.to_int32(tf.argmax(predictions, axis=-1))
  outputs = tf.squeeze(outputs, axis=[-1, -2])

  # Convert the labels to a [batch_size, 1, sequence_length] tensor.
  labels = tf.squeeze(labels, axis=[-1, -2])
  labels = tf.expand_dims(labels, axis=1)

  score, _, _, _ = get_sari(inputs, outputs, labels)
  return score, tf.constant(1.0)