add test for espnet2 quantization

chintu619 · May 10, 2022 · a48423f · a48423f
1 parent acb24c8
commit a48423f
Showing 1 changed file with 18 additions and 0 deletions.
diff --git a/test/espnet2/bin/test_asr_inference.py b/test/espnet2/bin/test_asr_inference.py
@@ -84,6 +84,24 @@ def test_Speech2Text(asr_config_file, lm_config_file):
         assert isinstance(hyp, Hypothesis)
 
 
+@pytest.mark.execution_timeout(5)
+def test_Speech2Text_quantized(asr_config_file, lm_config_file):
+    speech2text = Speech2Text(
+        asr_train_config=asr_config_file, 
+        lm_train_config=lm_config_file, 
+        beam_size=1,
+        quantize_asr_model=True, 
+        quantize_lm=True,
+    )
+    speech = np.random.randn(100000)
+    results = speech2text(speech)
+    for text, token, token_int, hyp in results:
+        assert isinstance(text, str)
+        assert isinstance(token[0], str)
+        assert isinstance(token_int[0], int)
+        assert isinstance(hyp, Hypothesis)
+
+
 @pytest.fixture()
 def asr_config_file_streaming(tmp_path: Path, token_list):
     # Write default configuration file