Encode chunk of input sequence. Args: x: Conformer input sequences. (B, T, D_block) pos_enc: Positional embedding sequences. (B, 2 * (T - 1), D_block) mask: Source mask. (B, T_2) left_context: Number of frames in left context. right
(
self,
x: torch.Tensor,
pos_enc: torch.Tensor,
mask: torch.Tensor,
chunk_size: int = 16,
left_context: int = 0,
right_context: int = 0,
)
| 851 | return x, mask, pos_enc |
| 852 | |
| 853 | def chunk_forward( |
| 854 | self, |
| 855 | x: torch.Tensor, |
| 856 | pos_enc: torch.Tensor, |
| 857 | mask: torch.Tensor, |
| 858 | chunk_size: int = 16, |
| 859 | left_context: int = 0, |
| 860 | right_context: int = 0, |
| 861 | ) -> Tuple[torch.Tensor, torch.Tensor]: |
| 862 | """Encode chunk of input sequence. |
| 863 | Args: |
| 864 | x: Conformer input sequences. (B, T, D_block) |
| 865 | pos_enc: Positional embedding sequences. (B, 2 * (T - 1), D_block) |
| 866 | mask: Source mask. (B, T_2) |
| 867 | left_context: Number of frames in left context. |
| 868 | right_context: Number of frames in right context. |
| 869 | Returns: |
| 870 | x: Conformer output sequences. (B, T, D_block) |
| 871 | pos_enc: Positional embedding sequences. (B, 2 * (T - 1), D_block) |
| 872 | """ |
| 873 | residual = x |
| 874 | |
| 875 | x = self.norm_macaron(x) |
| 876 | x = residual + self.feed_forward_scale * self.feed_forward_macaron(x) |
| 877 | |
| 878 | residual = x |
| 879 | x = self.norm_self_att(x) |
| 880 | if left_context > 0: |
| 881 | key = torch.cat([self.cache[0], x], dim=1) |
| 882 | else: |
| 883 | key = x |
| 884 | val = key |
| 885 | |
| 886 | if right_context > 0: |
| 887 | att_cache = key[:, -(left_context + right_context) : -right_context, :] |
| 888 | else: |
| 889 | att_cache = key[:, -left_context:, :] |
| 890 | x = residual + self.self_att( |
| 891 | x, |
| 892 | key, |
| 893 | val, |
| 894 | pos_enc, |
| 895 | mask, |
| 896 | left_context=left_context, |
| 897 | ) |
| 898 | |
| 899 | residual = x |
| 900 | x = self.norm_conv(x) |
| 901 | x, conv_cache = self.conv_mod(x, cache=self.cache[1], right_context=right_context) |
| 902 | x = residual + x |
| 903 | residual = x |
| 904 | |
| 905 | x = self.norm_feed_forward(x) |
| 906 | x = residual + self.feed_forward_scale * self.feed_forward(x) |
| 907 | |
| 908 | x = self.norm_final(x) |
| 909 | self.cache = [att_cache, conv_cache] |
| 910 |