| | |
|---|
| | | 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() |
|---|
