def test_values_as_expected(self):
    """Test whether metrics computation produces expected values."""
    probs = np.ones((self.batch_size, self.seq_len, self.num_items),
                    dtype=np.float32) / self.num_items

    classes = np.ones((self.batch_size, self.seq_len, self.num_items),
                      dtype=np.int32)
    targets = np.ones((self.batch_size, self.seq_len), dtype=np.int32)
    weights = np.random.randint(2, size=(self.batch_size, self.seq_len))

    # Check case where all classes are targets
    loss, avg_prob, denominator = metric_utils.compute_loss_and_prob_from_probs_with_duplicates(
        probs,
        classes,
        targets,
        weights,
    )

    self.assertAlmostEqual(loss / denominator, 0.0, 4)
    self.assertAlmostEqual(avg_prob / denominator, 1.0, 4)
    self.assertAlmostEqual(denominator, weights.sum(), 4)

    # Check case where no classes are targets
    targets = np.zeros((self.batch_size, self.seq_len), dtype=np.int32)
    loss, avg_prob, denominator = metric_utils.compute_loss_and_prob_from_probs_with_duplicates(
        probs,
        classes,
        targets,
        weights,
    )

    self.assertAlmostEqual(avg_prob / denominator, 0.0, 4)

    # Check random cases
    classes = np.random.randint(
        self.num_classes, size=(self.batch_size, self.seq_len, self.num_items))
    targets = np.random.randint(
        self.num_classes, size=(self.batch_size, self.seq_len))

    loss, avg_prob, denominator = metric_utils.compute_loss_and_prob_from_probs_with_duplicates(
        probs,
        classes,
        targets,
        weights,
    )

    correct_probs = (classes == np.expand_dims(targets, axis=-1)) * probs
    expected_avg_prob = (
        correct_probs * np.expand_dims(weights, axis=-1)).sum() / weights.sum()

    self.assertAlmostEqual(avg_prob / denominator, expected_avg_prob, 4)
Test whether metrics computation produces expected values.
language/mentionmemory/utils/metric_utils_test.py