text: import math
  import torch
  from torch import nn
  import torch.nn.functional as F
  
  ######################################################################
  # ``PositionalEncoding`` module injects 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.
  #
  
  
  class PositionalEncoding(nn.Module):
      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):
          x = x + self.pe[:x.size(0), :]
          return self.dropout(x)
  
  
  def generate_subsequent_mask(tgt_sz, src_sz):
      mask = (torch.triu(torch.ones(src_sz, tgt_sz)) == 1).transpose(0, 1)
      mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
      return mask
  
  
  def generate_square_subsequent_mask(sz):
      r"""Generate a square mask for the sequence. The masked positions are filled with float('-inf').
          Unmasked positions are filled with float(0.0).
      """
      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
  
  
  class TransformerMT(nn.Module):
      def __init__(self, src_vocab_size, tgt_vocab_size, d_model, nhead, num_encoder_layers, num_decoder_layers, dim_feedforward, max_seq_length, pos_dropout, trans_dropout):
          super().__init__()
          self.d_model = d_model
          self.embed_src = nn.Embedding(src_vocab_size, d_model)
          self.embed_tgt = nn.Embedding(tgt_vocab_size, d_model)
          self.pos_enc = PositionalEncoding(d_model, pos_dropout, max_seq_length)
  
          self.transformer = nn.Transformer(d_model, nhead, num_encoder_layers, num_decoder_layers, dim_feedforward, trans_dropout)
          self.fc = nn.Linear(d_model, tgt_vocab_size)
  
      def forward(self, src, tgt, src_key_padding_mask, tgt_key_padding_mask, memory_key_padding_mask, tgt_mask):
          # src = rearrange(src, 'n s -> s n')
          # tgt = rearrange(tgt, 'n t -> t n')
          if src_key_padding_mask is not None:
              src_key_padding_mask = src_key_padding_mask.transpose(0, 1)
          if tgt_key_padding_mask is not None:
              tgt_key_padding_mask = tgt_key_padding_mask.transpose(0, 1)
          if memory_key_padding_mask is not None:
              memory_key_padding_mask = memory_key_padding_mask.transpose(0, 1)
          src = self.pos_enc(self.embed_src(src) * math.sqrt(self.d_model))
          tgt = self.pos_enc(self.embed_tgt(tgt) * math.sqrt(self.d_model))
  
          output = self.transformer(src, tgt, tgt_mask=tgt_mask, src_key_padding_mask=src_key_padding_mask,
                                    tgt_key_padding_mask=tgt_key_padding_mask, memory_key_padding_mask=memory_key_padding_mask)
          # output = rearrange(output, 't n e -> n t e')
          # output = output.transpose(0, 1)
          output = self.fc(output)
          return F.softmax(output, dim=2)
  