import torch


def get_best_non_unk_predictions(logits: torch.Tensor, unk_id):
    """
    For every sample, returns the prediction with the highest score. If this is the <unk> token, then the prediction
    with the second highest score will be returned instead
    :param logits: batch_size x num_predict x num_subtokens x vocab_size
    :param unk_id: ID of the <unk> token
    """
    top_2_predictions = logits.topk(2, -1).indices
    best_non_unk_predictions = top_2_predictions[:, :, :, 0]
    idx_first_unk = best_non_unk_predictions == unk_id
    best_non_unk_predictions[idx_first_unk] = top_2_predictions[:, :, :, 1][idx_first_unk]
    return best_non_unk_predictions