Init cache. Args: cache: State cache dict for streaming inference. **kwargs: Additional keyword arguments.
(self, cache: dict = None, **kwargs)
| 692 | return results |
| 693 | |
| 694 | def init_cache(self, cache: dict = None, **kwargs): |
| 695 | """Init cache. |
| 696 | |
| 697 | Args: |
| 698 | cache: State cache dict for streaming inference. |
| 699 | **kwargs: Additional keyword arguments. |
| 700 | """ |
| 701 | if cache is None: |
| 702 | cache = {} |
| 703 | device = kwargs.get("device", "cuda") |
| 704 | |
| 705 | chunk_size = kwargs.get("chunk_size", [0, 10, 5]) |
| 706 | encoder_chunk_look_back = kwargs.get("encoder_chunk_look_back", 0) |
| 707 | decoder_chunk_look_back = kwargs.get("decoder_chunk_look_back", 0) |
| 708 | batch_size = 1 |
| 709 | |
| 710 | enc_output_size = kwargs["encoder_conf"]["output_size"] |
| 711 | feats_dims = kwargs["frontend_conf"]["n_mels"] * kwargs["frontend_conf"]["lfr_m"] |
| 712 | |
| 713 | cache_encoder = { |
| 714 | "start_idx": 0, |
| 715 | "cif_hidden": torch.zeros((batch_size, 1, enc_output_size)).to(device=device), |
| 716 | "cif_alphas": torch.zeros((batch_size, 1)).to(device=device), |
| 717 | "chunk_size": chunk_size, |
| 718 | "encoder_chunk_look_back": encoder_chunk_look_back, |
| 719 | "last_chunk": False, |
| 720 | "opt": None, |
| 721 | "feats": torch.zeros((batch_size, chunk_size[0] + chunk_size[2], feats_dims)).to( |
| 722 | device=device |
| 723 | ), |
| 724 | "tail_chunk": False, |
| 725 | } |
| 726 | cache["encoder"] = cache_encoder |
| 727 | |
| 728 | cache_decoder = { |
| 729 | "decode_fsmn": None, |
| 730 | "decoder_chunk_look_back": decoder_chunk_look_back, |
| 731 | "opt": None, |
| 732 | "chunk_size": chunk_size, |
| 733 | } |
| 734 | cache["decoder"] = cache_decoder |
| 735 | cache["frontend"] = {} |
| 736 | |
| 737 | cache["prev_samples"] = torch.empty(0) |
| 738 | |
| 739 | return cache |
| 740 | |
| 741 | def inference( |
| 742 | self, |