import torch


def accuracy(pred, y, regression=True, n_classes=3):
    """Accuracy of a prediction pred vs. ground truth y
    Arguments:
    - pred: a vector of sentiment predictions between 0 and 1
    - y: a vector of sentiment ground truths between 0 and 1
    - regression: whether the model output is a regression of the sentiment
                  or the class scores
    - n_classes: The number of classes in the model output

    Returns:
    Accuracy of the predictions when rounded to the nearest label (0/0.5/1)"""
    if regression:
        labels = torch.round(pred * (n_classes - 1)).type(torch.int)
    else:
        labels = torch.argmax(pred, dim=-1).type(torch.int)
    labels = labels.squeeze()
    y = (y * (n_classes - 1)).type(torch.int)
    return torch.mean((labels == y).float())