import torch


def count_correct(labels, outputs):
    """ Count correctly reconstructed tokens (out of masked tokens) """
    predictions = torch.argmax(outputs, dim=-1)
    mask = labels != -100
    correct = torch.eq(predictions, labels)[mask].sum().item()
    total = mask.sum().item()
    return correct, total