Method setUpClass

tests/test_text_decoder.py:82–154

(cls)

class TestWhisperTextDecoder(argmaxtools_test_utils.CoreMLTestsMixin, unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.test_output_names = TEST_OUTPUT_NAMES
        cls.test_cache_dir = TEST_CACHE_DIR
        cls.model_name = "TextDecoder"

        if not TEST_TOKEN_TIMESTAMPS:
            cls.test_output_names.pop(cls.test_output_names.index("alignment_heads_weights"))

        # Original model
        orig_torch_model = load_whisper_model(TEST_WHISPER_VERSION, TEST_TORCH_DTYPE)
        cls.orig_torch_model = (
            orig_torch_model.model.decoder.to(TEST_DEV).to(TEST_TORCH_DTYPE).eval()
        )

        # Base test model
        cls.test_torch_model = text_decoder.WhisperTextDecoder(
            cls.orig_torch_model.config
        )
        cls.test_torch_model.load_state_dict(cls.orig_torch_model.state_dict())
        cls.test_torch_model = (
            cls.test_torch_model.to(TEST_DEV).to(TEST_TORCH_DTYPE).eval()
        )
        cls.gen_cfg = orig_torch_model.generation_config

        if TEST_TOKEN_TIMESTAMPS:
            cls.test_torch_model.configure_for_token_timestamps(cls.gen_cfg)

        # Elaboration: I/O and architecture config
        cfg = cls.orig_torch_model.config
        cls.cfg = dict(
            n_heads=cfg.decoder_attention_heads,
            n_layers=cfg.decoder_layers,
            embed_dim=cfg.d_model,
            batch_size=1,
            vocab_size=cfg.vocab_size,
            enc_seq_len=cfg.max_source_positions,
            dec_kv_seq_len=TEST_DEC_KV_SEQ_LEN or cfg.max_target_positions,
        )

        cls.cfg["active_dec_kv_seq_len"] = random.randint(1, cls.cfg["dec_kv_seq_len"])
        logger.info(
            f"Decoding token at index {cls.cfg['active_dec_kv_seq_len']} "
            f"(max={cls.cfg['dec_kv_seq_len']}) "
        )

        (
            cls.test_torch_inputs,
            cls.orig_inputs,
        ) = test_utils._prepare_test_inputs_for_decoder(**cls.cfg)

        # Do casting and device placement per test config
        cls.test_torch_inputs = {k: place(v) for k, v in cls.test_torch_inputs.items()}
        cls.orig_inputs = {
            k: place(v)
            if isinstance(v, torch.Tensor)
            else [(place(vi[0]), place(vi[1])) for vi in v]
            for k, v in cls.orig_inputs.items()
        }