Method test_text_input

axlearn/audio/input_asr_test.py:218–249

(
        self,
        expected: Sequence[dict[str, Any]],
        text: str,
        max_len: int,
        eos_id: int,
        input_key: str = "text",
        truncate: bool = False,
    )

        ),
    )
    def test_text_input(
        self,
        expected: Sequence[dict[str, Any]],
        text: str,
        max_len: int,
        eos_id: int,
        input_key: str = "text",
        truncate: bool = False,
    ):
        spm_file = os.path.join(tokenizers_dir, "librispeech_unigram_1024.model")
        vocab_cfg = config_for_class(seqio.SentencePieceVocabulary).set(
            sentencepiece_model_file=spm_file
        )
        vocab = vocab_cfg.instantiate()
        processor = input_tf_data.chain(
            input_asr.text_input(
                max_len=max_len,
                vocab=vocab_cfg,
                input_key=input_key,
                truncate=truncate,
                eos_id=eos_id,
            ),
            input_asr.make_autoregressive_inputs(vocab=vocab_cfg, bos_id=eos_id),
        )
        source = input_fake.fake_source(examples=[{input_key: text}], is_training=False)
        actual = list(processor(source()))
        for ex in expected:
            for k, v in ex.items():
                if k in ["input_ids", "target_labels"]:
                    v = [vocab.tokenizer.PieceToId(x) if isinstance(x, str) else x for x in v]
                ex[k] = tf.constant(v)
        tf.nest.map_structure(self.assertAllEqual, expected, actual)


class SpeechTextInputTest(parameterized.TestCase, tf.test.TestCase):