import torch


def get_prediction(prediction_scores, maximize_word, maximize_id, normalize):
  """Get the prediction for the current input.

  Args:
    prediction_scores: The scores that the model has output.
    maximize_word: The word id for which the activation should be maximized.
    maximize_id: The token id that we wish to be maximized for.
    normalize: Whether to normalize the prediction value to sum to one.

  Returns:
    prediction_score: Current prediction score for said word/token combination.
  """
  if normalize:
    return torch.nn.functional.softmax(
        prediction_scores[0][maximize_word], -1)[maximize_id]
  else:
    return prediction_scores[0][maximize_word][maximize_id]