Repository: huggingface/transformers
File: tests/generation/test_utils.py (Python)

Code before:
        self.original_matches = [1, 1, 1, 1, 1]
        self.candidate_generator.matches = self.original_matches
        self.num_matches = 2
        if sklearn_available:
            self.candidate_generator.update_candidate_strategy(self.input_ids, None, self.num_matches)
            self.assertEqual(self.candidate_generator.matches, [1, 1, 1, 1, 1, 1, 1, 0])
            self.assertEqual(self.candidate_generator.probs, [0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2])
            self.assertEqual(self.candidate_generator.assistant_generation_config.assistant_confidence_threshold, 0.3)
        else:
            self.assert_no_sklearn()

    @parameterized.expand([(is_sklearn_available(),), (False,)])
    def test_update_candidate_strategy_with_matches_1(self, sklearn_available):
        self.original_matches = [1, 1, 1, 1, 1]
        self.candidate_generator.matches = self.original_matches
        self.num_matches = 1
        if sklearn_available:
            self.candidate_generator.update_candidate_strategy(self.input_ids, None, self.num_matches)
            self.assertEqual(self.candidate_generator.matches, [1, 1, 1, 1, 1, 1, 0])
            self.assertEqual(self.candidate_generator.probs, [0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3])
            self.assertEqual(self.candidate_generator.assistant_generation_config.assistant_confidence_threshold, 0.4)
        else:
            self.assert_no_sklearn()


def has_similar_generate_outputs(output_1, output_2, atol=1e-3, rtol=1e-3) -> bool:
    """
    Returns a boolean indicating whether a pair of generate outputs are similar. Two `generate` call outputs are
    considered similar in the following situations:
    1. The sequences are the same
    2. The sequences are different, but the scores up to (and including) the first mismatch are nearly identical

    Args:
        output_1 (`GenerateOutput`): The first `generate` call output.
        output_2 (`GenerateOutput`): The second `generate` call output.
        atol (`float`, *optional*, defaults to 1e-5): The absolute tolerance for the scores.
        rtol (`float`, *optional*, defaults to 1e-5): The relative tolerance for the scores.

    Returns:
        A boolean indicating whether the two generate outputs are similar.
    """
    # scores doesn't include data regarding decoder input tokens
    decoder_input_length = output_1.sequences.shape[1] - len(output_1.scores)
    output_matches = output_1.sequences == output_2.sequences
    has_matching_outputs = output_matches.all()
    has_matching_scores = None
    if not has_matching_outputs:
        for batch_idx in range(output_1.sequences.shape[0]):
            batch_matches = output_matches[batch_idx]
            if batch_matches.all():
                continue

Review comment (suggestion):
Not feeling super confident about this, it's likely due to moe model's using `grouped_mm` which is inherently a bit flaky

It would be nicer to isolate this to moes only

Code after:
        if sklearn_available:
            self.candidate_generator.update_candidate_strategy(self.input_ids, None, self.num_matches)
            self.assertEqual(self.candidate_generator.matches, [1, 0, 1, 0, 1, 1, 1, 1, 0])
            self.assertEqual(self.candidate_generator.probs, [0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1])
            self.assertEqual(self.candidate_generator.assistant_generation_config.assistant_confidence_threshold, 0.2)
        else:
            self.assert_no_sklearn()

    @parameterized.expand([(is_sklearn_available(),), (False,)])
    def test_update_candidate_strategy_with_matches_4(self, sklearn_available):
        self.original_matches = [1, 1, 1, 1, 1]
        self.candidate_generator.matches = self.original_matches
        self.num_matches = 4
        if sklearn_available:
            self.candidate_generator.update_candidate_strategy(self.input_ids, None, self.num_matches)
            self.assertEqual(self.candidate_generator.matches, [1, 1, 1, 1, 1, 1, 1, 1, 1])
            self.assertEqual(self.candidate_generator.probs, [0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1])
            self.assertEqual(self.candidate_generator.assistant_generation_config.assistant_confidence_threshold, 0.4)
        else:
            self.assert_no_sklearn()

    @parameterized.expand([(is_sklearn_available(),), (False,)])
    def test_update_candidate_strategy_with_matches_3(self, sklearn_available):
        self.original_matches = [1, 1, 1, 1, 1]
        self.candidate_generator.matches = self.original_matches
        self.num_matches = 3
        if sklearn_available:
            self.candidate_generator.update_candidate_strategy(self.input_ids, None, self.num_matches)
            self.assertEqual(self.candidate_generator.matches, [1, 1, 1, 1, 1, 1, 1, 1, 0])
            self.assertEqual(self.candidate_generator.probs, [0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1])
            self.assertEqual(self.candidate_generator.assistant_generation_config.assistant_confidence_threshold, 0.2)
        else:
            self.assert_no_sklearn()

    @parameterized.expand([(is_sklearn_available(),), (False,)])
    def test_update_candidate_strategy_with_matches_2(self, sklearn_available):
        self.original_matches = [1, 1, 1, 1, 1]
        self.candidate_generator.matches = self.original_matches
        self.num_matches = 2
        if sklearn_available:
            self.candidate_generator.update_candidate_strategy(self.input_ids, None, self.num_matches)
            self.assertEqual(self.candidate_generator.matches, [1, 1, 1, 1, 1, 1, 1, 0])
            self.assertEqual(self.candidate_generator.probs, [0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2])
            self.assertEqual(self.candidate_generator.assistant_generation_config.assistant_confidence_threshold, 0.3)
        else:
            self.assert_no_sklearn()

    @parameterized.expand([(is_sklearn_available(),), (False,)])
    def test_update_candidate_strategy_with_matches_1(self, sklearn_available):
        self.original_matches = [1, 1, 1, 1, 1]
        self.candidate_generator.matches = self.original_matches