Method test_postprocess

axlearn/audio/decoder_asr_test.py:879–946

(self)

        self.assertLess(beam_search_outputs.scores[3, 3], 0.5 * NEG_INF)

    def test_postprocess(self):
        input_dim, vocab_size = 3, 8
        cfg = CTCDecoderModel.default_config().set(input_dim=input_dim, vocab_size=vocab_size)

        # Initialize layer parameters.
        layer: CTCDecoderModel = cfg.set(name="test").instantiate(parent=None)
        layer_params = layer.initialize_parameters_recursively(jax.random.PRNGKey(123))

        sequences = jnp.array(
            [
                [[3, 0, 5, 0, 0, 6, 6, 2, 0, 0, 0], [5, 4, 4, 5, 5, 0, 0, 4, 3, 0, 0]],
                [[4, 4, 4, 0, 4, 6, 4, 4, 2, 0, 0], [3, 5, 0, 0, 0, 6, 2, 2, 2, 2, 0]],
            ],
        )
        paddings = jnp.array(
            [
                [0, 0, 0, 0, 0, 0, 0, 0, 1, 1],
                [0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
            ]
        ).astype(jnp.bool)
        scores = jnp.array(
            [
                [[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10], [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]],
                [[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10], [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]],
            ]
        )

        @jax.jit
        def jit_postprocess(sequences, paddings, scores):
            outputs, _ = F(
                layer,
                inputs=dict(sequences=sequences, paddings=paddings, scores=scores),
                is_training=True,
                prng_key=jax.random.PRNGKey(456),
                state=layer_params,
                method="_postprocess_outputs",
            )
            return outputs

        outputs: DecodeOutputs = jit_postprocess(sequences, paddings, scores)
        self.assertNestedEqual(
            outputs.raw_sequences,
            jnp.array(
                [
                    [[3, 0, 5, 0, 0, 6, 6, 2, 0, 0], [5, 4, 4, 5, 5, 0, 0, 4, 0, 0]],
                    [[4, 4, 4, 0, 4, 6, 4, 4, 2, 0], [3, 5, 0, 0, 0, 6, 2, 2, 2, 0]],
                ]
            ),
        )
        self.assertNestedEqual(
            outputs.sequences,
            jnp.array(
                [
                    [[3, 5, 6, 2, 0, 0, 0, 0, 0, 0], [5, 4, 5, 4, 0, 0, 0, 0, 0, 0]],
                    [[4, 4, 6, 4, 2, 0, 0, 0, 0, 0], [3, 5, 6, 2, 0, 0, 0, 0, 0, 0]],
                ]
            ),
        )