Spaces:
Runtime error
Runtime error
| import math | |
| import torch | |
| import torch.nn as nn | |
| from tencentpretrain.utils.constants import * | |
| class SinusoidalposEmbedding(nn.Module): | |
| """Sinusoidal positional encoding for non-recurrent neural networks. | |
| Implementation based on "Attention Is All You Need" | |
| :cite:`DBLP:journals/corr/VaswaniSPUJGKP17` | |
| Args: | |
| dropout (float): dropout parameter | |
| dim (int): embedding size | |
| """ | |
| def __init__(self, args, _): | |
| super(SinusoidalposEmbedding, self).__init__() | |
| if "speech" in args.embedding: | |
| self.max_seq_length = max(args.max_seq_length, args.max_audio_frames) | |
| self.arrange_sincos_cross = False | |
| else: | |
| self.max_seq_length = args.max_seq_length | |
| self.arrange_sincos_cross = True | |
| self.emb_size = args.emb_size | |
| half_dim = self.emb_size // 2 | |
| value = math.log(10000) / (half_dim - 1) | |
| half_exp = torch.exp(torch.arange(half_dim, dtype=torch.float) * -value) | |
| half_mat = torch.arange(self.max_seq_length, dtype=torch.float).unsqueeze( | |
| 1 | |
| ) * half_exp.unsqueeze(0) | |
| if not self.arrange_sincos_cross: #Same as the implementation of huggingface/transformers, tensor2tensor | |
| self.emb = torch.cat([torch.sin(half_mat), torch.cos(half_mat)], dim=1).view( | |
| self.max_seq_length, -1 | |
| ) | |
| else: #Implementation based on "Attention Is All You Need" | |
| self.emb = torch.zeros(self.max_seq_length, args.emb_size) | |
| self.emb[:, 0::2] = torch.sin(half_mat) | |
| self.emb[:, 1::2] = torch.cos(half_mat) | |
| if self.emb_size % 2 == 1: | |
| # zero pad | |
| self.emb = torch.cat([self.emb, torch.zeros(self.max_seq_length, 1)], dim=1) | |
| self.emb[args.tokenizer.vocab.get(PAD_TOKEN), :] = 0 | |
| def forward(self, src, seg): | |
| """Embed inputs. | |
| Args: | |
| emb (FloatTensor): Sequence of word vectors | |
| ``(batch_size, seq_len, self.dim)`` | |
| step (int or NoneType): If stepwise (``seq_len = 1``), use | |
| the encoding for this position. | |
| """ | |
| if seg is not None: | |
| batch_size, seq_length = seg.size() | |
| device = seg.device | |
| no_pad_num = seg.sum(dim=-1) | |
| else: | |
| batch_size, seq_length = src.size() | |
| device = src.device | |
| no_pad_num = (src != 0).sum(dim=-1) | |
| emb = torch.zeros(batch_size, seq_length, self.emb_size) | |
| for i in range(batch_size): | |
| emb[i, :no_pad_num[i], :] = self.emb[2: no_pad_num[i]+2] | |
| return emb.to(device) | |