Method test_greedy_decode

axlearn/audio/decoder_asr_test.py:567–666

(self, num_decodes, vocab_size, blank_id, logits_modifier)

        logits_modifier=[top_k_logits(1), config_for_function(top_k_logits).set(k=1)],
    )
    def test_greedy_decode(self, num_decodes, vocab_size, blank_id, logits_modifier):
        cfg: CTCDecoderModel.Config = CTCDecoderModel.default_config().set(
            input_dim=6,
            vocab_size=vocab_size,
            blank_id=blank_id,
        )

        # Initialize layer parameters.
        layer: CTCDecoderModel = cfg.set(name="test").instantiate(parent=None)
        prng_key = jax.random.PRNGKey(123)
        decode_key, predict_key, init_key, input_key = jax.random.split(prng_key, num=4)
        layer_params = layer.initialize_parameters_recursively(init_key)

        batch_size, max_seq_len = 4, 10
        seq_len = jnp.array([10, 7, 5, 8])
        # [batch_size, max_seq_len, dim].
        inputs = jax.random.normal(input_key, [batch_size, max_seq_len, cfg.input_dim]) * 1000
        # [batch_size, max_seq_len].
        paddings = jnp.arange(max_seq_len) >= seq_len[:, None]

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

        sample_decode_outputs: DecodeOutputs = jit_method(
            dict(input_batch=dict(inputs=inputs, paddings=paddings)),
            prng_key=decode_key,
            method="sample_decode",
            modify_logits=True,
            num_decodes=num_decodes,
        )

        greedy_decode_outputs: DecodeOutputs = jit_method(
            dict(input_batch=dict(inputs=inputs, paddings=paddings)),
            prng_key=decode_key,
            method="greedy_decode",
        )

        # Should be equivalent to taking the top logit of each output.
        # [batch_size, max_seq_len, vocab_size].
        logits = jit_method(
            dict(input_batch=dict(inputs=inputs, paddings=paddings)),
            prng_key=predict_key,
            method="predict",
        )
        log_probs = jax.nn.log_softmax(logits, axis=-1)