# 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.
"""PyTorch BERT model."""

from __future__ import absolute_import, division, print_function, unicode_literals

import copy
import json
import logging
import math
import os
import sys
from io import open

import torch
from torch import nn
from torch.nn import CrossEntropyLoss

from relogic.utils.file_utils import cached_path
from relogic.logickit.inference.modeling_utils import prune_linear_layer

logger = logging.getLogger(__name__)

PRETRAINED_MODEL_ARCHIVE_MAP = {
  'bert-base-uncased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-uncased-pytorch_model.bin",
  'bert-large-uncased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-pytorch_model.bin",
  'bert-base-cased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-cased-pytorch_model.bin",
  'bert-large-cased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-pytorch_model.bin",
  'bert-base-multilingual-uncased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-multilingual-uncased-pytorch_model.bin",
  'bert-base-multilingual-cased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-multilingual-cased-pytorch_model.bin",
  'bert-base-chinese': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-chinese-pytorch_model.bin",
  'bert-base-german-cased': "https://int-deepset-models-bert.s3.eu-central-1.amazonaws.com/pytorch/bert-base-german-cased-pytorch_model.bin",
  'bert-large-uncased-whole-word-masking': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-whole-word-masking-pytorch_model.bin",
  'bert-large-cased-whole-word-masking': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-whole-word-masking-pytorch_model.bin",
  'bert-large-uncased-whole-word-masking-finetuned-squad': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-whole-word-masking-finetuned-squad-pytorch_model.bin",
  'bert-large-cased-whole-word-masking-finetuned-squad': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-whole-word-masking-finetuned-squad-pytorch_model.bin",
  'bert-base-cased-finetuned-mrpc': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-cased-finetuned-mrpc-pytorch_model.bin",
}
PRETRAINED_CONFIG_ARCHIVE_MAP = {
  'bert-base-uncased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-uncased-config.json",
  'bert-large-uncased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-config.json",
  'bert-base-cased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-cased-config.json",
  'bert-large-cased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-config.json",
  'bert-base-multilingual-uncased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-multilingual-uncased-config.json",
  'bert-base-multilingual-cased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-multilingual-cased-config.json",
  'bert-base-chinese': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-chinese-config.json",
  'bert-base-german-cased': "https://int-deepset-models-bert.s3.eu-central-1.amazonaws.com/pytorch/bert-base-german-cased-config.json",
  'bert-large-uncased-whole-word-masking': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-whole-word-masking-config.json",
  'bert-large-cased-whole-word-masking': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-whole-word-masking-config.json",
  'bert-large-uncased-whole-word-masking-finetuned-squad': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-whole-word-masking-finetuned-squad-config.json",
  'bert-large-cased-whole-word-masking-finetuned-squad': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-whole-word-masking-finetuned-squad-config.json",
  'bert-base-cased-finetuned-mrpc': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-cased-finetuned-mrpc-config.json",
}
CONFIG_NAME = 'bert_config.json'
WEIGHTS_NAME = 'pytorch_model.bin'
TF_WEIGHTS_NAME = 'model.ckpt'


def load_tf_weights_in_bert(model, tf_checkpoint_path):
  """ Load tf checkpoints in a pytorch model
  """
  try:
    import re
    import numpy as np
    import tensorflow as tf
  except ImportError:
    print("Loading a TensorFlow models in PyTorch, requires TensorFlow to be installed. Please see "
          "https://www.tensorflow.org/install/ for installation instructions.")
    raise
  tf_path = os.path.abspath(tf_checkpoint_path)
  print("Converting TensorFlow checkpoint from {}".format(tf_path))
  # Load weights from TF model
  init_vars = tf.train.list_variables(tf_path)
  names = []
  arrays = []
  for name, shape in init_vars:
    print("Loading TF weight {} with shape {}".format(name, shape))
    array = tf.train.load_variable(tf_path, name)
    names.append(name)
    arrays.append(array)

  for name, array in zip(names, arrays):
    name = name.split('/')
    # adam_v and adam_m are variables used in AdamWeightDecayOptimizer to calculated m and v
    # which are not required for using pretrained model
    if any(n in ["adam_v", "adam_m", "global_step"] for n in name):
      print("Skipping {}".format("/".join(name)))
      continue
    pointer = model
    for m_name in name:
      if re.fullmatch(r'[A-Za-z]+_\d+', m_name):
        l = re.split(r'_(\d+)', m_name)
      else:
        l = [m_name]
      if l[0] == 'kernel' or l[0] == 'gamma':
        pointer = getattr(pointer, 'weight')
      elif l[0] == 'output_bias' or l[0] == 'beta':
        pointer = getattr(pointer, 'bias')
      elif l[0] == 'output_weights':
        pointer = getattr(pointer, 'weight')
      elif l[0] == 'squad':
        pointer = getattr(pointer, 'classifier')
      else:
        try:
          pointer = getattr(pointer, l[0])
        except AttributeError:
          print("Skipping {}".format("/".join(name)))
          continue
      if len(l) >= 2:
        num = int(l[1])
        pointer = pointer[num]
    if m_name[-11:] == '_embeddings':
      pointer = getattr(pointer, 'weight')
    elif m_name == 'kernel':
      array = np.transpose(array)
    try:
      assert pointer.shape == array.shape
    except AssertionError as e:
      e.args += (pointer.shape, array.shape)
      raise
    print("Initialize PyTorch weight {}".format(name))
    pointer.data = torch.from_numpy(array)
  return model


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)))


def swish(x):
  return x * torch.sigmoid(x)


ACT2FN = {"gelu": gelu, "relu": torch.nn.functional.relu, "swish": swish}


class BertConfig(object):
  """Configuration class to store the configuration of a `BertModel`.
  """

  def __init__(self,
               vocab_size_or_config_json_file,
               hidden_size=768,
               num_hidden_layers=12,
               num_attention_heads=12,
               intermediate_size=3072,
               hidden_act="gelu",
               hidden_dropout_prob=0.1,
               attention_probs_dropout_prob=0.1,
               max_position_embeddings=512,
               type_vocab_size=2,
               initializer_range=0.02,
               layer_norm_eps=1e-12):
    """Constructs BertConfig.

    Args:
        vocab_size_or_config_json_file: Vocabulary size of `inputs_ids` in `BertModel`.
        hidden_size: Size of the encoder layers and the pooler layer.
        num_hidden_layers: Number of hidden layers in the Transformer encoder.
        num_attention_heads: Number of attention heads for each attention layer in
            the Transformer encoder.
        intermediate_size: The size of the "intermediate" (i.e., feed-forward)
            layer in the Transformer encoder.
        hidden_act: The non-linear activation function (function or string) in the
            encoder and pooler. If string, "gelu", "relu" and "swish" are supported.
        hidden_dropout_prob: The dropout probabilitiy for all fully connected
            layers in the embeddings, encoder, and pooler.
        attention_probs_dropout_prob: The dropout ratio for the attention
            probabilities.
        max_position_embeddings: The maximum sequence length that this model might
            ever be used with. Typically set this to something large just in case
            (e.g., 512 or 1024 or 2048).
        type_vocab_size: The vocabulary size of the `token_type_ids` passed into
            `BertModel`.
        initializer_range: The sttdev of the truncated_normal_initializer for
            initializing all weight matrices.
        layer_norm_eps: The epsilon used by LayerNorm.
    """
    if isinstance(vocab_size_or_config_json_file, str) or (sys.version_info[0] == 2
                    and isinstance(vocab_size_or_config_json_file, unicode)):
      with open(vocab_size_or_config_json_file, "r", encoding='utf-8') as reader:
        json_config = json.loads(reader.read())
      for key, value in json_config.items():
        self.__dict__[key] = value
    elif isinstance(vocab_size_or_config_json_file, int):
      self.vocab_size = vocab_size_or_config_json_file
      self.hidden_size = hidden_size
      self.num_hidden_layers = num_hidden_layers
      self.num_attention_heads = num_attention_heads
      self.hidden_act = hidden_act
      self.intermediate_size = intermediate_size
      self.hidden_dropout_prob = hidden_dropout_prob
      self.attention_probs_dropout_prob = attention_probs_dropout_prob
      self.max_position_embeddings = max_position_embeddings
      self.type_vocab_size = type_vocab_size
      self.initializer_range = initializer_range
      self.layer_norm_eps = layer_norm_eps
    else:
      raise ValueError("First argument must be either a vocabulary size (int)"
                       "or the path to a pretrained model config file (str)")

  @classmethod
  def from_dict(cls, json_object):
    """Constructs a `BertConfig` from a Python dictionary of parameters."""
    config = BertConfig(vocab_size_or_config_json_file=-1)
    for key, value in json_object.items():
      config.__dict__[key] = value
    return config

  @classmethod
  def from_json_file(cls, json_file):
    """Constructs a `BertConfig` from a json file of parameters."""
    with open(json_file, "r", encoding='utf-8') as reader:
      text = reader.read()
    return cls.from_dict(json.loads(text))

  def __repr__(self):
    return str(self.to_json_string())

  def to_dict(self):
    """Serializes this instance to a Python dictionary."""
    output = copy.deepcopy(self.__dict__)
    return output

  def to_json_string(self):
    """Serializes this instance to a JSON string."""
    return json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n"

  def to_json_file(self, json_file_path):
    """ Save this instance to a json file."""
    with open(json_file_path, "w", encoding='utf-8') as writer:
      writer.write(self.to_json_string())


try:
  from apex.normalization.fused_layer_norm import FusedLayerNorm as BertLayerNorm
except ImportError:
  logger.info("Better speed can be achieved with apex installed from https://www.github.com/nvidia/apex.")
  class BertLayerNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-12):
      """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 BertEmbeddings(nn.Module):
  """Construct the embeddings from word, position and token_type embeddings.
  """

  def __init__(self, config):
    super(BertEmbeddings, self).__init__()
    self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=0)
    self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
    self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)

    # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
    # any TensorFlow checkpoint file
    self.LayerNorm = BertLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
    self.dropout = nn.Dropout(config.hidden_dropout_prob)

  def forward(self, input_ids, token_type_ids=None):
    seq_length = input_ids.size(1)
    position_ids = torch.arange(seq_length, dtype=torch.long, device=input_ids.device)
    position_ids = position_ids.unsqueeze(0).expand_as(input_ids)
    if token_type_ids is None:
      token_type_ids = torch.zeros_like(input_ids)

    words_embeddings = self.word_embeddings(input_ids)
    position_embeddings = self.position_embeddings(position_ids)
    token_type_embeddings = self.token_type_embeddings(token_type_ids)

    embeddings = words_embeddings + position_embeddings + token_type_embeddings
    embeddings = self.LayerNorm(embeddings)
    embeddings = self.dropout(embeddings)
    return embeddings


class BertSelfAttention(nn.Module):
  def __init__(self, config, output_attentions=False, keep_multihead_output=False):
    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.output_attentions = output_attentions
    self.keep_multihead_output = keep_multihead_output
    self.multihead_output = None

    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)

    # self.token_level_attention_mask_embedding = nn.Embedding(2, self.all_head_size)

  def transpose_for_scores(self, x):
    new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
    x = x.view(*new_x_shape)
    return x.permute(0, 2, 1, 3)

  def forward(self, hidden_states, attention_mask, head_mask=None, token_level_attention_mask=None):
    """
    :param token_level_attention_mask: Token level attention mask is to implement dependency tree injection
    :return:
    """
    mixed_query_layer = self.query(hidden_states)
    mixed_key_layer = self.key(hidden_states)
    mixed_value_layer = self.value(hidden_states)

    query_layer = self.transpose_for_scores(mixed_query_layer)
    # batch_size, num_heads,sent_length, head_size
    key_layer = self.transpose_for_scores(mixed_key_layer)
    value_layer = self.transpose_for_scores(mixed_value_layer)

    if token_level_attention_mask is not None:
      # query  * key    = (W_q * h) * (W_k * h + W_embeding * m_k[i, j])
      # token_level_attention_mask = token_level_attention_mask.unsqueeze(1).repeat(1, self.num_attention_heads, 1, 1)
      # batch_size, sent_length, sent_length
      token_level_attention_mask_embed = self.token_level_attention_mask_embedding(token_level_attention_mask)
      # batch_size, sent_length, sent_length, dim
      new_shape = token_level_attention_mask_embed.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
      token_level_attention_mask_embed = token_level_attention_mask_embed.view(*new_shape).permute(0, 3, 1, 2, 4)
      # batch_size, num_heads, sent_length, sent_length, dim


    # 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))

    if token_level_attention_mask is not None:
      # A = Q' * K + Q' * M
      # Q (batch_size, num_heads, sent_length, head_size)
      # M (batch_size, num_heads, sent_length, sent_length, head_size)
      token_level_attention_matrix = torch.matmul(
        query_layer.unsqueeze(3), token_level_attention_mask_embed.transpose(-1, -2)).squeeze(3)
      # A = [length, length]
      # Q = [q_1, q_2, q_3]
      # K = [k_1, k_2, k_3]
      # a_11 = q_1' * (k_1 + embed[1, 1])
      # a_12 = q_1' * (k_2 + embed[1, 2])
      # A_1 = q_1 * embed[1,:]
      # A = Q' * (K + embed) = Q' K + Q' * embed
      # Q (batch_size, num_heads, sent_length, 1, head_size)
      # M (batch_size, num_heads, sent_length, head_size, sent_length)
      # A (batch_size, num_heads, sent_length, sent_length)
      attention_scores = attention_scores + token_level_attention_matrix

    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)

    # if token_level_attention_mask is not None:
    #   print("Using token level attention")
    #   # attention_probs (batch_size, head_num, length, length)
    #   # token_level_attention_mask (batch_size, length, length)
    #   token_level_attention_mask = token_level_attention_mask.to(dtype=next(self.parameters()).dtype)
    #   # expand to (batch, head_num, length, length)
    #   token_level_attention_mask = token_level_attention_mask.unsqueeze(1).repeat(1, attention_probs.size(1) ,1, 1)
    #   attention_probs = nn.Softmax(dim=-1)(attention_probs + token_level_attention_mask)

    # 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)

    # Mask heads if we want to
    if head_mask is not None:
      attention_probs = attention_probs * head_mask

    context_layer = torch.matmul(attention_probs, value_layer)
    # (batch, n, seq_len, head_dim)
    if self.keep_multihead_output:
      self.multihead_output = context_layer
      self.multihead_output.retain_grad()
      # TODO: Why do we need to keep gradients in internal layer?

    # Notice here multihead_output is same as context_layer, which is the same as
    # the self_output.

    context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
    # (batch, seq_len, n, head_dim)
    new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
    context_layer = context_layer.view(*new_context_layer_shape)
    # (batch, seq_len, n * head_dim)
    if self.output_attentions:
      return attention_probs, context_layer
    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=config.layer_norm_eps)
    self.dropout = nn.Dropout(config.hidden_dropout_prob)

  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, output_attentions=False, keep_multihead_output=False):
    super(BertAttention, self).__init__()
    self.output_attentions = output_attentions

    self.self = BertSelfAttention(config, output_attentions=output_attentions,
                                          keep_multihead_output=keep_multihead_output)
    self.output = BertSelfOutput(config)

  def prune_heads(self, heads):
    if len(heads) == 0:
      return
    mask = torch.ones(self.self.num_attention_heads, self.self.attention_head_size)
    for head in heads:
      mask[head] = 0
    mask = mask.view(-1).contiguous().eq(1)
    index = torch.arange(len(mask))[mask].long()
    # Prune linear layers
    self.self.query = prune_linear_layer(self.self.query, index)
    self.self.key = prune_linear_layer(self.self.key, index)
    self.self.value = prune_linear_layer(self.self.value, index)
    self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
    # Update hyper params
    self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
    self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads

  def forward(self, input_tensor, attention_mask, head_mask=None, token_level_attention_mask=None):
    self_output = self.self(input_tensor, attention_mask, head_mask, token_level_attention_mask=token_level_attention_mask)
    # (batch, seq_len, n * head_dim)
    if self.output_attentions:
      attentions, self_output = self_output
    attention_output = self.output(self_output, input_tensor)

    if self.output_attentions:
      return attentions, attention_output
    return attention_output
    # This part is different from current implementation of huggingface
    # https://github.com/huggingface/pytorch-pretrained-BERT/blob/master/pytorch_pretrained_bert/modeling.py#L434
    # They only return attention_output, here I also return self_output.

    # More on this. I decided not to return self_output because self_output == multihead_output
    # And this has been saved in the class `BertSelfAttention`.
    # This can be aggregate later.


class BertIntermediate(nn.Module):
  def __init__(self, config):
    super(BertIntermediate, self).__init__()
    self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
    if isinstance(config.hidden_act, str) or (sys.version_info[0] == 2 and isinstance(config.hidden_act, unicode)):
      self.intermediate_act_fn = ACT2FN[config.hidden_act]
    else:
      self.intermediate_act_fn = config.hidden_act

  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=config.layer_norm_eps)
    self.dropout = nn.Dropout(config.hidden_dropout_prob)

  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):
  """
  Return layer output and intermediate output of multi_head_repr
  """
  def __init__(self, config, output_attentions=False, keep_multihead_output=False):
    super(BertLayer, self).__init__()
    self.output_attentions = output_attentions
    self.attention = BertAttention(config, output_attentions=output_attentions,
                                           keep_multihead_output=keep_multihead_output)
    self.intermediate = BertIntermediate(config)
    self.output = BertOutput(config)

  def forward(self, hidden_states, attention_mask, head_mask=None, token_level_attention_mask=None):

    attention_output = self.attention(hidden_states, attention_mask, head_mask,
                                    token_level_attention_mask=token_level_attention_mask)
    if self.output_attentions:
      attentions, attention_output = attention_output
    intermediate_output = self.intermediate(attention_output)
    layer_output = self.output(intermediate_output, attention_output)

    if self.output_attentions:
      return attentions, layer_output
    return layer_output


class BertEncoder(nn.Module):
  def __init__(self, config, output_attentions=False, keep_multihead_output=False):
    super(BertEncoder, self).__init__()
    self.output_attentions = output_attentions
    layer = BertLayer(config, output_attentions=output_attentions,
                              keep_multihead_output=keep_multihead_output)
    self.layer = nn.ModuleList([copy.deepcopy(layer) for _ in range(config.num_hidden_layers)])

  def forward(self,
              hidden_states,
              attention_mask,
              token_level_attention_mask=None,
              output_all_encoded_layers=True,
              head_mask=None,
              output_final_multi_head_repr=False,
              selected_non_final_layers=None):
    """
    `token_level_attention_mask`
    `output_all_encoder_layers` is not suggested to be used because it take up large amount of
    unnecessary gpu memory.
    `selected_non_final_layers` is designed to solve above problem.

    `head_mask` and `token_level_attention_mask` must be a list with `layer num` items
    """
    if selected_non_final_layers is None:
      selected_non_final_layers = []
    else:
      assert isinstance(selected_non_final_layers, list), "selected_non_final_layers should be a list"

    all_encoder_layers = []
    all_attentions = []

    for idx, layer_module in enumerate(self.layer):
      hidden_states = layer_module(hidden_states, attention_mask,
                                                    head_mask[idx], token_level_attention_mask[idx])
      if self.output_attentions:
        attentions, hidden_states = hidden_states
        all_attentions.append(attentions)
      if output_all_encoded_layers:
        all_encoder_layers.append(hidden_states)
      if idx in selected_non_final_layers:
        all_encoder_layers.append(hidden_states)
    if not output_all_encoded_layers:
      all_encoder_layers.append(hidden_states)
    if self.output_attentions:
      return all_attentions, all_encoder_layers

    return all_encoder_layers



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)
    if isinstance(config.hidden_act, str) or (sys.version_info[0] == 2 and isinstance(config.hidden_act, unicode)):
      self.transform_act_fn = ACT2FN[config.hidden_act]
    else:
      self.transform_act_fn = config.hidden_act
    self.LayerNorm = BertLayerNorm(config.hidden_size, eps=config.layer_norm_eps)

  def forward(self, hidden_states):
    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, bert_model_embedding_weights):
    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(bert_model_embedding_weights.size(1),
                             bert_model_embedding_weights.size(0),
                             bias=False)
    self.decoder.weight = bert_model_embedding_weights
    self.bias = nn.Parameter(torch.zeros(bert_model_embedding_weights.size(0)))

  def forward(self, hidden_states):
    hidden_states = self.transform(hidden_states)
    hidden_states = self.decoder(hidden_states) + self.bias
    return hidden_states


class BertOnlyMLMHead(nn.Module):
  def __init__(self, config, bert_model_embedding_weights):
    super(BertOnlyMLMHead, self).__init__()
    self.predictions = BertLMPredictionHead(config, bert_model_embedding_weights)

  def forward(self, sequence_output):
    prediction_scores = self.predictions(sequence_output)
    return prediction_scores


class BertOnlyNSPHead(nn.Module):
  def __init__(self, config):
    super(BertOnlyNSPHead, self).__init__()
    self.seq_relationship = nn.Linear(config.hidden_size, 2)

  def forward(self, pooled_output):
    seq_relationship_score = self.seq_relationship(pooled_output)
    return seq_relationship_score


class BertPreTrainingHeads(nn.Module):
  def __init__(self, config, bert_model_embedding_weights):
    super(BertPreTrainingHeads, self).__init__()
    self.predictions = BertLMPredictionHead(config, bert_model_embedding_weights)
    self.seq_relationship = nn.Linear(config.hidden_size, 2)

  def forward(self, sequence_output, pooled_output):
    prediction_scores = self.predictions(sequence_output)
    seq_relationship_score = self.seq_relationship(pooled_output)
    return prediction_scores, seq_relationship_score


class BertPreTrainedModel(nn.Module):
  """ An abstract class to handle weights initialization and
      a simple interface for dowloading and loading pretrained models.
  """

  def __init__(self, config, *inputs, **kwargs):
    super(BertPreTrainedModel, self).__init__()
    if not isinstance(config, BertConfig):
      raise ValueError(
        "Parameter config in `{}(config)` should be an instance of class `BertConfig`. "
        "To create a model from a Google pretrained model use "
        "`model = {}.from_pretrained(PRETRAINED_MODEL_NAME)`".format(
          self.__class__.__name__, self.__class__.__name__
        ))
    self.config = config

  def init_bert_weights(self, module):
    """ Initialize the weights.
    """
    if isinstance(module, (nn.Linear, nn.Embedding)):
      # Slightly different from the TF version which uses truncated_normal for initialization
      # cf https://github.com/pytorch/pytorch/pull/5617
      module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
    elif isinstance(module, BertLayerNorm):
      module.bias.data.zero_()
      module.weight.data.fill_(1.0)
    if isinstance(module, nn.Linear) and module.bias is not None:
      module.bias.data.zero_()

  @classmethod
  def from_pretrained(cls, pretrained_model_name_or_path, *inputs, **kwargs):
    """
    Instantiate a BertPreTrainedModel from a pre-trained model file or a pytorch state dict.
    Download and cache the pre-trained model file if needed.
    Params:
      pretrained_model_name_or_path: either:
        - a str with the name of a pre-trained model to load selected in the list of:
          . `bert-base-uncased`
          . `bert-large-uncased`
          . `bert-base-cased`
          . `bert-large-cased`
          . `bert-base-multilingual-uncased`
          . `bert-base-multilingual-cased`
          . `bert-base-chinese`
          . `bert-base-german-cased`
          . `bert-large-uncased-whole-word-masking`
          . `bert-large-cased-whole-word-masking`
        - a path or url to a pretrained model archive containing:
          . `bert_config.json` a configuration file for the model
          . `pytorch_model.bin` a PyTorch dump of a BertForPreTraining instance
        - a path or url to a pretrained model archive containing:
          . `bert_config.json` a configuration file for the model
          . `model.chkpt` a TensorFlow checkpoint
    from_tf: should we load the weights from a locally saved TensorFlow checkpoint
    cache_dir: an optional path to a folder in which the pre-trained models will be cached.
    state_dict: an optional state dictionnary (collections.OrderedDict object) to use instead of Google pre-trained models
    *inputs, **kwargs: additional input for the specific Bert class
        (ex: num_labels for BertForSequenceClassification)
    """
    state_dict = kwargs.get('state_dict', None)
    kwargs.pop('state_dict', None)
    cache_dir = kwargs.get('cache_dir', None)
    kwargs.pop('cache_dir', None)
    from_tf = kwargs.get('from_tf', False)
    kwargs.pop('from_tf', None)

    if pretrained_model_name_or_path in PRETRAINED_MODEL_ARCHIVE_MAP:
      archive_file = PRETRAINED_MODEL_ARCHIVE_MAP[pretrained_model_name_or_path]
      config_file = PRETRAINED_CONFIG_ARCHIVE_MAP[pretrained_model_name_or_path]
    else:
      archive_file = os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)
      config_file = os.path.join(pretrained_model_name_or_path, CONFIG_NAME)

    # redirect to the cache, if necessary
    try:
      resolved_archive_file = cached_path(archive_file, cache_dir=cache_dir)
    except EnvironmentError:
      if pretrained_model_name_or_path in PRETRAINED_MODEL_ARCHIVE_MAP:
        logger.error(
          "Couldn't reach server at '{}' to download pretrained weights.".format(
            archive_file))
      else:
        logger.error(
          "Model name '{}' was not found in model name list ({}). "
          "We assumed '{}' was a path or url but couldn't find any file "
          "associated to this path or url.".format(
            pretrained_model_name_or_path,
            ', '.join(PRETRAINED_MODEL_ARCHIVE_MAP.keys()),
            archive_file))
      return None

    try:
      resolved_config_file = cached_path(config_file, cache_dir=cache_dir)
    except EnvironmentError:
      if pretrained_model_name_or_path in PRETRAINED_CONFIG_ARCHIVE_MAP:
        logger.error(
          "Couldn't reach server at '{}' to download pretrained model configuration file.".format(
            config_file))
      else:
        logger.error(
          "Model name '{}' was not found in model name list ({}). "
          "We assumed '{}' was a path or url but couldn't find any file "
          "associated to this path or url.".format(
            pretrained_model_name_or_path,
            ', '.join(PRETRAINED_CONFIG_ARCHIVE_MAP.keys()),
            config_file))
      return None

    if resolved_archive_file == archive_file and resolved_config_file == config_file:
      logger.info("loading weights file {}".format(archive_file))
      logger.info("loading configuration file {}".format(config_file))
    else:
      logger.info("loading weights file {} from cache at {}".format(
        archive_file, resolved_archive_file))
      logger.info("loading configuration file {} from cache at {}".format(
        config_file, resolved_config_file))

    # Load config
    config = BertConfig.from_json_file(resolved_config_file)
    logger.info("Model config {}".format(config))
    # Instantiate model.
    model = cls(config, *inputs, **kwargs)
    if state_dict is None and not from_tf:
      state_dict = torch.load(resolved_archive_file, map_location='cpu')
    if from_tf:
      # Directly load from a TensorFlow checkpoint
      return load_tf_weights_in_bert(model, os.path.join(cache_dir, WEIGHTS_NAME))
      # TODO: THIS IS THE ERROR FROM SOURCE CODE. JUST A QUICK FIX WITHOUT VERIFICATION
    # Load from a PyTorch state_dict
    old_keys = []
    new_keys = []
    for key in state_dict.keys():
      new_key = None
      if 'gamma' in key:
        new_key = key.replace('gamma', 'weight')
      if 'beta' in key:
        new_key = key.replace('beta', 'bias')
      if new_key:
        old_keys.append(key)
        new_keys.append(new_key)
    for old_key, new_key in zip(old_keys, new_keys):
      state_dict[new_key] = state_dict.pop(old_key)

    missing_keys = []
    unexpected_keys = []
    error_msgs = []
    # copy state_dict so _load_from_state_dict can modify it
    metadata = getattr(state_dict, '_metadata', None)
    state_dict = state_dict.copy()
    if metadata is not None:
      state_dict._metadata = metadata

    def load(module, prefix=''):
      local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
      module._load_from_state_dict(
        state_dict, prefix, local_metadata, True, missing_keys, unexpected_keys, error_msgs)
      for name, child in module._modules.items():
        if child is not None:
          load(child, prefix + name + '.')

    start_prefix = ''
    if not hasattr(model, 'bert') and any(s.startswith('bert.') for s in state_dict.keys()):
      start_prefix = 'bert.'
    load(model, prefix=start_prefix)
    if len(missing_keys) > 0:
      logger.info("Weights of {} not initialized from pretrained model: {}".format(
        model.__class__.__name__, missing_keys))
    if len(unexpected_keys) > 0:
      logger.info("Weights from pretrained model not used in {}: {}".format(
        model.__class__.__name__, unexpected_keys))
    if len(error_msgs) > 0:
      raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
        model.__class__.__name__, "\n\t".join(error_msgs)))
    return model


class BertModel(BertPreTrainedModel):
  """BERT model ("Bidirectional Embedding Representations from a Transformer").
  Params:
      `config`: a BertConfig class instance with the configuration to build a new model
      `output_attentions`: If True, also output attentions weights computed by the model at each layer. Default: False
      `keep_multihead_output`: If True, saves output of the multi-head attention module with its gradient.
          This can be used to compute head importance metrics. Default: False
  Inputs:
      `input_ids`: a torch.LongTensor of shape [batch_size, sequence_length]
          with the word token indices in the vocabulary(see the tokens preprocessing logic in the scripts
          `extract_features.py`, `run_classifier.py` and `run_squad.py`)
      `token_type_ids`: an optional torch.LongTensor of shape [batch_size, sequence_length] with the token
          types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to
          a `sentence B` token (see BERT paper for more details).
      `attention_mask`: an optional torch.LongTensor of shape [batch_size, sequence_length] with indices
          selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max
          input sequence length in the current batch. It's the mask that we typically use for attention when
          a batch has varying length sentences.
      `output_all_encoded_layers`: boolean which controls the content of the `encoded_layers` output as described below. Default: `True`.
      `head_mask`: an optional torch.Tensor of shape [num_heads] or [num_layers, num_heads] with indices between 0 and 1.
          It's a mask to be used to nullify some heads of the transformer. 1.0 => head is fully masked, 0.0 => head is not masked.
  Outputs: Tuple of (encoded_layers, pooled_output)
      `encoded_layers`: controled by `output_all_encoded_layers` argument:
          - `output_all_encoded_layers=True`: outputs a list of the full sequences of encoded-hidden-states at the end
              of each attention block (i.e. 12 full sequences for BERT-base, 24 for BERT-large), each
              encoded-hidden-state is a torch.FloatTensor of size [batch_size, sequence_length, hidden_size],
          - `output_all_encoded_layers=False`: outputs only the full sequence of hidden-states corresponding
              to the last attention block of shape [batch_size, sequence_length, hidden_size],
      `pooled_output`: a torch.FloatTensor of size [batch_size, hidden_size] which is the output of a
          classifier pretrained on top of the hidden state associated to the first character of the
          input (`CLS`) to train on the Next-Sentence task (see BERT's paper).
  Example usage:
  ```python
  # Already been converted into WordPiece token ids
  input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
  input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])
  token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]])
  config = modeling.BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
      num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)
  model = modeling.BertModel(config=config)
  all_encoder_layers, pooled_output = model(input_ids, token_type_ids, input_mask)
  ```
  """

  def __init__(self, config, output_attentions=False, keep_multihead_output=False):
    super(BertModel, self).__init__(config)
    self.output_attentions = output_attentions

    self.embeddings = BertEmbeddings(config)
    self.encoder = BertEncoder(config, output_attentions=output_attentions,
                                       keep_multihead_output=keep_multihead_output)
    self.pooler = BertPooler(config)
    self.apply(self.init_bert_weights)

  def prune_heads(self, heads_to_prune):
    """ Prunes heads of the model.
        heads_to_prune: dict of {layer_num: list of heads to prune in this layer}
    """
    for layer, heads in heads_to_prune.items():
      self.encoder.layer[layer].attention.prune_heads(heads)

  def get_multihead_outputs(self):
    """ Gather all multi-head outputs.
        Return: list (layers) of multihead module outputs with gradients
    """
    return [layer.attention.self.multihead_output for layer in self.encoder.layer]

  def forward(self,
              input_ids,
              token_type_ids=None,
              attention_mask=None,
              head_mask=None,
              token_level_attention_mask=None,
              output_all_encoded_layers=True,
              selected_non_final_layers=None):
    if attention_mask is None:
      attention_mask = torch.ones_like(input_ids)
    if token_type_ids is None:
      token_type_ids = torch.zeros_like(input_ids)

    # We create a 3D attention mask from a 2D tensor mask.
    # Sizes are [batch_size, 1, 1, to_seq_length]
    # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
    # this attention mask is more simple than the triangular masking of causal attention
    # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
    extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)

    # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
    # masked positions, this operation will create a tensor which is 0.0 for
    # positions we want to attend and -10000.0 for masked positions.
    # Since we are adding it to the raw scores before the softmax, this is
    # effectively the same as removing these entirely.
    extended_attention_mask = extended_attention_mask.to(dtype=next(self.parameters()).dtype)  # fp16 compatibility
    extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0

    # Prepare head mask if needed
    # 1.0 in head_mask indicate we keep the head
    # attention_probs has shape bsz x n_heads x N x N
    # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
    # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
    if head_mask is not None:
      if head_mask.dim() == 1:
        head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
        head_mask = head_mask.expand_as(self.config.num_hidden_layers, -1, -1, -1, -1)
      elif head_mask.dim() == 2:
        head_mask = head_mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1)  # We can specify head_mask for each layer
      head_mask = head_mask.to(dtype=next(self.parameters()).dtype)  # switch to fload if need + fp16 compatibility
    else:
      head_mask = [None] * self.config.num_hidden_layers

    if token_level_attention_mask is None:
      token_level_attention_mask = [None] * self.config.num_hidden_layers


    embedding_output = self.embeddings(input_ids, token_type_ids)
    encoded_layers = self.encoder(embedding_output,
                                  extended_attention_mask,
                                  output_all_encoded_layers=output_all_encoded_layers,
                                  head_mask=head_mask,
                                  selected_non_final_layers=selected_non_final_layers,
                                  token_level_attention_mask=token_level_attention_mask)
    if self.output_attentions:
      all_attentions, encoded_layers = encoded_layers

    sequence_output = encoded_layers[-1]
    pooled_output = self.pooler(sequence_output)
    if not output_all_encoded_layers and selected_non_final_layers is None:
      encoded_layers = encoded_layers[-1]
    if self.output_attentions:
      return all_attentions, encoded_layers, pooled_output

    return encoded_layers, pooled_output
      # (batch, seq_len, n * head_dim)

class BertForPreTraining(BertPreTrainedModel):
  """BERT model with pre-training heads.
  This module comprises the BERT model followed by the two pre-training heads:
      - the masked language modeling head, and
      - the next sentence classification head.
  Params:
      `config`: a BertConfig class instance with the configuration to build a new model
      `output_attentions`: If True, also output attentions weights computed by the model at each layer. Default: False
      `keep_multihead_output`: If True, saves output of the multi-head attention module with its gradient.
          This can be used to compute head importance metrics. Default: False
  Inputs:
      `input_ids`: a torch.LongTensor of shape [batch_size, sequence_length]
          with the word token indices in the vocabulary(see the tokens preprocessing logic in the scripts
          `extract_features.py`, `run_classifier.py` and `run_squad.py`)
      `token_type_ids`: an optional torch.LongTensor of shape [batch_size, sequence_length] with the token
          types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to
          a `sentence B` token (see BERT paper for more details).
      `attention_mask`: an optional torch.LongTensor of shape [batch_size, sequence_length] with indices
          selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max
          input sequence length in the current batch. It's the mask that we typically use for attention when
          a batch has varying length sentences.
      `masked_lm_labels`: optional masked language modeling labels: torch.LongTensor of shape [batch_size, sequence_length]
          with indices selected in [-1, 0, ..., vocab_size]. All labels set to -1 are ignored (masked), the loss
          is only computed for the labels set in [0, ..., vocab_size]
      `next_sentence_label`: optional next sentence classification loss: torch.LongTensor of shape [batch_size]
          with indices selected in [0, 1].
          0 => next sentence is the continuation, 1 => next sentence is a random sentence.
      `head_mask`: an optional torch.Tensor of shape [num_heads] or [num_layers, num_heads] with indices between 0 and 1.
          It's a mask to be used to nullify some heads of the transformer. 1.0 => head is fully masked, 0.0 => head is not masked.
  Outputs:
      if `masked_lm_labels` and `next_sentence_label` are not `None`:
          Outputs the total_loss which is the sum of the masked language modeling loss and the next
          sentence classification loss.
      if `masked_lm_labels` or `next_sentence_label` is `None`:
          Outputs a tuple comprising
          - the masked language modeling logits of shape [batch_size, sequence_length, vocab_size], and
          - the next sentence classification logits of shape [batch_size, 2].
  Example usage:
  ```python
  # Already been converted into WordPiece token ids
  input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
  input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])
  token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]])
  config = BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
      num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)
  model = BertForPreTraining(config)
  masked_lm_logits_scores, seq_relationship_logits = model(input_ids, token_type_ids, input_mask)
  ```
  """

  def __init__(self, config, output_attentions=False, keep_multihead_output=False):
    super(BertForPreTraining, self).__init__(config)
    self.output_attentions = output_attentions
    self.bert = BertModel(config, output_attentions=output_attentions,
                          keep_multihead_output=keep_multihead_output)
    self.cls = BertPreTrainingHeads(config, self.bert.embeddings.word_embeddings.weight)
    self.apply(self.init_bert_weights)

  def forward(self, input_ids, token_type_ids=None, attention_mask=None, masked_lm_labels=None,
              next_sentence_label=None, head_mask=None):
    outputs = self.bert(input_ids, token_type_ids, attention_mask,
                        output_all_encoded_layers=False, head_mask=head_mask)
    if self.output_attentions:
      all_attentions, sequence_output, pooled_output = outputs
    else:
      sequence_output, pooled_output = outputs
    prediction_scores, seq_relationship_score = self.cls(sequence_output, pooled_output)

    if masked_lm_labels is not None and next_sentence_label is not None:
      loss_fct = CrossEntropyLoss(ignore_index=-1)
      masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), masked_lm_labels.view(-1))
      next_sentence_loss = loss_fct(seq_relationship_score.view(-1, 2), next_sentence_label.view(-1))
      total_loss = masked_lm_loss + next_sentence_loss
      return total_loss
    elif self.output_attentions:
      return all_attentions, prediction_scores, seq_relationship_score
    return prediction_scores, seq_relationship_score


class BertForMaskedLM(BertPreTrainedModel):
  """BERT model with the masked language modeling head.
  This module comprises the BERT model followed by the masked language modeling head.
  Params:
      `config`: a BertConfig class instance with the configuration to build a new model
      `output_attentions`: If True, also output attentions weights computed by the model at each layer. Default: False
      `keep_multihead_output`: If True, saves output of the multi-head attention module with its gradient.
          This can be used to compute head importance metrics. Default: False
  Inputs:
      `input_ids`: a torch.LongTensor of shape [batch_size, sequence_length]
          with the word token indices in the vocabulary(see the tokens preprocessing logic in the scripts
          `extract_features.py`, `run_classifier.py` and `run_squad.py`)
      `token_type_ids`: an optional torch.LongTensor of shape [batch_size, sequence_length] with the token
          types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to
          a `sentence B` token (see BERT paper for more details).
      `attention_mask`: an optional torch.LongTensor of shape [batch_size, sequence_length] with indices
          selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max
          input sequence length in the current batch. It's the mask that we typically use for attention when
          a batch has varying length sentences.
      `masked_lm_labels`: masked language modeling labels: torch.LongTensor of shape [batch_size, sequence_length]
          with indices selected in [-1, 0, ..., vocab_size]. All labels set to -1 are ignored (masked), the loss
          is only computed for the labels set in [0, ..., vocab_size]
      `head_mask`: an optional torch.Tensor of shape [num_heads] or [num_layers, num_heads] with indices between 0 and 1.
          It's a mask to be used to nullify some heads of the transformer. 1.0 => head is fully masked, 0.0 => head is not masked.
  Outputs:
      if `masked_lm_labels` is  not `None`:
          Outputs the masked language modeling loss.
      if `masked_lm_labels` is `None`:
          Outputs the masked language modeling logits of shape [batch_size, sequence_length, vocab_size].
  Example usage:
  ```python
  # Already been converted into WordPiece token ids
  input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
  input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])
  token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]])
  config = BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
      num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)
  model = BertForMaskedLM(config)
  masked_lm_logits_scores = model(input_ids, token_type_ids, input_mask)
  ```
  """

  def __init__(self, config, output_attentions=False, keep_multihead_output=False):
    super(BertForMaskedLM, self).__init__(config)
    self.output_attentions = output_attentions
    self.bert = BertModel(config, output_attentions=output_attentions,
                          keep_multihead_output=keep_multihead_output)
    self.cls = BertOnlyMLMHead(config, self.bert.embeddings.word_embeddings.weight)
    self.apply(self.init_bert_weights)

  def forward(self, input_ids, token_type_ids=None, attention_mask=None, masked_lm_labels=None, head_mask=None):
    outputs = self.bert(input_ids, token_type_ids, attention_mask,
                        output_all_encoded_layers=False,
                        head_mask=head_mask)
    if self.output_attentions:
      all_attentions, sequence_output, _ = outputs
    else:
      sequence_output, _ = outputs
    prediction_scores = self.cls(sequence_output)

    if masked_lm_labels is not None:
      loss_fct = CrossEntropyLoss(ignore_index=-1)
      masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), masked_lm_labels.view(-1))
      return masked_lm_loss
    elif self.output_attentions:
      return all_attentions, prediction_scores
    return prediction_scores


class BertForNextSentencePrediction(BertPreTrainedModel):
  """BERT model with next sentence prediction head.
  This module comprises the BERT model followed by the next sentence classification head.
  Params:
      `config`: a BertConfig class instance with the configuration to build a new model
      `output_attentions`: If True, also output attentions weights computed by the model at each layer. Default: False
      `keep_multihead_output`: If True, saves output of the multi-head attention module with its gradient.
          This can be used to compute head importance metrics. Default: False
  Inputs:
      `input_ids`: a torch.LongTensor of shape [batch_size, sequence_length]
          with the word token indices in the vocabulary(see the tokens preprocessing logic in the scripts
          `extract_features.py`, `run_classifier.py` and `run_squad.py`)
      `token_type_ids`: an optional torch.LongTensor of shape [batch_size, sequence_length] with the token
          types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to
          a `sentence B` token (see BERT paper for more details).
      `attention_mask`: an optional torch.LongTensor of shape [batch_size, sequence_length] with indices
          selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max
          input sequence length in the current batch. It's the mask that we typically use for attention when
          a batch has varying length sentences.
      `next_sentence_label`: next sentence classification loss: torch.LongTensor of shape [batch_size]
          with indices selected in [0, 1].
          0 => next sentence is the continuation, 1 => next sentence is a random sentence.
      `head_mask`: an optional torch.Tensor of shape [num_heads] or [num_layers, num_heads] with indices between 0 and 1.
          It's a mask to be used to nullify some heads of the transformer. 1.0 => head is fully masked, 0.0 => head is not masked.
  Outputs:
      if `next_sentence_label` is not `None`:
          Outputs the total_loss which is the sum of the masked language modeling loss and the next
          sentence classification loss.
      if `next_sentence_label` is `None`:
          Outputs the next sentence classification logits of shape [batch_size, 2].
  Example usage:
  ```python
  # Already been converted into WordPiece token ids
  input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
  input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])
  token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]])
  config = BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
      num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)
  model = BertForNextSentencePrediction(config)
  seq_relationship_logits = model(input_ids, token_type_ids, input_mask)
  ```
  """

  def __init__(self, config, output_attentions=False, keep_multihead_output=False):
    super(BertForNextSentencePrediction, self).__init__(config)
    self.output_attentions = output_attentions
    self.bert = BertModel(config, output_attentions=output_attentions,
                          keep_multihead_output=keep_multihead_output)
    self.cls = BertOnlyNSPHead(config)
    self.apply(self.init_bert_weights)

  def forward(self, input_ids, token_type_ids=None, attention_mask=None, next_sentence_label=None, head_mask=None):
    outputs = self.bert(input_ids, token_type_ids, attention_mask,
                        output_all_encoded_layers=False,
                        head_mask=head_mask)
    if self.output_attentions:
      all_attentions, _, pooled_output = outputs
    else:
      _, pooled_output = outputs
    seq_relationship_score = self.cls(pooled_output)

    if next_sentence_label is not None:
      loss_fct = CrossEntropyLoss(ignore_index=-1)
      next_sentence_loss = loss_fct(seq_relationship_score.view(-1, 2), next_sentence_label.view(-1))
      return next_sentence_loss
    elif self.output_attentions:
      return all_attentions, seq_relationship_score
    return seq_relationship_score


class BertForSequenceClassification(BertPreTrainedModel):
  """BERT model for classification.
  This module is composed of the BERT model with a linear layer on top of
  the pooled output.
  Params:
      `config`: a BertConfig class instance with the configuration to build a new model
      `output_attentions`: If True, also output attentions weights computed by the model at each layer. Default: False
      `keep_multihead_output`: If True, saves output of the multi-head attention module with its gradient.
          This can be used to compute head importance metrics. Default: False
      `num_labels`: the number of classes for the classifier. Default = 2.
  Inputs:
      `input_ids`: a torch.LongTensor of shape [batch_size, sequence_length]
          with the word token indices in the vocabulary. Items in the batch should begin with the special "CLS" token. (see the tokens preprocessing logic in the scripts
          `extract_features.py`, `run_classifier.py` and `run_squad.py`)
      `token_type_ids`: an optional torch.LongTensor of shape [batch_size, sequence_length] with the token
          types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to
          a `sentence B` token (see BERT paper for more details).
      `attention_mask`: an optional torch.LongTensor of shape [batch_size, sequence_length] with indices
          selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max
          input sequence length in the current batch. It's the mask that we typically use for attention when
          a batch has varying length sentences.
      `labels`: labels for the classification output: torch.LongTensor of shape [batch_size]
          with indices selected in [0, ..., num_labels].
      `head_mask`: an optional torch.Tensor of shape [num_heads] or [num_layers, num_heads] with indices between 0 and 1.
          It's a mask to be used to nullify some heads of the transformer. 1.0 => head is fully masked, 0.0 => head is not masked.
  Outputs:
      if `labels` is not `None`:
          Outputs the CrossEntropy classification loss of the output with the labels.
      if `labels` is `None`:
          Outputs the classification logits of shape [batch_size, num_labels].
  Example usage:
  ```python
  # Already been converted into WordPiece token ids
  input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
  input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])
  token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]])
  config = BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
      num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)
  num_labels = 2
  model = BertForSequenceClassification(config, num_labels)
  logits = model(input_ids, token_type_ids, input_mask)
  ```
  """

  def __init__(self, config, num_labels=2, output_attentions=False, keep_multihead_output=False):
    super(BertForSequenceClassification, self).__init__(config)
    self.output_attentions = output_attentions
    self.num_labels = num_labels
    self.bert = BertModel(config, output_attentions=output_attentions,
                          keep_multihead_output=keep_multihead_output)
    self.dropout = nn.Dropout(config.hidden_dropout_prob)
    self.classifier = nn.Linear(config.hidden_size, num_labels)
    self.apply(self.init_bert_weights)

  def forward(self, input_ids, token_type_ids=None, attention_mask=None, labels=None, head_mask=None):
    outputs = self.bert(input_ids, token_type_ids, attention_mask, output_all_encoded_layers=False, head_mask=head_mask)
    if self.output_attentions:
      all_attentions, _, pooled_output = outputs
    else:
      _, pooled_output = outputs
    pooled_output = self.dropout(pooled_output)
    logits = self.classifier(pooled_output)

    if labels is not None:
      loss_fct = CrossEntropyLoss()
      loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
      return loss
    elif self.output_attentions:
      return all_attentions, logits
    return logits


class BertForMultipleChoice(BertPreTrainedModel):
  """BERT model for multiple choice tasks.
  This module is composed of the BERT model with a linear layer on top of
  the pooled output.
  Params:
      `config`: a BertConfig class instance with the configuration to build a new model
      `output_attentions`: If True, also output attentions weights computed by the model at each layer. Default: False
      `keep_multihead_output`: If True, saves output of the multi-head attention module with its gradient.
          This can be used to compute head importance metrics. Default: False
      `num_choices`: the number of classes for the classifier. Default = 2.
  Inputs:
      `input_ids`: a torch.LongTensor of shape [batch_size, num_choices, sequence_length]
          with the word token indices in the vocabulary(see the tokens preprocessing logic in the scripts
          `extract_features.py`, `run_classifier.py` and `run_squad.py`)
      `token_type_ids`: an optional torch.LongTensor of shape [batch_size, num_choices, sequence_length]
          with the token types indices selected in [0, 1]. Type 0 corresponds to a `sentence A`
          and type 1 corresponds to a `sentence B` token (see BERT paper for more details).
      `attention_mask`: an optional torch.LongTensor of shape [batch_size, num_choices, sequence_length] with indices
          selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max
          input sequence length in the current batch. It's the mask that we typically use for attention when
          a batch has varying length sentences.
      `labels`: labels for the classification output: torch.LongTensor of shape [batch_size]
          with indices selected in [0, ..., num_choices].
      `head_mask`: an optional torch.Tensor of shape [num_heads] or [num_layers, num_heads] with indices between 0 and 1.
          It's a mask to be used to nullify some heads of the transformer. 1.0 => head is fully masked, 0.0 => head is not masked.
  Outputs:
      if `labels` is not `None`:
          Outputs the CrossEntropy classification loss of the output with the labels.
      if `labels` is `None`:
          Outputs the classification logits of shape [batch_size, num_labels].
  Example usage:
  ```python
  # Already been converted into WordPiece token ids
  input_ids = torch.LongTensor([[[31, 51, 99], [15, 5, 0]], [[12, 16, 42], [14, 28, 57]]])
  input_mask = torch.LongTensor([[[1, 1, 1], [1, 1, 0]],[[1,1,0], [1, 0, 0]]])
  token_type_ids = torch.LongTensor([[[0, 0, 1], [0, 1, 0]],[[0, 1, 1], [0, 0, 1]]])
  config = BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
      num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)
  num_choices = 2
  model = BertForMultipleChoice(config, num_choices)
  logits = model(input_ids, token_type_ids, input_mask)
  ```
  """

  def __init__(self, config, num_choices=2, output_attentions=False, keep_multihead_output=False):
    super(BertForMultipleChoice, self).__init__(config)
    self.output_attentions = output_attentions
    self.num_choices = num_choices
    self.bert = BertModel(config, output_attentions=output_attentions,
                          keep_multihead_output=keep_multihead_output)
    self.dropout = nn.Dropout(config.hidden_dropout_prob)
    self.classifier = nn.Linear(config.hidden_size, 1)
    self.apply(self.init_bert_weights)

  def forward(self, input_ids, token_type_ids=None, attention_mask=None, labels=None, head_mask=None):
    flat_input_ids = input_ids.view(-1, input_ids.size(-1))
    flat_token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None
    flat_attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None
    outputs = self.bert(flat_input_ids, flat_token_type_ids, flat_attention_mask, output_all_encoded_layers=False,
                        head_mask=head_mask)
    if self.output_attentions:
      all_attentions, _, pooled_output = outputs
    else:
      _, pooled_output = outputs
    pooled_output = self.dropout(pooled_output)
    logits = self.classifier(pooled_output)
    reshaped_logits = logits.view(-1, self.num_choices)

    if labels is not None:
      loss_fct = CrossEntropyLoss()
      loss = loss_fct(reshaped_logits, labels)
      return loss
    elif self.output_attentions:
      return all_attentions, reshaped_logits
    return reshaped_logits


class BertForTokenClassification(BertPreTrainedModel):
  """BERT model for token-level classification.
  This module is composed of the BERT model with a linear layer on top of
  the full hidden state of the last layer.
  Params:
      `config`: a BertConfig class instance with the configuration to build a new model
      `output_attentions`: If True, also output attentions weights computed by the model at each layer. Default: False
      `keep_multihead_output`: If True, saves output of the multi-head attention module with its gradient.
          This can be used to compute head importance metrics. Default: False
      `num_labels`: the number of classes for the classifier. Default = 2.
  Inputs:
      `input_ids`: a torch.LongTensor of shape [batch_size, sequence_length]
          with the word token indices in the vocabulary(see the tokens preprocessing logic in the scripts
          `extract_features.py`, `run_classifier.py` and `run_squad.py`)
      `token_type_ids`: an optional torch.LongTensor of shape [batch_size, sequence_length] with the token
          types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to
          a `sentence B` token (see BERT paper for more details).
      `attention_mask`: an optional torch.LongTensor of shape [batch_size, sequence_length] with indices
          selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max
          input sequence length in the current batch. It's the mask that we typically use for attention when
          a batch has varying length sentences.
      `labels`: labels for the classification output: torch.LongTensor of shape [batch_size, sequence_length]
          with indices selected in [0, ..., num_labels].
      `head_mask`: an optional torch.Tensor of shape [num_heads] or [num_layers, num_heads] with indices between 0 and 1.
          It's a mask to be used to nullify some heads of the transformer. 1.0 => head is fully masked, 0.0 => head is not masked.
  Outputs:
      if `labels` is not `None`:
          Outputs the CrossEntropy classification loss of the output with the labels.
      if `labels` is `None`:
          Outputs the classification logits of shape [batch_size, sequence_length, num_labels].
  Example usage:
  ```python
  # Already been converted into WordPiece token ids
  input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
  input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])
  token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]])
  config = BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
      num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)
  num_labels = 2
  model = BertForTokenClassification(config, num_labels)
  logits = model(input_ids, token_type_ids, input_mask)
  ```
  """

  def __init__(self, config, num_labels=2, output_attentions=False, keep_multihead_output=False):
    super(BertForTokenClassification, self).__init__(config)
    self.output_attentions = output_attentions
    self.num_labels = num_labels
    self.bert = BertModel(config, output_attentions=output_attentions,
                          keep_multihead_output=keep_multihead_output)
    self.dropout = nn.Dropout(config.hidden_dropout_prob)
    self.classifier = nn.Linear(config.hidden_size, num_labels)
    self.apply(self.init_bert_weights)

  def forward(self, input_ids, token_type_ids=None, attention_mask=None, labels=None, head_mask=None):
    outputs = self.bert(input_ids, token_type_ids, attention_mask, output_all_encoded_layers=False, head_mask=head_mask)
    if self.output_attentions:
      all_attentions, sequence_output, _ = outputs
    else:
      sequence_output, _ = outputs
    sequence_output = self.dropout(sequence_output)
    logits = self.classifier(sequence_output)

    if labels is not None:
      loss_fct = CrossEntropyLoss()
      # Only keep active parts of the loss
      if attention_mask is not None:
        active_loss = attention_mask.view(-1) == 1
        active_logits = logits.view(-1, self.num_labels)[active_loss]
        active_labels = labels.view(-1)[active_loss]
        loss = loss_fct(active_logits, active_labels)
      else:
        loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
      return loss
    elif self.output_attentions:
      return all_attentions, logits
    return logits


class BertForQuestionAnswering(BertPreTrainedModel):
  """BERT model for Question Answering (span extraction).
  This module is composed of the BERT model with a linear layer on top of
  the sequence output that computes start_logits and end_logits
  Params:
      `config`: a BertConfig class instance with the configuration to build a new model
      `output_attentions`: If True, also output attentions weights computed by the model at each layer. Default: False
      `keep_multihead_output`: If True, saves output of the multi-head attention module with its gradient.
          This can be used to compute head importance metrics. Default: False
  Inputs:
      `input_ids`: a torch.LongTensor of shape [batch_size, sequence_length]
          with the word token indices in the vocabulary(see the tokens preprocessing logic in the scripts
          `extract_features.py`, `run_classifier.py` and `run_squad.py`)
      `token_type_ids`: an optional torch.LongTensor of shape [batch_size, sequence_length] with the token
          types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to
          a `sentence B` token (see BERT paper for more details).
      `attention_mask`: an optional torch.LongTensor of shape [batch_size, sequence_length] with indices
          selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max
          input sequence length in the current batch. It's the mask that we typically use for attention when
          a batch has varying length sentences.
      `start_positions`: position of the first token for the labeled span: torch.LongTensor of shape [batch_size].
          Positions are clamped to the length of the sequence and position outside of the sequence are not taken
          into account for computing the loss.
      `end_positions`: position of the last token for the labeled span: torch.LongTensor of shape [batch_size].
          Positions are clamped to the length of the sequence and position outside of the sequence are not taken
          into account for computing the loss.
      `head_mask`: an optional torch.Tensor of shape [num_heads] or [num_layers, num_heads] with indices between 0 and 1.
          It's a mask to be used to nullify some heads of the transformer. 1.0 => head is fully masked, 0.0 => head is not masked.
  Outputs:
      if `start_positions` and `end_positions` are not `None`:
          Outputs the total_loss which is the sum of the CrossEntropy loss for the start and end token positions.
      if `start_positions` or `end_positions` is `None`:
          Outputs a tuple of start_logits, end_logits which are the logits respectively for the start and end
          position tokens of shape [batch_size, sequence_length].
  Example usage:
  ```python
  # Already been converted into WordPiece token ids
  input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
  input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])
  token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]])
  config = BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
      num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)
  model = BertForQuestionAnswering(config)
  start_logits, end_logits = model(input_ids, token_type_ids, input_mask)
  ```
  """

  def __init__(self, config, output_attentions=False, keep_multihead_output=False):
    super(BertForQuestionAnswering, self).__init__(config)
    self.output_attentions = output_attentions
    self.bert = BertModel(config, output_attentions=output_attentions,
                          keep_multihead_output=keep_multihead_output)
    self.qa_outputs = nn.Linear(config.hidden_size, 2)
    self.apply(self.init_bert_weights)

  def forward(self, input_ids, token_type_ids=None, attention_mask=None, start_positions=None,
              end_positions=None, head_mask=None):
    outputs = self.bert(input_ids, token_type_ids, attention_mask,
                        output_all_encoded_layers=False,
                        head_mask=head_mask)
    if self.output_attentions:
      all_attentions, sequence_output, _ = outputs
    else:
      sequence_output, _ = outputs
    logits = self.qa_outputs(sequence_output)
    start_logits, end_logits = logits.split(1, dim=-1)
    start_logits = start_logits.squeeze(-1)
    end_logits = end_logits.squeeze(-1)

    if start_positions is not None and end_positions is not None:
      # If we are on multi-GPU, split add a dimension
      if len(start_positions.size()) > 1:
        start_positions = start_positions.squeeze(-1)
      if len(end_positions.size()) > 1:
        end_positions = end_positions.squeeze(-1)
      # sometimes the start/end positions are outside our model inputs, we ignore these terms
      ignored_index = start_logits.size(1)
      start_positions.clamp_(0, ignored_index)
      end_positions.clamp_(0, ignored_index)

      loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
      start_loss = loss_fct(start_logits, start_positions)
      end_loss = loss_fct(end_logits, end_positions)
      total_loss = (start_loss + end_loss) / 2
      return total_loss
    elif self.output_attentions:
      return all_attentions, start_logits, end_logits
    return start_logits, end_logits