text: """
  BERT layers from the huggingface implementation
  (https://github.com/huggingface/transformers)
  """
  # coding=utf-8
  # Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
  # Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  import logging
  import math
  import torch
  from torch import nn
  
  logger = logging.getLogger(__name__)
  
  
  def gelu(x):
      """Implementation of the gelu activation function.
          For information: OpenAI GPT's gelu is slightly different (and gives slightly different results):
          0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))
          Also see https://arxiv.org/abs/1606.08415
      """
      return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))
  
  
  class BertSelfAttention(nn.Module):
      def __init__(self, config):
          super(BertSelfAttention, self).__init__()
          if config.hidden_size % config.num_attention_heads != 0:
              raise ValueError(
                  "The hidden size (%d) is not a multiple of the number of attention "
                  "heads (%d)" % (config.hidden_size, config.num_attention_heads))
          self.num_attention_heads = config.num_attention_heads
          self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
          self.all_head_size = self.num_attention_heads * self.attention_head_size
  
          self.query = nn.Linear(config.hidden_size, self.all_head_size)
          self.key = nn.Linear(config.hidden_size, self.all_head_size)
          self.value = nn.Linear(config.hidden_size, self.all_head_size)
  
          self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
  
      def transpose_for_scores(self, x):
          new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)  # (N, L, nh, dh)
          x = x.view(*new_x_shape)
          return x.permute(0, 2, 1, 3)  # (N, nh, L, dh)
  
      def forward(self, query_states, key_states, value_states, attention_mask):
          """
          Args:
              query_states: (N, Lq, D)
              key_states: (N, L, D)
              value_states: (N, L, D)
              attention_mask: (N, Lq, L)
  
          Returns:
  
          """
  
          # only need to mask the dimension where the softmax (last dim) is applied, as another dim (second last)
          # will be ignored in future computation anyway
          attention_mask = (1 - attention_mask.unsqueeze(1)) * -10000.  # (N, 1, Lq, L)
          mixed_query_layer = self.query(query_states)
          mixed_key_layer = self.key(key_states)
          mixed_value_layer = self.value(value_states)
  
          query_layer = self.transpose_for_scores(mixed_query_layer)  # (N, nh, Lq, dh)
          key_layer = self.transpose_for_scores(mixed_key_layer)  # (N, nh, L, dh)
          value_layer = self.transpose_for_scores(mixed_value_layer)  # (N, nh, L, dh)
  
          # Take the dot product between "query" and "key" to get the raw attention scores.
          attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))  # (N, nh, Lq, L)
          attention_scores = attention_scores / math.sqrt(self.attention_head_size)
          # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
          attention_scores = attention_scores + attention_mask
  
          # Normalize the attention scores to probabilities.
          attention_probs = nn.Softmax(dim=-1)(attention_scores)
  
          # This is actually dropping out entire tokens to attend to, which might
          # seem a bit unusual, but is taken from the original Transformer paper.
          attention_probs = self.dropout(attention_probs)
  
          context_layer = torch.matmul(attention_probs, value_layer)
          context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
          new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
          context_layer = context_layer.view(*new_context_layer_shape)
          return context_layer
  
  
  class BertSelfOutput(nn.Module):
      def __init__(self, config):
          super(BertSelfOutput, self).__init__()
          self.dense = nn.Linear(config.hidden_size, config.hidden_size)
          self.LayerNorm = BertLayerNorm(config.hidden_size, eps=1e-12)
          self.dropout = nn.Dropout(config.dropout)
  
      def forward(self, hidden_states, input_tensor):
          hidden_states = self.dense(hidden_states)
          hidden_states = self.dropout(hidden_states)
          hidden_states = self.LayerNorm(hidden_states + input_tensor)
          return hidden_states
  
  
  class BertAttention(nn.Module):
      def __init__(self, config):
          super(BertAttention, self).__init__()
          self.self = BertSelfAttention(config)
          self.output = BertSelfOutput(config)
  
      def forward(self, input_tensor, attention_mask):
          self_output = self.self(input_tensor, input_tensor, input_tensor, attention_mask)
          attention_output = self.output(self_output, input_tensor)
          return attention_output
  
  
  class BertIntermediate(nn.Module):
      def __init__(self, config):
          super(BertIntermediate, self).__init__()
          self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
          self.intermediate_act_fn = gelu
  
      def forward(self, hidden_states):
          hidden_states = self.dense(hidden_states)
          hidden_states = self.intermediate_act_fn(hidden_states)
          return hidden_states
  
  
  class BertOutput(nn.Module):
      def __init__(self, config):
          super(BertOutput, self).__init__()
          self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
          self.LayerNorm = BertLayerNorm(config.hidden_size, eps=1e-12)
          self.dropout = nn.Dropout(config.dropout)
  
      def forward(self, hidden_states, input_tensor):
          hidden_states = self.dense(hidden_states)
          hidden_states = self.dropout(hidden_states)
          hidden_states = self.LayerNorm(hidden_states + input_tensor)
          return hidden_states
  
  
  class BertLayer(nn.Module):
      def __init__(self, config):
          super(BertLayer, self).__init__()
          self.attention = BertAttention(config)
          self.intermediate = BertIntermediate(config)
          self.output = BertOutput(config)
  
      def forward(self, hidden_states, attention_mask, diagonal_mask=True):
          # print("h : {}, mask : {}".format(hidden_states.shape, attention_mask.shape))
  
          self_attention_mask = attention_mask.unsqueeze(1)
  
  
          if diagonal_mask:  # mask subsequent words
              max_len = hidden_states.size(1)
              self_attention_mask = self_attention_mask * \
                                    torch.tril(self_attention_mask.new_ones(max_len, max_len), diagonal=0)
  
          attention_output = self.attention(hidden_states, self_attention_mask)
          intermediate_output = self.intermediate(attention_output)
          layer_output = self.output(intermediate_output, attention_output)
          return layer_output
  
  
  class BertPooler(nn.Module):
      def __init__(self, config):
          super(BertPooler, self).__init__()
          self.dense = nn.Linear(config.hidden_size, config.hidden_size)
          self.activation = nn.Tanh()
  
      def forward(self, hidden_states):
          # We "pool" the model by simply taking the hidden state corresponding
          # to the first token.
          first_token_tensor = hidden_states[:, 0]
          pooled_output = self.dense(first_token_tensor)
          pooled_output = self.activation(pooled_output)
          return pooled_output
  
  
  class BertPredictionHeadTransform(nn.Module):
      def __init__(self, config):
          super(BertPredictionHeadTransform, self).__init__()
          self.dense = nn.Linear(config.hidden_size, config.hidden_size)
          self.transform_act_fn = gelu
          self.LayerNorm = BertLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
  
      def forward(self, hidden_states):
          """(N, L, D)"""
          hidden_states = self.dense(hidden_states)
          hidden_states = self.transform_act_fn(hidden_states)
          hidden_states = self.LayerNorm(hidden_states)
          return hidden_states
  
  
  class BertLMPredictionHead(nn.Module):
      def __init__(self, config):
          super(BertLMPredictionHead, self).__init__()
          self.transform = BertPredictionHeadTransform(config)
  
          # The output weights are the same as the input embeddings, but there is
          # an output-only bias for each token.
          self.decoder = nn.Linear(config.hidden_size,
                                   config.vocab_size,
                                   bias=False)
          self.bias = nn.Parameter(torch.zeros(config.vocab_size))
          self.softmax = nn.Softmax(dim=1)
  
      def forward(self, hidden_states):
          hidden_states = self.transform(hidden_states)
          hidden_states = self.decoder(hidden_states) + self.bias
          hidden_states = self.softmax(hidden_states)
  
          return hidden_states
  
  
  class BertLayerNorm(nn.Module):
      def __init__(self, hidden_size, eps=1e-12):
          super(BertLayerNorm, self).__init__()
          """
          Construct a layernorm module in the TF style (epsilon inside the square root).
          """
          super(BertLayerNorm, self).__init__()
          self.weight = nn.Parameter(torch.ones(hidden_size))
          self.bias = nn.Parameter(torch.zeros(hidden_size))
          self.variance_epsilon = eps
  
      def forward(self, x):
          u = x.mean(-1, keepdim=True)
          s = (x - u).pow(2).mean(-1, keepdim=True)
          x = (x - u) / torch.sqrt(s + self.variance_epsilon)
          return self.weight * x + self.bias
  
  
  class PositionEncoding(nn.Module):
      """
      Add positional information to input tensor.
      """
  
      def __init__(self, n_filters=128, max_len=500):
          """
          :param n_filters: same with input hidden size
          :param max_len: maximum sequence length
          """
          super(PositionEncoding, self).__init__()
          # Compute the positional encodings once in log space.
          pe = torch.zeros(max_len, n_filters)  # (L, D)
          position = torch.arange(0, max_len).float().unsqueeze(1)
          div_term = torch.exp(torch.arange(0, n_filters, 2).float() * - (math.log(10000.0) / n_filters))
          pe[:, 0::2] = torch.sin(position * div_term)
          pe[:, 1::2] = torch.cos(position * div_term)
          self.register_buffer("pe", pe)  # buffer is a tensor, not a variable, (L, D)
  
      def forward(self, x):
          """
          :Input: (*, L, D)
          :Output: (*, L, D) the same size as input
          """
          pe = self.pe.data[:x.size(-2), :]  # (#x.size(-2), n_filters)
          extra_dim = len(x.size()) - 2
          for _ in range(extra_dim):
              pe = pe.unsqueeze(0)
          x = x + pe
          return x