def get_sari(source_ids, prediction_ids, target_ids, max_gram_size=4):
  """Computes the SARI scores from the given source, prediction and targets.

  Args:
    source_ids: A 2D tf.Tensor of size (batch_size , sequence_length)
    prediction_ids: A 2D tf.Tensor of size (batch_size, sequence_length)
    target_ids: A 3D tf.Tensor of size (batch_size, number_of_targets,
        sequence_length)
    max_gram_size: int. largest n-gram size we care about (e.g. 3 for unigrams,
        bigrams, and trigrams)

  Returns:
    A 4-tuple of 1D float Tensors of size (batch_size) for the SARI score and
        the keep, addition and deletion scores.
  """

  def get_sari_numpy(source_ids, prediction_ids, target_ids):
    """Iterate over elements in the batch and call the SARI function."""
    sari_scores = []
    keep_scores = []
    add_scores = []
    deletion_scores = []
    # Iterate over elements in the batch.
    for source_ids_i, prediction_ids_i, target_ids_i in zip(
        source_ids, prediction_ids, target_ids):
      sari, keep, add, deletion = get_sari_score(
          source_ids_i, prediction_ids_i, target_ids_i, max_gram_size,
          BETA_FOR_SARI_DELETION_F_MEASURE)
      sari_scores.append(sari)
      keep_scores.append(keep)
      add_scores.append(add)
      deletion_scores.append(deletion)
    return (np.asarray(sari_scores), np.asarray(keep_scores),
            np.asarray(add_scores), np.asarray(deletion_scores))

  sari, keep, add, deletion = tf.py_func(
      get_sari_numpy,
      [source_ids, prediction_ids, target_ids],
      [tf.float64, tf.float64, tf.float64, tf.float64])
  return sari, keep, add, deletion