Method test_beam_search_decode_no_prefix

axlearn/common/decoding_test.py:562–607

(self, pad_id: int)

    @parameterized.parameters([0, 1])
    def test_beam_search_decode_no_prefix(self, pad_id: int):
        beam_size = 2

        def token_to_scores(ids, cache):  # pylint: disable=unused-argument
            # Use id 2 then 3 for batch element 0 and id 3 then 2 for element 1.
            log_probs = np.repeat(
                np.expand_dims(
                    np.array(
                        [[-1e5, -1e6, -0.37, -1.17], [-1e5, -1e6, -1.17, -0.37]], dtype=np.float32
                    ),
                    axis=1,
                ),
                [beam_size],
                axis=1,
            )
            log_probs = decoding.flatten_decoding_dim(log_probs)
            return log_probs, {}

        # No prefix is passed.
        inputs = np.full((2, 5), pad_id, dtype=np.int32)
        beam_search_output = decoding.beam_search_decode(
            inputs=inputs,
            time_step=decoding.infer_initial_time_step(inputs, pad_id=pad_id),
            cache={},
            tokens_to_scores=token_to_scores,
            eos_id=EOS_ID,
            num_decodes=beam_size,
            pad_id=pad_id,
        )

        # Beam search scores are in the descending order.
        self.assertTrue(np.all(np.diff(beam_search_output.scores) <= 0))
        self.assertTrue(np.all(np.diff(beam_search_output.live_scores) <= 0))
        np.testing.assert_array_equal(
            beam_search_output.sequences,
            np.array(
                [
                    [[1, pad_id, pad_id, pad_id, pad_id], [2, 1, pad_id, pad_id, pad_id]],
                    [[1, pad_id, pad_id, pad_id, pad_id], [3, 1, pad_id, pad_id, pad_id]],
                ],
            ),
        )
        np.testing.assert_array_equal(
            beam_search_output.live_sequences,
            np.array([[[2, 2, 2, 2, 2], [2, 2, 2, 2, 3]], [[3, 3, 3, 3, 3], [3, 3, 3, 3, 2]]]),
        )

    @parameterized.product(prefix_length=[1, 3, 6, 7], low=[0, 1])
    def test_beam_init_with_inputs(self, prefix_length: int, low: int):