text: import math
  
  import torch
  from torch import nn
  from torch.nn import functional as F
  
  class AudioTransformer(nn.Module):
      def __init__(self, d_model, nhead, dim_feedforward, num_layers, num_classes, dropout=0.5):
          super(AudioTransformer, self).__init__()
          self.src_mask = None
          self.pos_encoder = PositionalEncoding(d_model, dropout)
          enc_layer = nn.TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout)
          self.encoder = nn.TransformerEncoder(enc_layer, num_layers, norm=None) #Todo add layer normalization
          #self.emb = nn.Embedding(ntoken, d_model)
          self.d_model = d_model
          self.classifier = nn.Linear(d_model, num_classes)
          self.init_weights()
  
      def _generate_square_subsequent_mask(self, sz):
          """
          Generate a mask matrix s.t. the tokens do not attend to future positions. It is as follows:
          [[  0  inf inf inf ],
           [  0   0  inf inf ],
           [  0   0   0  inf ],
           [  0   0   0   0  ]]
          :param sz: matrix dimension
          :return: mask matrix
          """
          mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
          mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
          return mask
  
      def init_weights(self):
          initrange = 0.1
          #self.emb.weight.data.uniform_(-initrange, initrange)
          #self.emb.bias.data.zero_()
          self.classifier.weight.data.uniform_(-initrange, initrange)
          self.classifier.bias.data.zero_()
  
      def forward(self, src, has_mask=False):
          """
          :param src: input
          :param has_mask: Put this as true if the tokens need to attend to future tokens
          """
          if has_mask:
              device = src.device
              if self.src_mask is None or self.src_mask.size(0) != len(src):
                  mask = self._generate_square_subsequent_mask(len(src)).to(device)
                  self.src_mask = mask
          else:
              self.src_mask = None
  
          src = self.encoder(src) * math.sqrt(self.d_model)
          src = self.pos_encoder(src)
          output = self.encoder(src, self.src_mask)
          output = self.classifier(output)
          # we are interested only in the first [cls] token
          output = output[:1,:,:]
          return F.log_softmax(output, dim=-1)
  
  
  class PositionalEncoding(nn.Module):
      r"""Inject some information about the relative or absolute position of the tokens
          in the sequence. The positional encodings have the same dimension as
          the embeddings, so that the two can be summed. Here, we use sine and cosine
          functions of different frequencies.
      .. math::
          \text{PosEncoder}(pos, 2i) = sin(pos/10000^(2i/d_model))
          \text{PosEncoder}(pos, 2i+1) = cos(pos/10000^(2i/d_model))
          \text{where pos is the word position and i is the embed idx)
      Args:
          d_model: the embed dim (required).
          dropout: the dropout value (default=0.1).
          max_len: the max. length of the incoming sequence (default=5000).
      Examples:
          >>> pos_encoder = PositionalEncoding(d_model)
      """
  
      def __init__(self, d_model, dropout=0.1, max_len=5000):
          super(PositionalEncoding, self).__init__()
          self.dropout = nn.Dropout(p=dropout)
  
          pe = torch.zeros(max_len, d_model)
          position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
          div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
          pe[:, 0::2] = torch.sin(position * div_term)
          pe[:, 1::2] = torch.cos(position * div_term)
          pe = pe.unsqueeze(0).transpose(0, 1)
          self.register_buffer('pe', pe)
  
      def forward(self, x):
          r"""Inputs of forward function
          Args:
              x: the sequence fed to the positional encoder model (required).
          Shape:
              x: [sequence length, batch size, embed dim]
              output: [sequence length, batch size, embed dim]
          Examples:
              >>> output = pos_encoder(x)
          """
  
          x = x + self.pe[:x.size(0), :]
          return self.dropout(x)
  
  class TransformerModel(nn.Module):
      """Container module with an encoder, a recurrent or transformer module, and a decoder."""
  
      def __init__(self, ntoken, ninp, nhead, nhid, nlayers, dropout=0.5):
          super(TransformerModel, self).__init__()
          try:
              from torch.nn import TransformerEncoder, TransformerEncoderLayer
          except:
              raise ImportError('TransformerEncoder module does not exist in PyTorch 1.1 or lower.')
          self.model_type = 'Transformer'
          self.src_mask = None
          self.pos_encoder = PositionalEncoding(ninp, dropout)
          encoder_layers = TransformerEncoderLayer(ninp, nhead, nhid, dropout)
          self.transformer_encoder = TransformerEncoder(encoder_layers, nlayers)
          self.encoder = nn.Embedding(ntoken, ninp)
          self.ninp = ninp
          self.decoder = nn.Linear(ninp, ntoken)
  
          self.init_weights()
  
      def _generate_square_subsequent_mask(self, sz):
          mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
          mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
          return mask
  
      def init_weights(self):
          initrange = 0.1
          self.encoder.weight.data.uniform_(-initrange, initrange)
          self.decoder.bias.data.zero_()
          self.decoder.weight.data.uniform_(-initrange, initrange)
  
      def forward(self, src, has_mask=True):
          if has_mask:
              device = src.device
              if self.src_mask is None or self.src_mask.size(0) != len(src):
                  mask = self._generate_square_subsequent_mask(len(src)).to(device)
                  self.src_mask = mask
          else:
              self.src_mask = None
  
          src = self.encoder(src) * math.sqrt(self.ninp)
          src = self.pos_encoder(src)
          output = self.transformer_encoder(src, self.src_mask)
          output = self.decoder(output)
          return F.log_softmax(output, dim=-1)