提交

kongdeqiang

5 天以前 28ccfbfc51068a663a80764e14074df5edf2b5ba

 tests/test_auto_model.py

@@ -22,7 +22,39 @@
        kwargs["spk_kwargs"] = {"cb_kwargs": {"merge_thr": merge_thr}}
        model = AutoModel(**kwargs)
        self.assertEqual(model.cb_model.model_config['merge_thr'], merge_thr)
        # res = model.generate(input="/test.wav", 
        # res = model.generate(input="/test.wav",
        #              batch_size_s=300)

    def test_progress_callback_called(self):
        class DummyModel:
            def __init__(self):
                self.param = torch.nn.Parameter(torch.zeros(1))

            def parameters(self):
                return iter([self.param])

            def eval(self):
                pass

            def inference(self, data_in=None, **kwargs):
                results = [{"text": str(d)} for d in data_in]
                return results, {"batch_data_time": 1}

        am = AutoModel.__new__(AutoModel)
        am.model = DummyModel()
        am.kwargs = {"batch_size": 2, "disable_pbar": True}

        progress = []

        res = AutoModel.inference(
            am,
            ["a", "b", "c"],
            progress_callback=lambda idx, total: progress.append((idx, total)),
        )

        self.assertEqual(len(progress), 2)
        self.assertEqual(progress, [(2, 3), (3, 3)])


if __name__ == '__main__':
    unittest.main()
    unittest.main()