Method test_decode

axlearn/audio/decoder_asr_test.py:1619–1695

(self)

        self.assertGreater(loss, 0.0)

    def test_decode(self):
        encoder_dim, decoder_dim, num_heads, vocab_size = 5, 16, 4, 20
        num_decodes = 5
        cfg: LASDecoderModel.Config = self._set_up_las(
            encoder_dim=encoder_dim,
            decoder_dim=decoder_dim,
            num_heads=num_heads,
            vocab_size=vocab_size,
        )
        pad_id = cfg.decoder.pad_token_id
        # Initialize layer parameters.
        layer: LASDecoderModel = cfg.set(name="test").instantiate(parent=None)
        prng_key = jax.random.PRNGKey(123)
        decode_key, init_key, input_key, prefix_key = jax.random.split(prng_key, num=4)
        layer_params = layer.initialize_parameters_recursively(init_key)

        batch_size, max_src_len, max_tgt_len = 3, 10, 8
        src_len = jnp.array([10, 7, 5])
        prefix_length = jnp.array([1, 3, 6])

        # [batch_size, max_seq_len, dim].
        inputs = jax.random.normal(input_key, [batch_size, max_src_len, cfg.input_dim]) * 1000
        # [batch_size, max_seq_len].
        paddings = jnp.arange(max_src_len) >= src_len[:, None]
        prefix = jax.random.randint(
            prefix_key,
            shape=[batch_size, max_tgt_len],
            # Prefix can consist of any tokens, including pad and eos.
            minval=0,
            maxval=vocab_size,
        )
        # Explicitly fill positions >= prefix_length with pad_id (-1).
        # Note that each batch example may have a different prefix length.
        # [batch_size, max_tgt_len].
        prefix_mask = sequence_mask(lengths=prefix_length, max_len=max_tgt_len)
        prefix = prefix * prefix_mask + pad_id * (1 - prefix_mask)

        @functools.partial(jax.jit, static_argnames=("method", "num_decodes", "logits_modifier"))
        def jit_method(inputs, prng_key, method, num_decodes, logits_modifier=None):
            if logits_modifier is not None:
                inputs["logits_modifier"] = logits_modifier
            outputs, _ = F(
                layer,
                inputs=dict(**inputs, max_decode_len=max_tgt_len, num_decodes=num_decodes),
                is_training=True,
                prng_key=prng_key,
                state=layer_params,
                method=method,
            )
            return outputs

        decode_inputs = dict(
            input_batch=dict(inputs=inputs, paddings=paddings, prefix=prefix),
        )

        # Beam search decode.
        beam_search_outputs: DecodeOutputs = jit_method(
            decode_inputs,