<filename>tests/metrics/test_nlp_metrics.py
# Copyright 2021 MosaicML. All Rights Reserved.

import math

import pytest
import torch
from torch.nn.functional import cross_entropy

from composer.metrics.nlp import BinaryF1Score, LanguageCrossEntropy, MaskedAccuracy


@pytest.mark.parametrize("ignore_index", [-100])
@pytest.mark.parametrize("num_classes", [2, 3, 4, 5])
@pytest.mark.xfail(reason="Test is flaky. See https://github.com/mosaicml/composer/issues/767.")
def test_masked_accuracy(ignore_index, num_classes):
    """Sanity check to make sure that masked accuracy has reasonable performance.

    Generates random targets and labels, and then ensures that the random targets and labels
    must hit at-chance accuracy.

    Args:
        batch_size (int): how many samples are in each batch
        ignore_index (Optional[int]): if present, the class index to ignore in accuracy calculations.
        num_classes (int): the number of classes in the classification task
    """
    batch_size = int(1e4)
    torchmetrics_masked_acc = MaskedAccuracy(ignore_index=ignore_index)
    # we're only testing binary accuracy -- expecteed accuracy should be 50%
    generated_preds = torch.rand((batch_size, num_classes))
    true_labels = torch.randint(low=0, high=num_classes - 1, size=(batch_size,))

    if ignore_index is not None:
        labels_mask = torch.rand((batch_size,))
        labels_mask[labels_mask > 0.8] = 1
        labels_mask[labels_mask <= 0.8] = 0
        labels_mask = labels_mask.bool()
        true_labels[labels_mask] = ignore_index

    true_labels = true_labels.float()
    generated_preds = generated_preds.float()

    torchmetrics_masked_acc.update(generated_preds, true_labels)
    final_acc = torchmetrics_masked_acc.compute()
    assert abs(final_acc - (1.0 / num_classes)) < 0.02


@pytest.mark.parametrize("ignore_index", [-100])
@pytest.mark.parametrize("batch_size", [1e2, 1e3])
@pytest.mark.parametrize("sequence_length", [128])
@pytest.mark.parametrize("num_classes", [2, 10])
@pytest.mark.parametrize("minibatch_size", [56, 256, 768])
def test_cross_entropy(batch_size: float, ignore_index: int, sequence_length: int, num_classes: int,
                       minibatch_size: int):
    """Sanity check to make sure that batched CrossEntropyLoss matches the expected performance.

    Generates a predicted distribution from a normal distribution, and a ground truth from a normal distribution.
    Verifies Cross Entropy Loss against the baseline performance.

    Args:
        batch_size (int): how many samples are in each batch
        ignore_index (Optional[int]): if present, the class index to ignore in accuracy calculations.
        sequence_length (int): the length of the generated sequence
        num_classes (int): the number of classes in the classification task
        minibatch_size (int): the minibatch size to simulate for model predictions
    """
    batch_size = int(batch_size)

    generated_preds = torch.randn((batch_size, sequence_length, num_classes))
    generated_true = torch.randint(low=0, high=num_classes, size=(batch_size, sequence_length))

    torchmetrics_xent = LanguageCrossEntropy(vocab_size=num_classes, dist_sync_on_step=False, ignore_index=ignore_index)

    if ignore_index is not None:
        labels_mask = torch.rand((batch_size, sequence_length))
        labels_mask[labels_mask > 0.8] = 1
        labels_mask[labels_mask <= 0.8] = 0
        labels_mask = labels_mask.bool()
        generated_true[labels_mask] = ignore_index

    num_batches = math.ceil(batch_size / minibatch_size)
    for batch_idx in range(num_batches):
        begin_idx = (batch_idx * minibatch_size)
        end_idx = ((batch_idx + 1) * minibatch_size)
        preds_subset = generated_preds[begin_idx:end_idx]
        true_subset = generated_true[begin_idx:end_idx]
        torchmetrics_xent.update(preds_subset, true_subset)

    torchmetrics_loss = torchmetrics_xent.compute()
    correct_loss = cross_entropy(generated_preds.view(-1, num_classes), generated_true.view(-1))
    assert torch.isclose(correct_loss, torchmetrics_loss)


@pytest.mark.parametrize("batch_size", [1e2, 1e3, 1e4])
@pytest.mark.parametrize("minibatch_size", [256, 768])
def test_binary_f1(batch_size, minibatch_size):
    """Sanity check to make sure that BinaryF1 TorchMetrics implementation matches the sklearn implementation.

    Generates a predicted set of labels, and a random set, and compares the resultant Binary F1 score.

    Args:
        batch_size (int): how many samples are in each batch
        minibatch_size (int): the minibatch size to simulate for model predictions
    """
    pytest.importorskip("sklearn", reason="sklearn is an optional dependency")
    from sklearn.metrics import f1_score

    batch_size = int(batch_size)

    generated_preds = torch.randn(size=(batch_size, 2))
    generated_true = torch.randint(low=0, high=2, size=(batch_size,))

    binary_f1 = BinaryF1Score()

    num_batches = math.ceil(batch_size / minibatch_size)
    for batch_idx in range(num_batches):
        begin_idx = (batch_idx * minibatch_size)
        end_idx = ((batch_idx + 1) * minibatch_size)
        preds_subset = generated_preds[begin_idx:end_idx]
        true_subset = generated_true[begin_idx:end_idx]
        binary_f1.update(preds_subset, true_subset)

    torchmetrics_f1 = binary_f1.compute()
    generated_preds = torch.argmax(generated_preds, dim=1)
    correct_f1 = f1_score(y_true=generated_true, y_pred=generated_preds)
    assert correct_f1 == torchmetrics_f1

<filename>tests/metrics/test_nlp_metrics.py
# Copyright 2021 MosaicML. All Rights Reserved.

import math

import pytest
import torch
from torch.nn.functional import cross_entropy

from composer.metrics.nlp import BinaryF1Score, LanguageCrossEntropy, MaskedAccuracy


@pytest.mark.parametrize("ignore_index", [-100])
@pytest.mark.parametrize("num_classes", [2, 3, 4, 5])
@pytest.mark.xfail(reason="Test is flaky. See https://github.com/mosaicml/composer/issues/767.")
def test_masked_accuracy(ignore_index, num_classes):
    """Sanity check to make sure that masked accuracy has reasonable performance.

    Generates random targets and labels, and then ensures that the random targets and labels
    must hit at-chance accuracy.

    Args:
        batch_size (int): how many samples are in each batch
        ignore_index (Optional[int]): if present, the class index to ignore in accuracy calculations.
        num_classes (int): the number of classes in the classification task
    """
    batch_size = int(1e4)
    torchmetrics_masked_acc = MaskedAccuracy(ignore_index=ignore_index)
    # we're only testing binary accuracy -- expecteed accuracy should be 50%
    generated_preds = torch.rand((batch_size, num_classes))
    true_labels = torch.randint(low=0, high=num_classes - 1, size=(batch_size,))

    if ignore_index is not None:
        labels_mask = torch.rand((batch_size,))
        labels_mask[labels_mask > 0.8] = 1
        labels_mask[labels_mask <= 0.8] = 0
        labels_mask = labels_mask.bool()
        true_labels[labels_mask] = ignore_index

    true_labels = true_labels.float()
    generated_preds = generated_preds.float()

    torchmetrics_masked_acc.update(generated_preds, true_labels)
    final_acc = torchmetrics_masked_acc.compute()
    assert abs(final_acc - (1.0 / num_classes)) < 0.02


@pytest.mark.parametrize("ignore_index", [-100])
@pytest.mark.parametrize("batch_size", [1e2, 1e3])
@pytest.mark.parametrize("sequence_length", [128])
@pytest.mark.parametrize("num_classes", [2, 10])
@pytest.mark.parametrize("minibatch_size", [56, 256, 768])
def test_cross_entropy(batch_size: float, ignore_index: int, sequence_length: int, num_classes: int,
                       minibatch_size: int):
    """Sanity check to make sure that batched CrossEntropyLoss matches the expected performance.

    Generates a predicted distribution from a normal distribution, and a ground truth from a normal distribution.
    Verifies Cross Entropy Loss against the baseline performance.

    Args:
        batch_size (int): how many samples are in each batch
        ignore_index (Optional[int]): if present, the class index to ignore in accuracy calculations.
        sequence_length (int): the length of the generated sequence
        num_classes (int): the number of classes in the classification task
        minibatch_size (int): the minibatch size to simulate for model predictions
    """
    batch_size = int(batch_size)

    generated_preds = torch.randn((batch_size, sequence_length, num_classes))
    generated_true = torch.randint(low=0, high=num_classes, size=(batch_size, sequence_length))

    torchmetrics_xent = LanguageCrossEntropy(vocab_size=num_classes, dist_sync_on_step=False, ignore_index=ignore_index)

    if ignore_index is not None:
        labels_mask = torch.rand((batch_size, sequence_length))
        labels_mask[labels_mask > 0.8] = 1
        labels_mask[labels_mask <= 0.8] = 0
        labels_mask = labels_mask.bool()
        generated_true[labels_mask] = ignore_index

    num_batches = math.ceil(batch_size / minibatch_size)
    for batch_idx in range(num_batches):
        begin_idx = (batch_idx * minibatch_size)
        end_idx = ((batch_idx + 1) * minibatch_size)
        preds_subset = generated_preds[begin_idx:end_idx]
        true_subset = generated_true[begin_idx:end_idx]
        torchmetrics_xent.update(preds_subset, true_subset)

    torchmetrics_loss = torchmetrics_xent.compute()
    correct_loss = cross_entropy(generated_preds.view(-1, num_classes), generated_true.view(-1))
    assert torch.isclose(correct_loss, torchmetrics_loss)


@pytest.mark.parametrize("batch_size", [1e2, 1e3, 1e4])
@pytest.mark.parametrize("minibatch_size", [256, 768])
def test_binary_f1(batch_size, minibatch_size):
    """Sanity check to make sure that BinaryF1 TorchMetrics implementation matches the sklearn implementation.

    Generates a predicted set of labels, and a random set, and compares the resultant Binary F1 score.

    Args:
        batch_size (int): how many samples are in each batch
        minibatch_size (int): the minibatch size to simulate for model predictions
    """
    pytest.importorskip("sklearn", reason="sklearn is an optional dependency")
    from sklearn.metrics import f1_score

    batch_size = int(batch_size)

    generated_preds = torch.randn(size=(batch_size, 2))
    generated_true = torch.randint(low=0, high=2, size=(batch_size,))

    binary_f1 = BinaryF1Score()

    num_batches = math.ceil(batch_size / minibatch_size)
    for batch_idx in range(num_batches):
        begin_idx = (batch_idx * minibatch_size)
        end_idx = ((batch_idx + 1) * minibatch_size)
        preds_subset = generated_preds[begin_idx:end_idx]
        true_subset = generated_true[begin_idx:end_idx]
        binary_f1.update(preds_subset, true_subset)

    torchmetrics_f1 = binary_f1.compute()
    generated_preds = torch.argmax(generated_preds, dim=1)
    correct_f1 = f1_score(y_true=generated_true, y_pred=generated_preds)
    assert correct_f1 == torchmetrics_f1

# Copyright 2021 MosaicML. All Rights Reserved. Sanity check to make sure that masked accuracy has reasonable performance.      Generates random targets and labels, and then ensures that the random targets and labels     must hit at-chance accuracy.      Args:         batch_size (int): how many samples are in each batch         ignore_index (Optional[int]): if present, the class index to ignore in accuracy calculations.         num_classes (int): the number of classes in the classification task # we're only testing binary accuracy -- expecteed accuracy should be 50% Sanity check to make sure that batched CrossEntropyLoss matches the expected performance.      Generates a predicted distribution from a normal distribution, and a ground truth from a normal distribution.     Verifies Cross Entropy Loss against the baseline performance.      Args:         batch_size (int): how many samples are in each batch         ignore_index (Optional[int]): if present, the class index to ignore in accuracy calculations.         sequence_length (int): the length of the generated sequence         num_classes (int): the number of classes in the classification task         minibatch_size (int): the minibatch size to simulate for model predictions Sanity check to make sure that BinaryF1 TorchMetrics implementation matches the sklearn implementation.      Generates a predicted set of labels, and a random set, and compares the resultant Binary F1 score.      Args:         batch_size (int): how many samples are in each batch         minibatch_size (int): the minibatch size to simulate for model predictions