# 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 logging
import math
import os
import sys

import torch
from torch import nn
from torch.nn import CrossEntropyLoss, MSELoss
import torch.nn.functional as F

from .modeling_utils import PreTrainedModel, prune_linear_layer
from .configuration_bert import BertConfig
from .file_utils import add_start_docstrings

# from GAT here
from torch_geometric.data import Data, Batch
from torch_geometric.nn import GCNConv, GATConv
from torch.nn import Parameter
from torch_scatter import scatter_add
from torch_geometric.nn.conv import MessagePassing
from torch_geometric.utils import add_remaining_self_loops, remove_self_loops, add_self_loops, softmax
from torch_geometric.nn.inits import glorot, zeros



import random
import numpy as np

INF = 1e30 # 定义正无穷

# set seed here(set the same with your run_xxx.py file)
# random.seed(42)
# np.random.seed(42)
# torch.manual_seed(42)
# torch.cuda.manual_seed(42)

logger = logging.getLogger(__name__)

BERT_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://s3.amazonaws.com/models.huggingface.co/bert/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",
    'bert-base-german-dbmdz-cased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-german-dbmdz-cased-pytorch_model.bin",
    'bert-base-german-dbmdz-uncased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-german-dbmdz-uncased-pytorch_model.bin",
}


def load_tf_weights_in_bert(model, config, tf_checkpoint_path):
    """ Load tf checkpoints in a pytorch model.
    """
    try:
        import re
        import numpy as np
        import tensorflow as tf
    except ImportError:
        logger.error("Loading a TensorFlow model 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)
    logger.info("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:
        logger.info("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):
            logger.info("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:
                    logger.info("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
        logger.info("Initialize PyTorch weight {}".format(name))
        pointer.data = torch.from_numpy(array)
    return model


def gelu(x):
    """ Original Implementation of the gelu activation function in Google Bert repo when initially created.
        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 gelu_new(x):
    """ Implementation of the gelu activation function currently in Google Bert repo (identical to OpenAI GPT).
        Also see https://arxiv.org/abs/1606.08415
    """
    return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))


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


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


BertLayerNorm = torch.nn.LayerNorm


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=None, token_type_ids=None, position_ids=None, inputs_embeds=None):
        if input_ids is not None:
            input_shape = input_ids.size()
        else:
            input_shape = inputs_embeds.size()[:-1]

        seq_length = input_shape[1]
        device = input_ids.device if input_ids is not None else inputs_embeds.device
        if position_ids is None:
            position_ids = torch.arange(seq_length, dtype=torch.long, device=device)
            position_ids = position_ids.unsqueeze(0).expand(input_shape)
        if token_type_ids is None:
            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)

        if inputs_embeds is None:
            inputs_embeds = self.word_embeddings(input_ids)
        position_embeddings = self.position_embeddings(position_ids)
        token_type_embeddings = self.token_type_embeddings(token_type_ids)

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


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.output_attentions = config.output_attentions

        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)
        x = x.view(*new_x_shape)
        return x.permute(0, 2, 1, 3)

    def forward(self, hidden_states, attention_mask=None, head_mask=None, encoder_hidden_states=None, encoder_attention_mask=None):
        mixed_query_layer = self.query(hidden_states)

        # If this is instantiated as a cross-attention module, the keys
        # and values come from an encoder; the attention mask needs to be
        # such that the encoder's padding tokens are not attended to.
        if encoder_hidden_states is not None:
            mixed_key_layer = self.key(encoder_hidden_states)
            mixed_value_layer = self.value(encoder_hidden_states)
            attention_mask = encoder_attention_mask
        else:
            mixed_key_layer = self.key(hidden_states)
            mixed_value_layer = self.value(hidden_states)

        query_layer = self.transpose_for_scores(mixed_query_layer)
        key_layer = self.transpose_for_scores(mixed_key_layer)
        value_layer = self.transpose_for_scores(mixed_value_layer)

        # 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))
        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
        if attention_mask is not None:
            # 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)

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

        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)

        outputs = (context_layer, attention_probs) if self.output_attentions else (context_layer,)
        return outputs


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):
        super(BertAttention, self).__init__()
        self.self = BertSelfAttention(config)
        self.output = BertSelfOutput(config)
        self.pruned_heads = set()

    def prune_heads(self, heads):
        if len(heads) == 0:
            return
        mask = torch.ones(self.self.num_attention_heads, self.self.attention_head_size)
        heads = set(heads) - self.pruned_heads  # Convert to set and emove already pruned heads
        for head in heads:
            # Compute how many pruned heads are before the head and move the index accordingly
            head = head - sum(1 if h < head else 0 for h in self.pruned_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 and store pruned heads
        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
        self.pruned_heads = self.pruned_heads.union(heads)

    def forward(self, hidden_states, attention_mask=None, head_mask=None, encoder_hidden_states=None, encoder_attention_mask=None):
        self_outputs = self.self(hidden_states, attention_mask, head_mask, encoder_hidden_states, encoder_attention_mask)
        attention_output = self.output(self_outputs[0], hidden_states)
        outputs = (attention_output,) + self_outputs[1:]  # add attentions if we output them
        return outputs


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):
    def __init__(self, config):
        super(BertLayer, self).__init__()
        self.attention = BertAttention(config)
        self.is_decoder = config.is_decoder
        if self.is_decoder:
            self.crossattention = BertAttention(config)
        self.intermediate = BertIntermediate(config)
        self.output = BertOutput(config)

    def forward(self, hidden_states, attention_mask=None, head_mask=None, encoder_hidden_states=None, encoder_attention_mask=None):
        self_attention_outputs = self.attention(hidden_states, attention_mask, head_mask)
        attention_output = self_attention_outputs[0]
        outputs = self_attention_outputs[1:]  # add self attentions if we output attention weights

        if self.is_decoder and encoder_hidden_states is not None:
            cross_attention_outputs = self.crossattention(attention_output, attention_mask, head_mask, encoder_hidden_states, encoder_attention_mask)
            attention_output = cross_attention_outputs[0]
            outputs = outputs + cross_attention_outputs[1:]  # add cross attentions if we output attention weights

        intermediate_output = self.intermediate(attention_output)
        layer_output = self.output(intermediate_output, attention_output)
        outputs = (layer_output,) + outputs
        return outputs


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

    def forward(self, hidden_states, attention_mask=None, head_mask=None, encoder_hidden_states=None, encoder_attention_mask=None):
        all_hidden_states = ()
        all_attentions = ()
        for i, layer_module in enumerate(self.layer):
            if self.output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            layer_outputs = layer_module(hidden_states, attention_mask, head_mask[i], encoder_hidden_states, encoder_attention_mask)
            hidden_states = layer_outputs[0]

            if self.output_attentions:
                all_attentions = all_attentions + (layer_outputs[1],)

        # Add last layer
        if self.output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        outputs = (hidden_states,)
        if self.output_hidden_states:
            outputs = outputs + (all_hidden_states,)
        if self.output_attentions:
            outputs = outputs + (all_attentions,)
        return outputs  # last-layer hidden state, (all hidden states), (all attentions)


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

    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):
        super(BertOnlyMLMHead, self).__init__()
        self.predictions = BertLMPredictionHead(config)

    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):
        super(BertPreTrainingHeads, self).__init__()
        self.predictions = BertLMPredictionHead(config)
        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(PreTrainedModel):
    """ An abstract class to handle weights initialization and
        a simple interface for dowloading and loading pretrained models.
    """
    config_class = BertConfig
    pretrained_model_archive_map = BERT_PRETRAINED_MODEL_ARCHIVE_MAP
    load_tf_weights = load_tf_weights_in_bert
    base_model_prefix = "bert"

    def _init_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_()


BERT_START_DOCSTRING = r"""    The BERT model was proposed in
    `BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding`_
    by <NAME>, <NAME>, <NAME> and <NAME>. It's a bidirectional transformer
    pre-trained using a combination of masked language modeling objective and next sentence prediction
    on a large corpus comprising the Toronto Book Corpus and Wikipedia.

    This model is a PyTorch `torch.nn.Module`_ sub-class. Use it as a regular PyTorch Module and
    refer to the PyTorch documentation for all matter related to general usage and behavior.

    .. _`BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding`:
        https://arxiv.org/abs/1810.04805

    .. _`torch.nn.Module`:
        https://pytorch.org/docs/stable/nn.html#module

    Parameters:
        config (:class:`~transformers.BertConfig`): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the configuration.
            Check out the :meth:`~transformers.PreTrainedModel.from_pretrained` method to load the model weights.
"""

BERT_INPUTS_DOCSTRING = r"""
    Inputs:
        **input_ids**: ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``:
            Indices of input sequence tokens in the vocabulary.
            To match pre-training, BERT input sequence should be formatted with [CLS] and [SEP] tokens as follows:

            (a) For sequence pairs:

                ``tokens:         [CLS] is this jack ##son ##ville ? [SEP] no it is not . [SEP]``

                ``token_type_ids:   0   0  0    0    0     0       0   0   1  1  1  1   1   1``

            (b) For single sequences:

                ``tokens:         [CLS] the dog is hairy . [SEP]``

                ``token_type_ids:   0   0   0   0  0     0   0``

            Bert is a model with absolute position embeddings so it's usually advised to pad the inputs on
            the right rather than the left.

            Indices can be obtained using :class:`transformers.BertTokenizer`.
            See :func:`transformers.PreTrainedTokenizer.encode` and
            :func:`transformers.PreTrainedTokenizer.convert_tokens_to_ids` for details.
        **attention_mask**: (`optional`) ``torch.FloatTensor`` of shape ``(batch_size, sequence_length)``:
            Mask to avoid performing attention on padding token indices.
            Mask values selected in ``[0, 1]``:
            ``1`` for tokens that are NOT MASKED, ``0`` for MASKED tokens.
        **token_type_ids**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``:
            Segment token indices to indicate first and second portions of the inputs.
            Indices are selected in ``[0, 1]``: ``0`` corresponds to a `sentence A` token, ``1``
            corresponds to a `sentence B` token
            (see `BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding`_ for more details).
        **position_ids**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``:
            Indices of positions of each input sequence tokens in the position embeddings.
            Selected in the range ``[0, config.max_position_embeddings - 1]``.
        **head_mask**: (`optional`) ``torch.FloatTensor`` of shape ``(num_heads,)`` or ``(num_layers, num_heads)``:
            Mask to nullify selected heads of the self-attention modules.
            Mask values selected in ``[0, 1]``:
            ``1`` indicates the head is **not masked**, ``0`` indicates the head is **masked**.
        **inputs_embeds**: (`optional`) ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, embedding_dim)``:
            Optionally, instead of passing ``input_ids`` you can choose to directly pass an embedded representation.
            This is useful if you want more control over how to convert `input_ids` indices into associated vectors
            than the model's internal embedding lookup matrix.
        **encoder_hidden_states**: (`optional`) ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, hidden_size)``:
            Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if the model
            is configured as a decoder.
        **encoder_attention_mask**: (`optional`) ``torch.FloatTensor`` of shape ``(batch_size, sequence_length)``:
            Mask to avoid performing attention on the padding token indices of the encoder input. This mask
            is used in the cross-attention if the model is configured as a decoder.
            Mask values selected in ``[0, 1]``:
            ``1`` for tokens that are NOT MASKED, ``0`` for MASKED tokens.
"""

@add_start_docstrings("The bare Bert Model transformer outputting raw hidden-states without any specific head on top.",
                      BERT_START_DOCSTRING, BERT_INPUTS_DOCSTRING)
class BertModel(BertPreTrainedModel):
    r"""
    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
        **last_hidden_state**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, hidden_size)``
            Sequence of hidden-states at the output of the last layer of the model.
        **pooler_output**: ``torch.FloatTensor`` of shape ``(batch_size, hidden_size)``
            Last layer hidden-state of the first token of the sequence (classification token)
            further processed by a Linear layer and a Tanh activation function. The Linear
            layer weights are trained from the next sentence prediction (classification)
            objective during Bert pretraining. This output is usually *not* a good summary
            of the semantic content of the input, you're often better with averaging or pooling
            the sequence of hidden-states for the whole input sequence.
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

    Examples::

        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        model = BertModel.from_pretrained('bert-base-uncased')
        input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute")).unsqueeze(0)  # Batch size 1
        outputs = model(input_ids)
        last_hidden_states = outputs[0]  # The last hidden-state is the first element of the output tuple

    """
    def __init__(self, config):
        super(BertModel, self).__init__(config)
        self.config = config

        self.embeddings = BertEmbeddings(config)
        self.encoder = BertEncoder(config)
        self.pooler = BertPooler(config)

        self.init_weights()

    def get_input_embeddings(self):
        return self.embeddings.word_embeddings

    def set_input_embeddings(self, value):
        self.embeddings.word_embeddings = value

    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}
            See base class PreTrainedModel
        """
        for layer, heads in heads_to_prune.items():
            self.encoder.layer[layer].attention.prune_heads(heads)

    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None,
                head_mask=None, inputs_embeds=None, encoder_hidden_states=None, encoder_attention_mask=None):
        """ Forward pass on the Model.

        The model can behave as an encoder (with only self-attention) as well
        as a decoder, in which case a layer of cross-attention is added between
        the self-attention layers, following the architecture described in `Attention is all you need`_ by <NAME>,
        <NAME>, <NAME>, <NAME>, <NAME>, <NAME>, <NAME> and <NAME>.

        To behave as an decoder the model needs to be initialized with the
        `is_decoder` argument of the configuration set to `True`; an
        `encoder_hidden_states` is expected as an input to the forward pass.

        .. _`Attention is all you need`:
            https://arxiv.org/abs/1706.03762

        """
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
            input_shape = input_ids.size()
        elif inputs_embeds is not None:
            input_shape = inputs_embeds.size()[:-1]
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        device = input_ids.device if input_ids is not None else inputs_embeds.device

        if attention_mask is None:
            attention_mask = torch.ones(input_shape, device=device)
        if encoder_attention_mask is None:
            encoder_attention_mask = torch.ones(input_shape, device=device)
        if token_type_ids is None:
            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)

        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
        # ourselves in which case we just need to make it broadcastable to all heads.
        if attention_mask.dim() == 3:
            extended_attention_mask = attention_mask[:, None, :, :]

        # Provided a padding mask of dimensions [batch_size, seq_length]
        # - if the model is a decoder, apply a causal mask in addition to the padding mask
        # - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length]
        if attention_mask.dim() == 2:
            if self.config.is_decoder:
                batch_size, seq_length = input_shape
                seq_ids = torch.arange(seq_length, device=device)
                causal_mask = seq_ids[None, None, :].repeat(batch_size, seq_length, 1) <= seq_ids[None, :, None]
                extended_attention_mask = causal_mask[:, None, :, :] * attention_mask[:, None, None, :]
            else:
                extended_attention_mask = attention_mask[:, None, None, :]

        # 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

        # If a 2D ou 3D attention mask is provided for the cross-attention
        # we need to make broadcastabe to [batch_size, num_heads, seq_length, seq_length]
        if encoder_attention_mask.dim() == 3:
            encoder_extended_attention_mask = encoder_attention_mask[:, None, :, :]
        if encoder_attention_mask.dim() == 2:
            encoder_extended_attention_mask = encoder_attention_mask[:, None, None, :]

        encoder_extended_attention_mask = encoder_extended_attention_mask.to(dtype=next(self.parameters()).dtype)  # fp16 compatibility
        encoder_extended_attention_mask = (1.0 - encoder_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(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

        embedding_output = self.embeddings(input_ids=input_ids, position_ids=position_ids, token_type_ids=token_type_ids, inputs_embeds=inputs_embeds)
        encoder_outputs = self.encoder(embedding_output,
                                       attention_mask=extended_attention_mask,
                                       head_mask=head_mask,
                                       encoder_hidden_states=encoder_hidden_states,
                                       encoder_attention_mask=encoder_extended_attention_mask)
        sequence_output = encoder_outputs[0]
        pooled_output = self.pooler(sequence_output)

        outputs = (sequence_output, pooled_output,) + encoder_outputs[1:]  # add hidden_states and attentions if they are here
        return outputs  # sequence_output, pooled_output, (hidden_states), (attentions)


@add_start_docstrings("""Bert Model with two heads on top as done during the pre-training:
                       a `masked language modeling` head and a `next sentence prediction (classification)` head. """,
                      BERT_START_DOCSTRING,
                      BERT_INPUTS_DOCSTRING)
class BertForPreTraining(BertPreTrainedModel):
    r"""
        **masked_lm_labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``:
            Labels for computing the masked language modeling loss.
            Indices should be in ``[-1, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring)
            Tokens with indices set to ``-1`` are ignored (masked), the loss is only computed for the tokens with labels
            in ``[0, ..., config.vocab_size]``
        **next_sentence_label**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:
            Labels for computing the next sequence prediction (classification) loss. Input should be a sequence pair (see ``input_ids`` docstring)
            Indices should be in ``[0, 1]``.
            ``0`` indicates sequence B is a continuation of sequence A,
            ``1`` indicates sequence B is a random sequence.

    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
        **loss**: (`optional`, returned when both ``masked_lm_labels`` and ``next_sentence_label`` are provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Total loss as the sum of the masked language modeling loss and the next sequence prediction (classification) loss.
        **prediction_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, config.vocab_size)``
            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
        **seq_relationship_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, 2)``
            Prediction scores of the next sequence prediction (classification) head (scores of True/False continuation before SoftMax).
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

    Examples::

        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        model = BertForPreTraining.from_pretrained('bert-base-uncased')
        input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute")).unsqueeze(0)  # Batch size 1
        outputs = model(input_ids)
        prediction_scores, seq_relationship_scores = outputs[:2]

    """
    def __init__(self, config):
        super(BertForPreTraining, self).__init__(config)

        self.bert = BertModel(config)
        self.cls = BertPreTrainingHeads(config)

        self.init_weights()

    def get_output_embeddings(self):
        return self.cls.predictions.decoder

    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None,
                masked_lm_labels=None, next_sentence_label=None):

        outputs = self.bert(input_ids,
                            attention_mask=attention_mask,
                            token_type_ids=token_type_ids,
                            position_ids=position_ids,
                            head_mask=head_mask,
                            inputs_embeds=inputs_embeds)

        sequence_output, pooled_output = outputs[:2]
        prediction_scores, seq_relationship_score = self.cls(sequence_output, pooled_output)

        outputs = (prediction_scores, seq_relationship_score,) + outputs[2:]  # add hidden states and attention if they are here

        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
            outputs = (total_loss,) + outputs

        return outputs  # (loss), prediction_scores, seq_relationship_score, (hidden_states), (attentions)


@add_start_docstrings("""Bert Model with a `language modeling` head on top. """,
                      BERT_START_DOCSTRING,
                      BERT_INPUTS_DOCSTRING)
class BertForMaskedLM(BertPreTrainedModel):
    r"""
        **masked_lm_labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``:
            Labels for computing the masked language modeling loss.
            Indices should be in ``[-1, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring)
            Tokens with indices set to ``-1`` are ignored (masked), the loss is only computed for the tokens with labels
            in ``[0, ..., config.vocab_size]``
        **lm_labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``:
            Labels for computing the left-to-right language modeling loss (next word prediction).
            Indices should be in ``[-1, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring)
            Tokens with indices set to ``-1`` are ignored (masked), the loss is only computed for the tokens with labels
            in ``[0, ..., config.vocab_size]``

    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
        **masked_lm_loss**: (`optional`, returned when ``masked_lm_labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Masked language modeling loss.
        **ltr_lm_loss**: (`optional`, returned when ``lm_labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Next token prediction loss.
        **prediction_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, config.vocab_size)``
            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

    Examples::

        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        model = BertForMaskedLM.from_pretrained('bert-base-uncased')
        input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute")).unsqueeze(0)  # Batch size 1
        outputs = model(input_ids, masked_lm_labels=input_ids)
        loss, prediction_scores = outputs[:2]

    """
    def __init__(self, config):
        super(BertForMaskedLM, self).__init__(config)

        self.bert = BertModel(config)
        self.cls = BertOnlyMLMHead(config)

        self.init_weights()

    def get_output_embeddings(self):
        return self.cls.predictions.decoder

    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None,
                masked_lm_labels=None, encoder_hidden_states=None, encoder_attention_mask=None, lm_labels=None, ):

        outputs = self.bert(input_ids,
                            attention_mask=attention_mask,
                            token_type_ids=token_type_ids,
                            position_ids=position_ids,
                            head_mask=head_mask,
                            inputs_embeds=inputs_embeds,
                            encoder_hidden_states=encoder_hidden_states,
                            encoder_attention_mask=encoder_attention_mask)

        sequence_output = outputs[0]
        prediction_scores = self.cls(sequence_output)

        outputs = (prediction_scores,) + outputs[2:]  # Add hidden states and attention if they are here

        # Although this may seem awkward, BertForMaskedLM supports two scenarios:
        # 1. If a tensor that contains the indices of masked labels is provided,
        #    the cross-entropy is the MLM cross-entropy that measures the likelihood
        #    of predictions for masked words.
        # 2. If `lm_labels` is provided we are in a causal scenario where we
        #    try to predict the next token for each input in the decoder.
        if masked_lm_labels is not None:
            loss_fct = CrossEntropyLoss(ignore_index=-1)  # -1 index = padding token
            masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), masked_lm_labels.view(-1))
            outputs = (masked_lm_loss,) + outputs

        if lm_labels is not None:
            # we are doing next-token prediction; shift prediction scores and input ids by one
            prediction_scores = prediction_scores[:, :-1, :].contiguous()
            lm_labels = lm_labels[:, 1:].contiguous()
            loss_fct = CrossEntropyLoss(ignore_index=-1)
            ltr_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), lm_labels.view(-1))
            outputs = (ltr_lm_loss,) + outputs

        return outputs  # (masked_lm_loss), (ltr_lm_loss), prediction_scores, (hidden_states), (attentions)


@add_start_docstrings("""Bert Model with a `next sentence prediction (classification)` head on top. """,
                      BERT_START_DOCSTRING,
                      BERT_INPUTS_DOCSTRING)
class BertForNextSentencePrediction(BertPreTrainedModel):
    r"""
        **next_sentence_label**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:
            Labels for computing the next sequence prediction (classification) loss. Input should be a sequence pair (see ``input_ids`` docstring)
            Indices should be in ``[0, 1]``.
            ``0`` indicates sequence B is a continuation of sequence A,
            ``1`` indicates sequence B is a random sequence.

    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
        **loss**: (`optional`, returned when ``next_sentence_label`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Next sequence prediction (classification) loss.
        **seq_relationship_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, 2)``
            Prediction scores of the next sequence prediction (classification) head (scores of True/False continuation before SoftMax).
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

    Examples::

        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        model = BertForNextSentencePrediction.from_pretrained('bert-base-uncased')
        input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute")).unsqueeze(0)  # Batch size 1
        outputs = model(input_ids)
        seq_relationship_scores = outputs[0]

    """
    def __init__(self, config):
        super(BertForNextSentencePrediction, self).__init__(config)

        self.bert = BertModel(config)
        self.cls = BertOnlyNSPHead(config)

        self.init_weights()

    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None,
                next_sentence_label=None):

        outputs = self.bert(input_ids,
                            attention_mask=attention_mask,
                            token_type_ids=token_type_ids,
                            position_ids=position_ids,
                            head_mask=head_mask,
                            inputs_embeds=inputs_embeds)

        pooled_output = outputs[1]

        seq_relationship_score = self.cls(pooled_output)

        outputs = (seq_relationship_score,) + outputs[2:]  # add hidden states and attention if they are here
        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))
            outputs = (next_sentence_loss,) + outputs

        return outputs  # (next_sentence_loss), seq_relationship_score, (hidden_states), (attentions)


@add_start_docstrings("""Bert Model transformer with a sequence classification/regression head on top (a linear layer on top of
                      the pooled output) e.g. for GLUE tasks. """,
                      BERT_START_DOCSTRING,
                      BERT_INPUTS_DOCSTRING)
class BertForSequenceClassification(BertPreTrainedModel):
    r"""
        **labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:
            Labels for computing the sequence classification/regression loss.
            Indices should be in ``[0, ..., config.num_labels - 1]``.
            If ``config.num_labels == 1`` a regression loss is computed (Mean-Square loss),
            If ``config.num_labels > 1`` a classification loss is computed (Cross-Entropy).

    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
        **loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Classification (or regression if config.num_labels==1) loss.
        **logits**: ``torch.FloatTensor`` of shape ``(batch_size, config.num_labels)``
            Classification (or regression if config.num_labels==1) scores (before SoftMax).
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

    Examples::

        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        model = BertForSequenceClassification.from_pretrained('bert-base-uncased')
        input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute")).unsqueeze(0)  # Batch size 1
        labels = torch.tensor([1]).unsqueeze(0)  # Batch size 1
        outputs = model(input_ids, labels=labels)
        loss, logits = outputs[:2]

    """
    def __init__(self, config):
        super(BertForSequenceClassification, self).__init__(config)
        self.num_labels = config.num_labels

        self.bert = BertModel(config)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.classifier = nn.Linear(config.hidden_size, self.config.num_labels)

        self.init_weights()

    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None,
                position_ids=None, head_mask=None, inputs_embeds=None, labels=None):

        outputs = self.bert(input_ids,
                            attention_mask=attention_mask,
                            token_type_ids=token_type_ids,
                            position_ids=position_ids,
                            head_mask=head_mask,
                            inputs_embeds=inputs_embeds)

        pooled_output = outputs[1]

        pooled_output = self.dropout(pooled_output)
        logits = self.classifier(pooled_output)

        outputs = (logits,) + outputs[2:]  # add hidden states and attention if they are here

        if labels is not None:
            if self.num_labels == 1:
                #  We are doing regression
                loss_fct = MSELoss()
                loss = loss_fct(logits.view(-1), labels.view(-1))
            else:
                loss_fct = CrossEntropyLoss()
                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
            outputs = (loss,) + outputs

        return outputs  # (loss), logits, (hidden_states), (attentions)


@add_start_docstrings("""Bert Model with a multiple choice classification head on top (a linear layer on top of
                      the pooled output and a softmax) e.g. for RocStories/SWAG tasks. """,
                      BERT_START_DOCSTRING,
                      BERT_INPUTS_DOCSTRING)
class BertForMultipleChoice(BertPreTrainedModel):
    r"""
        **labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:
            Labels for computing the multiple choice classification loss.
            Indices should be in ``[0, ..., num_choices]`` where `num_choices` is the size of the second dimension
            of the input tensors. (see `input_ids` above)

    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
        **loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Classification loss.
        **classification_scores**: ``torch.FloatTensor`` of shape ``(batch_size, num_choices)`` where `num_choices` is the size of the second dimension
            of the input tensors. (see `input_ids` above).
            Classification scores (before SoftMax).
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

    Examples::

        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        model = BertForMultipleChoice.from_pretrained('bert-base-uncased')
        choices = ["Hello, my dog is cute", "Hello, my cat is amazing"]
        input_ids = torch.tensor([tokenizer.encode(s) for s in choices]).unsqueeze(0)  # Batch size 1, 2 choices
        labels = torch.tensor(1).unsqueeze(0)  # Batch size 1
        outputs = model(input_ids, labels=labels)
        loss, classification_scores = outputs[:2]

    """
    def __init__(self, config):
        super(BertForMultipleChoice, self).__init__(config)

        self.bert = BertModel(config)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.classifier = nn.Linear(config.hidden_size, 1)

        self.init_weights()

    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None,
                position_ids=None, head_mask=None, inputs_embeds=None, labels=None):
        num_choices = input_ids.shape[1]

        input_ids = input_ids.view(-1, input_ids.size(-1))
        attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None
        token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None
        position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None

        outputs = self.bert(input_ids,
                            attention_mask=attention_mask,
                            token_type_ids=token_type_ids,
                            position_ids=position_ids,
                            head_mask=head_mask,
                            inputs_embeds=inputs_embeds)

        pooled_output = outputs[1]

        pooled_output = self.dropout(pooled_output)
        logits = self.classifier(pooled_output)
        reshaped_logits = logits.view(-1, num_choices)

        outputs = (reshaped_logits,) + outputs[2:]  # add hidden states and attention if they are here

        if labels is not None:
            loss_fct = CrossEntropyLoss()
            loss = loss_fct(reshaped_logits, labels)
            outputs = (loss,) + outputs

        return outputs  # (loss), reshaped_logits, (hidden_states), (attentions)


@add_start_docstrings("""Bert Model with a token classification head on top (a linear layer on top of
                      the hidden-states output) e.g. for Named-Entity-Recognition (NER) tasks. """,
                      BERT_START_DOCSTRING,
                      BERT_INPUTS_DOCSTRING)
class BertForTokenClassification(BertPreTrainedModel):
    r"""
        **labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``:
            Labels for computing the token classification loss.
            Indices should be in ``[0, ..., config.num_labels - 1]``.

    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
        **loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Classification loss.
        **scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, config.num_labels)``
            Classification scores (before SoftMax).
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

    Examples::

        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        model = BertForTokenClassification.from_pretrained('bert-base-uncased')
        input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute")).unsqueeze(0)  # Batch size 1
        labels = torch.tensor([1] * input_ids.size(1)).unsqueeze(0)  # Batch size 1
        outputs = model(input_ids, labels=labels)
        loss, scores = outputs[:2]

    """
    def __init__(self, config):
        super(BertForTokenClassification, self).__init__(config)
        self.num_labels = config.num_labels

        self.bert = BertModel(config)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.classifier = nn.Linear(config.hidden_size, config.num_labels)

        self.init_weights()

    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None,
                position_ids=None, head_mask=None, inputs_embeds=None, labels=None):

        outputs = self.bert(input_ids,
                            attention_mask=attention_mask,
                            token_type_ids=token_type_ids,
                            position_ids=position_ids,
                            head_mask=head_mask,
                            inputs_embeds=inputs_embeds)

        sequence_output = outputs[0]

        sequence_output = self.dropout(sequence_output)
        logits = self.classifier(sequence_output)

        outputs = (logits,) + outputs[2:]  # add hidden states and attention if they are here
        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 # [batch_size * seq_len]
                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))
            outputs = (loss,) + outputs

        return outputs  # (loss), scores, (hidden_states), (attentions)


@add_start_docstrings("""Bert Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear layers on top of
                      the hidden-states output to compute `span start logits` and `span end logits`). """,
                      BERT_START_DOCSTRING,
                      BERT_INPUTS_DOCSTRING)
class BertForQuestionAnswering(BertPreTrainedModel):
    r"""
        **start_positions**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:
            Labels for position (index) of the start of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`).
            Position outside of the sequence are not taken into account for computing the loss.
        **end_positions**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:
            Labels for position (index) of the end of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`).
            Position outside of the sequence are not taken into account for computing the loss.

    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
        **loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Total span extraction loss is the sum of a Cross-Entropy for the start and end positions.
        **start_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length,)``
            Span-start scores (before SoftMax).
        **end_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length,)``
            Span-end scores (before SoftMax).
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

    Examples::

        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        model = BertForQuestionAnswering.from_pretrained('bert-large-uncased-whole-word-masking-finetuned-squad')
        question, text = "Who was <NAME>?", "<NAME> was a nice puppet"
        input_text = "[CLS] " + question + " [SEP] " + text + " [SEP]"
        input_ids = tokenizer.encode(input_text)
        token_type_ids = [0 if i <= input_ids.index(102) else 1 for i in range(len(input_ids))] 
        start_scores, end_scores = model(torch.tensor([input_ids]), token_type_ids=torch.tensor([token_type_ids]))
        all_tokens = tokenizer.convert_ids_to_tokens(input_ids)  
        print(' '.join(all_tokens[torch.argmax(start_scores) : torch.argmax(end_scores)+1]))
        # a nice puppet


    """
    def __init__(self, config):
        super(BertForQuestionAnswering, self).__init__(config)
        self.num_labels = config.num_labels

        self.bert = BertModel(config)
        self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)

        self.init_weights()

    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None,
                start_positions=None, end_positions=None):

        outputs = self.bert(input_ids,
                            attention_mask=attention_mask,
                            token_type_ids=token_type_ids,
                            position_ids=position_ids,
                            head_mask=head_mask,
                            inputs_embeds=inputs_embeds)

        sequence_output = outputs[0]

        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)

        outputs = (start_logits, end_logits,) + outputs[2:]
        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
            outputs = (total_loss,) + outputs

        return outputs  # (loss), start_logits, end_logits, (hidden_states), (attentions)


###############################################################################################################################################
# below is authored by yyhaker
class FSNetForQuestionAnswering(BertPreTrainedModel):
    r"""
        **start_positions**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:
            Labels for position (index) of the start of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`).
            Position outside of the sequence are not taken into account for computing the loss.
        **end_positions**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:
            Labels for position (index) of the end of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`).
            Position outside of the sequence are not taken into account for computing the loss.

    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
        **loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Total span extraction loss is the sum of a Cross-Entropy for the start and end positions.
        **start_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length,)``
            Span-start scores (before SoftMax).
        **end_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length,)``
            Span-end scores (before SoftMax).
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

    Examples::

        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        model = BertForQuestionAnswering.from_pretrained('bert-large-uncased-whole-word-masking-finetuned-squad')
        question, text = "Who was <NAME>?", "<NAME> was a nice puppet"
        input_text = "[CLS] " + question + " [SEP] " + text + " [SEP]"
        input_ids = tokenizer.encode(input_text)
        token_type_ids = [0 if i <= input_ids.index(102) else 1 for i in range(len(input_ids))] 
        start_scores, end_scores = model(torch.tensor([input_ids]), token_type_ids=torch.tensor([token_type_ids]))
        all_tokens = tokenizer.convert_ids_to_tokens(input_ids)  
        print(' '.join(all_tokens[torch.argmax(start_scores) : torch.argmax(end_scores)+1]))
        # a nice puppet
    """
    def __init__(self, config, args, concept_embedding_mat):
        super(FSNetForQuestionAnswering, self).__init__(config, args, concept_embedding_mat)
        self.num_labels = config.num_labels
        
        '''1st Layer: BERT Encoding Layer'''
        self.bert = BertModel(config)

        '''2st Layer: Knowledge Integration Layer'''
        self.memory_embs = torch.nn.Embedding.from_pretrained(torch.Tensor(concept_embedding_mat), freeze=args.freeze_concept)

        concept_vocab_size = concept_embedding_mat.shape[0]
        concept_dim = concept_embedding_mat.shape[1]

        # init a sentinel vector
        # self.sentinel_vec = torch.nn.Parameter(torch.rand(concept_dim), requires_grad=False)

        # init a sentinel vector
        # elf.sentinel_vec = torch.nn.Parameter(torch.zeros(concept_dim), requires_grad=True)

        self.relavance_layer = torch.nn.Bilinear(config.hidden_size, concept_dim, 1)

        self.trilinear = torch.nn.Linear(3 * (config.hidden_size + concept_dim), 1)

        self.qa_outputs = torch.nn.Linear(6 * (config.hidden_size + concept_dim), config.num_labels)

        # self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)

        self.init_weights()
    
    def self_matching(self, U, mask):
        '''self-matching layer.
            U, (b, len, d)
            mask, (b, len)
        '''
        b, len_, d = U.size()

        # direct info interaction
        ui = U.unsqueeze(2) # [b, len, 1, d]
        uj = U.unsqueeze(1) # [b, 1, len, d]
        dot = ui * uj  # [b, len, len, d]
        ui_expand = ui.expand(b, len_, len_, d).contiguous() # [b, len, len, d] 
        uj_expand = uj.expand(b, len_, len_, d).contiguous() # [b, len, len, d]
        R = self.trilinear(torch.cat((ui_expand, uj_expand, dot), -1)).squeeze(-1) # [b, len, len, 3 * d] -> [b, len, len, 1] -> [b, len, len]
        mask_expand = mask.unsqueeze(-1).expand(b, len_, len_).contiguous() # [b, len, len]
        R = INF * (mask_expand.float() - 1.0) + R   # mask here
        A = torch.nn.functional.softmax(R, -1) # [b, len, len]
        V = torch.bmm(A, U) # [b, len, d]

        # indirect info interaction
        A_hat = torch.bmm(A, A)  # [b, len, len]
        V_hat = torch.bmm(A_hat, U) # [b, len, d]
        O = torch.cat((U, V, U - V, U * V, V_hat, U - V_hat), -1) # [b, len, 6*d]
        return O


    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None,
                start_positions=None, end_positions=None, concept_ids=None, concept_masks=None):
        
        '''1st Layer: Bert Layer'''
        outputs = self.bert(input_ids,
                            attention_mask=attention_mask,
                            token_type_ids=token_type_ids,
                            position_ids=position_ids,
                            head_mask=head_mask,
                            inputs_embeds=inputs_embeds)

        sequence_output = outputs[0] # [batch_size, seq_length, hidden_size]
        batch_size, seq_length, hidden_size = sequence_output.size()

        '''2nd Layer: Knowledge Integration Layer'''
        concept_emb = self.memory_embs(concept_ids) # [batch_size, seq_length, num_concepts, cpnet_dim]
        num_concepts = concept_emb.size()[2]
        cpnet_dim = concept_emb.size()[3]
        
        # add sentinel vector
        # sentinel_vec = self.sentinel_vec.to(concept_emb.device)
        # sentinel_vec_expand = sentinel_vec.expand(batch_size, seq_length, cpnet_dim).unsqueeze(2).contiguous()  # [batch_size, seq_length, 1, cpnet_dim]
        
        # concept_emb_add = torch.cat((concept_emb, sentinel_vec_expand), 2)  # [batch_size, seq_length, num_concepts + 1, cpnet_dim]

        sequence_output_expand = sequence_output.unsqueeze(2).expand(
            batch_size, seq_length, num_concepts, hidden_size).contiguous() # [batch_size, seq_length, num_concepts, hidden_size]

        similarity_weight = self.relavance_layer(sequence_output_expand, concept_emb).squeeze(-1) # [batch_size, seq_length, num_concepts]
        
        # calc masked attention
        similarity_weight = INF * (concept_masks - 1.0) + similarity_weight
        similarity_weight_softmax = torch.nn.functional.softmax(similarity_weight, -1) # [batch_size, seq_length, num_concepts]

        concept_relavance = torch.bmm(similarity_weight_softmax.unsqueeze(2).reshape(batch_size * seq_length, 1, -1),
            concept_emb.reshape(batch_size * seq_length, num_concepts, -1))  # [batch_size * seq_length, 1, cpnet_dim]
        concept_relavance = concept_relavance.squeeze(1).reshape(batch_size, seq_length, -1) # [batch_size, seq_length, cpnet_dim]
        
        # mask and cat representation
        attention_mask_expand = attention_mask.unsqueeze(-1).expand(batch_size, seq_length, cpnet_dim).contiguous() # [batch_size, seq_length, cpnet_dim]
        concept_relavance = attention_mask_expand.float() * concept_relavance
        sequence_kg_output = torch.cat((sequence_output, concept_relavance), dim=-1) # [batch_size, seq_length, hidden_size + cpnet_dim]

        '''3nd Layer: Self-Matching Layer'''
        sequence_kg_output = self.self_matching(sequence_kg_output, attention_mask)

        # residule connection
        # sequence_kg_output = torch.cat((sequence_output, sequence_kg_output), dim=-1)

        '''4nd Layer: Output Layer'''
        logits = self.qa_outputs(sequence_kg_output)
        start_logits, end_logits = logits.split(1, dim=-1)
        start_logits = start_logits.squeeze(-1)
        end_logits = end_logits.squeeze(-1)

        outputs = (start_logits, end_logits,) + outputs[2:]
        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
            outputs = (total_loss,) + outputs

        return outputs  # (loss), start_logits, end_logits, (hidden_states), (attentions)


# here is FSNET_ner for QA
class FSNet_ner_ForQuestionAnswering(BertPreTrainedModel):
    r"""
        **start_positions**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:
            Labels for position (index) of the start of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`).
            Position outside of the sequence are not taken into account for computing the loss.
        **end_positions**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:
            Labels for position (index) of the end of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`).
            Position outside of the sequence are not taken into account for computing the loss.

    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
        **loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Total span extraction loss is the sum of a Cross-Entropy for the start and end positions.
        **start_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length,)``
            Span-start scores (before SoftMax).
        **end_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length,)``
            Span-end scores (before SoftMax).
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

    Examples::

        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        model = BertForQuestionAnswering.from_pretrained('bert-large-uncased-whole-word-masking-finetuned-squad')
        question, text = "Who was <NAME>?", "<NAME> was a nice puppet"
        input_text = "[CLS] " + question + " [SEP] " + text + " [SEP]"
        input_ids = tokenizer.encode(input_text)
        token_type_ids = [0 if i <= input_ids.index(102) else 1 for i in range(len(input_ids))] 
        start_scores, end_scores = model(torch.tensor([input_ids]), token_type_ids=torch.tensor([token_type_ids]))
        all_tokens = tokenizer.convert_ids_to_tokens(input_ids)  
        print(' '.join(all_tokens[torch.argmax(start_scores) : torch.argmax(end_scores)+1]))
        # a nice puppet
    """
    def __init__(self, config, args, concept_embedding_mat):
        super(FSNet_ner_ForQuestionAnswering, self).__init__(config, args, concept_embedding_mat)
        self.num_labels = config.num_labels
        
        '''1st Layer: BERT Encoding Layer'''
        self.bert = BertModel(config)

        '''2st Layer: Knowledge Integration Layer'''
        self.memory_embs = torch.nn.Embedding.from_pretrained(torch.Tensor(concept_embedding_mat), freeze=args.freeze_concept)

        concept_vocab_size = concept_embedding_mat.shape[0]
        concept_dim = concept_embedding_mat.shape[1]

        self.relavance_layer = torch.nn.Bilinear(config.hidden_size, concept_dim, 1)

        self.trilinear = torch.nn.Linear(3 * (config.hidden_size + concept_dim), 1)

        self.qa_outputs = torch.nn.Linear(6 * (config.hidden_size + concept_dim), config.num_labels)

        # for Named Entity task
        self.tag_num_labels = 4
        self.classifier = nn.Linear(config.hidden_size, self.tag_num_labels)

        self.init_weights()
    
    def self_matching(self, U, mask):
        '''self-matching layer.
            U, (b, len, d)
            mask, (b, len)
        '''
        b, len_, d = U.size()

        # direct info interaction
        ui = U.unsqueeze(2) # [b, len, 1, d]
        uj = U.unsqueeze(1) # [b, 1, len, d]
        dot = ui * uj  # [b, len, len, d]
        ui_expand = ui.expand(b, len_, len_, d).contiguous() # [b, len, len, d] 
        uj_expand = uj.expand(b, len_, len_, d).contiguous() # [b, len, len, d]
        R = self.trilinear(torch.cat((ui_expand, uj_expand, dot), -1)).squeeze(-1) # [b, len, len, 3 * d] -> [b, len, len, 1] -> [b, len, len]
        mask_expand = mask.unsqueeze(-1).expand(b, len_, len_).contiguous() # [b, len, len]
        R = INF * (mask_expand.float() - 1.0) + R   # mask here
        A = torch.nn.functional.softmax(R, -1) # [b, len, len]
        V = torch.bmm(A, U) # [b, len, d]

        # indirect info interaction
        A_hat = torch.bmm(A, A)  # [b, len, len]
        V_hat = torch.bmm(A_hat, U) # [b, len, d]
        O = torch.cat((U, V, U - V, U * V, V_hat, U - V_hat), -1) # [b, len, 6*d]
        return O


    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None,
                start_positions=None, end_positions=None, concept_ids=None, concept_masks=None, labels_ids=None):
        
        '''1st Layer: Bert Layer'''
        outputs = self.bert(input_ids,
                            attention_mask=attention_mask,
                            token_type_ids=token_type_ids,
                            position_ids=position_ids,
                            head_mask=head_mask,
                            inputs_embeds=inputs_embeds)

        sequence_output = outputs[0] # [batch_size, seq_length, hidden_size]
        batch_size, seq_length, hidden_size = sequence_output.size()

        # here has two tasks: (1) Name Entity recognition (2) knowledge Enhanced QA
        # task (1) below

        # sequence_output_for_tag = self.dropout(sequence_output)
        logits_tag = self.classifier(sequence_output)
        if labels_ids 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 # [batch_size * seq_len]
                active_logits = logits_tag.view(-1, self.tag_num_labels)[active_loss]
                active_labels = labels_ids.view(-1)[active_loss]
                tag_loss = loss_fct(active_logits, active_labels)
            else:
                tag_loss = loss_fct(logits_tag.view(-1, self.tag_num_labels), labels_ids.view(-1))

        # task (2) below

        '''2nd Layer: Knowledge Integration Layer'''
        concept_emb = self.memory_embs(concept_ids) # [batch_size, seq_length, num_concepts, cpnet_dim]
        num_concepts = concept_emb.size()[2]
        cpnet_dim = concept_emb.size()[3]
        
        # add sentinel vector
        # sentinel_vec = self.sentinel_vec.to(concept_emb.device)
        # sentinel_vec_expand = sentinel_vec.expand(batch_size, seq_length, cpnet_dim).unsqueeze(2).contiguous()  # [batch_size, seq_length, 1, cpnet_dim]
        
        # concept_emb_add = torch.cat((concept_emb, sentinel_vec_expand), 2)  # [batch_size, seq_length, num_concepts + 1, cpnet_dim]

        sequence_output_expand = sequence_output.unsqueeze(2).expand(
            batch_size, seq_length, num_concepts, hidden_size).contiguous() # [batch_size, seq_length, num_concepts, hidden_size]

        similarity_weight = self.relavance_layer(sequence_output_expand, concept_emb).squeeze(-1) # [batch_size, seq_length, num_concepts]
        
        # calc masked attention
        similarity_weight = INF * (concept_masks - 1.0) + similarity_weight
        similarity_weight_softmax = torch.nn.functional.softmax(similarity_weight, -1) # [batch_size, seq_length, num_concepts]

        concept_relavance = torch.bmm(similarity_weight_softmax.unsqueeze(2).reshape(batch_size * seq_length, 1, -1),
            concept_emb.reshape(batch_size * seq_length, num_concepts, -1))  # [batch_size * seq_length, 1, cpnet_dim]
        concept_relavance = concept_relavance.squeeze(1).reshape(batch_size, seq_length, -1) # [batch_size, seq_length, cpnet_dim]
        
        # mask and cat representation
        attention_mask_expand = attention_mask.unsqueeze(-1).expand(batch_size, seq_length, cpnet_dim).contiguous() # [batch_size, seq_length, cpnet_dim]
        concept_relavance = attention_mask_expand.float() * concept_relavance
        sequence_kg_output = torch.cat((sequence_output, concept_relavance), dim=-1) # [batch_size, seq_length, hidden_size + cpnet_dim]

        '''3nd Layer: Self-Matching Layer'''
        sequence_kg_output = self.self_matching(sequence_kg_output, attention_mask)

        # residule connection
        # sequence_kg_output = torch.cat((sequence_output, sequence_kg_output), dim=-1)

        '''4nd Layer: Output Layer'''
        logits = self.qa_outputs(sequence_kg_output)
        start_logits, end_logits = logits.split(1, dim=-1)
        start_logits = start_logits.squeeze(-1)
        end_logits = end_logits.squeeze(-1)

        outputs = (start_logits, end_logits,) + outputs[2:]
        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
            # add multi-task loss
            total_loss += tag_loss
            outputs = (total_loss,) + outputs

        return outputs  # (loss), start_logits, end_logits, (hidden_states), (attentions)


# here KT-NET support (wordnet and nell)
class KTNetForQuestionAnswering_both(BertPreTrainedModel):
    r"""
        **start_positions**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:
            Labels for position (index) of the start of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`).
            Position outside of the sequence are not taken into account for computing the loss.
        **end_positions**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:
            Labels for position (index) of the end of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`).
            Position outside of the sequence are not taken into account for computing the loss.

    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
        **loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Total span extraction loss is the sum of a Cross-Entropy for the start and end positions.
        **start_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length,)``
            Span-start scores (before SoftMax).
        **end_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length,)``
            Span-end scores (before SoftMax).
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

    Examples::

        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        model = BertForQuestionAnswering.from_pretrained('bert-large-uncased-whole-word-masking-finetuned-squad')
        question, text = "Who was <NAME>?", "<NAME> was a nice puppet"
        input_text = "[CLS] " + question + " [SEP] " + text + " [SEP]"
        input_ids = tokenizer.encode(input_text)
        token_type_ids = [0 if i <= input_ids.index(102) else 1 for i in range(len(input_ids))] 
        start_scores, end_scores = model(torch.tensor([input_ids]), token_type_ids=torch.tensor([token_type_ids]))
        all_tokens = tokenizer.convert_ids_to_tokens(input_ids)  
        print(' '.join(all_tokens[torch.argmax(start_scores) : torch.argmax(end_scores)+1]))
        # a nice puppet
    """
    def __init__(self, config, wn_concept_embedding_mat, nell_concept_embedding_mat):
        super(KTNetForQuestionAnswering_both, self).__init__(config, wn_concept_embedding_mat, nell_concept_embedding_mat)
        self.num_labels = config.num_labels
        
        '''1st Layer: BERT Encoding Layer'''
        self.bert = BertModel(config)

        '''2st Layer: Knowledge Integration Layer'''
        self.wn_memory_embs = torch.nn.Embedding.from_pretrained(torch.Tensor(wn_concept_embedding_mat), freeze=True)
        self.nell_memory_embs = torch.nn.Embedding.from_pretrained(torch.Tensor(nell_concept_embedding_mat), freeze=True)

        wn_concept_vocab_size = wn_concept_embedding_mat.shape[0]
        wn_concept_dim = wn_concept_embedding_mat.shape[1]

        nell_concept_vocab_size = nell_concept_embedding_mat.shape[0]
        nell_concept_dim = wn_concept_embedding_mat.shape[1]

        assert wn_concept_dim == nell_concept_dim
        self.cpnet_dim = wn_concept_dim

        # init a sentinel vector
        self.sentinel_vec = torch.nn.Parameter(torch.rand(self.cpnet_dim), requires_grad=False)

        self.wn_relavance_layer = torch.nn.Bilinear(config.hidden_size, self.cpnet_dim, 1)
        self.nell_relavance_layer = torch.nn.Bilinear(config.hidden_size, self.cpnet_dim, 1)

        self.trilinear = torch.nn.Linear(3 * (config.hidden_size + 2 * self.cpnet_dim), 1)

        self.qa_outputs = torch.nn.Linear(6 * (config.hidden_size + 2 * self.cpnet_dim), config.num_labels)

        # self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)

        self.init_weights()
    
    def self_matching(self, U, mask):
        '''self-matching layer.
            U, (b, len, d)
            mask, (b, len)
        '''
        b, len_, d = U.size()

        # direct info interaction
        ui = U.unsqueeze(2) # [b, len, 1, d]
        uj = U.unsqueeze(1) # [b, 1, len, d]
        dot = ui * uj  # [b, len, len, d]
        ui_expand = ui.expand(b, len_, len_, d).contiguous() # [b, len, len, d] 
        uj_expand = uj.expand(b, len_, len_, d).contiguous() # [b, len, len, d]
        R = self.trilinear(torch.cat((ui_expand, uj_expand, dot), -1)).squeeze(-1) # [b, len, len, 3 * d] -> [b, len, len, 1] -> [b, len, len]
        mask_expand = mask.unsqueeze(-1).expand(b, len_, len_).contiguous() # [b, len, len]
        R = INF * (mask_expand.float() - 1.0) + R   # mask here
        A = torch.nn.functional.softmax(R, -1) # [b, len, len]
        V = torch.bmm(A, U) # [b, len, d]

        # indirect info interaction
        A_hat = torch.bmm(A, A)  # [b, len, len]
        V_hat = torch.bmm(A_hat, U) # [b, len, d]
        O = torch.cat((U, V, U - V, U * V, V_hat, U - V_hat), -1) # [b, len, 6*d]
        return O
    
    def knowledge_integration(self, concept_emb, sequence_output, sentinel_vec_expand, attention_mask_expand, relavance_layer):
        '''knowledge integration layer.
            concept_emb: [batch_size, seq_length, num_concepts, cpnet_dim],
            sequence_output: [batch_size, seq_length, hidden_size],
            sentinel_vec_expand: [batch_size, seq_length, 1, cpnet_dim],
            attention_mask_expand: [batch_size, seq_length, cpnet_dim]
        '''
        # get size
        batch_size, seq_length, num_concepts, cpnet_dim = concept_emb.size()
        _, _, hidden_size = sequence_output.size()
        # cat emb and sentinel
        concept_emb_add = torch.cat((concept_emb, sentinel_vec_expand), 2)  # [batch_size, seq_length, num_concepts + 1, cpnet_dim]
        # calc similarity weight
        sequence_output_expand = sequence_output.unsqueeze(2).expand(batch_size, seq_length, num_concepts + 1, hidden_size).contiguous() # [batch_size, seq_length, num_concepts + 1, hidden_size]
        similarity_weight = relavance_layer(sequence_output_expand, concept_emb_add).squeeze(-1) # [batch_size, seq_length, num_concepts + 1]
        similarity_weight_softmax = torch.nn.functional.softmax(similarity_weight, -1)
        # attention result
        concept_relavance = torch.bmm(similarity_weight_softmax.unsqueeze(2).reshape(batch_size * seq_length, 1, -1),
            concept_emb_add.reshape(batch_size*seq_length, num_concepts + 1, -1))  # [batch_size * seq_length, 1, cpnet_dim]
        concept_relavance = concept_relavance.squeeze(1).reshape(batch_size, seq_length, -1) # [batch_size, seq_length, cpnet_dim]
        # mask the concept
        concept_relavance = attention_mask_expand.float() * concept_relavance # [batch_size, seq_length, cpnet_dim]
        return concept_relavance

    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None,
                start_positions=None, end_positions=None, wn_concept_ids=None, nell_concept_ids=None):
        
        '''1st Layer: Bert Layer'''
        outputs = self.bert(input_ids,
                            attention_mask=attention_mask,
                            token_type_ids=token_type_ids,
                            position_ids=position_ids,
                            head_mask=head_mask,
                            inputs_embeds=inputs_embeds)

        sequence_output = outputs[0] # [batch_size, seq_length, hidden_size]
        batch_size, seq_length, hidden_size = sequence_output.size()

        '''2nd Layer: Knowledge Integration Layer'''
        # get concept emb
        wn_concept_emb = self.wn_memory_embs(wn_concept_ids) # [batch_size, seq_length, num_concepts, cpnet_dim]
        nell_concept_emb = self.nell_memory_embs(nell_concept_ids) # [batch_size, seq_length, num_concepts, cpnet_dim]

        # add sentinel vector
        sentinel_vec = self.sentinel_vec.to(sequence_output.device)
        sentinel_vec_expand = sentinel_vec.expand(batch_size, seq_length, self.cpnet_dim).unsqueeze(2).contiguous()  # [batch_size, seq_length, 1, cpnet_dim]
        
        # mask and cat representation
        attention_mask_expand = attention_mask.unsqueeze(-1).expand(batch_size, seq_length, self.cpnet_dim).contiguous() # [batch_size, seq_length, cpnet_dim]
        
        # knowledge integration
        wn_concept_relavance = self.knowledge_integration(wn_concept_emb, sequence_output, sentinel_vec_expand, attention_mask_expand, self.wn_relavance_layer)
        nell_concept_relavance = self.knowledge_integration(nell_concept_emb, sequence_output, sentinel_vec_expand, attention_mask_expand, self.nell_relavance_layer)

        sequence_output = torch.cat((sequence_output, wn_concept_relavance, nell_concept_relavance), dim=-1) # [batch_size, seq_length, hidden_size + cpnet_dim * 2]

        '''3nd Layer: Self-Matching Layer'''
        sequence_output = self.self_matching(sequence_output, attention_mask) # [batch_size, seq_length, 6 * (hidden_size + cpnet_dim * 2)]

        '''4nd Layer: Output Layer'''
        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)

        outputs = (start_logits, end_logits,) + outputs[2:]
        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
            outputs = (total_loss,) + outputs

        return outputs  # (loss), start_logits, end_logits, (hidden_states), (attentions)


# SKG-NET for QA models here
class SKGNetForQuestionAnswering(BertPreTrainedModel):
    r"""
        **start_positions**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:
            Labels for position (index) of the start of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`).
            Position outside of the sequence are not taken into account for computing the loss.
        **end_positions**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:
            Labels for position (index) of the end of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`).
            Position outside of the sequence are not taken into account for computing the loss.

    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
        **loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Total span extraction loss is the sum of a Cross-Entropy for the start and end positions.
        **start_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length,)``
            Span-start scores (before SoftMax).
        **end_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length,)``
            Span-end scores (before SoftMax).
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

    Examples::

        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        model = BertForQuestionAnswering.from_pretrained('bert-large-uncased-whole-word-masking-finetuned-squad')
        question, text = "Who was <NAME>?", "<NAME> was a nice puppet"
        input_text = "[CLS] " + question + " [SEP] " + text + " [SEP]"
        input_ids = tokenizer.encode(input_text)
        token_type_ids = [0 if i <= input_ids.index(102) else 1 for i in range(len(input_ids))] 
        start_scores, end_scores = model(torch.tensor([input_ids]), token_type_ids=torch.tensor([token_type_ids]))
        all_tokens = tokenizer.convert_ids_to_tokens(input_ids)  
        print(' '.join(all_tokens[torch.argmax(start_scores) : torch.argmax(end_scores)+1]))
        # a nice puppet
    """
    def __init__(self, config, device, args, entity_emb_mat, relation_emb_mat):
        super(SKGNetForQuestionAnswering, self).__init__(config, device, args, entity_emb_mat, relation_emb_mat)
        self.num_labels = config.num_labels
        self.device = device
        self.gnn_layers = args.gnn_layers
        
        '''1st Layer: BERT Encoding Layer'''
        self.bert = BertModel(config)

        '''2st Layer: Knowledge subgraph Layer'''
        self.entity_embs = torch.nn.Embedding.from_pretrained(torch.Tensor(entity_emb_mat), freeze=True)
        self.relation_embs = torch.nn.Embedding.from_pretrained(torch.Tensor(relation_emb_mat), freeze=True)

        concept_dim = entity_emb_mat.shape[1]

        '''3st Layer: Graph Attention Layer'''
        self.conv = GATConv(concept_dim, concept_dim, heads=8, dropout=config.hidden_dropout_prob)
        self.reflect_layer = torch.nn.Linear(8 * concept_dim, config.hidden_size)

        '''4st Layer: Output Layer'''
        self.gate_layer = torch.nn.Linear(config.hidden_size * 2, 1)
        self.qa_outputs = torch.nn.Linear(config.hidden_size, config.num_labels)

        # self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)

        self.init_weights()
    
    def to_device(self, tensor):
        '''move tensor to device'''
        if tensor is not None:
            return tensor.to(self.device)
        else:
            return None

    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None,
                start_positions=None, end_positions=None, core_entity_ids=None, graph_data_dict=None):
        # get sub graph info, and combine batch sub-graph
        # print("graph ids lens: {}".format(len(example_graph_info)))
        # nodes_ids_list = torch.tensor([node_id[0] for node_id in example_graph_info["nodes_ids"]]).numpy().tolist()
        # core_entity_ids_list = core_entity_ids[0].numpy().tolist()
        # core_entity_index_list = []
        # for entity in core_entity_ids_list:
        #     if entity in nodes_ids_list:
        #         idx = nodes_ids_list.index(entity)
        #     else:
        #         k = entity
        #         idx = 0
        #     core_entity_index_list.append(idx)

        # nodes_ids_list = sorted(list(set(nodes_ids_list)))
        # sort_nodes_ids_list = list(set(nodes_ids_list))
        # graph_data_dict = graph_data_dict
            
        # core_entity_index_list = [nodes_ids_list.index(entity) for entity in core_entity_ids_list]
        # core_entity_index = torch.tensor(core_entity_index_list, dtype=torch.long)
        # core_entity_index = self.to_device(core_entity_index)

        # set all tensors to device
        input_ids = self.to_device(input_ids)
        attention_mask = self.to_device(attention_mask)
        token_type_ids = self.to_device(token_type_ids)
        position_ids = self.to_device(position_ids)
        head_mask = self.to_device(head_mask)
        inputs_embeds = self.to_device(inputs_embeds)
        start_positions = self.to_device(start_positions)
        end_positions = self.to_device(end_positions)
        core_entity_ids = self.to_device(core_entity_ids)
        
        '''1st Layer: Bert Layer'''
        outputs = self.bert(input_ids,
                            attention_mask=attention_mask,
                            token_type_ids=token_type_ids,
                            position_ids=position_ids,
                            head_mask=head_mask,
                            inputs_embeds=inputs_embeds)

        sequence_output = outputs[0] # [batch_size, seq_length, hidden_size]
        batch_size, seq_length, hidden_size = sequence_output.size()
 
        '''2nd Layer: Knowledge subgraph Layer'''
        # # get sub graph info, and combine batch sub-graph
        # x_ids = torch.tensor([node_id[0] for node_id in example_graph_info["nodes_ids"]], dtype=torch.long)
        # x_ids = self.to_device(x_ids)
        # x = self.entity_embs(x_ids)  # [num_nodes, node_features]
        # x = self.to_device(x)

        # edges_ids = torch.tensor(example_graph_info["edges_ids"], dtype=torch.long) # [2, num_edges]
        # edges_ids = self.to_device(edges_ids)

        # edges_attr_ids = torch.tensor([edge_attr_id[0] for edge_attr_id in example_graph_info["edges_attr_ids"]], dtype=torch.long)
        # edges_attr_ids = self.to_device(edges_attr_ids)
        # edges_attr = self.relation_embs(edges_attr_ids) # [num_edges, edge_features]

        # graph = Data(x=x, edge_index=edges_ids.t().contiguous())

        nodes_features = graph_data_dict["nodes_features"] # [b, nodes_nums, feature_dim]
        edges_index = graph_data_dict["edges_index"] # [b, edges_nums, 2]
        edges_attr = graph_data_dict["edges_attr"] # [b, edges_nums, feature_dim]

        graph_list = []
        for i in range(batch_size):
            nodes_features_i = nodes_features[i]
            edges_index_i = edges_index[i]
            edges_attr_i = edges_attr[i]
            graph = Data(x=nodes_features_i, edge_index=edges_index_i, edge_attr=edges_attr_i).to(self.device)
            graph_list.append(graph)
        graph_batch = Batch.from_data_list(graph_list)
        
        '''3nd Layer: Graph Attention Layer'''
        out = graph_batch.x
        edge_emb = graph_batch.edge_attr
        # out = self.conv(graph_batch.x, graph_batch.edge_index)
        for i in range(self.gnn_layers):
            out = self.conv(out, graph_batch.edge_index, edge_emb)
        
        out_dim = out.size()[-1]
        out = out.reshape(batch_size, -1, out_dim)    # [batch_size, num_nodes, node_features]

        # find updated core entity representation
        core_entity_emb_list = []
        for i in range(batch_size):
            core_entity_emb = torch.index_select(out[i], 0, core_entity_ids[i])
            core_entity_emb_list.append(core_entity_emb)
        core_entity_embs = torch.stack(core_entity_emb_list).to(self.device) # [batch_size, seq_len, node_features]
        # core_entity_embs = torch.tensor(core_entity_emb_list).to(self.device) 

        # concat bert output and kg output
        kg_emb_output = self.reflect_layer(core_entity_embs)
        gate_W = torch.sigmoid(self.gate_layer(torch.cat((sequence_output, kg_emb_output), dim=-1)))
        sequence_kg_output = gate_W * sequence_output + (1 - gate_W) * kg_emb_output


        '''4nd Layer: Output Layer'''
        logits = self.qa_outputs(sequence_kg_output)
        start_logits, end_logits = logits.split(1, dim=-1)
        start_logits = start_logits.squeeze(-1)
        end_logits = end_logits.squeeze(-1)

        outputs = (start_logits, end_logits,) + outputs[2:]
        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
            outputs = (total_loss,) + outputs

        return outputs  # (loss), start_logits, end_logits, (hidden_states), (attentions)


# Graph Attention Networks here
class GATConv(MessagePassing):
    r"""The graph attentional operator from the `"Graph Attention Networks"
    <https://arxiv.org/abs/1710.10903>`_ paper
    .. math::
        \mathbf{x}^{\prime}_i = \alpha_{i,i}\mathbf{\Theta}\mathbf{x}_{i} +
        \sum_{j \in \mathcal{N}(i)} \alpha_{i,j}\mathbf{\Theta}\mathbf{x}_{j},
    where the attention coefficients :math:`\alpha_{i,j}` are computed as
    .. math::
        \alpha_{i,j} =
        \frac{
        \exp\left(\mathrm{LeakyReLU}\left(\mathbf{a}^{\top}
        [\mathbf{\Theta}\mathbf{x}_i \, \Vert \, \mathbf{\Theta}\mathbf{x}_j]
        \right)\right)}
        {\sum_{k \in \mathcal{N}(i) \cup \{ i \}}
        \exp\left(\mathrm{LeakyReLU}\left(\mathbf{a}^{\top}
        [\mathbf{\Theta}\mathbf{x}_i \, \Vert \, \mathbf{\Theta}\mathbf{x}_k]
        \right)\right)}.
    Args:
        in_channels (int): Size of each input sample.
        out_channels (int): Size of each output sample.
        heads (int, optional): Number of multi-head-attentions.
            (default: :obj:`1`)
        concat (bool, optional): If set to :obj:`False`, the multi-head
            attentions are averaged instead of concatenated.
            (default: :obj:`True`)
        negative_slope (float, optional): LeakyReLU angle of the negative
            slope. (default: :obj:`0.2`)
        dropout (float, optional): Dropout probability of the normalized
            attention coefficients which exposes each node to a stochastically
            sampled neighborhood during training. (default: :obj:`0`)
        bias (bool, optional): If set to :obj:`False`, the layer will not learn
            an additive bias. (default: :obj:`True`)
        **kwargs (optional): Additional arguments of
            :class:`torch_geometric.nn.conv.MessagePassing`.
    """
    def __init__(self, in_channels, out_channels, heads=1, concat=True,
                 negative_slope=0.2, dropout=0, bias=True, **kwargs):
        super(GATConv, self).__init__(aggr='add', **kwargs)

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.heads = heads
        self.concat = concat
        self.negative_slope = negative_slope
        self.dropout = dropout

        self.weight_h = Parameter(torch.Tensor(in_channels,
                                             heads * out_channels))
        self.weight_t = Parameter(torch.Tensor(in_channels,
                                             heads * out_channels))
        self.weight_r = Parameter(torch.Tensor(in_channels,
                                             heads * out_channels))
        
        self.att = Parameter(torch.Tensor(1, heads, out_channels))

        if bias and concat:
            self.bias = Parameter(torch.Tensor(heads * out_channels))
        elif bias and not concat:
            self.bias = Parameter(torch.Tensor(out_channels))
        else:
            self.register_parameter('bias', None)

        self.reset_parameters()

    def reset_parameters(self):
        glorot(self.weight_h)
        glorot(self.weight_t)
        glorot(self.weight_r)
        glorot(self.att)
        zeros(self.bias)

    def forward(self, x, edge_index, edge_attr):
        """"""
        # x : [N, in_channels]
        # edge_index: [2, E]
        # edge_attr: [E, attr_dim]
        
        # if torch.is_tensor(x):
        #     edge_index, _ = remove_self_loops(edge_index)
        #     edge_index, _ = add_self_loops(edge_index,
        #                                    num_nodes=x.size(self.node_dim))

        # edge feature
        edge_attr = edge_attr.unsqueeze(-1) if edge_attr.dim() == 1 else edge_attr

        # if torch.is_tensor(x):
        #     x = torch.matmul(x, self.weight)
        # else:
        #     x = (None if x[0] is None else torch.matmul(x[0], self.weight),
        #          None if x[1] is None else torch.matmul(x[1], self.weight))

        return self.propagate(edge_index, x=x, edge_attr=edge_attr)

    def message(self, edge_index_i, x_i, x_j, edge_attr):
        # Compute attention coefficients.
        x_i = torch.matmul(x_i, self.weight_t) # [N, num_head * out_channels]
        x_j = torch.matmul(x_j, self.weight_h) # [N, num_head * out_channels]
        e_ij = torch.matmul(edge_attr, self.weight_r) # [N, num_head * out_channels]

        x_i = x_i.view(-1, self.heads, self.out_channels)   # [N, num_heads, out_channels]
        x_j = x_j.view(-1, self.heads, self.out_channels)   # [N, num_heads, out_channels]
        e_ij = e_ij.view(-1, self.heads, self.out_channels) # [N, num_heads, out_channels]

        # alpha = (torch.cat([x_i, x_j], dim=-1) * self.att).sum(dim=-1)

        alpha = ((e_ij * F.leaky_relu(x_i + x_j, self.negative_slope)) * self.att).sum(-1)  # [N, num_heads]

        # alpha = F.leaky_relu(alpha, self.negative_slope)
        alpha = softmax(alpha, edge_index_i)

        # Sample attention coefficients stochastically.
        alpha = F.dropout(alpha, p=self.dropout, training=self.training)

        return x_j * alpha.view(-1, self.heads, 1)

    def update(self, aggr_out):
        if self.concat is True:
            aggr_out = aggr_out.view(-1, self.heads * self.out_channels)
        else:
            aggr_out = aggr_out.mean(dim=1)

        if self.bias is not None:
            aggr_out = aggr_out + self.bias
        return aggr_out

    def __repr__(self):
        return '{}({}, {}, heads={})'.format(self.__class__.__name__,
                                             self.in_channels,
                                             self.out_channels, self.heads)
# 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 logging
import math
import os
import sys

import torch
from torch import nn
from torch.nn import CrossEntropyLoss, MSELoss
import torch.nn.functional as F

from .modeling_utils import PreTrainedModel, prune_linear_layer
from .configuration_bert import BertConfig
from .file_utils import add_start_docstrings

# from GAT here
from torch_geometric.data import Data, Batch
from torch_geometric.nn import GCNConv, GATConv
from torch.nn import Parameter
from torch_scatter import scatter_add
from torch_geometric.nn.conv import MessagePassing
from torch_geometric.utils import add_remaining_self_loops, remove_self_loops, add_self_loops, softmax
from torch_geometric.nn.inits import glorot, zeros



import random
import numpy as np

INF = 1e30 # 定义正无穷

# set seed here(set the same with your run_xxx.py file)
# random.seed(42)
# np.random.seed(42)
# torch.manual_seed(42)
# torch.cuda.manual_seed(42)

logger = logging.getLogger(__name__)

BERT_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://s3.amazonaws.com/models.huggingface.co/bert/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",
    'bert-base-german-dbmdz-cased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-german-dbmdz-cased-pytorch_model.bin",
    'bert-base-german-dbmdz-uncased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-german-dbmdz-uncased-pytorch_model.bin",
}


def load_tf_weights_in_bert(model, config, tf_checkpoint_path):
    """ Load tf checkpoints in a pytorch model.
    """
    try:
        import re
        import numpy as np
        import tensorflow as tf
    except ImportError:
        logger.error("Loading a TensorFlow model 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)
    logger.info("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:
        logger.info("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):
            logger.info("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:
                    logger.info("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
        logger.info("Initialize PyTorch weight {}".format(name))
        pointer.data = torch.from_numpy(array)
    return model


def gelu(x):
    """ Original Implementation of the gelu activation function in Google Bert repo when initially created.
        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 gelu_new(x):
    """ Implementation of the gelu activation function currently in Google Bert repo (identical to OpenAI GPT).
        Also see https://arxiv.org/abs/1606.08415
    """
    return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))


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


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


BertLayerNorm = torch.nn.LayerNorm


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=None, token_type_ids=None, position_ids=None, inputs_embeds=None):
        if input_ids is not None:
            input_shape = input_ids.size()
        else:
            input_shape = inputs_embeds.size()[:-1]

        seq_length = input_shape[1]
        device = input_ids.device if input_ids is not None else inputs_embeds.device
        if position_ids is None:
            position_ids = torch.arange(seq_length, dtype=torch.long, device=device)
            position_ids = position_ids.unsqueeze(0).expand(input_shape)
        if token_type_ids is None:
            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)

        if inputs_embeds is None:
            inputs_embeds = self.word_embeddings(input_ids)
        position_embeddings = self.position_embeddings(position_ids)
        token_type_embeddings = self.token_type_embeddings(token_type_ids)

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


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.output_attentions = config.output_attentions

        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)
        x = x.view(*new_x_shape)
        return x.permute(0, 2, 1, 3)

    def forward(self, hidden_states, attention_mask=None, head_mask=None, encoder_hidden_states=None, encoder_attention_mask=None):
        mixed_query_layer = self.query(hidden_states)

        # If this is instantiated as a cross-attention module, the keys
        # and values come from an encoder; the attention mask needs to be
        # such that the encoder's padding tokens are not attended to.
        if encoder_hidden_states is not None:
            mixed_key_layer = self.key(encoder_hidden_states)
            mixed_value_layer = self.value(encoder_hidden_states)
            attention_mask = encoder_attention_mask
        else:
            mixed_key_layer = self.key(hidden_states)
            mixed_value_layer = self.value(hidden_states)

        query_layer = self.transpose_for_scores(mixed_query_layer)
        key_layer = self.transpose_for_scores(mixed_key_layer)
        value_layer = self.transpose_for_scores(mixed_value_layer)

        # 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))
        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
        if attention_mask is not None:
            # 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)

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

        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)

        outputs = (context_layer, attention_probs) if self.output_attentions else (context_layer,)
        return outputs


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):
        super(BertAttention, self).__init__()
        self.self = BertSelfAttention(config)
        self.output = BertSelfOutput(config)
        self.pruned_heads = set()

    def prune_heads(self, heads):
        if len(heads) == 0:
            return
        mask = torch.ones(self.self.num_attention_heads, self.self.attention_head_size)
        heads = set(heads) - self.pruned_heads  # Convert to set and emove already pruned heads
        for head in heads:
            # Compute how many pruned heads are before the head and move the index accordingly
            head = head - sum(1 if h < head else 0 for h in self.pruned_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 and store pruned heads
        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
        self.pruned_heads = self.pruned_heads.union(heads)

    def forward(self, hidden_states, attention_mask=None, head_mask=None, encoder_hidden_states=None, encoder_attention_mask=None):
        self_outputs = self.self(hidden_states, attention_mask, head_mask, encoder_hidden_states, encoder_attention_mask)
        attention_output = self.output(self_outputs[0], hidden_states)
        outputs = (attention_output,) + self_outputs[1:]  # add attentions if we output them
        return outputs


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):
    def __init__(self, config):
        super(BertLayer, self).__init__()
        self.attention = BertAttention(config)
        self.is_decoder = config.is_decoder
        if self.is_decoder:
            self.crossattention = BertAttention(config)
        self.intermediate = BertIntermediate(config)
        self.output = BertOutput(config)

    def forward(self, hidden_states, attention_mask=None, head_mask=None, encoder_hidden_states=None, encoder_attention_mask=None):
        self_attention_outputs = self.attention(hidden_states, attention_mask, head_mask)
        attention_output = self_attention_outputs[0]
        outputs = self_attention_outputs[1:]  # add self attentions if we output attention weights

        if self.is_decoder and encoder_hidden_states is not None:
            cross_attention_outputs = self.crossattention(attention_output, attention_mask, head_mask, encoder_hidden_states, encoder_attention_mask)
            attention_output = cross_attention_outputs[0]
            outputs = outputs + cross_attention_outputs[1:]  # add cross attentions if we output attention weights

        intermediate_output = self.intermediate(attention_output)
        layer_output = self.output(intermediate_output, attention_output)
        outputs = (layer_output,) + outputs
        return outputs


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

    def forward(self, hidden_states, attention_mask=None, head_mask=None, encoder_hidden_states=None, encoder_attention_mask=None):
        all_hidden_states = ()
        all_attentions = ()
        for i, layer_module in enumerate(self.layer):
            if self.output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            layer_outputs = layer_module(hidden_states, attention_mask, head_mask[i], encoder_hidden_states, encoder_attention_mask)
            hidden_states = layer_outputs[0]

            if self.output_attentions:
                all_attentions = all_attentions + (layer_outputs[1],)

        # Add last layer
        if self.output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        outputs = (hidden_states,)
        if self.output_hidden_states:
            outputs = outputs + (all_hidden_states,)
        if self.output_attentions:
            outputs = outputs + (all_attentions,)
        return outputs  # last-layer hidden state, (all hidden states), (all attentions)


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

    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):
        super(BertOnlyMLMHead, self).__init__()
        self.predictions = BertLMPredictionHead(config)

    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):
        super(BertPreTrainingHeads, self).__init__()
        self.predictions = BertLMPredictionHead(config)
        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(PreTrainedModel):
    """ An abstract class to handle weights initialization and
        a simple interface for dowloading and loading pretrained models.
    """
    config_class = BertConfig
    pretrained_model_archive_map = BERT_PRETRAINED_MODEL_ARCHIVE_MAP
    load_tf_weights = load_tf_weights_in_bert
    base_model_prefix = "bert"

    def _init_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_()


BERT_START_DOCSTRING = r"""    The BERT model was proposed in
    `BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding`_
    by <NAME>, <NAME>, <NAME> and <NAME>. It's a bidirectional transformer
    pre-trained using a combination of masked language modeling objective and next sentence prediction
    on a large corpus comprising the Toronto Book Corpus and Wikipedia.

    This model is a PyTorch `torch.nn.Module`_ sub-class. Use it as a regular PyTorch Module and
    refer to the PyTorch documentation for all matter related to general usage and behavior.

    .. _`BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding`:
        https://arxiv.org/abs/1810.04805

    .. _`torch.nn.Module`:
        https://pytorch.org/docs/stable/nn.html#module

    Parameters:
        config (:class:`~transformers.BertConfig`): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the configuration.
            Check out the :meth:`~transformers.PreTrainedModel.from_pretrained` method to load the model weights.
"""

BERT_INPUTS_DOCSTRING = r"""
    Inputs:
        **input_ids**: ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``:
            Indices of input sequence tokens in the vocabulary.
            To match pre-training, BERT input sequence should be formatted with [CLS] and [SEP] tokens as follows:

            (a) For sequence pairs:

                ``tokens:         [CLS] is this jack ##son ##ville ? [SEP] no it is not . [SEP]``

                ``token_type_ids:   0   0  0    0    0     0       0   0   1  1  1  1   1   1``

            (b) For single sequences:

                ``tokens:         [CLS] the dog is hairy . [SEP]``

                ``token_type_ids:   0   0   0   0  0     0   0``

            Bert is a model with absolute position embeddings so it's usually advised to pad the inputs on
            the right rather than the left.

            Indices can be obtained using :class:`transformers.BertTokenizer`.
            See :func:`transformers.PreTrainedTokenizer.encode` and
            :func:`transformers.PreTrainedTokenizer.convert_tokens_to_ids` for details.
        **attention_mask**: (`optional`) ``torch.FloatTensor`` of shape ``(batch_size, sequence_length)``:
            Mask to avoid performing attention on padding token indices.
            Mask values selected in ``[0, 1]``:
            ``1`` for tokens that are NOT MASKED, ``0`` for MASKED tokens.
        **token_type_ids**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``:
            Segment token indices to indicate first and second portions of the inputs.
            Indices are selected in ``[0, 1]``: ``0`` corresponds to a `sentence A` token, ``1``
            corresponds to a `sentence B` token
            (see `BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding`_ for more details).
        **position_ids**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``:
            Indices of positions of each input sequence tokens in the position embeddings.
            Selected in the range ``[0, config.max_position_embeddings - 1]``.
        **head_mask**: (`optional`) ``torch.FloatTensor`` of shape ``(num_heads,)`` or ``(num_layers, num_heads)``:
            Mask to nullify selected heads of the self-attention modules.
            Mask values selected in ``[0, 1]``:
            ``1`` indicates the head is **not masked**, ``0`` indicates the head is **masked**.
        **inputs_embeds**: (`optional`) ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, embedding_dim)``:
            Optionally, instead of passing ``input_ids`` you can choose to directly pass an embedded representation.
            This is useful if you want more control over how to convert `input_ids` indices into associated vectors
            than the model's internal embedding lookup matrix.
        **encoder_hidden_states**: (`optional`) ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, hidden_size)``:
            Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if the model
            is configured as a decoder.
        **encoder_attention_mask**: (`optional`) ``torch.FloatTensor`` of shape ``(batch_size, sequence_length)``:
            Mask to avoid performing attention on the padding token indices of the encoder input. This mask
            is used in the cross-attention if the model is configured as a decoder.
            Mask values selected in ``[0, 1]``:
            ``1`` for tokens that are NOT MASKED, ``0`` for MASKED tokens.
"""

@add_start_docstrings("The bare Bert Model transformer outputting raw hidden-states without any specific head on top.",
                      BERT_START_DOCSTRING, BERT_INPUTS_DOCSTRING)
class BertModel(BertPreTrainedModel):
    r"""
    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
        **last_hidden_state**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, hidden_size)``
            Sequence of hidden-states at the output of the last layer of the model.
        **pooler_output**: ``torch.FloatTensor`` of shape ``(batch_size, hidden_size)``
            Last layer hidden-state of the first token of the sequence (classification token)
            further processed by a Linear layer and a Tanh activation function. The Linear
            layer weights are trained from the next sentence prediction (classification)
            objective during Bert pretraining. This output is usually *not* a good summary
            of the semantic content of the input, you're often better with averaging or pooling
            the sequence of hidden-states for the whole input sequence.
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

    Examples::

        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        model = BertModel.from_pretrained('bert-base-uncased')
        input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute")).unsqueeze(0)  # Batch size 1
        outputs = model(input_ids)
        last_hidden_states = outputs[0]  # The last hidden-state is the first element of the output tuple

    """
    def __init__(self, config):
        super(BertModel, self).__init__(config)
        self.config = config

        self.embeddings = BertEmbeddings(config)
        self.encoder = BertEncoder(config)
        self.pooler = BertPooler(config)

        self.init_weights()

    def get_input_embeddings(self):
        return self.embeddings.word_embeddings

    def set_input_embeddings(self, value):
        self.embeddings.word_embeddings = value

    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}
            See base class PreTrainedModel
        """
        for layer, heads in heads_to_prune.items():
            self.encoder.layer[layer].attention.prune_heads(heads)

    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None,
                head_mask=None, inputs_embeds=None, encoder_hidden_states=None, encoder_attention_mask=None):
        """ Forward pass on the Model.

        The model can behave as an encoder (with only self-attention) as well
        as a decoder, in which case a layer of cross-attention is added between
        the self-attention layers, following the architecture described in `Attention is all you need`_ by <NAME>,
        <NAME>, <NAME>, <NAME>, <NAME>, <NAME>, <NAME> and <NAME>.

        To behave as an decoder the model needs to be initialized with the
        `is_decoder` argument of the configuration set to `True`; an
        `encoder_hidden_states` is expected as an input to the forward pass.

        .. _`Attention is all you need`:
            https://arxiv.org/abs/1706.03762

        """
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
            input_shape = input_ids.size()
        elif inputs_embeds is not None:
            input_shape = inputs_embeds.size()[:-1]
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        device = input_ids.device if input_ids is not None else inputs_embeds.device

        if attention_mask is None:
            attention_mask = torch.ones(input_shape, device=device)
        if encoder_attention_mask is None:
            encoder_attention_mask = torch.ones(input_shape, device=device)
        if token_type_ids is None:
            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)

        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
        # ourselves in which case we just need to make it broadcastable to all heads.
        if attention_mask.dim() == 3:
            extended_attention_mask = attention_mask[:, None, :, :]

        # Provided a padding mask of dimensions [batch_size, seq_length]
        # - if the model is a decoder, apply a causal mask in addition to the padding mask
        # - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length]
        if attention_mask.dim() == 2:
            if self.config.is_decoder:
                batch_size, seq_length = input_shape
                seq_ids = torch.arange(seq_length, device=device)
                causal_mask = seq_ids[None, None, :].repeat(batch_size, seq_length, 1) <= seq_ids[None, :, None]
                extended_attention_mask = causal_mask[:, None, :, :] * attention_mask[:, None, None, :]
            else:
                extended_attention_mask = attention_mask[:, None, None, :]

        # 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

        # If a 2D ou 3D attention mask is provided for the cross-attention
        # we need to make broadcastabe to [batch_size, num_heads, seq_length, seq_length]
        if encoder_attention_mask.dim() == 3:
            encoder_extended_attention_mask = encoder_attention_mask[:, None, :, :]
        if encoder_attention_mask.dim() == 2:
            encoder_extended_attention_mask = encoder_attention_mask[:, None, None, :]

        encoder_extended_attention_mask = encoder_extended_attention_mask.to(dtype=next(self.parameters()).dtype)  # fp16 compatibility
        encoder_extended_attention_mask = (1.0 - encoder_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(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

        embedding_output = self.embeddings(input_ids=input_ids, position_ids=position_ids, token_type_ids=token_type_ids, inputs_embeds=inputs_embeds)
        encoder_outputs = self.encoder(embedding_output,
                                       attention_mask=extended_attention_mask,
                                       head_mask=head_mask,
                                       encoder_hidden_states=encoder_hidden_states,
                                       encoder_attention_mask=encoder_extended_attention_mask)
        sequence_output = encoder_outputs[0]
        pooled_output = self.pooler(sequence_output)

        outputs = (sequence_output, pooled_output,) + encoder_outputs[1:]  # add hidden_states and attentions if they are here
        return outputs  # sequence_output, pooled_output, (hidden_states), (attentions)


@add_start_docstrings("""Bert Model with two heads on top as done during the pre-training:
                       a `masked language modeling` head and a `next sentence prediction (classification)` head. """,
                      BERT_START_DOCSTRING,
                      BERT_INPUTS_DOCSTRING)
class BertForPreTraining(BertPreTrainedModel):
    r"""
        **masked_lm_labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``:
            Labels for computing the masked language modeling loss.
            Indices should be in ``[-1, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring)
            Tokens with indices set to ``-1`` are ignored (masked), the loss is only computed for the tokens with labels
            in ``[0, ..., config.vocab_size]``
        **next_sentence_label**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:
            Labels for computing the next sequence prediction (classification) loss. Input should be a sequence pair (see ``input_ids`` docstring)
            Indices should be in ``[0, 1]``.
            ``0`` indicates sequence B is a continuation of sequence A,
            ``1`` indicates sequence B is a random sequence.

    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
        **loss**: (`optional`, returned when both ``masked_lm_labels`` and ``next_sentence_label`` are provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Total loss as the sum of the masked language modeling loss and the next sequence prediction (classification) loss.
        **prediction_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, config.vocab_size)``
            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
        **seq_relationship_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, 2)``
            Prediction scores of the next sequence prediction (classification) head (scores of True/False continuation before SoftMax).
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

    Examples::

        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        model = BertForPreTraining.from_pretrained('bert-base-uncased')
        input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute")).unsqueeze(0)  # Batch size 1
        outputs = model(input_ids)
        prediction_scores, seq_relationship_scores = outputs[:2]

    """
    def __init__(self, config):
        super(BertForPreTraining, self).__init__(config)

        self.bert = BertModel(config)
        self.cls = BertPreTrainingHeads(config)

        self.init_weights()

    def get_output_embeddings(self):
        return self.cls.predictions.decoder

    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None,
                masked_lm_labels=None, next_sentence_label=None):

        outputs = self.bert(input_ids,
                            attention_mask=attention_mask,
                            token_type_ids=token_type_ids,
                            position_ids=position_ids,
                            head_mask=head_mask,
                            inputs_embeds=inputs_embeds)

        sequence_output, pooled_output = outputs[:2]
        prediction_scores, seq_relationship_score = self.cls(sequence_output, pooled_output)

        outputs = (prediction_scores, seq_relationship_score,) + outputs[2:]  # add hidden states and attention if they are here

        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
            outputs = (total_loss,) + outputs

        return outputs  # (loss), prediction_scores, seq_relationship_score, (hidden_states), (attentions)


@add_start_docstrings("""Bert Model with a `language modeling` head on top. """,
                      BERT_START_DOCSTRING,
                      BERT_INPUTS_DOCSTRING)
class BertForMaskedLM(BertPreTrainedModel):
    r"""
        **masked_lm_labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``:
            Labels for computing the masked language modeling loss.
            Indices should be in ``[-1, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring)
            Tokens with indices set to ``-1`` are ignored (masked), the loss is only computed for the tokens with labels
            in ``[0, ..., config.vocab_size]``
        **lm_labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``:
            Labels for computing the left-to-right language modeling loss (next word prediction).
            Indices should be in ``[-1, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring)
            Tokens with indices set to ``-1`` are ignored (masked), the loss is only computed for the tokens with labels
            in ``[0, ..., config.vocab_size]``

    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
        **masked_lm_loss**: (`optional`, returned when ``masked_lm_labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Masked language modeling loss.
        **ltr_lm_loss**: (`optional`, returned when ``lm_labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Next token prediction loss.
        **prediction_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, config.vocab_size)``
            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

    Examples::

        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        model = BertForMaskedLM.from_pretrained('bert-base-uncased')
        input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute")).unsqueeze(0)  # Batch size 1
        outputs = model(input_ids, masked_lm_labels=input_ids)
        loss, prediction_scores = outputs[:2]

    """
    def __init__(self, config):
        super(BertForMaskedLM, self).__init__(config)

        self.bert = BertModel(config)
        self.cls = BertOnlyMLMHead(config)

        self.init_weights()

    def get_output_embeddings(self):
        return self.cls.predictions.decoder

    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None,
                masked_lm_labels=None, encoder_hidden_states=None, encoder_attention_mask=None, lm_labels=None, ):

        outputs = self.bert(input_ids,
                            attention_mask=attention_mask,
                            token_type_ids=token_type_ids,
                            position_ids=position_ids,
                            head_mask=head_mask,
                            inputs_embeds=inputs_embeds,
                            encoder_hidden_states=encoder_hidden_states,
                            encoder_attention_mask=encoder_attention_mask)

        sequence_output = outputs[0]
        prediction_scores = self.cls(sequence_output)

        outputs = (prediction_scores,) + outputs[2:]  # Add hidden states and attention if they are here

        # Although this may seem awkward, BertForMaskedLM supports two scenarios:
        # 1. If a tensor that contains the indices of masked labels is provided,
        #    the cross-entropy is the MLM cross-entropy that measures the likelihood
        #    of predictions for masked words.
        # 2. If `lm_labels` is provided we are in a causal scenario where we
        #    try to predict the next token for each input in the decoder.
        if masked_lm_labels is not None:
            loss_fct = CrossEntropyLoss(ignore_index=-1)  # -1 index = padding token
            masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), masked_lm_labels.view(-1))
            outputs = (masked_lm_loss,) + outputs

        if lm_labels is not None:
            # we are doing next-token prediction; shift prediction scores and input ids by one
            prediction_scores = prediction_scores[:, :-1, :].contiguous()
            lm_labels = lm_labels[:, 1:].contiguous()
            loss_fct = CrossEntropyLoss(ignore_index=-1)
            ltr_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), lm_labels.view(-1))
            outputs = (ltr_lm_loss,) + outputs

        return outputs  # (masked_lm_loss), (ltr_lm_loss), prediction_scores, (hidden_states), (attentions)


@add_start_docstrings("""Bert Model with a `next sentence prediction (classification)` head on top. """,
                      BERT_START_DOCSTRING,
                      BERT_INPUTS_DOCSTRING)
class BertForNextSentencePrediction(BertPreTrainedModel):
    r"""
        **next_sentence_label**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:
            Labels for computing the next sequence prediction (classification) loss. Input should be a sequence pair (see ``input_ids`` docstring)
            Indices should be in ``[0, 1]``.
            ``0`` indicates sequence B is a continuation of sequence A,
            ``1`` indicates sequence B is a random sequence.

    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
        **loss**: (`optional`, returned when ``next_sentence_label`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Next sequence prediction (classification) loss.
        **seq_relationship_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, 2)``
            Prediction scores of the next sequence prediction (classification) head (scores of True/False continuation before SoftMax).
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

    Examples::

        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        model = BertForNextSentencePrediction.from_pretrained('bert-base-uncased')
        input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute")).unsqueeze(0)  # Batch size 1
        outputs = model(input_ids)
        seq_relationship_scores = outputs[0]

    """
    def __init__(self, config):
        super(BertForNextSentencePrediction, self).__init__(config)

        self.bert = BertModel(config)
        self.cls = BertOnlyNSPHead(config)

        self.init_weights()

    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None,
                next_sentence_label=None):

        outputs = self.bert(input_ids,
                            attention_mask=attention_mask,
                            token_type_ids=token_type_ids,
                            position_ids=position_ids,
                            head_mask=head_mask,
                            inputs_embeds=inputs_embeds)

        pooled_output = outputs[1]

        seq_relationship_score = self.cls(pooled_output)

        outputs = (seq_relationship_score,) + outputs[2:]  # add hidden states and attention if they are here
        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))
            outputs = (next_sentence_loss,) + outputs

        return outputs  # (next_sentence_loss), seq_relationship_score, (hidden_states), (attentions)


@add_start_docstrings("""Bert Model transformer with a sequence classification/regression head on top (a linear layer on top of
                      the pooled output) e.g. for GLUE tasks. """,
                      BERT_START_DOCSTRING,
                      BERT_INPUTS_DOCSTRING)
class BertForSequenceClassification(BertPreTrainedModel):
    r"""
        **labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:
            Labels for computing the sequence classification/regression loss.
            Indices should be in ``[0, ..., config.num_labels - 1]``.
            If ``config.num_labels == 1`` a regression loss is computed (Mean-Square loss),
            If ``config.num_labels > 1`` a classification loss is computed (Cross-Entropy).

    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
        **loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Classification (or regression if config.num_labels==1) loss.
        **logits**: ``torch.FloatTensor`` of shape ``(batch_size, config.num_labels)``
            Classification (or regression if config.num_labels==1) scores (before SoftMax).
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

    Examples::

        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        model = BertForSequenceClassification.from_pretrained('bert-base-uncased')
        input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute")).unsqueeze(0)  # Batch size 1
        labels = torch.tensor([1]).unsqueeze(0)  # Batch size 1
        outputs = model(input_ids, labels=labels)
        loss, logits = outputs[:2]

    """
    def __init__(self, config):
        super(BertForSequenceClassification, self).__init__(config)
        self.num_labels = config.num_labels

        self.bert = BertModel(config)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.classifier = nn.Linear(config.hidden_size, self.config.num_labels)

        self.init_weights()

    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None,
                position_ids=None, head_mask=None, inputs_embeds=None, labels=None):

        outputs = self.bert(input_ids,
                            attention_mask=attention_mask,
                            token_type_ids=token_type_ids,
                            position_ids=position_ids,
                            head_mask=head_mask,
                            inputs_embeds=inputs_embeds)

        pooled_output = outputs[1]

        pooled_output = self.dropout(pooled_output)
        logits = self.classifier(pooled_output)

        outputs = (logits,) + outputs[2:]  # add hidden states and attention if they are here

        if labels is not None:
            if self.num_labels == 1:
                #  We are doing regression
                loss_fct = MSELoss()
                loss = loss_fct(logits.view(-1), labels.view(-1))
            else:
                loss_fct = CrossEntropyLoss()
                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
            outputs = (loss,) + outputs

        return outputs  # (loss), logits, (hidden_states), (attentions)


@add_start_docstrings("""Bert Model with a multiple choice classification head on top (a linear layer on top of
                      the pooled output and a softmax) e.g. for RocStories/SWAG tasks. """,
                      BERT_START_DOCSTRING,
                      BERT_INPUTS_DOCSTRING)
class BertForMultipleChoice(BertPreTrainedModel):
    r"""
        **labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:
            Labels for computing the multiple choice classification loss.
            Indices should be in ``[0, ..., num_choices]`` where `num_choices` is the size of the second dimension
            of the input tensors. (see `input_ids` above)

    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
        **loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Classification loss.
        **classification_scores**: ``torch.FloatTensor`` of shape ``(batch_size, num_choices)`` where `num_choices` is the size of the second dimension
            of the input tensors. (see `input_ids` above).
            Classification scores (before SoftMax).
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

    Examples::

        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        model = BertForMultipleChoice.from_pretrained('bert-base-uncased')
        choices = ["Hello, my dog is cute", "Hello, my cat is amazing"]
        input_ids = torch.tensor([tokenizer.encode(s) for s in choices]).unsqueeze(0)  # Batch size 1, 2 choices
        labels = torch.tensor(1).unsqueeze(0)  # Batch size 1
        outputs = model(input_ids, labels=labels)
        loss, classification_scores = outputs[:2]

    """
    def __init__(self, config):
        super(BertForMultipleChoice, self).__init__(config)

        self.bert = BertModel(config)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.classifier = nn.Linear(config.hidden_size, 1)

        self.init_weights()

    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None,
                position_ids=None, head_mask=None, inputs_embeds=None, labels=None):
        num_choices = input_ids.shape[1]

        input_ids = input_ids.view(-1, input_ids.size(-1))
        attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None
        token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None
        position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None

        outputs = self.bert(input_ids,
                            attention_mask=attention_mask,
                            token_type_ids=token_type_ids,
                            position_ids=position_ids,
                            head_mask=head_mask,
                            inputs_embeds=inputs_embeds)

        pooled_output = outputs[1]

        pooled_output = self.dropout(pooled_output)
        logits = self.classifier(pooled_output)
        reshaped_logits = logits.view(-1, num_choices)

        outputs = (reshaped_logits,) + outputs[2:]  # add hidden states and attention if they are here

        if labels is not None:
            loss_fct = CrossEntropyLoss()
            loss = loss_fct(reshaped_logits, labels)
            outputs = (loss,) + outputs

        return outputs  # (loss), reshaped_logits, (hidden_states), (attentions)


@add_start_docstrings("""Bert Model with a token classification head on top (a linear layer on top of
                      the hidden-states output) e.g. for Named-Entity-Recognition (NER) tasks. """,
                      BERT_START_DOCSTRING,
                      BERT_INPUTS_DOCSTRING)
class BertForTokenClassification(BertPreTrainedModel):
    r"""
        **labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``:
            Labels for computing the token classification loss.
            Indices should be in ``[0, ..., config.num_labels - 1]``.

    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
        **loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Classification loss.
        **scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, config.num_labels)``
            Classification scores (before SoftMax).
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

    Examples::

        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        model = BertForTokenClassification.from_pretrained('bert-base-uncased')
        input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute")).unsqueeze(0)  # Batch size 1
        labels = torch.tensor([1] * input_ids.size(1)).unsqueeze(0)  # Batch size 1
        outputs = model(input_ids, labels=labels)
        loss, scores = outputs[:2]

    """
    def __init__(self, config):
        super(BertForTokenClassification, self).__init__(config)
        self.num_labels = config.num_labels

        self.bert = BertModel(config)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.classifier = nn.Linear(config.hidden_size, config.num_labels)

        self.init_weights()

    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None,
                position_ids=None, head_mask=None, inputs_embeds=None, labels=None):

        outputs = self.bert(input_ids,
                            attention_mask=attention_mask,
                            token_type_ids=token_type_ids,
                            position_ids=position_ids,
                            head_mask=head_mask,
                            inputs_embeds=inputs_embeds)

        sequence_output = outputs[0]

        sequence_output = self.dropout(sequence_output)
        logits = self.classifier(sequence_output)

        outputs = (logits,) + outputs[2:]  # add hidden states and attention if they are here
        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 # [batch_size * seq_len]
                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))
            outputs = (loss,) + outputs

        return outputs  # (loss), scores, (hidden_states), (attentions)


@add_start_docstrings("""Bert Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear layers on top of
                      the hidden-states output to compute `span start logits` and `span end logits`). """,
                      BERT_START_DOCSTRING,
                      BERT_INPUTS_DOCSTRING)
class BertForQuestionAnswering(BertPreTrainedModel):
    r"""
        **start_positions**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:
            Labels for position (index) of the start of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`).
            Position outside of the sequence are not taken into account for computing the loss.
        **end_positions**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:
            Labels for position (index) of the end of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`).
            Position outside of the sequence are not taken into account for computing the loss.

    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
        **loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Total span extraction loss is the sum of a Cross-Entropy for the start and end positions.
        **start_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length,)``
            Span-start scores (before SoftMax).
        **end_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length,)``
            Span-end scores (before SoftMax).
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

    Examples::

        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        model = BertForQuestionAnswering.from_pretrained('bert-large-uncased-whole-word-masking-finetuned-squad')
        question, text = "Who was <NAME>?", "<NAME> was a nice puppet"
        input_text = "[CLS] " + question + " [SEP] " + text + " [SEP]"
        input_ids = tokenizer.encode(input_text)
        token_type_ids = [0 if i <= input_ids.index(102) else 1 for i in range(len(input_ids))] 
        start_scores, end_scores = model(torch.tensor([input_ids]), token_type_ids=torch.tensor([token_type_ids]))
        all_tokens = tokenizer.convert_ids_to_tokens(input_ids)  
        print(' '.join(all_tokens[torch.argmax(start_scores) : torch.argmax(end_scores)+1]))
        # a nice puppet


    """
    def __init__(self, config):
        super(BertForQuestionAnswering, self).__init__(config)
        self.num_labels = config.num_labels

        self.bert = BertModel(config)
        self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)

        self.init_weights()

    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None,
                start_positions=None, end_positions=None):

        outputs = self.bert(input_ids,
                            attention_mask=attention_mask,
                            token_type_ids=token_type_ids,
                            position_ids=position_ids,
                            head_mask=head_mask,
                            inputs_embeds=inputs_embeds)

        sequence_output = outputs[0]

        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)

        outputs = (start_logits, end_logits,) + outputs[2:]
        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
            outputs = (total_loss,) + outputs

        return outputs  # (loss), start_logits, end_logits, (hidden_states), (attentions)


###############################################################################################################################################
# below is authored by yyhaker
class FSNetForQuestionAnswering(BertPreTrainedModel):
    r"""
        **start_positions**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:
            Labels for position (index) of the start of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`).
            Position outside of the sequence are not taken into account for computing the loss.
        **end_positions**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:
            Labels for position (index) of the end of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`).
            Position outside of the sequence are not taken into account for computing the loss.

    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
        **loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Total span extraction loss is the sum of a Cross-Entropy for the start and end positions.
        **start_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length,)``
            Span-start scores (before SoftMax).
        **end_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length,)``
            Span-end scores (before SoftMax).
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

    Examples::

        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        model = BertForQuestionAnswering.from_pretrained('bert-large-uncased-whole-word-masking-finetuned-squad')
        question, text = "Who was <NAME>?", "<NAME> was a nice puppet"
        input_text = "[CLS] " + question + " [SEP] " + text + " [SEP]"
        input_ids = tokenizer.encode(input_text)
        token_type_ids = [0 if i <= input_ids.index(102) else 1 for i in range(len(input_ids))] 
        start_scores, end_scores = model(torch.tensor([input_ids]), token_type_ids=torch.tensor([token_type_ids]))
        all_tokens = tokenizer.convert_ids_to_tokens(input_ids)  
        print(' '.join(all_tokens[torch.argmax(start_scores) : torch.argmax(end_scores)+1]))
        # a nice puppet
    """
    def __init__(self, config, args, concept_embedding_mat):
        super(FSNetForQuestionAnswering, self).__init__(config, args, concept_embedding_mat)
        self.num_labels = config.num_labels
        
        '''1st Layer: BERT Encoding Layer'''
        self.bert = BertModel(config)

        '''2st Layer: Knowledge Integration Layer'''
        self.memory_embs = torch.nn.Embedding.from_pretrained(torch.Tensor(concept_embedding_mat), freeze=args.freeze_concept)

        concept_vocab_size = concept_embedding_mat.shape[0]
        concept_dim = concept_embedding_mat.shape[1]

        # init a sentinel vector
        # self.sentinel_vec = torch.nn.Parameter(torch.rand(concept_dim), requires_grad=False)

        # init a sentinel vector
        # elf.sentinel_vec = torch.nn.Parameter(torch.zeros(concept_dim), requires_grad=True)

        self.relavance_layer = torch.nn.Bilinear(config.hidden_size, concept_dim, 1)

        self.trilinear = torch.nn.Linear(3 * (config.hidden_size + concept_dim), 1)

        self.qa_outputs = torch.nn.Linear(6 * (config.hidden_size + concept_dim), config.num_labels)

        # self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)

        self.init_weights()
    
    def self_matching(self, U, mask):
        '''self-matching layer.
            U, (b, len, d)
            mask, (b, len)
        '''
        b, len_, d = U.size()

        # direct info interaction
        ui = U.unsqueeze(2) # [b, len, 1, d]
        uj = U.unsqueeze(1) # [b, 1, len, d]
        dot = ui * uj  # [b, len, len, d]
        ui_expand = ui.expand(b, len_, len_, d).contiguous() # [b, len, len, d] 
        uj_expand = uj.expand(b, len_, len_, d).contiguous() # [b, len, len, d]
        R = self.trilinear(torch.cat((ui_expand, uj_expand, dot), -1)).squeeze(-1) # [b, len, len, 3 * d] -> [b, len, len, 1] -> [b, len, len]
        mask_expand = mask.unsqueeze(-1).expand(b, len_, len_).contiguous() # [b, len, len]
        R = INF * (mask_expand.float() - 1.0) + R   # mask here
        A = torch.nn.functional.softmax(R, -1) # [b, len, len]
        V = torch.bmm(A, U) # [b, len, d]

        # indirect info interaction
        A_hat = torch.bmm(A, A)  # [b, len, len]
        V_hat = torch.bmm(A_hat, U) # [b, len, d]
        O = torch.cat((U, V, U - V, U * V, V_hat, U - V_hat), -1) # [b, len, 6*d]
        return O


    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None,
                start_positions=None, end_positions=None, concept_ids=None, concept_masks=None):
        
        '''1st Layer: Bert Layer'''
        outputs = self.bert(input_ids,
                            attention_mask=attention_mask,
                            token_type_ids=token_type_ids,
                            position_ids=position_ids,
                            head_mask=head_mask,
                            inputs_embeds=inputs_embeds)

        sequence_output = outputs[0] # [batch_size, seq_length, hidden_size]
        batch_size, seq_length, hidden_size = sequence_output.size()

        '''2nd Layer: Knowledge Integration Layer'''
        concept_emb = self.memory_embs(concept_ids) # [batch_size, seq_length, num_concepts, cpnet_dim]
        num_concepts = concept_emb.size()[2]
        cpnet_dim = concept_emb.size()[3]
        
        # add sentinel vector
        # sentinel_vec = self.sentinel_vec.to(concept_emb.device)
        # sentinel_vec_expand = sentinel_vec.expand(batch_size, seq_length, cpnet_dim).unsqueeze(2).contiguous()  # [batch_size, seq_length, 1, cpnet_dim]
        
        # concept_emb_add = torch.cat((concept_emb, sentinel_vec_expand), 2)  # [batch_size, seq_length, num_concepts + 1, cpnet_dim]

        sequence_output_expand = sequence_output.unsqueeze(2).expand(
            batch_size, seq_length, num_concepts, hidden_size).contiguous() # [batch_size, seq_length, num_concepts, hidden_size]

        similarity_weight = self.relavance_layer(sequence_output_expand, concept_emb).squeeze(-1) # [batch_size, seq_length, num_concepts]
        
        # calc masked attention
        similarity_weight = INF * (concept_masks - 1.0) + similarity_weight
        similarity_weight_softmax = torch.nn.functional.softmax(similarity_weight, -1) # [batch_size, seq_length, num_concepts]

        concept_relavance = torch.bmm(similarity_weight_softmax.unsqueeze(2).reshape(batch_size * seq_length, 1, -1),
            concept_emb.reshape(batch_size * seq_length, num_concepts, -1))  # [batch_size * seq_length, 1, cpnet_dim]
        concept_relavance = concept_relavance.squeeze(1).reshape(batch_size, seq_length, -1) # [batch_size, seq_length, cpnet_dim]
        
        # mask and cat representation
        attention_mask_expand = attention_mask.unsqueeze(-1).expand(batch_size, seq_length, cpnet_dim).contiguous() # [batch_size, seq_length, cpnet_dim]
        concept_relavance = attention_mask_expand.float() * concept_relavance
        sequence_kg_output = torch.cat((sequence_output, concept_relavance), dim=-1) # [batch_size, seq_length, hidden_size + cpnet_dim]

        '''3nd Layer: Self-Matching Layer'''
        sequence_kg_output = self.self_matching(sequence_kg_output, attention_mask)

        # residule connection
        # sequence_kg_output = torch.cat((sequence_output, sequence_kg_output), dim=-1)

        '''4nd Layer: Output Layer'''
        logits = self.qa_outputs(sequence_kg_output)
        start_logits, end_logits = logits.split(1, dim=-1)
        start_logits = start_logits.squeeze(-1)
        end_logits = end_logits.squeeze(-1)

        outputs = (start_logits, end_logits,) + outputs[2:]
        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
            outputs = (total_loss,) + outputs

        return outputs  # (loss), start_logits, end_logits, (hidden_states), (attentions)


# here is FSNET_ner for QA
class FSNet_ner_ForQuestionAnswering(BertPreTrainedModel):
    r"""
        **start_positions**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:
            Labels for position (index) of the start of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`).
            Position outside of the sequence are not taken into account for computing the loss.
        **end_positions**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:
            Labels for position (index) of the end of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`).
            Position outside of the sequence are not taken into account for computing the loss.

    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
        **loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Total span extraction loss is the sum of a Cross-Entropy for the start and end positions.
        **start_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length,)``
            Span-start scores (before SoftMax).
        **end_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length,)``
            Span-end scores (before SoftMax).
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

    Examples::

        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        model = BertForQuestionAnswering.from_pretrained('bert-large-uncased-whole-word-masking-finetuned-squad')
        question, text = "Who was <NAME>?", "<NAME> was a nice puppet"
        input_text = "[CLS] " + question + " [SEP] " + text + " [SEP]"
        input_ids = tokenizer.encode(input_text)
        token_type_ids = [0 if i <= input_ids.index(102) else 1 for i in range(len(input_ids))] 
        start_scores, end_scores = model(torch.tensor([input_ids]), token_type_ids=torch.tensor([token_type_ids]))
        all_tokens = tokenizer.convert_ids_to_tokens(input_ids)  
        print(' '.join(all_tokens[torch.argmax(start_scores) : torch.argmax(end_scores)+1]))
        # a nice puppet
    """
    def __init__(self, config, args, concept_embedding_mat):
        super(FSNet_ner_ForQuestionAnswering, self).__init__(config, args, concept_embedding_mat)
        self.num_labels = config.num_labels
        
        '''1st Layer: BERT Encoding Layer'''
        self.bert = BertModel(config)

        '''2st Layer: Knowledge Integration Layer'''
        self.memory_embs = torch.nn.Embedding.from_pretrained(torch.Tensor(concept_embedding_mat), freeze=args.freeze_concept)

        concept_vocab_size = concept_embedding_mat.shape[0]
        concept_dim = concept_embedding_mat.shape[1]

        self.relavance_layer = torch.nn.Bilinear(config.hidden_size, concept_dim, 1)

        self.trilinear = torch.nn.Linear(3 * (config.hidden_size + concept_dim), 1)

        self.qa_outputs = torch.nn.Linear(6 * (config.hidden_size + concept_dim), config.num_labels)

        # for Named Entity task
        self.tag_num_labels = 4
        self.classifier = nn.Linear(config.hidden_size, self.tag_num_labels)

        self.init_weights()
    
    def self_matching(self, U, mask):
        '''self-matching layer.
            U, (b, len, d)
            mask, (b, len)
        '''
        b, len_, d = U.size()

        # direct info interaction
        ui = U.unsqueeze(2) # [b, len, 1, d]
        uj = U.unsqueeze(1) # [b, 1, len, d]
        dot = ui * uj  # [b, len, len, d]
        ui_expand = ui.expand(b, len_, len_, d).contiguous() # [b, len, len, d] 
        uj_expand = uj.expand(b, len_, len_, d).contiguous() # [b, len, len, d]
        R = self.trilinear(torch.cat((ui_expand, uj_expand, dot), -1)).squeeze(-1) # [b, len, len, 3 * d] -> [b, len, len, 1] -> [b, len, len]
        mask_expand = mask.unsqueeze(-1).expand(b, len_, len_).contiguous() # [b, len, len]
        R = INF * (mask_expand.float() - 1.0) + R   # mask here
        A = torch.nn.functional.softmax(R, -1) # [b, len, len]
        V = torch.bmm(A, U) # [b, len, d]

        # indirect info interaction
        A_hat = torch.bmm(A, A)  # [b, len, len]
        V_hat = torch.bmm(A_hat, U) # [b, len, d]
        O = torch.cat((U, V, U - V, U * V, V_hat, U - V_hat), -1) # [b, len, 6*d]
        return O


    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None,
                start_positions=None, end_positions=None, concept_ids=None, concept_masks=None, labels_ids=None):
        
        '''1st Layer: Bert Layer'''
        outputs = self.bert(input_ids,
                            attention_mask=attention_mask,
                            token_type_ids=token_type_ids,
                            position_ids=position_ids,
                            head_mask=head_mask,
                            inputs_embeds=inputs_embeds)

        sequence_output = outputs[0] # [batch_size, seq_length, hidden_size]
        batch_size, seq_length, hidden_size = sequence_output.size()

        # here has two tasks: (1) Name Entity recognition (2) knowledge Enhanced QA
        # task (1) below

        # sequence_output_for_tag = self.dropout(sequence_output)
        logits_tag = self.classifier(sequence_output)
        if labels_ids 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 # [batch_size * seq_len]
                active_logits = logits_tag.view(-1, self.tag_num_labels)[active_loss]
                active_labels = labels_ids.view(-1)[active_loss]
                tag_loss = loss_fct(active_logits, active_labels)
            else:
                tag_loss = loss_fct(logits_tag.view(-1, self.tag_num_labels), labels_ids.view(-1))

        # task (2) below

        '''2nd Layer: Knowledge Integration Layer'''
        concept_emb = self.memory_embs(concept_ids) # [batch_size, seq_length, num_concepts, cpnet_dim]
        num_concepts = concept_emb.size()[2]
        cpnet_dim = concept_emb.size()[3]
        
        # add sentinel vector
        # sentinel_vec = self.sentinel_vec.to(concept_emb.device)
        # sentinel_vec_expand = sentinel_vec.expand(batch_size, seq_length, cpnet_dim).unsqueeze(2).contiguous()  # [batch_size, seq_length, 1, cpnet_dim]
        
        # concept_emb_add = torch.cat((concept_emb, sentinel_vec_expand), 2)  # [batch_size, seq_length, num_concepts + 1, cpnet_dim]

        sequence_output_expand = sequence_output.unsqueeze(2).expand(
            batch_size, seq_length, num_concepts, hidden_size).contiguous() # [batch_size, seq_length, num_concepts, hidden_size]

        similarity_weight = self.relavance_layer(sequence_output_expand, concept_emb).squeeze(-1) # [batch_size, seq_length, num_concepts]
        
        # calc masked attention
        similarity_weight = INF * (concept_masks - 1.0) + similarity_weight
        similarity_weight_softmax = torch.nn.functional.softmax(similarity_weight, -1) # [batch_size, seq_length, num_concepts]

        concept_relavance = torch.bmm(similarity_weight_softmax.unsqueeze(2).reshape(batch_size * seq_length, 1, -1),
            concept_emb.reshape(batch_size * seq_length, num_concepts, -1))  # [batch_size * seq_length, 1, cpnet_dim]
        concept_relavance = concept_relavance.squeeze(1).reshape(batch_size, seq_length, -1) # [batch_size, seq_length, cpnet_dim]
        
        # mask and cat representation
        attention_mask_expand = attention_mask.unsqueeze(-1).expand(batch_size, seq_length, cpnet_dim).contiguous() # [batch_size, seq_length, cpnet_dim]
        concept_relavance = attention_mask_expand.float() * concept_relavance
        sequence_kg_output = torch.cat((sequence_output, concept_relavance), dim=-1) # [batch_size, seq_length, hidden_size + cpnet_dim]

        '''3nd Layer: Self-Matching Layer'''
        sequence_kg_output = self.self_matching(sequence_kg_output, attention_mask)

        # residule connection
        # sequence_kg_output = torch.cat((sequence_output, sequence_kg_output), dim=-1)

        '''4nd Layer: Output Layer'''
        logits = self.qa_outputs(sequence_kg_output)
        start_logits, end_logits = logits.split(1, dim=-1)
        start_logits = start_logits.squeeze(-1)
        end_logits = end_logits.squeeze(-1)

        outputs = (start_logits, end_logits,) + outputs[2:]
        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
            # add multi-task loss
            total_loss += tag_loss
            outputs = (total_loss,) + outputs

        return outputs  # (loss), start_logits, end_logits, (hidden_states), (attentions)


# here KT-NET support (wordnet and nell)
class KTNetForQuestionAnswering_both(BertPreTrainedModel):
    r"""
        **start_positions**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:
            Labels for position (index) of the start of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`).
            Position outside of the sequence are not taken into account for computing the loss.
        **end_positions**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:
            Labels for position (index) of the end of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`).
            Position outside of the sequence are not taken into account for computing the loss.

    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
        **loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Total span extraction loss is the sum of a Cross-Entropy for the start and end positions.
        **start_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length,)``
            Span-start scores (before SoftMax).
        **end_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length,)``
            Span-end scores (before SoftMax).
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

    Examples::

        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        model = BertForQuestionAnswering.from_pretrained('bert-large-uncased-whole-word-masking-finetuned-squad')
        question, text = "Who was <NAME>?", "<NAME> was a nice puppet"
        input_text = "[CLS] " + question + " [SEP] " + text + " [SEP]"
        input_ids = tokenizer.encode(input_text)
        token_type_ids = [0 if i <= input_ids.index(102) else 1 for i in range(len(input_ids))] 
        start_scores, end_scores = model(torch.tensor([input_ids]), token_type_ids=torch.tensor([token_type_ids]))
        all_tokens = tokenizer.convert_ids_to_tokens(input_ids)  
        print(' '.join(all_tokens[torch.argmax(start_scores) : torch.argmax(end_scores)+1]))
        # a nice puppet
    """
    def __init__(self, config, wn_concept_embedding_mat, nell_concept_embedding_mat):
        super(KTNetForQuestionAnswering_both, self).__init__(config, wn_concept_embedding_mat, nell_concept_embedding_mat)
        self.num_labels = config.num_labels
        
        '''1st Layer: BERT Encoding Layer'''
        self.bert = BertModel(config)

        '''2st Layer: Knowledge Integration Layer'''
        self.wn_memory_embs = torch.nn.Embedding.from_pretrained(torch.Tensor(wn_concept_embedding_mat), freeze=True)
        self.nell_memory_embs = torch.nn.Embedding.from_pretrained(torch.Tensor(nell_concept_embedding_mat), freeze=True)

        wn_concept_vocab_size = wn_concept_embedding_mat.shape[0]
        wn_concept_dim = wn_concept_embedding_mat.shape[1]

        nell_concept_vocab_size = nell_concept_embedding_mat.shape[0]
        nell_concept_dim = wn_concept_embedding_mat.shape[1]

        assert wn_concept_dim == nell_concept_dim
        self.cpnet_dim = wn_concept_dim

        # init a sentinel vector
        self.sentinel_vec = torch.nn.Parameter(torch.rand(self.cpnet_dim), requires_grad=False)

        self.wn_relavance_layer = torch.nn.Bilinear(config.hidden_size, self.cpnet_dim, 1)
        self.nell_relavance_layer = torch.nn.Bilinear(config.hidden_size, self.cpnet_dim, 1)

        self.trilinear = torch.nn.Linear(3 * (config.hidden_size + 2 * self.cpnet_dim), 1)

        self.qa_outputs = torch.nn.Linear(6 * (config.hidden_size + 2 * self.cpnet_dim), config.num_labels)

        # self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)

        self.init_weights()
    
    def self_matching(self, U, mask):
        '''self-matching layer.
            U, (b, len, d)
            mask, (b, len)
        '''
        b, len_, d = U.size()

        # direct info interaction
        ui = U.unsqueeze(2) # [b, len, 1, d]
        uj = U.unsqueeze(1) # [b, 1, len, d]
        dot = ui * uj  # [b, len, len, d]
        ui_expand = ui.expand(b, len_, len_, d).contiguous() # [b, len, len, d] 
        uj_expand = uj.expand(b, len_, len_, d).contiguous() # [b, len, len, d]
        R = self.trilinear(torch.cat((ui_expand, uj_expand, dot), -1)).squeeze(-1) # [b, len, len, 3 * d] -> [b, len, len, 1] -> [b, len, len]
        mask_expand = mask.unsqueeze(-1).expand(b, len_, len_).contiguous() # [b, len, len]
        R = INF * (mask_expand.float() - 1.0) + R   # mask here
        A = torch.nn.functional.softmax(R, -1) # [b, len, len]
        V = torch.bmm(A, U) # [b, len, d]

        # indirect info interaction
        A_hat = torch.bmm(A, A)  # [b, len, len]
        V_hat = torch.bmm(A_hat, U) # [b, len, d]
        O = torch.cat((U, V, U - V, U * V, V_hat, U - V_hat), -1) # [b, len, 6*d]
        return O
    
    def knowledge_integration(self, concept_emb, sequence_output, sentinel_vec_expand, attention_mask_expand, relavance_layer):
        '''knowledge integration layer.
            concept_emb: [batch_size, seq_length, num_concepts, cpnet_dim],
            sequence_output: [batch_size, seq_length, hidden_size],
            sentinel_vec_expand: [batch_size, seq_length, 1, cpnet_dim],
            attention_mask_expand: [batch_size, seq_length, cpnet_dim]
        '''
        # get size
        batch_size, seq_length, num_concepts, cpnet_dim = concept_emb.size()
        _, _, hidden_size = sequence_output.size()
        # cat emb and sentinel
        concept_emb_add = torch.cat((concept_emb, sentinel_vec_expand), 2)  # [batch_size, seq_length, num_concepts + 1, cpnet_dim]
        # calc similarity weight
        sequence_output_expand = sequence_output.unsqueeze(2).expand(batch_size, seq_length, num_concepts + 1, hidden_size).contiguous() # [batch_size, seq_length, num_concepts + 1, hidden_size]
        similarity_weight = relavance_layer(sequence_output_expand, concept_emb_add).squeeze(-1) # [batch_size, seq_length, num_concepts + 1]
        similarity_weight_softmax = torch.nn.functional.softmax(similarity_weight, -1)
        # attention result
        concept_relavance = torch.bmm(similarity_weight_softmax.unsqueeze(2).reshape(batch_size * seq_length, 1, -1),
            concept_emb_add.reshape(batch_size*seq_length, num_concepts + 1, -1))  # [batch_size * seq_length, 1, cpnet_dim]
        concept_relavance = concept_relavance.squeeze(1).reshape(batch_size, seq_length, -1) # [batch_size, seq_length, cpnet_dim]
        # mask the concept
        concept_relavance = attention_mask_expand.float() * concept_relavance # [batch_size, seq_length, cpnet_dim]
        return concept_relavance

    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None,
                start_positions=None, end_positions=None, wn_concept_ids=None, nell_concept_ids=None):
        
        '''1st Layer: Bert Layer'''
        outputs = self.bert(input_ids,
                            attention_mask=attention_mask,
                            token_type_ids=token_type_ids,
                            position_ids=position_ids,
                            head_mask=head_mask,
                            inputs_embeds=inputs_embeds)

        sequence_output = outputs[0] # [batch_size, seq_length, hidden_size]
        batch_size, seq_length, hidden_size = sequence_output.size()

        '''2nd Layer: Knowledge Integration Layer'''
        # get concept emb
        wn_concept_emb = self.wn_memory_embs(wn_concept_ids) # [batch_size, seq_length, num_concepts, cpnet_dim]
        nell_concept_emb = self.nell_memory_embs(nell_concept_ids) # [batch_size, seq_length, num_concepts, cpnet_dim]

        # add sentinel vector
        sentinel_vec = self.sentinel_vec.to(sequence_output.device)
        sentinel_vec_expand = sentinel_vec.expand(batch_size, seq_length, self.cpnet_dim).unsqueeze(2).contiguous()  # [batch_size, seq_length, 1, cpnet_dim]
        
        # mask and cat representation
        attention_mask_expand = attention_mask.unsqueeze(-1).expand(batch_size, seq_length, self.cpnet_dim).contiguous() # [batch_size, seq_length, cpnet_dim]
        
        # knowledge integration
        wn_concept_relavance = self.knowledge_integration(wn_concept_emb, sequence_output, sentinel_vec_expand, attention_mask_expand, self.wn_relavance_layer)
        nell_concept_relavance = self.knowledge_integration(nell_concept_emb, sequence_output, sentinel_vec_expand, attention_mask_expand, self.nell_relavance_layer)

        sequence_output = torch.cat((sequence_output, wn_concept_relavance, nell_concept_relavance), dim=-1) # [batch_size, seq_length, hidden_size + cpnet_dim * 2]

        '''3nd Layer: Self-Matching Layer'''
        sequence_output = self.self_matching(sequence_output, attention_mask) # [batch_size, seq_length, 6 * (hidden_size + cpnet_dim * 2)]

        '''4nd Layer: Output Layer'''
        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)

        outputs = (start_logits, end_logits,) + outputs[2:]
        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
            outputs = (total_loss,) + outputs

        return outputs  # (loss), start_logits, end_logits, (hidden_states), (attentions)


# SKG-NET for QA models here
class SKGNetForQuestionAnswering(BertPreTrainedModel):
    r"""
        **start_positions**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:
            Labels for position (index) of the start of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`).
            Position outside of the sequence are not taken into account for computing the loss.
        **end_positions**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:
            Labels for position (index) of the end of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`).
            Position outside of the sequence are not taken into account for computing the loss.

    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
        **loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Total span extraction loss is the sum of a Cross-Entropy for the start and end positions.
        **start_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length,)``
            Span-start scores (before SoftMax).
        **end_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length,)``
            Span-end scores (before SoftMax).
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

    Examples::

        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        model = BertForQuestionAnswering.from_pretrained('bert-large-uncased-whole-word-masking-finetuned-squad')
        question, text = "Who was <NAME>?", "<NAME> was a nice puppet"
        input_text = "[CLS] " + question + " [SEP] " + text + " [SEP]"
        input_ids = tokenizer.encode(input_text)
        token_type_ids = [0 if i <= input_ids.index(102) else 1 for i in range(len(input_ids))] 
        start_scores, end_scores = model(torch.tensor([input_ids]), token_type_ids=torch.tensor([token_type_ids]))
        all_tokens = tokenizer.convert_ids_to_tokens(input_ids)  
        print(' '.join(all_tokens[torch.argmax(start_scores) : torch.argmax(end_scores)+1]))
        # a nice puppet
    """
    def __init__(self, config, device, args, entity_emb_mat, relation_emb_mat):
        super(SKGNetForQuestionAnswering, self).__init__(config, device, args, entity_emb_mat, relation_emb_mat)
        self.num_labels = config.num_labels
        self.device = device
        self.gnn_layers = args.gnn_layers
        
        '''1st Layer: BERT Encoding Layer'''
        self.bert = BertModel(config)

        '''2st Layer: Knowledge subgraph Layer'''
        self.entity_embs = torch.nn.Embedding.from_pretrained(torch.Tensor(entity_emb_mat), freeze=True)
        self.relation_embs = torch.nn.Embedding.from_pretrained(torch.Tensor(relation_emb_mat), freeze=True)

        concept_dim = entity_emb_mat.shape[1]

        '''3st Layer: Graph Attention Layer'''
        self.conv = GATConv(concept_dim, concept_dim, heads=8, dropout=config.hidden_dropout_prob)
        self.reflect_layer = torch.nn.Linear(8 * concept_dim, config.hidden_size)

        '''4st Layer: Output Layer'''
        self.gate_layer = torch.nn.Linear(config.hidden_size * 2, 1)
        self.qa_outputs = torch.nn.Linear(config.hidden_size, config.num_labels)

        # self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)

        self.init_weights()
    
    def to_device(self, tensor):
        '''move tensor to device'''
        if tensor is not None:
            return tensor.to(self.device)
        else:
            return None

    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None,
                start_positions=None, end_positions=None, core_entity_ids=None, graph_data_dict=None):
        # get sub graph info, and combine batch sub-graph
        # print("graph ids lens: {}".format(len(example_graph_info)))
        # nodes_ids_list = torch.tensor([node_id[0] for node_id in example_graph_info["nodes_ids"]]).numpy().tolist()
        # core_entity_ids_list = core_entity_ids[0].numpy().tolist()
        # core_entity_index_list = []
        # for entity in core_entity_ids_list:
        #     if entity in nodes_ids_list:
        #         idx = nodes_ids_list.index(entity)
        #     else:
        #         k = entity
        #         idx = 0
        #     core_entity_index_list.append(idx)

        # nodes_ids_list = sorted(list(set(nodes_ids_list)))
        # sort_nodes_ids_list = list(set(nodes_ids_list))
        # graph_data_dict = graph_data_dict
            
        # core_entity_index_list = [nodes_ids_list.index(entity) for entity in core_entity_ids_list]
        # core_entity_index = torch.tensor(core_entity_index_list, dtype=torch.long)
        # core_entity_index = self.to_device(core_entity_index)

        # set all tensors to device
        input_ids = self.to_device(input_ids)
        attention_mask = self.to_device(attention_mask)
        token_type_ids = self.to_device(token_type_ids)
        position_ids = self.to_device(position_ids)
        head_mask = self.to_device(head_mask)
        inputs_embeds = self.to_device(inputs_embeds)
        start_positions = self.to_device(start_positions)
        end_positions = self.to_device(end_positions)
        core_entity_ids = self.to_device(core_entity_ids)
        
        '''1st Layer: Bert Layer'''
        outputs = self.bert(input_ids,
                            attention_mask=attention_mask,
                            token_type_ids=token_type_ids,
                            position_ids=position_ids,
                            head_mask=head_mask,
                            inputs_embeds=inputs_embeds)

        sequence_output = outputs[0] # [batch_size, seq_length, hidden_size]
        batch_size, seq_length, hidden_size = sequence_output.size()
 
        '''2nd Layer: Knowledge subgraph Layer'''
        # # get sub graph info, and combine batch sub-graph
        # x_ids = torch.tensor([node_id[0] for node_id in example_graph_info["nodes_ids"]], dtype=torch.long)
        # x_ids = self.to_device(x_ids)
        # x = self.entity_embs(x_ids)  # [num_nodes, node_features]
        # x = self.to_device(x)

        # edges_ids = torch.tensor(example_graph_info["edges_ids"], dtype=torch.long) # [2, num_edges]
        # edges_ids = self.to_device(edges_ids)

        # edges_attr_ids = torch.tensor([edge_attr_id[0] for edge_attr_id in example_graph_info["edges_attr_ids"]], dtype=torch.long)
        # edges_attr_ids = self.to_device(edges_attr_ids)
        # edges_attr = self.relation_embs(edges_attr_ids) # [num_edges, edge_features]

        # graph = Data(x=x, edge_index=edges_ids.t().contiguous())

        nodes_features = graph_data_dict["nodes_features"] # [b, nodes_nums, feature_dim]
        edges_index = graph_data_dict["edges_index"] # [b, edges_nums, 2]
        edges_attr = graph_data_dict["edges_attr"] # [b, edges_nums, feature_dim]

        graph_list = []
        for i in range(batch_size):
            nodes_features_i = nodes_features[i]
            edges_index_i = edges_index[i]
            edges_attr_i = edges_attr[i]
            graph = Data(x=nodes_features_i, edge_index=edges_index_i, edge_attr=edges_attr_i).to(self.device)
            graph_list.append(graph)
        graph_batch = Batch.from_data_list(graph_list)
        
        '''3nd Layer: Graph Attention Layer'''
        out = graph_batch.x
        edge_emb = graph_batch.edge_attr
        # out = self.conv(graph_batch.x, graph_batch.edge_index)
        for i in range(self.gnn_layers):
            out = self.conv(out, graph_batch.edge_index, edge_emb)
        
        out_dim = out.size()[-1]
        out = out.reshape(batch_size, -1, out_dim)    # [batch_size, num_nodes, node_features]

        # find updated core entity representation
        core_entity_emb_list = []
        for i in range(batch_size):
            core_entity_emb = torch.index_select(out[i], 0, core_entity_ids[i])
            core_entity_emb_list.append(core_entity_emb)
        core_entity_embs = torch.stack(core_entity_emb_list).to(self.device) # [batch_size, seq_len, node_features]
        # core_entity_embs = torch.tensor(core_entity_emb_list).to(self.device) 

        # concat bert output and kg output
        kg_emb_output = self.reflect_layer(core_entity_embs)
        gate_W = torch.sigmoid(self.gate_layer(torch.cat((sequence_output, kg_emb_output), dim=-1)))
        sequence_kg_output = gate_W * sequence_output + (1 - gate_W) * kg_emb_output


        '''4nd Layer: Output Layer'''
        logits = self.qa_outputs(sequence_kg_output)
        start_logits, end_logits = logits.split(1, dim=-1)
        start_logits = start_logits.squeeze(-1)
        end_logits = end_logits.squeeze(-1)

        outputs = (start_logits, end_logits,) + outputs[2:]
        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
            outputs = (total_loss,) + outputs

        return outputs  # (loss), start_logits, end_logits, (hidden_states), (attentions)


# Graph Attention Networks here
class GATConv(MessagePassing):
    r"""The graph attentional operator from the `"Graph Attention Networks"
    <https://arxiv.org/abs/1710.10903>`_ paper
    .. math::
        \mathbf{x}^{\prime}_i = \alpha_{i,i}\mathbf{\Theta}\mathbf{x}_{i} +
        \sum_{j \in \mathcal{N}(i)} \alpha_{i,j}\mathbf{\Theta}\mathbf{x}_{j},
    where the attention coefficients :math:`\alpha_{i,j}` are computed as
    .. math::
        \alpha_{i,j} =
        \frac{
        \exp\left(\mathrm{LeakyReLU}\left(\mathbf{a}^{\top}
        [\mathbf{\Theta}\mathbf{x}_i \, \Vert \, \mathbf{\Theta}\mathbf{x}_j]
        \right)\right)}
        {\sum_{k \in \mathcal{N}(i) \cup \{ i \}}
        \exp\left(\mathrm{LeakyReLU}\left(\mathbf{a}^{\top}
        [\mathbf{\Theta}\mathbf{x}_i \, \Vert \, \mathbf{\Theta}\mathbf{x}_k]
        \right)\right)}.
    Args:
        in_channels (int): Size of each input sample.
        out_channels (int): Size of each output sample.
        heads (int, optional): Number of multi-head-attentions.
            (default: :obj:`1`)
        concat (bool, optional): If set to :obj:`False`, the multi-head
            attentions are averaged instead of concatenated.
            (default: :obj:`True`)
        negative_slope (float, optional): LeakyReLU angle of the negative
            slope. (default: :obj:`0.2`)
        dropout (float, optional): Dropout probability of the normalized
            attention coefficients which exposes each node to a stochastically
            sampled neighborhood during training. (default: :obj:`0`)
        bias (bool, optional): If set to :obj:`False`, the layer will not learn
            an additive bias. (default: :obj:`True`)
        **kwargs (optional): Additional arguments of
            :class:`torch_geometric.nn.conv.MessagePassing`.
    """
    def __init__(self, in_channels, out_channels, heads=1, concat=True,
                 negative_slope=0.2, dropout=0, bias=True, **kwargs):
        super(GATConv, self).__init__(aggr='add', **kwargs)

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.heads = heads
        self.concat = concat
        self.negative_slope = negative_slope
        self.dropout = dropout

        self.weight_h = Parameter(torch.Tensor(in_channels,
                                             heads * out_channels))
        self.weight_t = Parameter(torch.Tensor(in_channels,
                                             heads * out_channels))
        self.weight_r = Parameter(torch.Tensor(in_channels,
                                             heads * out_channels))
        
        self.att = Parameter(torch.Tensor(1, heads, out_channels))

        if bias and concat:
            self.bias = Parameter(torch.Tensor(heads * out_channels))
        elif bias and not concat:
            self.bias = Parameter(torch.Tensor(out_channels))
        else:
            self.register_parameter('bias', None)

        self.reset_parameters()

    def reset_parameters(self):
        glorot(self.weight_h)
        glorot(self.weight_t)
        glorot(self.weight_r)
        glorot(self.att)
        zeros(self.bias)

    def forward(self, x, edge_index, edge_attr):
        """"""
        # x : [N, in_channels]
        # edge_index: [2, E]
        # edge_attr: [E, attr_dim]
        
        # if torch.is_tensor(x):
        #     edge_index, _ = remove_self_loops(edge_index)
        #     edge_index, _ = add_self_loops(edge_index,
        #                                    num_nodes=x.size(self.node_dim))

        # edge feature
        edge_attr = edge_attr.unsqueeze(-1) if edge_attr.dim() == 1 else edge_attr

        # if torch.is_tensor(x):
        #     x = torch.matmul(x, self.weight)
        # else:
        #     x = (None if x[0] is None else torch.matmul(x[0], self.weight),
        #          None if x[1] is None else torch.matmul(x[1], self.weight))

        return self.propagate(edge_index, x=x, edge_attr=edge_attr)

    def message(self, edge_index_i, x_i, x_j, edge_attr):
        # Compute attention coefficients.
        x_i = torch.matmul(x_i, self.weight_t) # [N, num_head * out_channels]
        x_j = torch.matmul(x_j, self.weight_h) # [N, num_head * out_channels]
        e_ij = torch.matmul(edge_attr, self.weight_r) # [N, num_head * out_channels]

        x_i = x_i.view(-1, self.heads, self.out_channels)   # [N, num_heads, out_channels]
        x_j = x_j.view(-1, self.heads, self.out_channels)   # [N, num_heads, out_channels]
        e_ij = e_ij.view(-1, self.heads, self.out_channels) # [N, num_heads, out_channels]

        # alpha = (torch.cat([x_i, x_j], dim=-1) * self.att).sum(dim=-1)

        alpha = ((e_ij * F.leaky_relu(x_i + x_j, self.negative_slope)) * self.att).sum(-1)  # [N, num_heads]

        # alpha = F.leaky_relu(alpha, self.negative_slope)
        alpha = softmax(alpha, edge_index_i)

        # Sample attention coefficients stochastically.
        alpha = F.dropout(alpha, p=self.dropout, training=self.training)

        return x_j * alpha.view(-1, self.heads, 1)

    def update(self, aggr_out):
        if self.concat is True:
            aggr_out = aggr_out.view(-1, self.heads * self.out_channels)
        else:
            aggr_out = aggr_out.mean(dim=1)

        if self.bias is not None:
            aggr_out = aggr_out + self.bias
        return aggr_out

    def __repr__(self):
        return '{}({}, {}, heads={})'.format(self.__class__.__name__,
                                             self.in_channels,
                                             self.out_channels, self.heads)
# 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 GAT here # 定义正无穷 # set seed here(set the same with your run_xxx.py file) # random.seed(42) # np.random.seed(42) # torch.manual_seed(42) # torch.cuda.manual_seed(42) Load tf checkpoints in a pytorch model. # Load weights from TF model # adam_v and adam_m are variables used in AdamWeightDecayOptimizer to calculated m and v # which are not required for using pretrained model Original Implementation of the gelu activation function in Google Bert repo when initially created.         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 Implementation of the gelu activation function currently in Google Bert repo (identical to OpenAI GPT).         Also see https://arxiv.org/abs/1606.08415 Construct the embeddings from word, position and token_type embeddings. # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load # any TensorFlow checkpoint file # If this is instantiated as a cross-attention module, the keys # and values come from an encoder; the attention mask needs to be # such that the encoder's padding tokens are not attended to. # Take the dot product between "query" and "key" to get the raw attention scores. # Apply the attention mask is (precomputed for all layers in BertModel forward() function) # Normalize the attention scores to probabilities. # This is actually dropping out entire tokens to attend to, which might # seem a bit unusual, but is taken from the original Transformer paper. # Mask heads if we want to # Convert to set and emove already pruned heads # Compute how many pruned heads are before the head and move the index accordingly # Prune linear layers # Update hyper params and store pruned heads # add attentions if we output them # add self attentions if we output attention weights # add cross attentions if we output attention weights # Add last layer # last-layer hidden state, (all hidden states), (all attentions) # We "pool" the model by simply taking the hidden state corresponding # to the first token. # The output weights are the same as the input embeddings, but there is # an output-only bias for each token. An abstract class to handle weights initialization and         a simple interface for dowloading and loading pretrained models. Initialize the weights # Slightly different from the TF version which uses truncated_normal for initialization # cf https://github.com/pytorch/pytorch/pull/5617 The BERT model was proposed in     `BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding`_     by <NAME>, <NAME>, <NAME> and <NAME>. It's a bidirectional transformer     pre-trained using a combination of masked language modeling objective and next sentence prediction     on a large corpus comprising the Toronto Book Corpus and Wikipedia.      This model is a PyTorch `torch.nn.Module`_ sub-class. Use it as a regular PyTorch Module and     refer to the PyTorch documentation for all matter related to general usage and behavior.      .. _`BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding`:         https://arxiv.org/abs/1810.04805      .. _`torch.nn.Module`:         https://pytorch.org/docs/stable/nn.html#module      Parameters:         config (:class:`~transformers.BertConfig`): Model configuration class with all the parameters of the model.             Initializing with a config file does not load the weights associated with the model, only the configuration.             Check out the :meth:`~transformers.PreTrainedModel.from_pretrained` method to load the model weights. Inputs:         **input_ids**: ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``:             Indices of input sequence tokens in the vocabulary.             To match pre-training, BERT input sequence should be formatted with [CLS] and [SEP] tokens as follows:              (a) For sequence pairs:                  ``tokens:         [CLS] is this jack ##son ##ville ? [SEP] no it is not . [SEP]``                  ``token_type_ids:   0   0  0    0    0     0       0   0   1  1  1  1   1   1``              (b) For single sequences:                  ``tokens:         [CLS] the dog is hairy . [SEP]``                  ``token_type_ids:   0   0   0   0  0     0   0``              Bert is a model with absolute position embeddings so it's usually advised to pad the inputs on             the right rather than the left.              Indices can be obtained using :class:`transformers.BertTokenizer`.             See :func:`transformers.PreTrainedTokenizer.encode` and             :func:`transformers.PreTrainedTokenizer.convert_tokens_to_ids` for details.         **attention_mask**: (`optional`) ``torch.FloatTensor`` of shape ``(batch_size, sequence_length)``:             Mask to avoid performing attention on padding token indices.             Mask values selected in ``[0, 1]``:             ``1`` for tokens that are NOT MASKED, ``0`` for MASKED tokens.         **token_type_ids**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``:             Segment token indices to indicate first and second portions of the inputs.             Indices are selected in ``[0, 1]``: ``0`` corresponds to a `sentence A` token, ``1``             corresponds to a `sentence B` token             (see `BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding`_ for more details).         **position_ids**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``:             Indices of positions of each input sequence tokens in the position embeddings.             Selected in the range ``[0, config.max_position_embeddings - 1]``.         **head_mask**: (`optional`) ``torch.FloatTensor`` of shape ``(num_heads,)`` or ``(num_layers, num_heads)``:             Mask to nullify selected heads of the self-attention modules.             Mask values selected in ``[0, 1]``:             ``1`` indicates the head is **not masked**, ``0`` indicates the head is **masked**.         **inputs_embeds**: (`optional`) ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, embedding_dim)``:             Optionally, instead of passing ``input_ids`` you can choose to directly pass an embedded representation.             This is useful if you want more control over how to convert `input_ids` indices into associated vectors             than the model's internal embedding lookup matrix.         **encoder_hidden_states**: (`optional`) ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, hidden_size)``:             Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if the model             is configured as a decoder.         **encoder_attention_mask**: (`optional`) ``torch.FloatTensor`` of shape ``(batch_size, sequence_length)``:             Mask to avoid performing attention on the padding token indices of the encoder input. This mask             is used in the cross-attention if the model is configured as a decoder.             Mask values selected in ``[0, 1]``:             ``1`` for tokens that are NOT MASKED, ``0`` for MASKED tokens. Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:         **last_hidden_state**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, hidden_size)``             Sequence of hidden-states at the output of the last layer of the model.         **pooler_output**: ``torch.FloatTensor`` of shape ``(batch_size, hidden_size)``             Last layer hidden-state of the first token of the sequence (classification token)             further processed by a Linear layer and a Tanh activation function. The Linear             layer weights are trained from the next sentence prediction (classification)             objective during Bert pretraining. This output is usually *not* a good summary             of the semantic content of the input, you're often better with averaging or pooling             the sequence of hidden-states for the whole input sequence.         **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)             list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)             of shape ``(batch_size, sequence_length, hidden_size)``:             Hidden-states of the model at the output of each layer plus the initial embedding outputs.         **attentions**: (`optional`, returned when ``config.output_attentions=True``)             list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:             Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.      Examples::          tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')         model = BertModel.from_pretrained('bert-base-uncased')         input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute")).unsqueeze(0)  # Batch size 1         outputs = model(input_ids)         last_hidden_states = outputs[0]  # The last hidden-state is the first element of the output tuple Prunes heads of the model.             heads_to_prune: dict of {layer_num: list of heads to prune in this layer}             See base class PreTrainedModel Forward pass on the Model.          The model can behave as an encoder (with only self-attention) as well         as a decoder, in which case a layer of cross-attention is added between         the self-attention layers, following the architecture described in `Attention is all you need`_ by <NAME>,         <NAME>, <NAME>, <NAME>, <NAME>, <NAME>, <NAME> and <NAME>.          To behave as an decoder the model needs to be initialized with the         `is_decoder` argument of the configuration set to `True`; an         `encoder_hidden_states` is expected as an input to the forward pass.          .. _`Attention is all you need`:             https://arxiv.org/abs/1706.03762 # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length] # ourselves in which case we just need to make it broadcastable to all heads. # Provided a padding mask of dimensions [batch_size, seq_length] # - if the model is a decoder, apply a causal mask in addition to the padding mask # - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length] # 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. # fp16 compatibility # If a 2D ou 3D attention mask is provided for the cross-attention # we need to make broadcastabe to [batch_size, num_heads, seq_length, seq_length] # fp16 compatibility # 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] # We can specify head_mask for each layer # switch to fload if need + fp16 compatibility # add hidden_states and attentions if they are here # sequence_output, pooled_output, (hidden_states), (attentions) Bert Model with two heads on top as done during the pre-training:                        a `masked language modeling` head and a `next sentence prediction (classification)` head. **masked_lm_labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``:             Labels for computing the masked language modeling loss.             Indices should be in ``[-1, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring)             Tokens with indices set to ``-1`` are ignored (masked), the loss is only computed for the tokens with labels             in ``[0, ..., config.vocab_size]``         **next_sentence_label**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:             Labels for computing the next sequence prediction (classification) loss. Input should be a sequence pair (see ``input_ids`` docstring)             Indices should be in ``[0, 1]``.             ``0`` indicates sequence B is a continuation of sequence A,             ``1`` indicates sequence B is a random sequence.      Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:         **loss**: (`optional`, returned when both ``masked_lm_labels`` and ``next_sentence_label`` are provided) ``torch.FloatTensor`` of shape ``(1,)``:             Total loss as the sum of the masked language modeling loss and the next sequence prediction (classification) loss.         **prediction_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, config.vocab_size)``             Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).         **seq_relationship_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, 2)``             Prediction scores of the next sequence prediction (classification) head (scores of True/False continuation before SoftMax).         **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)             list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)             of shape ``(batch_size, sequence_length, hidden_size)``:             Hidden-states of the model at the output of each layer plus the initial embedding outputs.         **attentions**: (`optional`, returned when ``config.output_attentions=True``)             list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:             Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.      Examples::          tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')         model = BertForPreTraining.from_pretrained('bert-base-uncased')         input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute")).unsqueeze(0)  # Batch size 1         outputs = model(input_ids)         prediction_scores, seq_relationship_scores = outputs[:2] # add hidden states and attention if they are here # (loss), prediction_scores, seq_relationship_score, (hidden_states), (attentions) Bert Model with a `language modeling` head on top. **masked_lm_labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``:             Labels for computing the masked language modeling loss.             Indices should be in ``[-1, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring)             Tokens with indices set to ``-1`` are ignored (masked), the loss is only computed for the tokens with labels             in ``[0, ..., config.vocab_size]``         **lm_labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``:             Labels for computing the left-to-right language modeling loss (next word prediction).             Indices should be in ``[-1, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring)             Tokens with indices set to ``-1`` are ignored (masked), the loss is only computed for the tokens with labels             in ``[0, ..., config.vocab_size]``      Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:         **masked_lm_loss**: (`optional`, returned when ``masked_lm_labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:             Masked language modeling loss.         **ltr_lm_loss**: (`optional`, returned when ``lm_labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:             Next token prediction loss.         **prediction_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, config.vocab_size)``             Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).         **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)             list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)             of shape ``(batch_size, sequence_length, hidden_size)``:             Hidden-states of the model at the output of each layer plus the initial embedding outputs.         **attentions**: (`optional`, returned when ``config.output_attentions=True``)             list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:             Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.      Examples::          tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')         model = BertForMaskedLM.from_pretrained('bert-base-uncased')         input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute")).unsqueeze(0)  # Batch size 1         outputs = model(input_ids, masked_lm_labels=input_ids)         loss, prediction_scores = outputs[:2] # Add hidden states and attention if they are here # Although this may seem awkward, BertForMaskedLM supports two scenarios: # 1. If a tensor that contains the indices of masked labels is provided, #    the cross-entropy is the MLM cross-entropy that measures the likelihood #    of predictions for masked words. # 2. If `lm_labels` is provided we are in a causal scenario where we #    try to predict the next token for each input in the decoder. # -1 index = padding token # we are doing next-token prediction; shift prediction scores and input ids by one # (masked_lm_loss), (ltr_lm_loss), prediction_scores, (hidden_states), (attentions) Bert Model with a `next sentence prediction (classification)` head on top. **next_sentence_label**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:             Labels for computing the next sequence prediction (classification) loss. Input should be a sequence pair (see ``input_ids`` docstring)             Indices should be in ``[0, 1]``.             ``0`` indicates sequence B is a continuation of sequence A,             ``1`` indicates sequence B is a random sequence.      Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:         **loss**: (`optional`, returned when ``next_sentence_label`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:             Next sequence prediction (classification) loss.         **seq_relationship_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, 2)``             Prediction scores of the next sequence prediction (classification) head (scores of True/False continuation before SoftMax).         **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)             list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)             of shape ``(batch_size, sequence_length, hidden_size)``:             Hidden-states of the model at the output of each layer plus the initial embedding outputs.         **attentions**: (`optional`, returned when ``config.output_attentions=True``)             list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:             Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.      Examples::          tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')         model = BertForNextSentencePrediction.from_pretrained('bert-base-uncased')         input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute")).unsqueeze(0)  # Batch size 1         outputs = model(input_ids)         seq_relationship_scores = outputs[0] # add hidden states and attention if they are here # (next_sentence_loss), seq_relationship_score, (hidden_states), (attentions) Bert Model transformer with a sequence classification/regression head on top (a linear layer on top of                       the pooled output) e.g. for GLUE tasks. **labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:             Labels for computing the sequence classification/regression loss.             Indices should be in ``[0, ..., config.num_labels - 1]``.             If ``config.num_labels == 1`` a regression loss is computed (Mean-Square loss),             If ``config.num_labels > 1`` a classification loss is computed (Cross-Entropy).      Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:         **loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:             Classification (or regression if config.num_labels==1) loss.         **logits**: ``torch.FloatTensor`` of shape ``(batch_size, config.num_labels)``             Classification (or regression if config.num_labels==1) scores (before SoftMax).         **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)             list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)             of shape ``(batch_size, sequence_length, hidden_size)``:             Hidden-states of the model at the output of each layer plus the initial embedding outputs.         **attentions**: (`optional`, returned when ``config.output_attentions=True``)             list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:             Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.      Examples::          tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')         model = BertForSequenceClassification.from_pretrained('bert-base-uncased')         input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute")).unsqueeze(0)  # Batch size 1         labels = torch.tensor([1]).unsqueeze(0)  # Batch size 1         outputs = model(input_ids, labels=labels)         loss, logits = outputs[:2] # add hidden states and attention if they are here #  We are doing regression # (loss), logits, (hidden_states), (attentions) Bert Model with a multiple choice classification head on top (a linear layer on top of                       the pooled output and a softmax) e.g. for RocStories/SWAG tasks. **labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:             Labels for computing the multiple choice classification loss.             Indices should be in ``[0, ..., num_choices]`` where `num_choices` is the size of the second dimension             of the input tensors. (see `input_ids` above)      Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:         **loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:             Classification loss.         **classification_scores**: ``torch.FloatTensor`` of shape ``(batch_size, num_choices)`` where `num_choices` is the size of the second dimension             of the input tensors. (see `input_ids` above).             Classification scores (before SoftMax).         **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)             list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)             of shape ``(batch_size, sequence_length, hidden_size)``:             Hidden-states of the model at the output of each layer plus the initial embedding outputs.         **attentions**: (`optional`, returned when ``config.output_attentions=True``)             list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:             Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.      Examples::          tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')         model = BertForMultipleChoice.from_pretrained('bert-base-uncased')         choices = ["Hello, my dog is cute", "Hello, my cat is amazing"]         input_ids = torch.tensor([tokenizer.encode(s) for s in choices]).unsqueeze(0)  # Batch size 1, 2 choices         labels = torch.tensor(1).unsqueeze(0)  # Batch size 1         outputs = model(input_ids, labels=labels)         loss, classification_scores = outputs[:2] # add hidden states and attention if they are here # (loss), reshaped_logits, (hidden_states), (attentions) Bert Model with a token classification head on top (a linear layer on top of                       the hidden-states output) e.g. for Named-Entity-Recognition (NER) tasks. **labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``:             Labels for computing the token classification loss.             Indices should be in ``[0, ..., config.num_labels - 1]``.      Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:         **loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:             Classification loss.         **scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, config.num_labels)``             Classification scores (before SoftMax).         **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)             list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)             of shape ``(batch_size, sequence_length, hidden_size)``:             Hidden-states of the model at the output of each layer plus the initial embedding outputs.         **attentions**: (`optional`, returned when ``config.output_attentions=True``)             list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:             Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.      Examples::          tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')         model = BertForTokenClassification.from_pretrained('bert-base-uncased')         input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute")).unsqueeze(0)  # Batch size 1         labels = torch.tensor([1] * input_ids.size(1)).unsqueeze(0)  # Batch size 1         outputs = model(input_ids, labels=labels)         loss, scores = outputs[:2] # add hidden states and attention if they are here # Only keep active parts of the loss # [batch_size * seq_len] # (loss), scores, (hidden_states), (attentions) Bert Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear layers on top of                       the hidden-states output to compute `span start logits` and `span end logits`). **start_positions**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:             Labels for position (index) of the start of the labelled span for computing the token classification loss.             Positions are clamped to the length of the sequence (`sequence_length`).             Position outside of the sequence are not taken into account for computing the loss.         **end_positions**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:             Labels for position (index) of the end of the labelled span for computing the token classification loss.             Positions are clamped to the length of the sequence (`sequence_length`).             Position outside of the sequence are not taken into account for computing the loss.      Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:         **loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:             Total span extraction loss is the sum of a Cross-Entropy for the start and end positions.         **start_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length,)``             Span-start scores (before SoftMax).         **end_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length,)``             Span-end scores (before SoftMax).         **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)             list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)             of shape ``(batch_size, sequence_length, hidden_size)``:             Hidden-states of the model at the output of each layer plus the initial embedding outputs.         **attentions**: (`optional`, returned when ``config.output_attentions=True``)             list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:             Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.      Examples::          tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')         model = BertForQuestionAnswering.from_pretrained('bert-large-uncased-whole-word-masking-finetuned-squad')         question, text = "Who was <NAME>?", "<NAME> was a nice puppet"         input_text = "[CLS] " + question + " [SEP] " + text + " [SEP]"         input_ids = tokenizer.encode(input_text)         token_type_ids = [0 if i <= input_ids.index(102) else 1 for i in range(len(input_ids))]          start_scores, end_scores = model(torch.tensor([input_ids]), token_type_ids=torch.tensor([token_type_ids]))         all_tokens = tokenizer.convert_ids_to_tokens(input_ids)           print(' '.join(all_tokens[torch.argmax(start_scores) : torch.argmax(end_scores)+1]))         # a nice puppet # If we are on multi-GPU, split add a dimension # sometimes the start/end positions are outside our model inputs, we ignore these terms # (loss), start_logits, end_logits, (hidden_states), (attentions) ############################################################################################################################################### # below is authored by yyhaker **start_positions**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:             Labels for position (index) of the start of the labelled span for computing the token classification loss.             Positions are clamped to the length of the sequence (`sequence_length`).             Position outside of the sequence are not taken into account for computing the loss.         **end_positions**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:             Labels for position (index) of the end of the labelled span for computing the token classification loss.             Positions are clamped to the length of the sequence (`sequence_length`).             Position outside of the sequence are not taken into account for computing the loss.      Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:         **loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:             Total span extraction loss is the sum of a Cross-Entropy for the start and end positions.         **start_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length,)``             Span-start scores (before SoftMax).         **end_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length,)``             Span-end scores (before SoftMax).         **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)             list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)             of shape ``(batch_size, sequence_length, hidden_size)``:             Hidden-states of the model at the output of each layer plus the initial embedding outputs.         **attentions**: (`optional`, returned when ``config.output_attentions=True``)             list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:             Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.      Examples::          tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')         model = BertForQuestionAnswering.from_pretrained('bert-large-uncased-whole-word-masking-finetuned-squad')         question, text = "Who was <NAME>?", "<NAME> was a nice puppet"         input_text = "[CLS] " + question + " [SEP] " + text + " [SEP]"         input_ids = tokenizer.encode(input_text)         token_type_ids = [0 if i <= input_ids.index(102) else 1 for i in range(len(input_ids))]          start_scores, end_scores = model(torch.tensor([input_ids]), token_type_ids=torch.tensor([token_type_ids]))         all_tokens = tokenizer.convert_ids_to_tokens(input_ids)           print(' '.join(all_tokens[torch.argmax(start_scores) : torch.argmax(end_scores)+1]))         # a nice puppet 1st Layer: BERT Encoding Layer 2st Layer: Knowledge Integration Layer # init a sentinel vector # self.sentinel_vec = torch.nn.Parameter(torch.rand(concept_dim), requires_grad=False) # init a sentinel vector # elf.sentinel_vec = torch.nn.Parameter(torch.zeros(concept_dim), requires_grad=True) # self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels) self-matching layer.             U, (b, len, d)             mask, (b, len) # direct info interaction # [b, len, 1, d] # [b, 1, len, d] # [b, len, len, d] # [b, len, len, d] # [b, len, len, d] # [b, len, len, 3 * d] -> [b, len, len, 1] -> [b, len, len] # [b, len, len] # mask here # [b, len, len] # [b, len, d] # indirect info interaction # [b, len, len] # [b, len, d] # [b, len, 6*d] 1st Layer: Bert Layer # [batch_size, seq_length, hidden_size] 2nd Layer: Knowledge Integration Layer # [batch_size, seq_length, num_concepts, cpnet_dim] # add sentinel vector # sentinel_vec = self.sentinel_vec.to(concept_emb.device) # sentinel_vec_expand = sentinel_vec.expand(batch_size, seq_length, cpnet_dim).unsqueeze(2).contiguous()  # [batch_size, seq_length, 1, cpnet_dim] # concept_emb_add = torch.cat((concept_emb, sentinel_vec_expand), 2)  # [batch_size, seq_length, num_concepts + 1, cpnet_dim] # [batch_size, seq_length, num_concepts, hidden_size] # [batch_size, seq_length, num_concepts] # calc masked attention # [batch_size, seq_length, num_concepts] # [batch_size * seq_length, 1, cpnet_dim] # [batch_size, seq_length, cpnet_dim] # mask and cat representation # [batch_size, seq_length, cpnet_dim] # [batch_size, seq_length, hidden_size + cpnet_dim] 3nd Layer: Self-Matching Layer # residule connection # sequence_kg_output = torch.cat((sequence_output, sequence_kg_output), dim=-1) 4nd Layer: Output Layer # If we are on multi-GPU, split add a dimension # sometimes the start/end positions are outside our model inputs, we ignore these terms # (loss), start_logits, end_logits, (hidden_states), (attentions) # here is FSNET_ner for QA **start_positions**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:             Labels for position (index) of the start of the labelled span for computing the token classification loss.             Positions are clamped to the length of the sequence (`sequence_length`).             Position outside of the sequence are not taken into account for computing the loss.         **end_positions**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:             Labels for position (index) of the end of the labelled span for computing the token classification loss.             Positions are clamped to the length of the sequence (`sequence_length`).             Position outside of the sequence are not taken into account for computing the loss.      Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:         **loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:             Total span extraction loss is the sum of a Cross-Entropy for the start and end positions.         **start_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length,)``             Span-start scores (before SoftMax).         **end_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length,)``             Span-end scores (before SoftMax).         **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)             list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)             of shape ``(batch_size, sequence_length, hidden_size)``:             Hidden-states of the model at the output of each layer plus the initial embedding outputs.         **attentions**: (`optional`, returned when ``config.output_attentions=True``)             list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:             Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.      Examples::          tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')         model = BertForQuestionAnswering.from_pretrained('bert-large-uncased-whole-word-masking-finetuned-squad')         question, text = "Who was <NAME>?", "<NAME> was a nice puppet"         input_text = "[CLS] " + question + " [SEP] " + text + " [SEP]"         input_ids = tokenizer.encode(input_text)         token_type_ids = [0 if i <= input_ids.index(102) else 1 for i in range(len(input_ids))]          start_scores, end_scores = model(torch.tensor([input_ids]), token_type_ids=torch.tensor([token_type_ids]))         all_tokens = tokenizer.convert_ids_to_tokens(input_ids)           print(' '.join(all_tokens[torch.argmax(start_scores) : torch.argmax(end_scores)+1]))         # a nice puppet 1st Layer: BERT Encoding Layer 2st Layer: Knowledge Integration Layer # for Named Entity task self-matching layer.             U, (b, len, d)             mask, (b, len) # direct info interaction # [b, len, 1, d] # [b, 1, len, d] # [b, len, len, d] # [b, len, len, d] # [b, len, len, d] # [b, len, len, 3 * d] -> [b, len, len, 1] -> [b, len, len] # [b, len, len] # mask here # [b, len, len] # [b, len, d] # indirect info interaction # [b, len, len] # [b, len, d] # [b, len, 6*d] 1st Layer: Bert Layer # [batch_size, seq_length, hidden_size] # here has two tasks: (1) Name Entity recognition (2) knowledge Enhanced QA # task (1) below # sequence_output_for_tag = self.dropout(sequence_output) # Only keep active parts of the loss # [batch_size * seq_len] # task (2) below 2nd Layer: Knowledge Integration Layer # [batch_size, seq_length, num_concepts, cpnet_dim] # add sentinel vector # sentinel_vec = self.sentinel_vec.to(concept_emb.device) # sentinel_vec_expand = sentinel_vec.expand(batch_size, seq_length, cpnet_dim).unsqueeze(2).contiguous()  # [batch_size, seq_length, 1, cpnet_dim] # concept_emb_add = torch.cat((concept_emb, sentinel_vec_expand), 2)  # [batch_size, seq_length, num_concepts + 1, cpnet_dim] # [batch_size, seq_length, num_concepts, hidden_size] # [batch_size, seq_length, num_concepts] # calc masked attention # [batch_size, seq_length, num_concepts] # [batch_size * seq_length, 1, cpnet_dim] # [batch_size, seq_length, cpnet_dim] # mask and cat representation # [batch_size, seq_length, cpnet_dim] # [batch_size, seq_length, hidden_size + cpnet_dim] 3nd Layer: Self-Matching Layer # residule connection # sequence_kg_output = torch.cat((sequence_output, sequence_kg_output), dim=-1) 4nd Layer: Output Layer # If we are on multi-GPU, split add a dimension # sometimes the start/end positions are outside our model inputs, we ignore these terms # add multi-task loss # (loss), start_logits, end_logits, (hidden_states), (attentions) # here KT-NET support (wordnet and nell) **start_positions**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:             Labels for position (index) of the start of the labelled span for computing the token classification loss.             Positions are clamped to the length of the sequence (`sequence_length`).             Position outside of the sequence are not taken into account for computing the loss.         **end_positions**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:             Labels for position (index) of the end of the labelled span for computing the token classification loss.             Positions are clamped to the length of the sequence (`sequence_length`).             Position outside of the sequence are not taken into account for computing the loss.      Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:         **loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:             Total span extraction loss is the sum of a Cross-Entropy for the start and end positions.         **start_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length,)``             Span-start scores (before SoftMax).         **end_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length,)``             Span-end scores (before SoftMax).         **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)             list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)             of shape ``(batch_size, sequence_length, hidden_size)``:             Hidden-states of the model at the output of each layer plus the initial embedding outputs.         **attentions**: (`optional`, returned when ``config.output_attentions=True``)             list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:             Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.      Examples::          tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')         model = BertForQuestionAnswering.from_pretrained('bert-large-uncased-whole-word-masking-finetuned-squad')         question, text = "Who was <NAME>?", "<NAME> was a nice puppet"         input_text = "[CLS] " + question + " [SEP] " + text + " [SEP]"         input_ids = tokenizer.encode(input_text)         token_type_ids = [0 if i <= input_ids.index(102) else 1 for i in range(len(input_ids))]          start_scores, end_scores = model(torch.tensor([input_ids]), token_type_ids=torch.tensor([token_type_ids]))         all_tokens = tokenizer.convert_ids_to_tokens(input_ids)           print(' '.join(all_tokens[torch.argmax(start_scores) : torch.argmax(end_scores)+1]))         # a nice puppet 1st Layer: BERT Encoding Layer 2st Layer: Knowledge Integration Layer # init a sentinel vector # self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels) self-matching layer.             U, (b, len, d)             mask, (b, len) # direct info interaction # [b, len, 1, d] # [b, 1, len, d] # [b, len, len, d] # [b, len, len, d] # [b, len, len, d] # [b, len, len, 3 * d] -> [b, len, len, 1] -> [b, len, len] # [b, len, len] # mask here # [b, len, len] # [b, len, d] # indirect info interaction # [b, len, len] # [b, len, d] # [b, len, 6*d] knowledge integration layer.             concept_emb: [batch_size, seq_length, num_concepts, cpnet_dim],             sequence_output: [batch_size, seq_length, hidden_size],             sentinel_vec_expand: [batch_size, seq_length, 1, cpnet_dim],             attention_mask_expand: [batch_size, seq_length, cpnet_dim] # get size # cat emb and sentinel # [batch_size, seq_length, num_concepts + 1, cpnet_dim] # calc similarity weight # [batch_size, seq_length, num_concepts + 1, hidden_size] # [batch_size, seq_length, num_concepts + 1] # attention result # [batch_size * seq_length, 1, cpnet_dim] # [batch_size, seq_length, cpnet_dim] # mask the concept # [batch_size, seq_length, cpnet_dim] 1st Layer: Bert Layer # [batch_size, seq_length, hidden_size] 2nd Layer: Knowledge Integration Layer # get concept emb # [batch_size, seq_length, num_concepts, cpnet_dim] # [batch_size, seq_length, num_concepts, cpnet_dim] # add sentinel vector # [batch_size, seq_length, 1, cpnet_dim] # mask and cat representation # [batch_size, seq_length, cpnet_dim] # knowledge integration # [batch_size, seq_length, hidden_size + cpnet_dim * 2] 3nd Layer: Self-Matching Layer # [batch_size, seq_length, 6 * (hidden_size + cpnet_dim * 2)] 4nd Layer: Output Layer # If we are on multi-GPU, split add a dimension # sometimes the start/end positions are outside our model inputs, we ignore these terms # (loss), start_logits, end_logits, (hidden_states), (attentions) # SKG-NET for QA models here **start_positions**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:             Labels for position (index) of the start of the labelled span for computing the token classification loss.             Positions are clamped to the length of the sequence (`sequence_length`).             Position outside of the sequence are not taken into account for computing the loss.         **end_positions**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:             Labels for position (index) of the end of the labelled span for computing the token classification loss.             Positions are clamped to the length of the sequence (`sequence_length`).             Position outside of the sequence are not taken into account for computing the loss.      Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:         **loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:             Total span extraction loss is the sum of a Cross-Entropy for the start and end positions.         **start_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length,)``             Span-start scores (before SoftMax).         **end_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length,)``             Span-end scores (before SoftMax).         **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)             list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)             of shape ``(batch_size, sequence_length, hidden_size)``:             Hidden-states of the model at the output of each layer plus the initial embedding outputs.         **attentions**: (`optional`, returned when ``config.output_attentions=True``)             list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:             Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.      Examples::          tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')         model = BertForQuestionAnswering.from_pretrained('bert-large-uncased-whole-word-masking-finetuned-squad')         question, text = "Who was <NAME>?", "<NAME> was a nice puppet"         input_text = "[CLS] " + question + " [SEP] " + text + " [SEP]"         input_ids = tokenizer.encode(input_text)         token_type_ids = [0 if i <= input_ids.index(102) else 1 for i in range(len(input_ids))]          start_scores, end_scores = model(torch.tensor([input_ids]), token_type_ids=torch.tensor([token_type_ids]))         all_tokens = tokenizer.convert_ids_to_tokens(input_ids)           print(' '.join(all_tokens[torch.argmax(start_scores) : torch.argmax(end_scores)+1]))         # a nice puppet 1st Layer: BERT Encoding Layer 2st Layer: Knowledge subgraph Layer 3st Layer: Graph Attention Layer 4st Layer: Output Layer # self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels) move tensor to device # get sub graph info, and combine batch sub-graph # print("graph ids lens: {}".format(len(example_graph_info))) # nodes_ids_list = torch.tensor([node_id[0] for node_id in example_graph_info["nodes_ids"]]).numpy().tolist() # core_entity_ids_list = core_entity_ids[0].numpy().tolist() # core_entity_index_list = [] # for entity in core_entity_ids_list: #     if entity in nodes_ids_list: #         idx = nodes_ids_list.index(entity) #     else: #         k = entity #         idx = 0 #     core_entity_index_list.append(idx) # nodes_ids_list = sorted(list(set(nodes_ids_list))) # sort_nodes_ids_list = list(set(nodes_ids_list)) # graph_data_dict = graph_data_dict # core_entity_index_list = [nodes_ids_list.index(entity) for entity in core_entity_ids_list] # core_entity_index = torch.tensor(core_entity_index_list, dtype=torch.long) # core_entity_index = self.to_device(core_entity_index) # set all tensors to device 1st Layer: Bert Layer # [batch_size, seq_length, hidden_size] 2nd Layer: Knowledge subgraph Layer # # get sub graph info, and combine batch sub-graph # x_ids = torch.tensor([node_id[0] for node_id in example_graph_info["nodes_ids"]], dtype=torch.long) # x_ids = self.to_device(x_ids) # x = self.entity_embs(x_ids)  # [num_nodes, node_features] # x = self.to_device(x) # edges_ids = torch.tensor(example_graph_info["edges_ids"], dtype=torch.long) # [2, num_edges] # edges_ids = self.to_device(edges_ids) # edges_attr_ids = torch.tensor([edge_attr_id[0] for edge_attr_id in example_graph_info["edges_attr_ids"]], dtype=torch.long) # edges_attr_ids = self.to_device(edges_attr_ids) # edges_attr = self.relation_embs(edges_attr_ids) # [num_edges, edge_features] # graph = Data(x=x, edge_index=edges_ids.t().contiguous()) # [b, nodes_nums, feature_dim] # [b, edges_nums, 2] # [b, edges_nums, feature_dim] 3nd Layer: Graph Attention Layer # out = self.conv(graph_batch.x, graph_batch.edge_index) # [batch_size, num_nodes, node_features] # find updated core entity representation # [batch_size, seq_len, node_features] # core_entity_embs = torch.tensor(core_entity_emb_list).to(self.device) # concat bert output and kg output 4nd Layer: Output Layer # If we are on multi-GPU, split add a dimension # sometimes the start/end positions are outside our model inputs, we ignore these terms # (loss), start_logits, end_logits, (hidden_states), (attentions) # Graph Attention Networks here The graph attentional operator from the `"Graph Attention Networks"     <https://arxiv.org/abs/1710.10903>`_ paper     .. math::         \mathbf{x}^{\prime}_i = \alpha_{i,i}\mathbf{\Theta}\mathbf{x}_{i} +         \sum_{j \in \mathcal{N}(i)} \alpha_{i,j}\mathbf{\Theta}\mathbf{x}_{j},     where the attention coefficients :math:`\alpha_{i,j}` are computed as     .. math::         \alpha_{i,j} =         \frac{         \exp\left(\mathrm{LeakyReLU}\left(\mathbf{a}^{\top}         [\mathbf{\Theta}\mathbf{x}_i \, \Vert \, \mathbf{\Theta}\mathbf{x}_j]         \right)\right)}         {\sum_{k \in \mathcal{N}(i) \cup \{ i \}}         \exp\left(\mathrm{LeakyReLU}\left(\mathbf{a}^{\top}         [\mathbf{\Theta}\mathbf{x}_i \, \Vert \, \mathbf{\Theta}\mathbf{x}_k]         \right)\right)}.     Args:         in_channels (int): Size of each input sample.         out_channels (int): Size of each output sample.         heads (int, optional): Number of multi-head-attentions.             (default: :obj:`1`)         concat (bool, optional): If set to :obj:`False`, the multi-head             attentions are averaged instead of concatenated.             (default: :obj:`True`)         negative_slope (float, optional): LeakyReLU angle of the negative             slope. (default: :obj:`0.2`)         dropout (float, optional): Dropout probability of the normalized             attention coefficients which exposes each node to a stochastically             sampled neighborhood during training. (default: :obj:`0`)         bias (bool, optional): If set to :obj:`False`, the layer will not learn             an additive bias. (default: :obj:`True`)         **kwargs (optional): Additional arguments of             :class:`torch_geometric.nn.conv.MessagePassing`. # x : [N, in_channels] # edge_index: [2, E] # edge_attr: [E, attr_dim] # if torch.is_tensor(x): #     edge_index, _ = remove_self_loops(edge_index) #     edge_index, _ = add_self_loops(edge_index, #                                    num_nodes=x.size(self.node_dim)) # edge feature # if torch.is_tensor(x): #     x = torch.matmul(x, self.weight) # else: #     x = (None if x[0] is None else torch.matmul(x[0], self.weight), #          None if x[1] is None else torch.matmul(x[1], self.weight)) # Compute attention coefficients. # [N, num_head * out_channels] # [N, num_head * out_channels] # [N, num_head * out_channels] # [N, num_heads, out_channels] # [N, num_heads, out_channels] # [N, num_heads, out_channels] # alpha = (torch.cat([x_i, x_j], dim=-1) * self.att).sum(dim=-1) # [N, num_heads] # alpha = F.leaky_relu(alpha, self.negative_slope) # Sample attention coefficients stochastically.