Spaces:
Sleeping
Sleeping
| # This module is from [WeNet](https://github.com/wenet-e2e/wenet). | |
| # ## Citations | |
| # ```bibtex | |
| # @inproceedings{yao2021wenet, | |
| # title={WeNet: Production oriented Streaming and Non-streaming End-to-End Speech Recognition Toolkit}, | |
| # author={Yao, Zhuoyuan and Wu, Di and Wang, Xiong and Zhang, Binbin and Yu, Fan and Yang, Chao and Peng, Zhendong and Chen, Xiaoyu and Xie, Lei and Lei, Xin}, | |
| # booktitle={Proc. Interspeech}, | |
| # year={2021}, | |
| # address={Brno, Czech Republic }, | |
| # organization={IEEE} | |
| # } | |
| # @article{zhang2022wenet, | |
| # title={WeNet 2.0: More Productive End-to-End Speech Recognition Toolkit}, | |
| # author={Zhang, Binbin and Wu, Di and Peng, Zhendong and Song, Xingchen and Yao, Zhuoyuan and Lv, Hang and Xie, Lei and Yang, Chao and Pan, Fuping and Niu, Jianwei}, | |
| # journal={arXiv preprint arXiv:2203.15455}, | |
| # year={2022} | |
| # } | |
| # | |
| """Decoder definition.""" | |
| from typing import Tuple, List, Optional | |
| import torch | |
| from modules.wenet_extractor.transformer.attention import MultiHeadedAttention | |
| from modules.wenet_extractor.transformer.decoder_layer import DecoderLayer | |
| from modules.wenet_extractor.transformer.embedding import PositionalEncoding | |
| from modules.wenet_extractor.transformer.embedding import NoPositionalEncoding | |
| from modules.wenet_extractor.transformer.positionwise_feed_forward import ( | |
| PositionwiseFeedForward, | |
| ) | |
| from modules.wenet_extractor.utils.mask import subsequent_mask, make_pad_mask | |
| class TransformerDecoder(torch.nn.Module): | |
| """Base class of Transfomer decoder module. | |
| Args: | |
| vocab_size: output dim | |
| encoder_output_size: dimension of attention | |
| attention_heads: the number of heads of multi head attention | |
| linear_units: the hidden units number of position-wise feedforward | |
| num_blocks: the number of decoder blocks | |
| dropout_rate: dropout rate | |
| self_attention_dropout_rate: dropout rate for attention | |
| input_layer: input layer type | |
| use_output_layer: whether to use output layer | |
| pos_enc_class: PositionalEncoding or ScaledPositionalEncoding | |
| normalize_before: | |
| True: use layer_norm before each sub-block of a layer. | |
| False: use layer_norm after each sub-block of a layer. | |
| src_attention: if false, encoder-decoder cross attention is not | |
| applied, such as CIF model | |
| """ | |
| def __init__( | |
| self, | |
| vocab_size: int, | |
| encoder_output_size: int, | |
| attention_heads: int = 4, | |
| linear_units: int = 2048, | |
| num_blocks: int = 6, | |
| dropout_rate: float = 0.1, | |
| positional_dropout_rate: float = 0.1, | |
| self_attention_dropout_rate: float = 0.0, | |
| src_attention_dropout_rate: float = 0.0, | |
| input_layer: str = "embed", | |
| use_output_layer: bool = True, | |
| normalize_before: bool = True, | |
| src_attention: bool = True, | |
| ): | |
| super().__init__() | |
| attention_dim = encoder_output_size | |
| if input_layer == "embed": | |
| self.embed = torch.nn.Sequential( | |
| torch.nn.Embedding(vocab_size, attention_dim), | |
| PositionalEncoding(attention_dim, positional_dropout_rate), | |
| ) | |
| elif input_layer == "none": | |
| self.embed = NoPositionalEncoding(attention_dim, positional_dropout_rate) | |
| else: | |
| raise ValueError(f"only 'embed' is supported: {input_layer}") | |
| self.normalize_before = normalize_before | |
| self.after_norm = torch.nn.LayerNorm(attention_dim, eps=1e-5) | |
| self.use_output_layer = use_output_layer | |
| self.output_layer = torch.nn.Linear(attention_dim, vocab_size) | |
| self.num_blocks = num_blocks | |
| self.decoders = torch.nn.ModuleList( | |
| [ | |
| DecoderLayer( | |
| attention_dim, | |
| MultiHeadedAttention( | |
| attention_heads, attention_dim, self_attention_dropout_rate | |
| ), | |
| MultiHeadedAttention( | |
| attention_heads, attention_dim, src_attention_dropout_rate | |
| ) | |
| if src_attention | |
| else None, | |
| PositionwiseFeedForward(attention_dim, linear_units, dropout_rate), | |
| dropout_rate, | |
| normalize_before, | |
| ) | |
| for _ in range(self.num_blocks) | |
| ] | |
| ) | |
| def forward( | |
| self, | |
| memory: torch.Tensor, | |
| memory_mask: torch.Tensor, | |
| ys_in_pad: torch.Tensor, | |
| ys_in_lens: torch.Tensor, | |
| r_ys_in_pad: torch.Tensor = torch.empty(0), | |
| reverse_weight: float = 0.0, | |
| ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: | |
| """Forward decoder. | |
| Args: | |
| memory: encoded memory, float32 (batch, maxlen_in, feat) | |
| memory_mask: encoder memory mask, (batch, 1, maxlen_in) | |
| ys_in_pad: padded input token ids, int64 (batch, maxlen_out) | |
| ys_in_lens: input lengths of this batch (batch) | |
| r_ys_in_pad: not used in transformer decoder, in order to unify api | |
| with bidirectional decoder | |
| reverse_weight: not used in transformer decoder, in order to unify | |
| api with bidirectional decode | |
| Returns: | |
| (tuple): tuple containing: | |
| x: decoded token score before softmax (batch, maxlen_out, | |
| vocab_size) if use_output_layer is True, | |
| torch.tensor(0.0), in order to unify api with bidirectional decoder | |
| olens: (batch, ) | |
| """ | |
| tgt = ys_in_pad | |
| maxlen = tgt.size(1) | |
| # tgt_mask: (B, 1, L) | |
| tgt_mask = ~make_pad_mask(ys_in_lens, maxlen).unsqueeze(1) | |
| tgt_mask = tgt_mask.to(tgt.device) | |
| # m: (1, L, L) | |
| m = subsequent_mask(tgt_mask.size(-1), device=tgt_mask.device).unsqueeze(0) | |
| # tgt_mask: (B, L, L) | |
| tgt_mask = tgt_mask & m | |
| x, _ = self.embed(tgt) | |
| for layer in self.decoders: | |
| x, tgt_mask, memory, memory_mask = layer(x, tgt_mask, memory, memory_mask) | |
| if self.normalize_before: | |
| x = self.after_norm(x) | |
| if self.use_output_layer: | |
| x = self.output_layer(x) | |
| olens = tgt_mask.sum(1) | |
| return x, torch.tensor(0.0), olens | |
| def forward_one_step( | |
| self, | |
| memory: torch.Tensor, | |
| memory_mask: torch.Tensor, | |
| tgt: torch.Tensor, | |
| tgt_mask: torch.Tensor, | |
| cache: Optional[List[torch.Tensor]] = None, | |
| ) -> Tuple[torch.Tensor, List[torch.Tensor]]: | |
| """Forward one step. | |
| This is only used for decoding. | |
| Args: | |
| memory: encoded memory, float32 (batch, maxlen_in, feat) | |
| memory_mask: encoded memory mask, (batch, 1, maxlen_in) | |
| tgt: input token ids, int64 (batch, maxlen_out) | |
| tgt_mask: input token mask, (batch, maxlen_out) | |
| dtype=torch.uint8 in PyTorch 1.2- | |
| dtype=torch.bool in PyTorch 1.2+ (include 1.2) | |
| cache: cached output list of (batch, max_time_out-1, size) | |
| Returns: | |
| y, cache: NN output value and cache per `self.decoders`. | |
| y.shape` is (batch, maxlen_out, token) | |
| """ | |
| x, _ = self.embed(tgt) | |
| new_cache = [] | |
| for i, decoder in enumerate(self.decoders): | |
| if cache is None: | |
| c = None | |
| else: | |
| c = cache[i] | |
| x, tgt_mask, memory, memory_mask = decoder( | |
| x, tgt_mask, memory, memory_mask, cache=c | |
| ) | |
| new_cache.append(x) | |
| if self.normalize_before: | |
| y = self.after_norm(x[:, -1]) | |
| else: | |
| y = x[:, -1] | |
| if self.use_output_layer: | |
| y = torch.log_softmax(self.output_layer(y), dim=-1) | |
| return y, new_cache | |
| class BiTransformerDecoder(torch.nn.Module): | |
| """Base class of Transfomer decoder module. | |
| Args: | |
| vocab_size: output dim | |
| encoder_output_size: dimension of attention | |
| attention_heads: the number of heads of multi head attention | |
| linear_units: the hidden units number of position-wise feedforward | |
| num_blocks: the number of decoder blocks | |
| r_num_blocks: the number of right to left decoder blocks | |
| dropout_rate: dropout rate | |
| self_attention_dropout_rate: dropout rate for attention | |
| input_layer: input layer type | |
| use_output_layer: whether to use output layer | |
| pos_enc_class: PositionalEncoding or ScaledPositionalEncoding | |
| normalize_before: | |
| True: use layer_norm before each sub-block of a layer. | |
| False: use layer_norm after each sub-block of a layer. | |
| """ | |
| def __init__( | |
| self, | |
| vocab_size: int, | |
| encoder_output_size: int, | |
| attention_heads: int = 4, | |
| linear_units: int = 2048, | |
| num_blocks: int = 6, | |
| r_num_blocks: int = 0, | |
| dropout_rate: float = 0.1, | |
| positional_dropout_rate: float = 0.1, | |
| self_attention_dropout_rate: float = 0.0, | |
| src_attention_dropout_rate: float = 0.0, | |
| input_layer: str = "embed", | |
| use_output_layer: bool = True, | |
| normalize_before: bool = True, | |
| ): | |
| super().__init__() | |
| self.left_decoder = TransformerDecoder( | |
| vocab_size, | |
| encoder_output_size, | |
| attention_heads, | |
| linear_units, | |
| num_blocks, | |
| dropout_rate, | |
| positional_dropout_rate, | |
| self_attention_dropout_rate, | |
| src_attention_dropout_rate, | |
| input_layer, | |
| use_output_layer, | |
| normalize_before, | |
| ) | |
| self.right_decoder = TransformerDecoder( | |
| vocab_size, | |
| encoder_output_size, | |
| attention_heads, | |
| linear_units, | |
| r_num_blocks, | |
| dropout_rate, | |
| positional_dropout_rate, | |
| self_attention_dropout_rate, | |
| src_attention_dropout_rate, | |
| input_layer, | |
| use_output_layer, | |
| normalize_before, | |
| ) | |
| def forward( | |
| self, | |
| memory: torch.Tensor, | |
| memory_mask: torch.Tensor, | |
| ys_in_pad: torch.Tensor, | |
| ys_in_lens: torch.Tensor, | |
| r_ys_in_pad: torch.Tensor, | |
| reverse_weight: float = 0.0, | |
| ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: | |
| """Forward decoder. | |
| Args: | |
| memory: encoded memory, float32 (batch, maxlen_in, feat) | |
| memory_mask: encoder memory mask, (batch, 1, maxlen_in) | |
| ys_in_pad: padded input token ids, int64 (batch, maxlen_out) | |
| ys_in_lens: input lengths of this batch (batch) | |
| r_ys_in_pad: padded input token ids, int64 (batch, maxlen_out), | |
| used for right to left decoder | |
| reverse_weight: used for right to left decoder | |
| Returns: | |
| (tuple): tuple containing: | |
| x: decoded token score before softmax (batch, maxlen_out, | |
| vocab_size) if use_output_layer is True, | |
| r_x: x: decoded token score (right to left decoder) | |
| before softmax (batch, maxlen_out, vocab_size) | |
| if use_output_layer is True, | |
| olens: (batch, ) | |
| """ | |
| l_x, _, olens = self.left_decoder(memory, memory_mask, ys_in_pad, ys_in_lens) | |
| r_x = torch.tensor(0.0) | |
| if reverse_weight > 0.0: | |
| r_x, _, olens = self.right_decoder( | |
| memory, memory_mask, r_ys_in_pad, ys_in_lens | |
| ) | |
| return l_x, r_x, olens | |
| def forward_one_step( | |
| self, | |
| memory: torch.Tensor, | |
| memory_mask: torch.Tensor, | |
| tgt: torch.Tensor, | |
| tgt_mask: torch.Tensor, | |
| cache: Optional[List[torch.Tensor]] = None, | |
| ) -> Tuple[torch.Tensor, List[torch.Tensor]]: | |
| """Forward one step. | |
| This is only used for decoding. | |
| Args: | |
| memory: encoded memory, float32 (batch, maxlen_in, feat) | |
| memory_mask: encoded memory mask, (batch, 1, maxlen_in) | |
| tgt: input token ids, int64 (batch, maxlen_out) | |
| tgt_mask: input token mask, (batch, maxlen_out) | |
| dtype=torch.uint8 in PyTorch 1.2- | |
| dtype=torch.bool in PyTorch 1.2+ (include 1.2) | |
| cache: cached output list of (batch, max_time_out-1, size) | |
| Returns: | |
| y, cache: NN output value and cache per `self.decoders`. | |
| y.shape` is (batch, maxlen_out, token) | |
| """ | |
| return self.left_decoder.forward_one_step( | |
| memory, memory_mask, tgt, tgt_mask, cache | |
| ) | |