Spaces:
Runtime error
Runtime error
| # -------------------------------------------------------- | |
| # ArTST: Arabic Text and Speech Transformer (https://arxiv.org/abs/2310.16621) | |
| # Github source: https://github.com/mbzuai-nlp/ArTST | |
| # Based on speecht5, fairseq and espnet code bases | |
| # https://github.com/microsoft/SpeechT5/tree/main/SpeechT5; https://github.com/pytorch/fairseq; https://github.com/espnet/espnet | |
| # -------------------------------------------------------- | |
| import contextlib | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from espnet.nets.pytorch_backend.tacotron2.decoder import Prenet as TacotronDecoderPrenet | |
| from espnet.nets.pytorch_backend.transformer.embedding import PositionalEncoding | |
| from espnet.nets.pytorch_backend.transformer.embedding import ScaledPositionalEncoding | |
| from espnet.nets.pytorch_backend.nets_utils import make_non_pad_mask | |
| class SpeechDecoderPrenet(nn.Module): | |
| """ | |
| Args: | |
| in_channels (int): the number of input channels | |
| mid_channels (int): the number of intermediate channels | |
| out_channels (int): the number of output channels | |
| kernel_sizes (List[int]): the kernel size for each convolutional layer | |
| """ | |
| def __init__( | |
| self, | |
| odim, | |
| args, | |
| ): | |
| super(SpeechDecoderPrenet, self).__init__() | |
| # define decoder prenet | |
| if args.dprenet_layers != 0: | |
| # decoder prenet | |
| decoder_input_layer = torch.nn.Sequential( | |
| TacotronDecoderPrenet( | |
| idim=odim, | |
| n_layers=args.dprenet_layers, | |
| n_units=args.dprenet_units, | |
| dropout_rate=args.dprenet_dropout_rate, | |
| ), | |
| torch.nn.Linear(args.dprenet_units, args.decoder_embed_dim), | |
| ) | |
| else: | |
| decoder_input_layer = "linear" | |
| pos_enc_class = ( | |
| ScaledPositionalEncoding if args.dec_use_scaled_pos_enc else PositionalEncoding | |
| ) | |
| if decoder_input_layer == "linear": | |
| self.decoder_prenet = torch.nn.Sequential( | |
| torch.nn.Linear(odim, args.decoder_embed_dim), | |
| torch.nn.LayerNorm(args.decoder_embed_dim), | |
| torch.nn.Dropout(args.transformer_dec_dropout_rate), | |
| torch.nn.ReLU(), | |
| pos_enc_class(args.decoder_embed_dim, args.transformer_dec_positional_dropout_rate), | |
| ) | |
| elif isinstance(decoder_input_layer, torch.nn.Module): | |
| self.decoder_prenet = torch.nn.Sequential( | |
| decoder_input_layer, pos_enc_class(args.decoder_embed_dim, args.transformer_dec_positional_dropout_rate, max_len=args.max_speech_positions) | |
| ) | |
| if args.spk_embed_integration_type == 'pre': | |
| self.spkembs_layer = torch.nn.Sequential( | |
| torch.nn.Linear(args.spk_embed_dim + args.decoder_embed_dim, args.decoder_embed_dim), torch.nn.ReLU() | |
| ) | |
| self.num_updates = 0 | |
| self.freeze_decoder_updates = args.freeze_decoder_updates | |
| def forward(self, prev_output_tokens, tgt_lengths_in=None, spkembs=None): | |
| ft = self.freeze_decoder_updates <= self.num_updates | |
| with torch.no_grad() if not ft else contextlib.ExitStack(): | |
| prev_output_tokens = self.decoder_prenet(prev_output_tokens) | |
| if spkembs is not None: | |
| spkembs = F.normalize(spkembs).unsqueeze(1).expand(-1, prev_output_tokens.size(1), -1) | |
| prev_output_tokens = self.spkembs_layer(torch.cat([prev_output_tokens, spkembs], dim=-1)) | |
| if tgt_lengths_in is not None: | |
| tgt_frames_mask = ~(self._source_mask(tgt_lengths_in).squeeze(1)) | |
| else: | |
| tgt_frames_mask = None | |
| return prev_output_tokens, tgt_frames_mask | |
| def _source_mask(self, ilens): | |
| """Make masks for self-attention. | |
| Args: | |
| ilens (LongTensor or List): Batch of lengths (B,). | |
| Returns: | |
| Tensor: Mask tensor for self-attention. | |
| dtype=torch.uint8 in PyTorch 1.2- | |
| dtype=torch.bool in PyTorch 1.2+ (including 1.2) | |
| Examples: | |
| >>> ilens = [5, 3] | |
| >>> self._source_mask(ilens) | |
| tensor([[[1, 1, 1, 1, 1], | |
| [[1, 1, 1, 0, 0]]], dtype=torch.uint8) | |
| """ | |
| x_masks = make_non_pad_mask(ilens).to(next(self.parameters()).device) | |
| return x_masks.unsqueeze(-2) | |
| def set_num_updates(self, num_updates): | |
| """Set the number of parameters updates.""" | |
| self.num_updates = num_updates | |