Method test_decode

axlearn/common/encoder_decoder_test.py:174–273

Test beam search and sample decoding from a randomly initialized model.

(
        self,
        stack_cfg: BaseStackedTransformerLayer.Config,
        num_decodes: int,
        prefix_length: utils.Tensor,
        method: Literal["sample_decode", "beam_search_decode"],
        pad_token_id: int,
    )

    )
    # pylint: disable-next=too-many-statements
    def test_decode(
        self,
        stack_cfg: BaseStackedTransformerLayer.Config,
        num_decodes: int,
        prefix_length: utils.Tensor,
        method: Literal["sample_decode", "beam_search_decode"],
        pad_token_id: int,
    ):
        """Test beam search and sample decoding from a randomly initialized model."""
        with jax.checking_leaks():
            batch_size, src_len, tgt_len, vocab_size = prefix_length.shape[0], 11, 10, 6
            bos_id = 1
            init_key, prefix_key, source_key, method_key = jax.random.split(
                jax.random.PRNGKey(0), num=4
            )

            if isinstance(stack_cfg, RepeatedTransformerLayer.Config):
                remat_spec = RematSpec(prevent_cse=False)
            else:
                remat_spec = None

            cfg = _model_config(
                vocab_size=vocab_size, source_len=src_len, target_len=tgt_len, remat_spec=remat_spec
            )
            cfg.encoder.pad_token_id = pad_token_id
            cfg.decoder.pad_token_id = pad_token_id
            model = cfg.set(name="test").instantiate(parent=None)
            params = model.initialize_parameters_recursively(init_key)

            prefix = jax.random.randint(
                prefix_key,
                shape=[batch_size, tgt_len],
                # Prefix can consist of any tokens, including pad and eos.
                minval=0,
                maxval=vocab_size,
            )
            # Explicitly fill positions >= prefix_length with pad_token_id.
            # Note that each batch example may have a different prefix length.
            # [batch_size, tgt_len].
            prefix_mask = utils.sequence_mask(lengths=prefix_length, max_len=tgt_len)
            prefix = prefix * prefix_mask + pad_token_id * (1 - prefix_mask)
            # Set last token to a non-pad token, to fix the prefix length.
            oh_indices = jax.nn.one_hot(prefix_length - 1, tgt_len, dtype=prefix.dtype)
            prefix = prefix * (1 - oh_indices) + bos_id * oh_indices

            source_ids = jax.random.randint(
                source_key, minval=1, maxval=vocab_size, shape=(batch_size, src_len)
            )
            source_mask = dummy_padding_mask(batch_size=batch_size, max_seq_len=src_len)
            source_ids = source_ids * source_mask + pad_token_id * (1 - source_mask)
            inputs = dict(
                input_batch=dict(prefix=prefix, source=dict(input_ids=source_ids)),
                max_sequence_length=tgt_len,
                num_decodes=num_decodes,
            )
            if method == "sample_decode":
                # Modify logits so that we will always sample the last token ID.
                inputs["logits_modifier"] = lambda logits: (