text: 
  from src.modules import MultiHeadAttention, PositionwiseFeedForward, LayerNorm, _gen_bias_mask, _gen_timing_signal
  import torch
  import torch.nn as nn
  import numpy as np
  import math
  
  
  class TransformerEncoder(nn.Module):
      """
      A Transformer Encoder module. 
      Inputs should be in the shape [batch_size, length, hidden_size]
      Outputs will have the shape [batch_size, length, hidden_size]
      Refer Fig.1 in https://arxiv.org/pdf/1706.03762.pdf
      """
      def __init__(self, input_size, hidden_size, num_layers, num_heads, dim_key, dim_value,
                   filter_size, max_length=100, input_dropout=0.0, layer_dropout=0.0, 
                   attn_dropout=0.0, relu_dropout=0.0, use_mask=False):
          """
          Parameters:
              input_size: Size of inputs
              hidden_size: Hidden size
              num_layers: Total layers in the Encoder
              num_heads: Number of attention heads
              total_key_depth: Size of last dimension of keys. Must be divisible by num_head
              total_value_depth: Size of last dimension of values. Must be divisible by num_head
              output_depth: Size last dimension of the final output
              filter_size: Hidden size of the middle layer in FFN
              max_length: Max sequence length (required for timing signal)
              input_dropout: Dropout just after embedding
              layer_dropout: Dropout for each layer
              attention_dropout: Dropout probability after attention (Should be non-zero only during training)
              relu_dropout: Dropout probability after relu in FFN (Should be non-zero only during training)
              use_mask: Set to True to turn on future value masking
          """
          
          super(TransformerEncoder, self).__init__()
          
          self.timing_signal = _gen_timing_signal(max_length, hidden_size)
          
          params =(hidden_size, 
                   dim_key or hidden_size,
                   dim_value or hidden_size,
                   filter_size, 
                   num_heads, 
                   _gen_bias_mask(max_length) if use_mask else None,
                   layer_dropout,
                   attn_dropout,
                   relu_dropout)
          
          self.embedding_proj = nn.Linear(input_size, hidden_size, bias=False)
          self.enc = nn.Sequential(*[EncoderLayer(*params) for l in range(num_layers)])
          
          self.layer_norm = LayerNorm(hidden_size)
          self.input_dropout = nn.Dropout(input_dropout)
          
      def forward(self, inputs):
          #Add input dropout
          x = self.input_dropout(inputs)
          
          # Project to hidden size
          x = self.embedding_proj(x)
          
          # Add timing signal
          x += self.timing_signal[:, :inputs.shape[1], :].type_as(inputs.data)
      
          y = self.enc(x)
          y = self.layer_norm(y)
          return y
  
  
  class EncoderLayer(nn.Module):
      """
      Represents one Encoder layer of the Transformer Encoder
      Refer Fig. 1 in https://arxiv.org/pdf/1706.03762.pdf
      NOTE: The layer normalization step has been moved to the input as per latest version of T2T
      """
      def __init__(self, hidden_size, total_key_depth, total_value_depth, filter_size, num_heads,
                   bias_mask=None, layer_dropout=0.0, attention_dropout=0.0, relu_dropout=0.0):
          """
          Parameters:
              hidden_size: Hidden size
              total_key_depth: Size of last dimension of keys. Must be divisible by num_head
              total_value_depth: Size of last dimension of values. Must be divisible by num_head
              output_depth: Size last dimension of the final output
              filter_size: Hidden size of the middle layer in FFN
              num_heads: Number of attention heads
              bias_mask: Masking tensor to prevent connections to future elements
              layer_dropout: Dropout for this layer
              attention_dropout: Dropout probability after attention (Should be non-zero only during training)
              relu_dropout: Dropout probability after relu in FFN (Should be non-zero only during training)
          """
          
          super(EncoderLayer, self).__init__()
          
          self.multi_head_attention = MultiHeadAttention(hidden_size, total_key_depth, total_value_depth, hidden_size, num_heads, bias_mask, attention_dropout)
          
          self.positionwise_feed_forward = PositionwiseFeedForward(hidden_size, filter_size, hidden_size, layer_config='cc', padding = 'both', dropout=relu_dropout)
          self.dropout = nn.Dropout(layer_dropout)
          self.layer_norm_mha = LayerNorm(hidden_size)
          self.layer_norm_ffn = LayerNorm(hidden_size)
          
      def forward(self, inputs):
          x = inputs
          
          # Layer Normalization
          x_norm = self.layer_norm_mha(x)
          
          # Multi-head attention
          y = self.multi_head_attention(x_norm, x_norm, x_norm)
          
          # Dropout and residual
          x = self.dropout(x + y)
          
          # Layer Normalization
          x_norm = self.layer_norm_ffn(x)
          
          # Positionwise Feedforward
          y = self.positionwise_feed_forward(x_norm)
          
          # Dropout and residual
          y = self.dropout(x + y)
          
          return y
  