# Copyright (c) Facebook, Inc. and its affiliates.

# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import copy
import json
import logging
import math
import os
import shutil
import tarfile
import tempfile
import sys
import numpy as np
from io import open

import torch
from torch import nn
from torch.nn import CrossEntropyLoss
import torch.nn.functional as F
from torch.nn.utils.weight_norm import weight_norm

from .utils import PreTrainedModel
import pdb
from typing import Any, Dict, Tuple
from transformers.configuration_utils import PretrainedConfig
from trainers.train_utils import render_order_heatmap
from models.heatmap_module import HeatMapOutput

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",
    "roberta-base": "https://s3.amazonaws.com/models.huggingface.co/bert/roberta-base-pytorch_model.bin",
    "roberta-large": "https://s3.amazonaws.com/models.huggingface.co/bert/roberta-large-pytorch_model.bin",
    "roberta-large-mnli": "https://s3.amazonaws.com/models.huggingface.co/bert/roberta-large-mnli-pytorch_model.bin",
}


WEIGHTS_NAME = "pytorch_model.bin"
TF2_WEIGHTS_NAME = "tf_model.h5"
TF_WEIGHTS_NAME = "model.ckpt"
CONFIG_NAME = "config.json"
MODEL_CARD_NAME = "modelcard.json"


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

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


def gelu(x):
    """Implementation of the gelu activation function.
        For information: OpenAI GPT's gelu is slightly different (and gives slightly different results):
        0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))
        Also see https://arxiv.org/abs/1606.08415
    """
    return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))


class GeLU(nn.Module):
    """Implementation of the gelu activation function.
        For information: OpenAI GPT's gelu is slightly different (and gives slightly different results):
        0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))
        Also see https://arxiv.org/abs/1606.08415
    """

    def __init__(self):
        super().__init__()

    def forward(self, x):
        return gelu(x)


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


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


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

    def __init__(
        self,
        vocab_size_or_config_json_file,
        hidden_size=768,
        num_hidden_layers=12,
        num_attention_heads=12,
        intermediate_size=3072,
        hidden_act="gelu",
        hidden_dropout_prob=0.1,
        attention_probs_dropout_prob=0.1,
        max_position_embeddings=512,
        type_vocab_size=2,
        initializer_range=0.02,
        v_feature_size=2048,
        v_target_size=1601,
        v_hidden_size=768,
        v_num_hidden_layers=3,
        v_num_attention_heads=12,
        v_intermediate_size=3072,
        bi_hidden_size=1024,
        bi_num_attention_heads=16,
        v_attention_probs_dropout_prob=0.1,
        v_hidden_act="gelu",
        v_hidden_dropout_prob=0.1,
        v_initializer_range=0.2,
        v_biattention_id=[0, 1],
        t_biattention_id=[10, 11],
        visual_target=0,
        fast_mode=False,
        fixed_v_layer=0,
        fixed_t_layer=0,
        in_batch_pairs=False,
        fusion_method="mul",
        dynamic_attention=False,
        with_coattention=True,
        objective=0,
        num_negative=128,
        model="bert",
        task_specific_tokens=False,
        visualization=False,
    ):

        """Constructs BertConfig.

        Args:
            vocab_size_or_config_json_file: Vocabulary size of `inputs_ids` in `BertModel`.
            hidden_size: Size of the encoder layers and the pooler layer.
            num_hidden_layers: Number of hidden layers in the Transformer encoder.
            num_attention_heads: Number of attention heads for each attention layer in
                the Transformer encoder.
            intermediate_size: The size of the "intermediate" (i.e., feed-forward)
                layer in the Transformer encoder.
            hidden_act: The non-linear activation function (function or string) in the
                encoder and pooler. If string, "gelu", "relu" and "swish" are supported.
            hidden_dropout_prob: The dropout probabilitiy for all fully connected
                layers in the embeddings, encoder, and pooler.
            attention_probs_dropout_prob: The dropout ratio for the attention
                probabilities.
            max_position_embeddings: The maximum sequence length that this model might
                ever be used with. Typically set this to something large just in case
                (e.g., 512 or 1024 or 2048).
            type_vocab_size: The vocabulary size of the `token_type_ids` passed into
                `BertModel`.
            initializer_range: The sttdev of the truncated_normal_initializer for
                initializing all weight matrices.
        """
        assert len(v_biattention_id) == len(t_biattention_id)
        assert max(v_biattention_id) < v_num_hidden_layers
        assert max(t_biattention_id) < num_hidden_layers

        if isinstance(vocab_size_or_config_json_file, str) or (
            sys.version_info[0] == 2
            and isinstance(vocab_size_or_config_json_file, unicode)
        ):
            with open(vocab_size_or_config_json_file, "r", encoding="utf-8") as reader:
                json_config = json.loads(reader.read())
            for key, value in json_config.items():
                self.__dict__[key] = value
        elif isinstance(vocab_size_or_config_json_file, int):
            self.vocab_size = vocab_size_or_config_json_file
            self.hidden_size = hidden_size
            self.num_hidden_layers = num_hidden_layers
            self.num_attention_heads = num_attention_heads
            self.hidden_act = hidden_act
            self.intermediate_size = intermediate_size
            self.hidden_dropout_prob = hidden_dropout_prob
            self.attention_probs_dropout_prob = attention_probs_dropout_prob
            self.max_position_embeddings = max_position_embeddings
            self.type_vocab_size = type_vocab_size
            self.initializer_range = initializer_range
            self.v_feature_size = v_feature_size
            self.v_hidden_size = v_hidden_size
            self.v_num_hidden_layers = v_num_hidden_layers
            self.v_num_attention_heads = v_num_attention_heads
            self.v_intermediate_size = v_intermediate_size
            self.v_attention_probs_dropout_prob = v_attention_probs_dropout_prob
            self.v_hidden_act = v_hidden_act
            self.v_hidden_dropout_prob = v_hidden_dropout_prob
            self.v_initializer_range = v_initializer_range
            self.v_biattention_id = v_biattention_id
            self.t_biattention_id = t_biattention_id
            self.v_target_size = v_target_size
            self.bi_hidden_size = bi_hidden_size
            self.bi_num_attention_heads = bi_num_attention_heads
            self.visual_target = visual_target
            self.fast_mode = fast_mode
            self.fixed_v_layer = fixed_v_layer
            self.fixed_t_layer = fixed_t_layer

            self.model = model
            self.in_batch_pairs = in_batch_pairs
            self.fusion_method = fusion_method
            self.dynamic_attention = dynamic_attention
            self.with_coattention = with_coattention
            self.objective = objective
            self.num_negative = num_negative
            self.task_specific_tokens = task_specific_tokens
            self.visualization = visualization
        else:
            raise ValueError(
                "First argument must be either a vocabulary size (int)"
                "or the path to a pretrained model config file (str)"
            )

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

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

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

    def to_diff_dict(self) -> Dict[str, Any]:
        """
        Removes all attributes from config which correspond to the default config attributes for better readability and
        serializes to a Python dictionary.

        Returns:
            :obj:`Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance,
        """
        config_dict = self.to_dict()

        # get the default config dict
        default_config_dict = PretrainedConfig().to_dict()

        # get class specific config dict
        class_config_dict = self.__class__(self.vocab_size).to_dict() if not self.is_composition else {}

        serializable_config_dict = {}

        # only serialize values that differ from the default config
        for key, value in config_dict.items():
            if (
                key not in default_config_dict
                or value != default_config_dict[key]
                or (key in class_config_dict and value != class_config_dict[key])
            ):
                serializable_config_dict[key] = value

        return serializable_config_dict

    def to_dict(self) -> Dict[str, Any]:
        """
        Serializes this instance to a Python dictionary.

        Returns:
            :obj:`Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance.
        """
        output = copy.deepcopy(self.__dict__)
        if hasattr(self.__class__, "model_type"):
            output["model_type"] = self.__class__.model_type
        return output

    def to_json_string(self, use_diff: bool = True) -> str:
        """
        Serializes this instance to a JSON string.

        Args:
            use_diff (:obj:`bool`, `optional`, defaults to :obj:`True`):
                If set to ``True``, only the difference between the config instance and the default
                ``PretrainedConfig()`` is serialized to JSON string.

        Returns:
            :obj:`str`: String containing all the attributes that make up this configuration instance in JSON format.
        """
        if use_diff is True:
            config_dict = self.to_diff_dict()
        else:
            config_dict = self.to_dict()
        return json.dumps(config_dict, indent=2, sort_keys=True) + "\n"

    def to_json_file(self, json_file_path: str, use_diff: bool = True):
        """
        Save this instance to a JSON file.

        Args:
            json_file_path (:obj:`str`):
                Path to the JSON file in which this configuration instance's parameters will be saved.
            use_diff (:obj:`bool`, `optional`, defaults to :obj:`True`):
                If set to ``True``, only the difference between the config instance and the default
                ``PretrainedConfig()`` is serialized to JSON file.
        """
        with open(json_file_path, "w", encoding="utf-8") as writer:
            writer.write(self.to_json_string(use_diff=use_diff))

    def save_pretrained(self, save_directory: str):
        """
        Save a configuration object to the directory ``save_directory``, so that it can be re-loaded using the
        :func:`~transformers.PretrainedConfig.from_pretrained` class method.

        Args:
            save_directory (:obj:`str`):
                Directory where the configuration JSON file will be saved (will be created if it does not exist).
        """
        if os.path.isfile(save_directory):
            raise AssertionError("Provided path ({}) should be a directory, not a file".format(save_directory))
        os.makedirs(save_directory, exist_ok=True)
        # If we save using the predefined names, we can load using `from_pretrained`
        output_config_file = os.path.join(save_directory, CONFIG_NAME)

        self.to_json_file(output_config_file, use_diff=True)
        logger.info("Configuration saved in {}".format(output_config_file))


try:
    # from apex.normalization.fused_layer_norm import FusedLayerNorm as BertLayerNorm
    import torch.nn.LayerNorm as BertLayerNorm
except ImportError:
    logger.info(
        "Better speed can be achieved with apex installed from https://www.github.com/nvidia/apex ."
    )

    class BertLayerNorm(nn.Module):
        def __init__(self, hidden_size, eps=1e-12):
            """Construct a layernorm module in the TF style (epsilon inside the square root).
            """
            super(BertLayerNorm, self).__init__()
            self.weight = nn.Parameter(torch.ones(hidden_size))
            self.bias = nn.Parameter(torch.zeros(hidden_size))
            self.variance_epsilon = eps

        def forward(self, x):
            u = x.mean(-1, keepdim=True)
            s = (x - u).pow(2).mean(-1, keepdim=True)
            x = (x - u) / torch.sqrt(s + self.variance_epsilon)
            return self.weight * x + self.bias


class BertEmbeddings(nn.Module):
    """Construct the embeddings from word, position and token_type embeddings.
    """

    def __init__(self, config):
        super(BertEmbeddings, self).__init__()

        self.task_specific_tokens = config.task_specific_tokens
        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=1e-12)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

        if self.task_specific_tokens:
            self.task_embeddings = nn.Embedding(20, config.hidden_size)

    def forward(self, input_ids, token_type_ids=None, task_ids=None, position_ids=None):

        seq_length = input_ids.size(1)
        position_ids = torch.arange(
            seq_length, dtype=torch.long, device=input_ids.device
        )
        position_ids = position_ids.unsqueeze(0).expand_as(input_ids)
        words_embeddings = self.word_embeddings(input_ids)
        position_embeddings = self.position_embeddings(position_ids)
        token_type_embeddings = self.token_type_embeddings(token_type_ids)
        embeddings = words_embeddings + position_embeddings + token_type_embeddings

        if self.task_specific_tokens:
            task_embeddings = self.task_embeddings(task_ids)
            embeddings = torch.cat(
                [embeddings[:, 0:1], task_embeddings, embeddings[:, 1:]], dim=1
            )

        embeddings = self.LayerNorm(embeddings)
        embeddings = self.dropout(embeddings)

        return embeddings


class RobertaEmbeddings(BertEmbeddings):
    """
    Same as BertEmbeddings with a tiny tweak for positional embeddings indexing.
    """

    def __init__(self, config):
        super(RobertaEmbeddings, self).__init__(config)
        self.padding_idx = 1

    def forward(self, input_ids, token_type_ids=None, position_ids=None):
        seq_length = input_ids.size(1)
        if position_ids is None:
            # Position numbers begin at padding_idx+1. Padding symbols are ignored.
            # cf. fairseq's `utils.make_positions`
            position_ids = torch.arange(
                self.padding_idx + 1,
                seq_length + self.padding_idx + 1,
                dtype=torch.long,
                device=input_ids.device,
            )
            position_ids = position_ids.unsqueeze(0).expand_as(input_ids)
        return super(RobertaEmbeddings, self).forward(
            input_ids, token_type_ids=token_type_ids, position_ids=position_ids
        )


class BertSelfAttention(nn.Module):
    def __init__(self, config):
        super(BertSelfAttention, self).__init__()
        if config.hidden_size % config.num_attention_heads != 0:
            raise ValueError(
                "The hidden size (%d) is not a multiple of the number of attention "
                "heads (%d)" % (config.hidden_size, config.num_attention_heads)
            )
        self.num_attention_heads = config.num_attention_heads
        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
        self.all_head_size = self.num_attention_heads * self.attention_head_size

        self.visualization = config.visualization

        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):
        mixed_query_layer = self.query(hidden_states)
        mixed_key_layer = self.key(hidden_states)
        mixed_value_layer = self.value(hidden_states)

        query_layer = self.transpose_for_scores(mixed_query_layer)
        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)
        # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
        attention_scores = attention_scores + attention_mask

        # Normalize the attention scores to probabilities.
        attention_probs = nn.Softmax(dim=-1)(attention_scores)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.dropout(attention_probs)

        context_layer = torch.matmul(attention_probs, value_layer)
        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
        context_layer = context_layer.view(*new_context_layer_shape)

        if self.visualization:
            attn_data = {
                "attn": attention_probs,
                "queries": query_layer,
                "keys": key_layer,
            }
        else:
            attn_data = None

        return context_layer, attn_data


class BertSelfOutput(nn.Module):
    def __init__(self, config):
        super(BertSelfOutput, self).__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.LayerNorm = BertLayerNorm(config.hidden_size, eps=1e-12)
        self.dropout = nn.Dropout(config.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)

    def forward(self, input_tensor, attention_mask):
        self_output, attention_probs = self.self(input_tensor, attention_mask)
        attention_output = self.output(self_output, input_tensor)
        return attention_output, attention_probs


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=1e-12)
        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.intermediate = BertIntermediate(config)
        self.output = BertOutput(config)

    def forward(self, hidden_states, attention_mask):
        attention_output, attention_probs = self.attention(
            hidden_states, attention_mask
        )
        intermediate_output = self.intermediate(attention_output)
        layer_output = self.output(intermediate_output, attention_output)
        return layer_output, attention_probs


class BertImageSelfAttention(nn.Module):
    def __init__(self, config):
        super(BertImageSelfAttention, self).__init__()
        if config.v_hidden_size % config.v_num_attention_heads != 0:
            raise ValueError(
                "The hidden size (%d) is not a multiple of the number of attention "
                "heads (%d)" % (config.v_hidden_size, config.v_num_attention_heads)
            )
        self.dynamic_attention = config.dynamic_attention
        self.num_attention_heads = config.v_num_attention_heads
        self.attention_head_size = int(
            config.v_hidden_size / config.v_num_attention_heads
        )

        self.visualization = config.visualization

        self.all_head_size = self.num_attention_heads * self.attention_head_size
        self.query = nn.Linear(config.v_hidden_size, self.all_head_size)
        self.key = nn.Linear(config.v_hidden_size, self.all_head_size)
        self.value = nn.Linear(config.v_hidden_size, self.all_head_size)

        if self.dynamic_attention:
            self.dyLinear_q = nn.Linear(config.hidden_size, self.all_head_size)
            self.dyLinear_k = nn.Linear(config.hidden_size, self.all_head_size)

        self.dropout = nn.Dropout(config.v_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, txt_embedding, txt_attention_mask):

        mixed_query_layer = self.query(hidden_states)
        mixed_key_layer = self.key(hidden_states)
        mixed_value_layer = self.value(hidden_states)

        if self.dynamic_attention:
            pool_embedding = (txt_embedding * txt_attention_mask).sum(1)
            pool_embedding = pool_embedding / txt_attention_mask.sum(1)

            # given pool embedding, Linear and Sigmoid layer.
            gate_q = 1 + torch.sigmoid(self.dyLinear_q(pool_embedding))
            gate_k = 1 + torch.sigmoid(self.dyLinear_k(pool_embedding))

            mixed_query_layer = mixed_query_layer * gate_q.unsqueeze(1)
            mixed_key_layer = mixed_key_layer * gate_k.unsqueeze(1)

        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)
        # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
        attention_scores = attention_scores + attention_mask

        # Normalize the attention scores to probabilities.
        attention_probs = nn.Softmax(dim=-1)(attention_scores)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.dropout(attention_probs)

        context_layer = torch.matmul(attention_probs, value_layer)
        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
        context_layer = context_layer.view(*new_context_layer_shape)

        if self.visualization:
            attn_data = {
                "attn": attention_probs,
                "queries": query_layer,
                "keys": key_layer,
            }
        else:
            attn_data = None

        return context_layer, attn_data


class BertImageSelfOutput(nn.Module):
    def __init__(self, config):
        super(BertImageSelfOutput, self).__init__()
        self.dense = nn.Linear(config.v_hidden_size, config.v_hidden_size)
        self.LayerNorm = BertLayerNorm(config.v_hidden_size, eps=1e-12)
        self.dropout = nn.Dropout(config.v_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 BertImageAttention(nn.Module):
    def __init__(self, config):
        super(BertImageAttention, self).__init__()
        self.self = BertImageSelfAttention(config)
        self.output = BertImageSelfOutput(config)

    def forward(self, input_tensor, attention_mask, txt_embedding, txt_attention_mask):
        self_output, attention_probs = self.self(
            input_tensor, attention_mask, txt_embedding, txt_attention_mask
        )
        attention_output = self.output(self_output, input_tensor)
        return attention_output, attention_probs


class BertImageIntermediate(nn.Module):
    def __init__(self, config):
        super(BertImageIntermediate, self).__init__()
        self.dense = nn.Linear(config.v_hidden_size, config.v_intermediate_size)
        if isinstance(config.v_hidden_act, str) or (
            sys.version_info[0] == 2 and isinstance(config.v_hidden_act, unicode)
        ):
            self.intermediate_act_fn = ACT2FN[config.v_hidden_act]
        else:
            self.intermediate_act_fn = config.v_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 BertImageOutput(nn.Module):
    def __init__(self, config):
        super(BertImageOutput, self).__init__()
        self.dense = nn.Linear(config.v_intermediate_size, config.v_hidden_size)
        self.LayerNorm = BertLayerNorm(config.v_hidden_size, eps=1e-12)
        self.dropout = nn.Dropout(config.v_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 BertImageLayer(nn.Module):
    def __init__(self, config):
        super(BertImageLayer, self).__init__()
        self.attention = BertImageAttention(config)
        self.intermediate = BertImageIntermediate(config)
        self.output = BertImageOutput(config)

    def forward(self, hidden_states, attention_mask, txt_embedding, txt_attention_mask):
        attention_output, attention_probs = self.attention(
            hidden_states, attention_mask, txt_embedding, txt_attention_mask
        )
        intermediate_output = self.intermediate(attention_output)
        layer_output = self.output(intermediate_output, attention_output)
        return layer_output, attention_probs


class BertBiAttention(nn.Module):
    def __init__(self, config):
        super(BertBiAttention, self).__init__()
        if config.bi_hidden_size % config.bi_num_attention_heads != 0:
            raise ValueError(
                "The hidden size (%d) is not a multiple of the number of attention "
                "heads (%d)" % (config.bi_hidden_size, config.bi_num_attention_heads)
            )

        self.visualization = config.visualization
        self.num_attention_heads = config.bi_num_attention_heads
        self.attention_head_size = int(
            config.bi_hidden_size / config.bi_num_attention_heads
        )
        self.all_head_size = self.num_attention_heads * self.attention_head_size

        # self.scale = nn.Linear(1, self.num_attention_heads, bias=False)
        # self.scale_act_fn = ACT2FN['relu']

        self.query1 = nn.Linear(config.v_hidden_size, self.all_head_size)
        self.key1 = nn.Linear(config.v_hidden_size, self.all_head_size)
        self.value1 = nn.Linear(config.v_hidden_size, self.all_head_size)
        # self.logit1 = nn.Linear(config.hidden_size, self.num_attention_heads)

        self.dropout1 = nn.Dropout(config.v_attention_probs_dropout_prob)

        self.query2 = nn.Linear(config.hidden_size, self.all_head_size)
        self.key2 = nn.Linear(config.hidden_size, self.all_head_size)
        self.value2 = nn.Linear(config.hidden_size, self.all_head_size)
        # self.logit2 = nn.Linear(config.hidden_size, self.num_attention_heads)

        self.dropout2 = 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,
        input_tensor1,
        attention_mask1,
        input_tensor2,
        attention_mask2,
        co_attention_mask=None,
        use_co_attention_mask=False,
    ):

        # for vision input.
        mixed_query_layer1 = self.query1(input_tensor1)
        mixed_key_layer1 = self.key1(input_tensor1)
        mixed_value_layer1 = self.value1(input_tensor1)
        # mixed_logit_layer1 = self.logit1(input_tensor1)

        query_layer1 = self.transpose_for_scores(mixed_query_layer1)
        key_layer1 = self.transpose_for_scores(mixed_key_layer1)
        value_layer1 = self.transpose_for_scores(mixed_value_layer1)
        # logit_layer1 = self.transpose_for_logits(mixed_logit_layer1)

        # for text input:
        mixed_query_layer2 = self.query2(input_tensor2)
        mixed_key_layer2 = self.key2(input_tensor2)
        mixed_value_layer2 = self.value2(input_tensor2)
        # mixed_logit_layer2 = self.logit2(input_tensor2)

        query_layer2 = self.transpose_for_scores(mixed_query_layer2)
        key_layer2 = self.transpose_for_scores(mixed_key_layer2)
        value_layer2 = self.transpose_for_scores(mixed_value_layer2)
        # logit_layer2 = self.transpose_for_logits(mixed_logit_layer2)

        # Take the dot product between "query2" and "key1" to get the raw attention scores for value 1.
        attention_scores1 = torch.matmul(query_layer2, key_layer1.transpose(-1, -2))
        attention_scores1 = attention_scores1 / math.sqrt(self.attention_head_size)
        attention_scores1 = attention_scores1 + attention_mask1
        # if use_co_attention_mask:
        # attention_scores1 = attention_scores1 + co_attention_mask.permute(0,1,3,2)

        # Normalize the attention scores to probabilities.
        attention_probs1 = nn.Softmax(dim=-1)(attention_scores1)

        # 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_probs1 = self.dropout1(attention_probs1)

        context_layer1 = torch.matmul(attention_probs1, value_layer1)
        context_layer1 = context_layer1.permute(0, 2, 1, 3).contiguous()
        new_context_layer_shape1 = context_layer1.size()[:-2] + (self.all_head_size,)
        context_layer1 = context_layer1.view(*new_context_layer_shape1)

        # Take the dot product between "query1" and "key2" to get the raw attention scores for value 2.
        attention_scores2 = torch.matmul(query_layer1, key_layer2.transpose(-1, -2))
        attention_scores2 = attention_scores2 / math.sqrt(self.attention_head_size)
        # Apply the attention mask is (precomputed for all layers in BertModel forward() function)

        # we can comment this line for single flow.
        attention_scores2 = attention_scores2 + attention_mask2
        # if use_co_attention_mask:
        # attention_scores2 = attention_scores2 + co_attention_mask

        # Normalize the attention scores to probabilities.
        attention_probs2 = nn.Softmax(dim=-1)(attention_scores2)

        # 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_probs2 = self.dropout2(attention_probs2)

        context_layer2 = torch.matmul(attention_probs2, value_layer2)
        context_layer2 = context_layer2.permute(0, 2, 1, 3).contiguous()
        new_context_layer_shape2 = context_layer2.size()[:-2] + (self.all_head_size,)
        context_layer2 = context_layer2.view(*new_context_layer_shape2)

        attn_data = None

        if self.visualization:
            attn_data = {
                "attn1": attention_probs1,
                "queries1": query_layer2,
                "keys1": key_layer1,
                "attn2": attention_probs2,
                "querues2": query_layer1,
                "keys2": key_layer2,
            }

        return context_layer1, context_layer2, attn_data


class BertBiOutput(nn.Module):
    def __init__(self, config):
        super(BertBiOutput, self).__init__()

        self.dense1 = nn.Linear(config.bi_hidden_size, config.v_hidden_size)
        self.LayerNorm1 = BertLayerNorm(config.v_hidden_size, eps=1e-12)
        self.dropout1 = nn.Dropout(config.v_hidden_dropout_prob)

        self.q_dense1 = nn.Linear(config.bi_hidden_size, config.v_hidden_size)
        self.q_dropout1 = nn.Dropout(config.v_hidden_dropout_prob)

        self.dense2 = nn.Linear(config.bi_hidden_size, config.hidden_size)
        self.LayerNorm2 = BertLayerNorm(config.hidden_size, eps=1e-12)
        self.dropout2 = nn.Dropout(config.hidden_dropout_prob)

        self.q_dense2 = nn.Linear(config.bi_hidden_size, config.hidden_size)
        self.q_dropout2 = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, hidden_states1, input_tensor1, hidden_states2, input_tensor2):

        context_state1 = self.dense1(hidden_states1)
        context_state1 = self.dropout1(context_state1)

        context_state2 = self.dense2(hidden_states2)
        context_state2 = self.dropout2(context_state2)

        hidden_states1 = self.LayerNorm1(context_state1 + input_tensor1)
        hidden_states2 = self.LayerNorm2(context_state2 + input_tensor2)

        return hidden_states1, hidden_states2


class BertConnectionLayer(nn.Module):
    def __init__(self, config):
        super(BertConnectionLayer, self).__init__()
        self.biattention = BertBiAttention(config)

        self.biOutput = BertBiOutput(config)

        self.v_intermediate = BertImageIntermediate(config)
        self.v_output = BertImageOutput(config)

        self.t_intermediate = BertIntermediate(config)
        self.t_output = BertOutput(config)

    def forward(
        self,
        input_tensor1,
        attention_mask1,
        input_tensor2,
        attention_mask2,
        co_attention_mask=None,
        use_co_attention_mask=False,
    ):

        bi_output1, bi_output2, co_attention_probs = self.biattention(
            input_tensor1,
            attention_mask1,
            input_tensor2,
            attention_mask2,
            co_attention_mask,
            use_co_attention_mask,
        )

        attention_output1, attention_output2 = self.biOutput(
            bi_output2, input_tensor1, bi_output1, input_tensor2
        )

        intermediate_output1 = self.v_intermediate(attention_output1)
        layer_output1 = self.v_output(intermediate_output1, attention_output1)

        intermediate_output2 = self.t_intermediate(attention_output2)
        layer_output2 = self.t_output(intermediate_output2, attention_output2)

        return layer_output1, layer_output2, co_attention_probs


class BertEncoder(nn.Module):
    def __init__(self, config):
        super(BertEncoder, self).__init__()

        # in the bert encoder, we need to extract three things here.
        # text bert layer: BertLayer
        # vision bert layer: BertImageLayer
        # Bi-Attention: Given the output of two bertlayer, perform bi-directional
        # attention and add on two layers.

        self.FAST_MODE = config.fast_mode
        self.with_coattention = config.with_coattention
        self.v_biattention_id = config.v_biattention_id
        self.t_biattention_id = config.t_biattention_id
        self.in_batch_pairs = config.in_batch_pairs
        self.fixed_t_layer = config.fixed_t_layer
        self.fixed_v_layer = config.fixed_v_layer
        layer = BertLayer(config)
        v_layer = BertImageLayer(config)
        connect_layer = BertConnectionLayer(config)

        self.layer = nn.ModuleList(
            [copy.deepcopy(layer) for _ in range(config.num_hidden_layers)]
        )
        self.v_layer = nn.ModuleList(
            [copy.deepcopy(v_layer) for _ in range(config.v_num_hidden_layers)]
        )
        self.c_layer = nn.ModuleList(
            [copy.deepcopy(connect_layer) for _ in range(len(config.v_biattention_id))]
        )

    def forward(
        self,
        txt_embedding,
        image_embedding,
        txt_attention_mask,
        txt_attention_mask2,
        image_attention_mask,
        co_attention_mask=None,
        output_all_encoded_layers=True,
        output_all_attention_masks=False,
    ):

        v_start = 0
        t_start = 0
        count = 0
        all_encoder_layers_t = []
        all_encoder_layers_v = []

        all_attention_mask_t = []
        all_attnetion_mask_v = []
        all_attention_mask_c = []

        batch_size, num_words, t_hidden_size = txt_embedding.size()
        _, num_regions, v_hidden_size = image_embedding.size()

        use_co_attention_mask = False
        for v_layer_id, t_layer_id in zip(self.v_biattention_id, self.t_biattention_id):

            v_end = v_layer_id
            t_end = t_layer_id

            assert self.fixed_t_layer <= t_end
            assert self.fixed_v_layer <= v_end

            for idx in range(t_start, self.fixed_t_layer):
                with torch.no_grad():
                    txt_embedding, txt_attention_probs = self.layer[idx](
                        txt_embedding, txt_attention_mask
                    )
                    t_start = self.fixed_t_layer
                    if output_all_attention_masks:
                        all_attention_mask_t.append(txt_attention_probs)

            for idx in range(t_start, t_end):
                txt_embedding, txt_attention_probs = self.layer[idx](
                    txt_embedding, txt_attention_mask
                )
                if output_all_attention_masks:
                    all_attention_mask_t.append(txt_attention_probs)

            for idx in range(v_start, self.fixed_v_layer):
                with torch.no_grad():
                    image_embedding, image_attention_probs = self.v_layer[idx](
                        image_embedding,
                        image_attention_mask,
                        txt_embedding,
                        txt_attention_mask2,
                    )
                    v_start = self.fixed_v_layer

                    if output_all_attention_masks:
                        all_attnetion_mask_v.append(image_attention_probs)

            for idx in range(v_start, v_end):
                image_embedding, image_attention_probs = self.v_layer[idx](
                    image_embedding,
                    image_attention_mask,
                    txt_embedding,
                    txt_attention_mask2,
                )

                if output_all_attention_masks:
                    all_attnetion_mask_v.append(image_attention_probs)

            if count == 0 and self.in_batch_pairs:
                # new batch size is the batch_size ^2
                image_embedding = (
                    image_embedding.unsqueeze(0)
                    .expand(batch_size, batch_size, num_regions, v_hidden_size)
                    .contiguous()
                    .view(batch_size * batch_size, num_regions, v_hidden_size)
                )
                image_attention_mask = (
                    image_attention_mask.unsqueeze(0)
                    .expand(batch_size, batch_size, 1, 1, num_regions)
                    .contiguous()
                    .view(batch_size * batch_size, 1, 1, num_regions)
                )

                txt_embedding = (
                    txt_embedding.unsqueeze(1)
                    .expand(batch_size, batch_size, num_words, t_hidden_size)
                    .contiguous()
                    .view(batch_size * batch_size, num_words, t_hidden_size)
                )
                txt_attention_mask = (
                    txt_attention_mask.unsqueeze(1)
                    .expand(batch_size, batch_size, 1, 1, num_words)
                    .contiguous()
                    .view(batch_size * batch_size, 1, 1, num_words)
                )
                co_attention_mask = (
                    co_attention_mask.unsqueeze(1)
                    .expand(batch_size, batch_size, 1, num_regions, num_words)
                    .contiguous()
                    .view(batch_size * batch_size, 1, num_regions, num_words)
                )

            if count == 0 and self.FAST_MODE:
                txt_embedding = txt_embedding.expand(
                    image_embedding.size(0),
                    txt_embedding.size(1),
                    txt_embedding.size(2),
                )
                txt_attention_mask = txt_attention_mask.expand(
                    image_embedding.size(0),
                    txt_attention_mask.size(1),
                    txt_attention_mask.size(2),
                    txt_attention_mask.size(3),
                )

            if self.with_coattention:
                # do the bi attention.
                image_embedding, txt_embedding, co_attention_probs = self.c_layer[
                    count
                ](
                    image_embedding,
                    image_attention_mask,
                    txt_embedding,
                    txt_attention_mask,
                    co_attention_mask,
                    use_co_attention_mask,
                )

                if output_all_attention_masks:
                    all_attention_mask_c.append(co_attention_probs)

            v_start = v_end
            t_start = t_end
            count += 1

            if output_all_encoded_layers:
                all_encoder_layers_t.append(txt_embedding)
                all_encoder_layers_v.append(image_embedding)

        for idx in range(v_start, len(self.v_layer)):
            image_embedding, image_attention_probs = self.v_layer[idx](
                image_embedding,
                image_attention_mask,
                txt_embedding,
                txt_attention_mask2,
            )

            if output_all_attention_masks:
                all_attnetion_mask_v.append(image_attention_probs)

        for idx in range(t_start, len(self.layer)):
            txt_embedding, txt_attention_probs = self.layer[idx](
                txt_embedding, txt_attention_mask
            )

            if output_all_attention_masks:
                all_attention_mask_t.append(txt_attention_probs)

        # add the end part to finish.
        if not output_all_encoded_layers:
            all_encoder_layers_t.append(txt_embedding)
            all_encoder_layers_v.append(image_embedding)

        return (
            all_encoder_layers_t,
            all_encoder_layers_v,
            (all_attention_mask_t, all_attnetion_mask_v, all_attention_mask_c),
        )


class BertTextPooler(nn.Module):
    def __init__(self, config):
        super(BertTextPooler, self).__init__()
        self.dense = nn.Linear(config.hidden_size, config.bi_hidden_size)
        self.activation = nn.ReLU()

    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 BertImagePooler(nn.Module):
    def __init__(self, config):
        super(BertImagePooler, self).__init__()
        self.dense = nn.Linear(config.v_hidden_size, config.bi_hidden_size)
        self.activation = nn.ReLU()

    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=1e-12)

    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 BertImgPredictionHeadTransform(nn.Module):
    def __init__(self, config):
        super(BertImgPredictionHeadTransform, self).__init__()
        self.dense = nn.Linear(config.v_hidden_size, config.v_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.v_hidden_act
        self.LayerNorm = BertLayerNorm(config.v_hidden_size, eps=1e-12)

    def forward(self, hidden_states):
        hidden_states = self.dense(hidden_states)
        hidden_states = self.transform_act_fn(hidden_states)
        hidden_states = self.LayerNorm(hidden_states)
        return hidden_states


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

        # The output weights are the same as the input embeddings, but there is
        # an output-only bias for each token.
        self.decoder = nn.Linear(
            bert_model_embedding_weights.size(1),
            bert_model_embedding_weights.size(0),
            bias=False,
        )
        self.decoder.weight = bert_model_embedding_weights
        self.bias = nn.Parameter(torch.zeros(bert_model_embedding_weights.size(0)))

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


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

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


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

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


class BertPreTrainingHeads(nn.Module):
    def __init__(self, config, bert_model_embedding_weights):
        super(BertPreTrainingHeads, self).__init__()
        self.predictions = BertLMPredictionHead(config, bert_model_embedding_weights)
        self.bi_seq_relationship = nn.Linear(config.bi_hidden_size, 2)
        self.imagePredictions = BertImagePredictionHead(config)
        self.fusion_method = config.fusion_method
        self.dropout = nn.Dropout(0.1)

    def forward(
        self, sequence_output_t, sequence_output_v, pooled_output_t, pooled_output_v
    ):

        if self.fusion_method == "sum":
            pooled_output = self.dropout(pooled_output_t + pooled_output_v)
        elif (self.fusion_method == "mul" or self.fusion_method == "text_only"
              or self.fusion_method == "img_only"):
            pooled_output = self.dropout(pooled_output_t * pooled_output_v)
        else:
            assert False

        prediction_scores_t = self.predictions(sequence_output_t)
        seq_relationship_score = self.bi_seq_relationship(pooled_output)
        prediction_scores_v = self.imagePredictions(sequence_output_v)

        return prediction_scores_t, prediction_scores_v, seq_relationship_score


class BertImagePredictionHead(nn.Module):
    def __init__(self, config):
        super(BertImagePredictionHead, self).__init__()
        self.transform = BertImgPredictionHeadTransform(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.v_hidden_size, config.v_target_size)

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


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__(self, *inputs, **kwargs):
        super(BertPreTrainedModel, self).__init__(*inputs, **kwargs)

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


class BertModel(BertPreTrainedModel):
    def __init__(self, config):
        super(BertModel, self).__init__(config)

        # initilize word embedding
        if config.model == "bert":
            self.embeddings = BertEmbeddings(config)
        elif config.model == "roberta":
            self.embeddings = RobertaEmbeddings(config)

        self.task_specific_tokens = config.task_specific_tokens

        # initlize the vision embedding
        self.v_embeddings = BertImageEmbeddings(config)

        self.encoder = BertEncoder(config)
        self.t_pooler = BertTextPooler(config)
        self.v_pooler = BertImagePooler(config)

        self.apply(self.init_weights)

    def forward(
        self,
        input_txt,
        input_imgs,
        image_loc=None,
        token_type_ids=None,
        attention_mask=None,
        image_attention_mask=None,
        co_attention_mask=None,
        task_ids=None,
        output_all_encoded_layers=False,
        output_all_attention_masks=False,
    ):
        if attention_mask is None:
            attention_mask = torch.ones_like(input_txt)
        if token_type_ids is None:
            token_type_ids = torch.zeros_like(input_txt)
        if image_attention_mask is None:
            image_attention_mask = torch.ones(
                input_imgs.size(0), input_imgs.size(1)
            ).type_as(input_txt)

        if self.task_specific_tokens:
            # extend the mask
            mask_tokens = input_txt.new().resize_(input_txt.size(0), 1).fill_(1)
            attention_mask = torch.cat([mask_tokens, attention_mask], dim=1)

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

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

        extended_attention_mask2 = extended_attention_mask2.to(
            dtype=next(self.parameters()).dtype
            # dtype=torch.float32,
        )  # fp16 compatibility

        extended_image_attention_mask = extended_image_attention_mask.to(
            dtype=next(self.parameters()).dtype
            # dtype=torch.float32,
        )  # fp16 compatibility
        extended_image_attention_mask = (1.0 - extended_image_attention_mask) * -10000.0

        if co_attention_mask is None:
            co_attention_mask = torch.zeros(
                input_txt.size(0), input_imgs.size(1), input_txt.size(1)
            ).type_as(extended_image_attention_mask)

        extended_co_attention_mask = co_attention_mask.unsqueeze(1)

        # extended_co_attention_mask = co_attention_mask.unsqueeze(-1)
        extended_co_attention_mask = extended_co_attention_mask * 5.0
        extended_co_attention_mask = extended_co_attention_mask.to(
            dtype=next(self.parameters()).dtype
            # dtype=torch.float32,
        )  # fp16 compatibility

        embedding_output = self.embeddings(input_txt, token_type_ids, task_ids)
        v_embedding_output = self.v_embeddings(input_imgs, image_loc)
        encoded_layers_t, encoded_layers_v, all_attention_mask = self.encoder(
            embedding_output,
            v_embedding_output,
            extended_attention_mask,
            extended_attention_mask2,
            extended_image_attention_mask,
            extended_co_attention_mask,
            output_all_encoded_layers=output_all_encoded_layers,
            output_all_attention_masks=output_all_attention_masks,
        )

        sequence_output_t = encoded_layers_t[-1]
        sequence_output_v = encoded_layers_v[-1]

        pooled_output_t = self.t_pooler(sequence_output_t)
        pooled_output_v = self.v_pooler(sequence_output_v)

        if not output_all_encoded_layers:
            encoded_layers_t = encoded_layers_t[-1]
            encoded_layers_v = encoded_layers_v[-1]

        return (
            encoded_layers_t,
            encoded_layers_v,
            pooled_output_t,
            pooled_output_v,
            all_attention_mask,
        )


class RobertaModel(BertModel):
    base_model_prefix = "roberta"

    def __init__(self, config):
        super(RobertaModel, self).__init__(config)

        self.apply(self.init_weights)


class BertImageEmbeddings(nn.Module):
    """Construct the embeddings from image, spatial location (omit now) and token_type embeddings.
    """

    def __init__(self, config):
        super(BertImageEmbeddings, self).__init__()

        self.image_embeddings = nn.Linear(config.v_feature_size, config.v_hidden_size)
        self.image_location_embeddings = nn.Linear(5, config.v_hidden_size)
        self.LayerNorm = BertLayerNorm(config.v_hidden_size, eps=1e-12)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, input_ids, input_loc=None):

        img_embeddings = self.image_embeddings(input_ids)
        if input_loc is not None:
            loc_embeddings = self.image_location_embeddings(input_loc)

        # TODO: we want to make the padding_idx == 0, however, with custom initilization, it seems it will have a bias.
        # Let's do masking for now
        if input_loc is not None:
            embeddings = self.LayerNorm(img_embeddings + loc_embeddings)
        else:
            embeddings = self.LayerNorm(img_embeddings)
        # embeddings = self.LayerNorm(img_embeddings+loc_embeddings)
        embeddings = self.dropout(embeddings)

        return embeddings


class BertForMultiModalPreTraining(BertPreTrainedModel):
    """BERT model with multi modal pre-training heads.
    """

    def __init__(self, config, vision_model=None):
        super(BertForMultiModalPreTraining, self).__init__(config)

        if config.model_type == "roberta":
            self.bert = RobertaModel(config)
        else:
            self.bert = BertModel(config)
        self.cls = BertPreTrainingHeads(
            config, self.bert.embeddings.word_embeddings.weight
        )

        self.apply(self.init_weights)
        self.visual_target = config.visual_target
        self.num_negative = config.num_negative
        self.loss_fct = CrossEntropyLoss(ignore_index=config.mlm_ignore_index)

        print("model's visual target is ", config.visual_target)

        if self.visual_target == 0:
            self.vis_criterion = nn.KLDivLoss(reduction="none")
        elif self.visual_target == 1:
            self.vis_criterion = nn.MSELoss(reduction="none")
        elif self.visual_target == 2:
            self.vis_criterion = CrossEntropyLoss()
        elif self.visual_target == 3:
            self.vis_criterion = CrossEntropyLoss(
                ignore_index=config.visual_mlm_ignore_index)

        self.tie_weights()

        self.vision_model = vision_model
        if self.vision_model is not None:
            # Remove the final FC layer.
            try:
                self.num_img_dim = self.vision_model.fc.in_features
                self.vision_model.fc = nn.Identity()
            except:
                # Detectron2 models
                self.num_img_dim = self.config.v_feature_size

        # TODO: Work on this.
        self.freeze_vision_model = config.freeze_vision_model
        self.multimodal_text_part = config.multimodal_text_part
        self.multimodal_img_part = config.multimodal_img_part
        self.visual_mlm_probability = config.visual_mlm_probability
        self.visual_mlm_ignore_index = config.visual_mlm_ignore_index
        self.swapping_based_nsp = config.swapping_based_nsp
        self.swapping_based_nsp_prob = 0.5
        self.sequence_based_nsp = config.sequence_based_nsp
        self.sequence_based_nsp_prob = 0.5
        self.visual_mlm = config.visual_mlm
        self.no_mlm = config.no_mlm
        self.cls_id = config.cls_id
        self.vilbert_paired_coattention = config.vilbert_paired_coattention
        
        if self.sequence_based_nsp:
            self.seq_based_relationship = nn.Linear(config.bi_hidden_size, 2)

        self.max_story_length = config.max_story_length

        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.include_num_img_regional_features = config.include_num_img_regional_features

        self.config = config

    def tie_weights(self):
        """ Make sure we are sharing the input and output embeddings.
            Export to TorchScript can't handle parameter sharing so we are cloning them instead.
        """
        self._tie_or_clone_weights(
            self.cls.predictions.decoder, self.bert.embeddings.word_embeddings
        )

    def forward(self, inputs):
        outputs = self._forward(**inputs)

        masked_lm_labels = inputs["masked_lm_labels"]
        next_sentence_label = inputs["next_sentence_label"] if "next_sentence_label" in inputs else None
        image_target = inputs["image_target"] if "image_target" in inputs else None

        if (
            masked_lm_labels is not None
            # and next_sentence_label is not None
            # and image_target is not None
        ):
            (
                masked_lm_loss,
                masked_img_loss,
                next_sentence_loss
            ) = outputs

            total_loss = 0
            if not self.no_mlm:
                total_loss += 0.2 * masked_lm_loss.unsqueeze(0)
            if masked_img_loss is not None and self.visual_mlm:
                total_loss += masked_img_loss.unsqueeze(0)
            if next_sentence_loss is not None and self.swapping_based_nsp:
                total_loss += next_sentence_loss.unsqueeze(0)

            return total_loss, None  # TODO: Change this.

            # return (
            #     masked_lm_loss.unsqueeze(0),
            #     masked_img_loss.unsqueeze(0) if masked_img_loss is not None else None,
            #     next_sentence_loss.unsqueeze(0) if next_sentence_loss is not None else None,
            # )
        else:
            (
                prediction_scores_t,
                prediction_scores_v,
                seq_relationship_score,
                all_attention_mask,
            ) = outputs
        
        return outputs

    def _forward(
        self,
        input_ids,
        images,
        image_loc=None,
        token_type_ids=None,
        attention_mask=None,
        image_attention_mask=None,
        masked_lm_labels=None,
        image_label=None,
        image_target=None,
        next_sentence_label=None,
        output_all_attention_masks=False,
        img_regional_features=None,
    ):

        bz, text_len = input_ids.size()
        cls_repr_pos = []
        for i in range(bz):
            cls_repr_pos_i = torch.nonzero(input_ids[i]==self.cls_id, as_tuple=False)
            cls_repr_pos.append(cls_repr_pos_i)

        if self.vision_model is not None and images.ndim > 3:
            # print(self.vision_model.layer4[0].conv1.weight);raise
            # Reshape from B x L x C x H x W to (B*L) x C x H x W
            bz, img_len, C, H, W = images.size()
            if self.multimodal_text_part:
                images = torch.zeros(images.size(0), 1,
                                     self.num_img_dim).type_as(images).float()
            else:
                images = torch.reshape(images, (bz*img_len, C, H, W)).float()
                images = self.vision_model(images)
                if type(images) == tuple:
                    images, img_regional_features = images
                    img_regional_features = torch.reshape(
                        img_regional_features,
                        (bz, img_len, self.include_num_img_regional_features,  # -1,
                        self.num_img_dim))
                if self.freeze_vision_model:
                    images = images.detach()
                images = torch.reshape(images, (bz, img_len, self.num_img_dim))

        if self.multimodal_img_part:
            input_ids = torch.zeros(input_ids.size(0), 1).type_as(input_ids)
            attention_mask = torch.zeros(input_ids.size(0), 1).type_as(attention_mask)
            # input_ids = torch.zeros(input_ids.size(0),
            #     input_ids.size(1)).type_as(input_ids)
            # attention_mask = torch.zeros(input_ids.size(0),
            #     input_ids.size(1)).type_as(attention_mask)
            # print(torch.sum(input_ids))
            # print(torch.sum(attention_mask));raise
            token_type_ids = None

        if images is not None:
            bz, img_len, D = images.size()
            
            if self.include_num_img_regional_features is not None:
                img_regional_features = img_regional_features.float()

            # Perform image masking
            if self.visual_target == 3 and self.visual_mlm:
                if self.include_num_img_regional_features is not None:
                    raise NotImplementedError("Not done yet!")
                for i in range(bz):
                    image_ = images[i]  # L x D
                    image_lenwise_sum = torch.sum(image_, dim=-1)
                    # print(image_lenwise_sum)
                    non_zero_images = image_lenwise_sum.nonzero().t()[0]
                    masked_img_indices_ = torch.bernoulli(torch.full(
                        non_zero_images.shape,
                        self.visual_mlm_probability)).bool()
                    label_ = image_target[i]
                    label_[~masked_img_indices_] = self.visual_mlm_ignore_index
                    indices_replaced = torch.bernoulli(torch.full(
                        label_.shape, 0.8)).bool() & masked_img_indices_
                    indices_random = torch.bernoulli(torch.full(label_.shape,
                        0.5)).bool() & masked_img_indices_ & ~indices_replaced
                    visual_mask_token = torch.zeros_like(image_)
                    image_[indices_replaced] = visual_mask_token[indices_replaced]
                    image_[indices_random] = images[i][indices_random]
                    new_image_lenwise_sum = torch.sum(image_, dim=-1)
                    images[i] = image_
                    image_target[i] = label_
                    # print(label_)
                    # print(new_image_lenwise_sum)
                pass
            pass

            # Perform swapping-based multimodal alignment objective.
            if self.swapping_based_nsp:
                swapping_based_nsp_labels = []
                images_if_swapped = torch.zeros(bz, img_len)
                for i in range(bz):
                    image_ = images[i]  # L x D
                    image_lenwise_sum = torch.sum(image_, dim=-1)
                    # TODO: Since our visual mask token is 0.
                    # non_zero_images = image_lenwise_sum.nonzero().t()[0]
                    non_zero_images = torch.nonzero(image_lenwise_sum, as_tuple=False).t()[0]

                    if len(non_zero_images) == 0:
                        swapping_based_nsp_labels.append(1)
                        continue

                    sample_batch_idx = i + 1
                    if i == bz - 1:
                       sample_batch_idx = 0
                    image_cands_ = images[sample_batch_idx]
                    image_cands_lenwise_sum = torch.sum(image_cands_, dim=-1)
                    # non_zero_image_cands_ = image_cands_lenwise_sum.nonzero().t()[0]
                    non_zero_image_cands_ = torch.nonzero(image_cands_lenwise_sum, as_tuple=False).t()[0]
                    if len(non_zero_image_cands_) == 0:
                        swapping_based_nsp_labels.append(1)
                        continue

                    non_zero_image_cands_ = non_zero_image_cands_.detach().cpu().numpy().astype(int)
                    non_zero_images = non_zero_images.detach().cpu().numpy().astype(int)

                    # TODO: Prevent swapping the already swapped images.
                    non_zero_image_cands_ = set(list(non_zero_image_cands_))
                    images_if_swapped_i = torch.nonzero(
                        images_if_swapped[
                            sample_batch_idx],
                            as_tuple=False).t()[0].detach().cpu().numpy().astype(int)
                    images_if_swapped_i = set(list(images_if_swapped_i))
                    non_zero_image_cands_ -= images_if_swapped_i
                    non_zero_image_cands_ = list(non_zero_image_cands_)
                    if len(non_zero_image_cands_) == 0:
                        swapping_based_nsp_labels.append(1)
                        continue

                    chose_index = np.random.choice(non_zero_image_cands_)
                    swapped_index = np.random.choice(non_zero_images)

                    # Probability of swapping.
                    if_swap = np.random.rand()
                    if if_swap > self.swapping_based_nsp_prob:
                        image_[swapped_index] = image_cands_[chose_index]
                        swapping_based_nsp_labels.append(0)
                        images_if_swapped[i][swapped_index] = 1
                        if self.include_num_img_regional_features is not None:
                            img_regional_features[i][swapped_index] = \
                                img_regional_features[
                                    sample_batch_idx][chose_index]
                    else:
                        swapping_based_nsp_labels.append(1)

                    images[i] = image_

                swapping_based_nsp_labels = torch.Tensor(
                    swapping_based_nsp_labels).type_as(image_target)

            # Perform sequence-based multimodal alignment objective.
            if self.sequence_based_nsp:
                sequence_based_nsp_labels = []
                images_if_sequenced = torch.zeros(bz, img_len)
                for i in range(bz):
                    image_ = images[i]  # L x D
                    image_lenwise_sum = torch.sum(image_, dim=-1)
                    # TODO: Since our visual mask token is 0.
                    # non_zero_images = image_lenwise_sum.nonzero().t()[0]
                    non_zero_images = torch.nonzero(image_lenwise_sum, as_tuple=False).t()[0]

                    if len(non_zero_images) == 0:
                        swapping_based_nsp_labels.append(1)
                        continue

                    # TODO: Prevent swapping based duplicattions.
                    if self.swapping_based_nsp:
                        all_images_on = torch.zeros(img_len)
                        all_images_on[non_zero_images] = 1
                        all_images_on -= images_if_swapped[i]
                        all_images_on_pos = torch.nonzero(all_images_on, as_tuple=False).t()[0]
                    else:
                        all_images_on_pos = non_zero_images
                    
                    all_images_on_pos = all_images_on_pos.detach().cpu().numpy().astype(int)
                    if len(all_images_on_pos) < 2:
                        sequence_based_nsp_labels.append(1)
                        continue

                    # Probability of swapping.
                    chosen_index1, chosen_index2 = np.random.choice(
                        all_images_on_pos, 2)
                    if_seq_swap = np.random.rand()
                    if if_seq_swap > self.sequence_based_nsp_prob:
                        image_idx1 = image_[chosen_index1]
                        image_[chosen_index1] = image_[chosen_index2]
                        image_[chosen_index2] = image_idx1
                        sequence_based_nsp_labels.append(0)
                        images_if_sequenced[i][chosen_index1] = 1
                        images_if_sequenced[i][chosen_index2] = 1
                        if self.include_num_img_regional_features is not None:
                            img_regional_features_idx1 = \
                                img_regional_features[i][chosen_index1]
                            img_regional_features[i][chosen_index1] = \
                                img_regional_features[i][chosen_index1]
                            img_regional_features[i][chosen_index2] = \
                                img_regional_features_idx1
                    else:
                        sequence_based_nsp_labels.append(1)

                    images[i] = image_

                sequence_based_nsp_labels = torch.Tensor(
                    sequence_based_nsp_labels).type_as(image_target)

        # TODO: Double check this part.
        co_attention_mask = None
        if self.vilbert_paired_coattention:
            paired_co_attention_mask = torch.zeros(
                input_ids.size(0), images.size(1), input_ids.size(1)
            ).type_as(images)
            if self.include_num_img_regional_features is not None:
                paired_co_attention_mask = torch.zeros(
                    input_ids.size(0),
                    images.size(1) + self.include_num_img_regional_features * images.size(1),
                    input_ids.size(1)
                ).type_as(images)
            for i in range(bz):
                cls_repr_pos_i = cls_repr_pos[i].squeeze().cpu().numpy()
                if len(cls_repr_pos_i) > self.max_story_length:
                    print(input_ids)
                    print(cls_repr_pos_i)
                    cls_repr_pos_i = cls_repr_pos_i[:self.max_story_length]
                for j in range(len(cls_repr_pos_i)):
                    # TODO: Currently the last sequence item is not
                    # truncated yet.
                    if j == len(cls_repr_pos_i) - 1:  # The last item.
                        start = cls_repr_pos_i[j]
                        end = input_ids.size(1)
                    else:
                        start = cls_repr_pos_i[j]
                        end = cls_repr_pos_i[j+1]
                    if self.include_num_img_regional_features is not None:
                        beta = self.include_num_img_regional_features
                        paired_co_attention_mask[i, j*beta+j:(j+1)*beta+j+1, start:end] = 1.0
                    else:
                        paired_co_attention_mask[i, j, start:end] = 1.0
            co_attention_mask = paired_co_attention_mask

        if self.vision_model is not None:
            # print(self.vision_model.layer4[0].conv1.weight);raise
            # Reshape from B x L x C x H x W to (B*L) x C x H x W
            bz, img_len, CC = images.size()
            if not self.multimodal_text_part:
                if self.include_num_img_regional_features is not None:
                    new_images = []
                    for i in range(bz):
                        curr_img = images[i]
                        curr_img_region_ = []
                        for j in range(img_len):
                            curr_seq_img = curr_img[j].unsqueeze(0)
                            curr_regional_features = img_regional_features[i, j, :, :]
                            curr_new_img = torch.cat([curr_seq_img, curr_regional_features])
                            curr_img_region_.append(curr_new_img)
                        curr_img_region_ = torch.cat(curr_img_region_)
                        new_images.append(curr_img_region_)
                    new_images = torch.stack(new_images)
                    images = new_images
                pass
            pass
        pass

        if (co_attention_mask is None and self.config.with_coattention
            and not self.vilbert_paired_coattention):
            co_attention_mask = torch.ones(
                input_ids.size(0), images.size(1), input_ids.size(1)
            ).type_as(images)

        # In this model, we first embed the images.
        sequence_output_t, sequence_output_v, pooled_output_t, pooled_output_v, all_attention_mask = self.bert(
            input_ids,
            images,
            image_loc,
            token_type_ids,
            attention_mask,
            image_attention_mask,
            co_attention_mask=co_attention_mask,
            output_all_encoded_layers=False,
            output_all_attention_masks=output_all_attention_masks,
        )

        prediction_scores_t, prediction_scores_v, seq_relationship_score = self.cls(
            sequence_output_t, sequence_output_v, pooled_output_t, pooled_output_v
        )

        masked_lm_loss = None
        masked_img_loss = None
        next_sentence_loss = None

        if (
            masked_lm_labels is not None
            # and next_sentence_label is not None
            # and image_target is not None
        ):
            if image_target is not None:
                if self.visual_target != 3:
                    prediction_scores_v = prediction_scores_v[:, 1:]
                if self.visual_target == 1:
                    img_loss = self.vis_criterion(prediction_scores_v, image_target)
                    masked_img_loss = torch.sum(
                        img_loss * (image_label == 1).unsqueeze(2).float()
                    ) / max(
                        torch.sum((image_label == 1).unsqueeze(2).expand_as(img_loss)), 1
                    )

                elif self.visual_target == 0:
                    img_loss = self.vis_criterion(
                        F.log_softmax(prediction_scores_v, dim=2), image_target
                    )

                    masked_img_loss = torch.sum(
                        img_loss * (image_label == 1).unsqueeze(2).float()
                    # ) / max(torch.sum((image_label == 1)), 0)
                    ) / max(torch.sum((image_label == 1)), 1) # FIXME: NaN issue
                    # print (image_label[0], torch.sum(image_label[0]), len(image_label[0]))
                    # print (masked_img_loss)
                    # import ipdb
                    # ipdb.set_trace()

                elif self.visual_target == 2:
                    # generate negative sampled index.
                    num_negative = self.num_negative
                    num_across_batch = int(self.num_negative * 0.7)
                    num_inside_batch = int(self.num_negative * 0.3)

                    batch_size, num_regions, _ = prediction_scores_v.size()
                    assert batch_size != 0
                    # random negative across batches.
                    row_across_index = input_ids.new(
                        batch_size, num_regions, num_across_batch
                    ).random_(0, batch_size)
                    col_across_index = input_ids.new(
                        batch_size, num_regions, num_across_batch
                    ).random_(0, num_regions)

                    for i in range(batch_size - 1):
                        row_across_index[i][row_across_index[i] == i] = batch_size - 1
                    final_across_index = row_across_index * num_regions + col_across_index

                    # random negative inside batches.
                    row_inside_index = input_ids.new(
                        batch_size, num_regions, num_inside_batch
                    ).zero_()
                    col_inside_index = input_ids.new(
                        batch_size, num_regions, num_inside_batch
                    ).random_(0, num_regions - 1)

                    for i in range(batch_size):
                        row_inside_index[i] = i
                    for i in range(num_regions - 1):
                        col_inside_index[:, i, :][col_inside_index[:, i, :] == i] = (
                            num_regions - 1
                        )
                    final_inside_index = row_inside_index * num_regions + col_inside_index

                    final_index = torch.cat((final_across_index, final_inside_index), dim=2)

                    # Let's first sample where we need to compute.
                    predict_v = prediction_scores_v[image_label == 1]
                    neg_index_v = final_index[image_label == 1]

                    flat_image_target = image_target.view(batch_size * num_regions, -1)
                    # we also need to append the target feature at the begining.
                    negative_v = flat_image_target[neg_index_v]
                    positive_v = image_target[image_label == 1]
                    sample_v = torch.cat((positive_v.unsqueeze(1), negative_v), dim=1)

                    # calculate the loss.
                    score = torch.bmm(sample_v, predict_v.unsqueeze(2)).squeeze(2)
                    if len(score) > 0:
                        masked_img_loss = self.vis_criterion(
                            score, input_ids.new(score.size(0)).zero_()
                        )
                    else:
                        masked_img_loss = None

                elif self.visual_target == 3 and self.visual_mlm:
                    masked_img_loss = self.vis_criterion(
                        prediction_scores_v.view(-1,
                            self.config.v_target_size), image_target.view(-1))
                    pass

            # masked_img_loss = torch.sum(img_loss) / (img_loss.shape[0] * img_loss.shape[1])
            masked_lm_loss = self.loss_fct(
                prediction_scores_t.view(-1, self.config.vocab_size),
                masked_lm_labels.view(-1),
            )

            if self.swapping_based_nsp or self.sequence_based_nsp:
                next_sentence_loss = 0
                if self.swapping_based_nsp:
                    next_sentence_loss += self.loss_fct(
                        seq_relationship_score.view(-1, 2), swapping_based_nsp_labels.view(-1)
                    )
                if self.sequence_based_nsp:
                    if self.config.fusion_method == "sum":
                        pooled_output = self.dropout(pooled_output_t + pooled_output_v)
                    elif self.config.fusion_method == "mul":
                        pooled_output = self.dropout(pooled_output_t * pooled_output_v)
                    else:
                        assert False
                    seq_relationship_score2 = self.seq_based_relationship(pooled_output)
                    next_sentence_loss += self.loss_fct(
                        seq_relationship_score2.view(-1, 2), sequence_based_nsp_labels.view(-1)
                    )
            elif next_sentence_label is not None:
                next_sentence_loss = self.loss_fct(
                    seq_relationship_score.view(-1, 2), next_sentence_label.view(-1)
                )

            return (
                masked_lm_loss.unsqueeze(0),
                masked_img_loss.unsqueeze(0) if masked_img_loss is not None else None,
                next_sentence_loss.unsqueeze(0) if next_sentence_loss is not None else None,
            )
        else:
            return (
                prediction_scores_t,
                prediction_scores_v,
                seq_relationship_score,
                all_attention_mask,
            )


class VILBertForVLTasks(BertPreTrainedModel):
    def __init__(self, config, num_labels, vision_model=None,
                 dropout_prob=0.1, default_gpu=True):

        if (config.fusion_method == "img_only"
            or config.fusion_method == "text_only"):
            config.with_coattention = False
            logging.warning("Not using co-attention due to: {}".format(config.fusion_method))

        super(VILBertForVLTasks, self).__init__(config)
        self.config = config
        self.num_labels = num_labels

        if config.model_type == "roberta":
            self.bert = RobertaModel(config)
        else:
            self.bert = BertModel(config)
        self.dropout = nn.Dropout(dropout_prob)
        self.cls = BertPreTrainingHeads(
            config, self.bert.embeddings.word_embeddings.weight
        )
        self.vil_prediction = SimpleClassifier(
            config.bi_hidden_size, config.bi_hidden_size * 2, 3129, 0.5
        )
        self.vil_prediction_gqa = SimpleClassifier(
            config.bi_hidden_size, config.bi_hidden_size * 2, 1533, 0.5
        )
        self.vil_binary_prediction = SimpleClassifier(
            config.bi_hidden_size * 2, config.bi_hidden_size * 2, 2, 0.5
        )
        self.vil_logit = nn.Linear(config.bi_hidden_size, 1)
        self.vil_tri_prediction = nn.Linear(
            config.bi_hidden_size, 3
        )  # for Visual Entailiment tasks
        # TODO: for sort
        self.vil_sort_prediction = SimpleClassifier(
            config.bi_hidden_size, config.bi_hidden_size * 2, 2, 0.5
        )
        print (self.vil_sort_prediction)
        self.vision_logit = nn.Linear(config.v_hidden_size, 1)
        self.linguistic_logit = nn.Linear(config.hidden_size, 1)
        self.fusion_method = config.fusion_method
        self.apply(self.init_weights)

        self.tie_weights()

        self.vision_model = vision_model
        if self.vision_model is not None:
            # Remove the final FC layer.
            try:
                self.num_img_dim = self.vision_model.fc.in_features
                self.vision_model.fc = nn.Identity()
            except:
                # Detectron2 models
                self.num_img_dim = self.config.v_feature_size

        self.vil_sort_loss = nn.CrossEntropyLoss()

        # TODO: Work on this.
        self.freeze_vision_model = config.freeze_vision_model
        self.multimodal_text_part = config.multimodal_text_part
        self.multimodal_img_part = config.multimodal_img_part
        self.vilbert_use_3way_logits = config.vilbert_use_3way_logits
        if self.vilbert_use_3way_logits:
            logging.info("Using logits from ALL modalities (3-way)!")
        self.cls_id = config.cls_id
        self.hierarchical_version = config.hierarchical_version
        
        # For the heatmap
        self.max_story_length = config.max_story_length
        if self.hierarchical_version != "v0":
            if self.vision_model is not None:
                if self.num_img_dim != 2048:
                    self.img_embd_up_proj = torch.nn.Linear(512, 2048)
                    print(self.img_embd_up_proj)
            self.heat_map_prediction = SimpleClassifier(
                config.bi_hidden_size, config.bi_hidden_size * 2,
                # self.max_story_length, 0.5
                5, 0.5
            )
            self.heat_map_prediction = torch.nn.Linear(config.bi_hidden_size, self.max_story_length)
            self.heat_map_sort_loss = torch.nn.MSELoss()
            # self.heat_map_sort_loss = torch.nn.SmoothL1Loss()
            # self.heat_map_sort_loss = torch.nn.BCELoss()
            # self.heat_map_sort_loss = torch.nn.BCEWithLogitsLoss()
            self.use_kl_div = False
            self.use_cls = False
            if self.use_kl_div:
                self.heat_map_sort_loss = torch.nn.KLDivLoss()
            if self.use_cls:
                self.heat_map_sort_loss = torch.nn.CrossEntropyLoss()
            print(self.heat_map_prediction)
            print(self.heat_map_sort_loss)

            self.additional_hl = False  # False
            if self.additional_hl:
                self.additional_hl_layer = BertLayer(config)
                print(self.additional_hl_layer)

        self.simple_img_classifier = config.simple_img_classifier
        if self.multimodal_img_part and self.simple_img_classifier:
            self.simple_img_prediction = SimpleClassifier(
                self.num_img_dim, config.bi_hidden_size * 2,
                config.bi_hidden_size, 0.2
            )

        self.vilbert_paired_coattention = config.vilbert_paired_coattention
        self.include_num_img_regional_features = config.include_num_img_regional_features
        self.swapping_based_nsp = False

        if "hl_include_objectives" not in config.__dict__:
            config.hl_include_objectives = []
            self.hl_include_objectives = []
        else:
            self.hl_include_objectives = config.hl_include_objectives
            logging.info("Aux using: {}".format(self.hl_include_objectives))
            if self.hl_include_objectives is not None:
                assert type(self.hl_include_objectives) == list
                if ("binary" in self.hl_include_objectives
                    or "pairwise" in self.hl_include_objectives):
                    self.hl_bin_pred_layer = SimpleClassifier(
                        # config.bi_hidden_size * num_labels,
                        # config.bi_hidden_size,
                        # num_labels, 0.5
                        config.bi_hidden_size * 1,
                        config.bi_hidden_size,
                        1, 0.5
                    )
                    self.hl_bin_pred_crit = torch.nn.CrossEntropyLoss()
                    self.hl_bin_sparse_prob = 1.5  # 1.0
                    if self.hl_bin_sparse_prob < 1.0:
                        self.hl_bin_sparse_pos = []
                if "binary_cross_modal" in self.hl_include_objectives:
                    self.hl_binx_pred_layer = SimpleClassifier(
                        # config.bi_hidden_size * num_labels,
                        # config.bi_hidden_size,
                        # num_labels, 0.5
                        config.bi_hidden_size * 1 + config.v_hidden_size,
                        config.bi_hidden_size,
                        2, 0.5
                    )
                    self.hl_binx_pred_crit = torch.nn.CrossEntropyLoss()
                    self.hl_binx_sparse_prob = 1.5  # 1.0
                    if self.hl_binx_sparse_prob < 1.0:
                        self.hl_binx_sparse_pos = []
                if "head" in self.hl_include_objectives:
                    self.hl_head_pred_layer = SimpleClassifier(
                        # config.bi_hidden_size * num_labels,
                        # config.bi_hidden_size,
                        # num_labels, 0.5
                        config.bi_hidden_size * 1,
                        config.bi_hidden_size,
                        1, 0.5
                    )
                    self.hl_head_pred_crit = torch.nn.CrossEntropyLoss()

                if "mlm" in self.hl_include_objectives:
                    self.mlm_loss_fct = CrossEntropyLoss(
                        ignore_index=config.mlm_ignore_index)

                if "itm" in self.hl_include_objectives:
                    if "swapping_based_nsp" in config.__dict__:
                        self.swapping_based_nsp = config.swapping_based_nsp
                        self.swapping_based_nsp_prob = 0.5
                        self.itm_loss_fct = CrossEntropyLoss()
                    else:
                        raise ValueError("No `swapping_based_nsp` in config.")
                else:
                    self.swapping_based_nsp = False

                if "cross_modal_dependence" in self.hl_include_objectives:
                    self.cross_modal_dependence_prediction = SimpleClassifier(
                        config.v_hidden_size, config.v_hidden_size * 1,
                        self.max_story_length, 0.5
                    )
                    self.cross_modal_dependence_loss = torch.nn.MSELoss()

                if "heatmap_module" in self.hl_include_objectives:
                    self.heatmap = HeatMapOutput(config)

        self.heatmap_late_fusion = False  # False
        if self.heatmap_late_fusion:
            self.heatmap_img_proj_layer = torch.nn.Linear(768, 1024)
            print(self.heatmap_img_proj_layer)

        if config.fusion_method == "img_only":
            self._img_only_inter_proj = torch.nn.Linear(
                config.v_hidden_size, config.bi_hidden_size)
            print(self._img_only_inter_proj)

    def tie_weights(self):
        """ Make sure we are sharing the input and output embeddings.
            Export to TorchScript can't handle parameter sharing so we are cloning them instead.
        """
        self._tie_or_clone_weights(
            self.cls.predictions.decoder, self.bert.embeddings.word_embeddings
        )

    def forward(self, inputs):
        (vil_prediction, vil_prediction_gqa, vil_logit, vil_binary_prediction,
         vil_tri_prediction, vision_prediction, vision_logit,
         linguistic_prediction, linguistic_logit, _,
         vil_sort_prediction, hl_aux_predictions,
         sequence_output_t, sequence_output_v) = self._forward(**inputs)

        if self.config.wrapper_model_type == "berson":
            if self.config.multimodal_loss:
                if self.config.include_num_img_regional_features:
                    raise NotImplementedError("Not done yet!")
                return (sequence_output_t, sequence_output_v, None)
            return (sequence_output_t, None)
        
        vil_sort_logits = vil_sort_prediction
        # Heatmap handlings.
        if self.hl_include_objectives is not None:
            if "heatmap_module" in self.hl_include_objectives:
                x = self.heatmap(inputs, sequence_output_t)
                return x

        if self.hierarchical_version != "v0":
            if not self.use_kl_div and not self.use_cls:
                # vil_sort_logits = torch.sigmoid(vil_sort_logits)
                # vil_sort_logits = torch.tanh(vil_sort_logits)
                # vil_sort_logits = F.softmax(vil_sort_logits)
                pass

        if "labels" in inputs:
            if self.hierarchical_version == "v0":
                loss = self.vil_sort_loss(vil_sort_logits, inputs["labels"])
            else:
                heatmap_labels = []

                if self.use_cls:
                    loss = 0

                for i in range(len(inputs["labels"])):
                    labels_i = inputs["labels"][i].detach().cpu().numpy()
                    heatmap_labels_i = render_order_heatmap(args=self.config,  # TODO: change this.
                                                            order_list=labels_i,
                                                            soft=False)
                    heatmap_labels.append(heatmap_labels_i)

                    if self.use_cls:
                        labels_i = inputs["labels"][i]
                        # print(vil_sort_logits[i].size(), labels_i.size())
                        loss_cls_curr = self.heat_map_sort_loss(vil_sort_logits[i], labels_i)
                        loss += loss_cls_curr

                heatmap_labels = torch.stack(heatmap_labels)
                heatmap_labels = heatmap_labels.type_as(vil_sort_logits)
                # log_softmax = torch.nn.LogSoftmax()
                # log_softmax = torch.nn.Softmax()
                if self.use_kl_div:
                    # print(heatmap_labels)
                    vil_sort_logits = torch.sigmoid(vil_sort_logits)
                    vil_sort_logits_soft = F.log_softmax(vil_sort_logits)
                    heatmap_labels_soft = heatmap_labels + 1e-12
                    heatmap_labels_soft = F.softmax(heatmap_labels_soft)
                    loss = torch.nn.KLDivLoss(size_average=True)(vil_sort_logits_soft, heatmap_labels_soft)
                    # print(vil_sort_logits_soft)
                    # print(heatmap_labels_soft)
                    # print(loss)

                elif self.use_cls:
                    pass
                    vil_sort_logits = torch.sigmoid(vil_sort_logits)

                else:
                    # FIXME Trying transpose.
                    # print(vil_sort_logits.size())
                    vil_sort_logits = torch.transpose(vil_sort_logits, 1, 2)
                    # print(vil_sort_logits.size());raise
                    # heatmap_labels = F.softmax(heatmap_labels)
                    vil_sort_logits = vil_sort_logits[:, :self.max_story_length, :self.max_story_length]

                    loss = self.heat_map_sort_loss(vil_sort_logits, heatmap_labels)
                    # vil_sort_logits = torch.sigmoid(vil_sort_logits)

                if self.hl_include_objectives is not None:
                    if len(self.hl_include_objectives) > 0:
                        for h in range(len(self.hl_include_objectives)):
                            hl_aux_objective = self.hl_include_objectives[h]
                            if "mlm_wo_loss" == hl_aux_objective:
                                pass
                            elif "mlm" == hl_aux_objective:
                                masked_lm_labels = inputs["masked_lm_labels"]
                                masked_lm_loss = self.mlm_loss_fct(
                                    linguistic_prediction.view(-1, self.config.vocab_size),
                                    masked_lm_labels.view(-1),
                                )
                                loss += 0.05 * masked_lm_loss
                            elif "itm" == hl_aux_objective:
                                swapping_based_nsp_loss = self.itm_loss_fct(
                                    vil_binary_prediction,
                                    self.swapping_based_nsp_labels
                                )
                                loss += 0.1 * swapping_based_nsp_loss
                            elif "head" == hl_aux_objective:
                                head_labels = inputs["labels"][:, 0]
                                hl_aux_head_loss = self.hl_head_pred_crit(
                                    hl_aux_predictions[h], head_labels)
                                loss += hl_aux_head_loss
                            elif ("pairwise" == hl_aux_objective
                                  or "binary" == hl_aux_objective):
                                hl_bin_predictions = hl_aux_predictions[h]
                                bz, seq_len = inputs["labels"].size()
                                for b in range(bz):
                                    label_curr = inputs["labels"][b]
                                    label_index = torch.argsort(label_curr)
                                    hl_bin_label = torch.zeros(
                                        hl_bin_predictions.size()[1]).type_as(
                                            inputs["labels"])

                                    bin_idx = 0
                                    if self.hl_bin_sparse_prob < 1.0:
                                        hl_bin_label = []
                                    for seq_i in range(seq_len):
                                        for seq_j in range(seq_i+1, seq_len):
                                            if label_index[seq_i] < label_index[seq_j]:
                                                if self.hl_bin_sparse_prob < 1.0:
                                                    hl_bin_label.append(1)
                                                else:
                                                    hl_bin_label[bin_idx]  = 1
                                            else:
                                                if self.hl_bin_sparse_prob < 1.0:
                                                    hl_bin_label.append(0)
                                            bin_idx += 1

                                    if self.hl_bin_sparse_prob < 1.0:
                                        hl_bin_label = torch.Tensor(hl_bin_label).type_as(inputs["labels"])
                                        # TODO: Change this!!!
                                        # hl_bin_label = hl_bin_label[:3]
                                        hl_bin_label = hl_bin_label[self.hl_bin_sparse_pos[b]]

                                    hl_aux_bin_loss_curr = self.hl_bin_pred_crit(hl_bin_predictions[b], hl_bin_label)
                                    loss += 0.1 * hl_aux_bin_loss_curr
                                        
                            elif "binary_cross_modal" == hl_aux_objective:
                                hl_binx_predictions = hl_aux_predictions[h]
                                bz, seq_len = inputs["labels"].size()
                                for b in range(bz):
                                    label_curr = inputs["labels"][b]
                                    label_index = torch.argsort(label_curr)
                                    hl_binx_label = torch.zeros(
                                        hl_binx_predictions.size()[1]).type_as(
                                            inputs["labels"])

                                    binx_idx = 0
                                    if self.hl_binx_sparse_prob < 1.0:
                                        hl_binx_label = []
                                    for seq_i in range(seq_len):
                                        for seq_j in range(seq_i+1, seq_len):
                                            if label_index[seq_i] < label_index[seq_j]:
                                                if self.hl_binx_sparse_prob < 1.0:
                                                    hl_binx_label.append(1)
                                                else:
                                                    hl_binx_label[binx_idx]  = 1
                                            else:
                                                if self.hl_binx_sparse_prob < 1.0:
                                                    hl_binx_label.append(0)
                                            binx_idx += 1

                                    if self.hl_binx_sparse_prob < 1.0:
                                        hl_binx_label = torch.Tensor(hl_binx_label).type_as(inputs["labels"])
                                        # TODO: Change this!!!
                                        # hl_bin_label = hl_bin_label[:3]
                                        hl_binx_label = hl_binx_label[self.hl_binx_sparse_pos[b]]

                                    hl_aux_binx_loss_curr = self.hl_binx_pred_crit(hl_binx_predictions[b], hl_binx_label)
                                    loss += hl_aux_binx_loss_curr
                                        
                                pass

                            elif "cross_modal_dependence" == hl_aux_objective:
                                hl_binx_predictions = hl_aux_predictions[h]
                                cross_modal_dependence_preds = torch.sigmoid(
                                    hl_binx_predictions)
                                cross_modal_dependence_logits = torch.transpose(
                                    cross_modal_dependence_preds, 1, 2)
                                cross_modal_dependence_loss = self.cross_modal_dependence_loss(
                                    heatmap_labels, cross_modal_dependence_logits)
                                loss += cross_modal_dependence_loss

                return loss, vil_sort_logits, heatmap_labels
            return loss, vil_sort_logits

        return (vil_sort_logits, )

    def _forward(
        self,
        input_ids,
        images,
        image_loc=None,
        token_type_ids=None,
        attention_mask=None,
        image_attention_mask=None,
        co_attention_mask=None,
        task_ids=None,
        output_all_encoded_layers=False,
        output_all_attention_masks=False,
        labels=None,
        img_regional_features=None,
        masked_lm_labels=None,
    ):

        if self.hl_include_objectives is not None:
            hl_aux_predictions = [None] * len(self.hl_include_objectives)
        else:
            hl_aux_predictions = None

        bz, text_len = input_ids.size()
        cls_repr_pos = []
        for i in range(bz):
            cls_repr_pos_i = torch.nonzero(input_ids[i]==self.cls_id, as_tuple=False)
            cls_repr_pos.append(cls_repr_pos_i)

        if self.vision_model is not None and images.ndim > 3:
            # print(self.vision_model.layer4[0].conv1.weight);raise
            # Reshape from B x L x C x H x W to (B*L) x C x H x W
            bz, img_len, C, H, W = images.size()
            img_len_for_whole_images = img_len
            if self.multimodal_text_part:
                images = torch.zeros(images.size(0), img_len,
                                     self.num_img_dim).type_as(images).float()
            else:
                images = torch.reshape(images, (bz*img_len, C, H, W)).float()
                images = self.vision_model(images)
                if type(images) == tuple:
                    images, img_regional_features = images
                    img_regional_features = torch.reshape(
                        img_regional_features,
                        (bz, img_len, self.include_num_img_regional_features,  # -1,
                        self.num_img_dim))
                if self.freeze_vision_model:
                    images = images.detach()
                images = torch.reshape(images, (bz, img_len, self.num_img_dim))
                # if self.hierarchical_version != "v0":

                # Perform swapping-based multimodal alignment objective.
                if self.swapping_based_nsp:
                    swapping_based_nsp_labels = []
                    images_if_swapped = torch.zeros(bz, img_len)
                    for i in range(bz):
                        image_ = images[i]  # L x D
                        image_lenwise_sum = torch.sum(image_, dim=-1)
                        # TODO: Since our visual mask token is 0.
                        # non_zero_images = image_lenwise_sum.nonzero().t()[0]
                        non_zero_images = torch.nonzero(image_lenwise_sum, as_tuple=False).t()[0]

                        if len(non_zero_images) == 0:
                            swapping_based_nsp_labels.append(1)
                            continue

                        sample_batch_idx = i + 1
                        if i == bz - 1:
                           sample_batch_idx = 0
                        image_cands_ = images[sample_batch_idx]
                        image_cands_lenwise_sum = torch.sum(image_cands_, dim=-1)
                        # non_zero_image_cands_ = image_cands_lenwise_sum.nonzero().t()[0]
                        non_zero_image_cands_ = torch.nonzero(image_cands_lenwise_sum, as_tuple=False).t()[0]
                        if len(non_zero_image_cands_) == 0:
                            swapping_based_nsp_labels.append(1)
                            continue

                        non_zero_image_cands_ = non_zero_image_cands_.detach().cpu().numpy().astype(int)
                        non_zero_images = non_zero_images.detach().cpu().numpy().astype(int)

                        # TODO: Prevent swapping the already swapped images.
                        non_zero_image_cands_ = set(list(non_zero_image_cands_))
                        images_if_swapped_i = torch.nonzero(
                            images_if_swapped[
                                sample_batch_idx],
                                as_tuple=False).t()[0].detach().cpu().numpy().astype(int)
                        images_if_swapped_i = set(list(images_if_swapped_i))
                        non_zero_image_cands_ -= images_if_swapped_i
                        non_zero_image_cands_ = list(non_zero_image_cands_)
                        if len(non_zero_image_cands_) == 0:
                            swapping_based_nsp_labels.append(1)
                            continue

                        chose_index = np.random.choice(non_zero_image_cands_)
                        swapped_index = np.random.choice(non_zero_images)

                        # Probability of swapping.
                        if_swap = np.random.rand()
                        if if_swap > self.swapping_based_nsp_prob:
                            image_[swapped_index] = image_cands_[chose_index]
                            swapping_based_nsp_labels.append(0)
                            images_if_swapped[i][swapped_index] = 1
                            if self.include_num_img_regional_features is not None:
                                img_regional_features[i][swapped_index] = \
                                    img_regional_features[
                                        sample_batch_idx][chose_index]
                        else:
                            swapping_based_nsp_labels.append(1)

                        images[i] = image_

                    self.swapping_based_nsp_labels = torch.Tensor(
                        swapping_based_nsp_labels).type_as(labels)

                if False:
                    if self.num_img_dim != 2048:
                        images = torch.reshape(images, (bz*img_len, self.num_img_dim))
                        images = self.img_embd_up_proj(images)
                        images = torch.reshape(images, (bz, img_len, 2048))
                if self.include_num_img_regional_features is not None:
                    img_regional_features = img_regional_features.float()
                    new_images = []
                    for i in range(bz):
                        curr_img = images[i]
                        curr_img_region_ = []
                        for j in range(img_len):
                            curr_seq_img = curr_img[j].unsqueeze(0)
                            curr_regional_features = img_regional_features[i, j, :, :]
                            curr_new_img = torch.cat([curr_seq_img, curr_regional_features])
                            curr_img_region_.append(curr_new_img)
                        curr_img_region_ = torch.cat(curr_img_region_)
                        new_images.append(curr_img_region_)
                    new_images = torch.stack(new_images)

        if self.multimodal_img_part:
            input_ids = torch.zeros(input_ids.size(0), 1).type_as(input_ids)
            attention_mask = torch.zeros(input_ids.size(0), 1).type_as(attention_mask)
            # input_ids = torch.zeros(input_ids.size(0),
            #     input_ids.size(1)).type_as(input_ids)
            # attention_mask = torch.zeros(input_ids.size(0),
            #     input_ids.size(1)).type_as(attention_mask)
            # aprint(torch.sum(input_ids))
            # print(torch.sum(attention_mask));raise
            token_type_ids = None

        if self.vilbert_paired_coattention:
            paired_co_attention_mask = torch.zeros(
                input_ids.size(0), images.size(1), input_ids.size(1)
            ).type_as(images)
            if self.include_num_img_regional_features is not None:
                paired_co_attention_mask = torch.zeros(
                    input_ids.size(0),
                    images.size(1) + self.include_num_img_regional_features * images.size(1),
                    input_ids.size(1)
                ).type_as(images)
            for i in range(bz):
                cls_repr_pos_i = cls_repr_pos[i].squeeze().cpu().numpy()
                for j in range(len(cls_repr_pos_i)):
                    # TODO: Currently the last sequence item is not
                    # truncated yet.
                    if j == len(cls_repr_pos_i) - 1:  # The last item.
                        start = cls_repr_pos_i[j]
                        end = input_ids.size(1)
                    else:
                        start = cls_repr_pos_i[j]
                        end = cls_repr_pos_i[j+1]
                    if self.include_num_img_regional_features is not None:
                        beta = self.include_num_img_regional_features
                        paired_co_attention_mask[i, j*beta+j:(j+1)*beta+j+1, start:end] = 1.0
                    else:
                        paired_co_attention_mask[i, j, start:end] = 1.0
            co_attention_mask = paired_co_attention_mask

        if self.include_num_img_regional_features is not None:
            images = new_images

        if (co_attention_mask is None and self.config.with_coattention
            and not self.vilbert_paired_coattention):
            co_attention_mask = torch.ones(
                input_ids.size(0), images.size(1), input_ids.size(1)
            ).type_as(images)

        # print(input_ids.size(), images.size());raise

        sequence_output_t, sequence_output_v, pooled_output_t, pooled_output_v, all_attention_mask = self.bert(
            input_ids,
            images,
            image_loc,
            token_type_ids,
            attention_mask,
            image_attention_mask,
            co_attention_mask,
            task_ids,
            output_all_encoded_layers=output_all_encoded_layers,
            output_all_attention_masks=output_all_attention_masks,
        )

        if image_attention_mask is None:
            image_attention_mask = torch.ones(
                images.size(0), images.size(1)
            ).type_as(input_ids)

        if self.multimodal_img_part and self.simple_img_classifier:
            images_ = torch.reshape(images, (bz*img_len, self.num_img_dim))
            images_ = self.simple_img_prediction(images_)
            images_ = torch.reshape(images_, (bz, img_len, self.config.bi_hidden_size))
            pooled_output_v = images_[:, 0, :]

        vil_prediction = 0
        vil_logit = 0
        vil_binary_prediction = 0
        vision_prediction = 0
        vision_logit = 0
        linguistic_prediction = 0
        linguistic_logit = 0
        vil_sort_prediction = 0

        linguistic_prediction, vision_prediction, vil_binary_prediction = self.cls(
            sequence_output_t, sequence_output_v, pooled_output_t, pooled_output_v
        )

        if self.fusion_method == "sum":
            pooled_output = self.dropout(pooled_output_t + pooled_output_v)
        elif self.fusion_method == "mul":
            pooled_output = self.dropout(pooled_output_t * pooled_output_v)
        elif self.fusion_method == "text_only":
            pooled_output = self.dropout(pooled_output_t)
        elif self.fusion_method == "img_only":
            pooled_output = self.dropout(pooled_output_v)
        else:
            assert False
        if self.multimodal_img_part and not self.config.with_coattention:
            pooled_output = pooled_output_v

        vil_prediction = self.vil_prediction(pooled_output)
        vil_prediction_gqa = self.vil_prediction_gqa(pooled_output)
        # if pooled_output.size(0) % 2 == 0 and not self.simple_img_classifier:
        #     vil_binary_prediction = self.vil_binary_prediction(
        #         pooled_output.view(-1, pooled_output.size(1) * 2)
        #     )
        vil_logit = self.vil_logit(pooled_output)
        vil_tri_prediction = self.vil_tri_prediction(pooled_output)
        vision_logit = self.vision_logit(self.dropout(sequence_output_v)) + (
            (1.0 - image_attention_mask) * -10000.0
        ).unsqueeze(2).to(dtype=next(self.parameters()).dtype)
        # ).unsqueeze(2).to(dtype=torch.float32)
        linguistic_logit = self.linguistic_logit(self.dropout(sequence_output_t))

        # Sorting task.
        if self.hierarchical_version != "v0":

            if self.hl_include_objectives is not None:
                if "head" in self.hl_include_objectives:
                    hl_aux_head_predictions = []
                if ("pairwise" in self.hl_include_objectives
                    or "binary" in self.hl_include_objectives):
                    hl_aux_bin_predictions = []
                if "binary_cross_modal" in self.hl_include_objectives:
                    hl_aux_binx_predictions = []
                if "cross_modal_dependence" in self.hl_include_objectives:
                    hl_aux_x_dep_predictions = []

            cls_heat_map = []
            if self.additional_hl:
                vil_add_hl_repr = []
                for i in range(bz):
                    cls_repr_pos_i = cls_repr_pos[i].squeeze()
                    vil_add_hl_repr_curr = sequence_output_t[i][cls_repr_pos_i]
                    vil_add_hl_repr.append(vil_add_hl_repr_curr)
                vil_add_hl_repr = torch.stack(vil_add_hl_repr)
                add_hl_attention_mask = torch.ones(
                    vil_add_hl_repr.size(0),
                    vil_add_hl_repr.size(1),
                    vil_add_hl_repr.size(1)
                ).type_as(images)
                add_hl_attention_mask = add_hl_attention_mask.unsqueeze(1)
                vil_add_hl_repr = self.additional_hl_layer(
                    vil_add_hl_repr,
                    attention_mask=add_hl_attention_mask)
                vil_add_hl_repr = vil_add_hl_repr[0]

            for i in range(bz):
                if not self.additional_hl:
                    cls_repr_pos_i = cls_repr_pos[i].squeeze()
                    cls_repr_i = sequence_output_t[i][cls_repr_pos_i]
                    if self.fusion_method == "img_only" or self.heatmap_late_fusion:
                        if self.include_num_img_regional_features is None:
                            img_reg_added = 0
                        else:
                            img_reg_added = self.include_num_img_regional_features
                        img_cls_repr_pos_i = list(range(0,
                            img_len_for_whole_images*img_reg_added+img_len_for_whole_images,
                            img_reg_added+1))
                        img_cls_repr_i = sequence_output_v[i][img_cls_repr_pos_i]
                        if self.fusion_method == "img_only":
                            img_cls_repr_i = self._img_only_inter_proj(img_cls_repr_i)
                            cls_repr_i = img_cls_repr_i
                else:
                    cls_repr_i = vil_add_hl_repr[i]
                # XXX: Try drop out here?
                # cls_repr_i = self.dropout(cls_repr_i)
                if self.heatmap_late_fusion:
                    # if self.include_num_img_regional_features:
                    #     raise NotImplementedError("Not done yet!")
                    # FIXME: Change this.
                    # img_cls_repr_i = self.heatmap_img_proj_layer(sequence_output_v[i])
                    img_cls_repr_i = self.heatmap_img_proj_layer(img_cls_repr_i)
                    # img_cls_repr_i = sequence_output_v[i]
                    if self.fusion_method == "sum":
                        cls_repr_i = cls_repr_i + img_cls_repr_i
                    elif self.fusion_method == "mul":
                        cls_repr_i = cls_repr_i * img_cls_repr_i
                    elif self.functional == "text_only":
                        cls_repr_i = cls_repr_i
                        raise
                    else:
                        assert False

                # TODO HL auxiliary predictions.
                if self.hl_include_objectives is not None:
                    if len(self.hl_include_objectives) > 0:
                        for hl_aux_objective in self.hl_include_objectives:
                            if "head" == hl_aux_objective:
                                hl_aux_head_prediction_curr = self.hl_head_pred_layer(cls_repr_i)
                                hl_aux_head_prediction_curr = hl_aux_head_prediction_curr.squeeze()
                                hl_aux_head_predictions.append(hl_aux_head_prediction_curr)
                            elif ("pairwise" == hl_aux_objective
                                  or "binary" == hl_aux_objective):
                                hl_aux_bin_predictions_tmp = []
                                for seq_i in range(len(cls_repr_pos_i)):
                                    for seq_j in range(seq_i+1, len(cls_repr_pos_i)):
                                        cls_repr_seq_i = cls_repr_i[seq_i]
                                        cls_repr_seq_j = cls_repr_i[seq_j]
                                        cls_repr_seq_ij = torch.stack(
                                            [cls_repr_seq_i, cls_repr_seq_j])
                                        hl_aux_bin_prediction_curr = self.hl_bin_pred_layer(cls_repr_seq_ij)
                                        hl_aux_bin_prediction_curr = hl_aux_bin_prediction_curr.squeeze()
                                        hl_aux_bin_predictions_tmp.append(hl_aux_bin_prediction_curr)
                                hl_aux_bin_predictions_tmp = torch.stack(hl_aux_bin_predictions_tmp)

                                if self.hl_bin_sparse_prob < 1.0:
                                    hl_bin_sparse_len = int(
                                        len(hl_aux_bin_predictions_tmp)
                                            * self.hl_bin_sparse_prob)
                                    hl_bin_sparse_pos_tmp = np.random.choice(
                                        np.arange(hl_aux_bin_predictions_tmp.size(0)),
                                                  hl_bin_sparse_len)
                                    # TODO: Temporary!!!
                                    hl_aux_bin_predictions_tmp = hl_aux_bin_predictions_tmp[
                                        hl_bin_sparse_pos_tmp]
                                    # hl_aux_bin_predictions_tmp = hl_aux_bin_predictions_tmp[:3]
                                    self.hl_bin_sparse_pos.append(hl_bin_sparse_pos_tmp)

                                hl_aux_bin_predictions.append(hl_aux_bin_predictions_tmp)

                            elif "binary_cross_modal" == hl_aux_objective:
                                hl_aux_binx_predictions_tmp = []
                                for seq_i in range(len(cls_repr_pos_i)):
                                    for seq_j in range(seq_i+1, len(cls_repr_pos_i)):
                                        # Text modality.
                                        cls_repr_seq_i = cls_repr_i[seq_i]
                                        # Image modality.
                                        if self.include_num_img_regional_features  is not None:
                                            img_seq_j = (1 + self.include_num_img_regional_features) * seq_j
                                            raise NotImplementedError("Not debugged yet!")
                                        else:
                                            img_seq_j = seq_j
                                        cls_repr_seq_j = sequence_output_v[i][img_seq_j]
                                        # Stacking.
                                        cls_repr_seq_ij = torch.cat(
                                            [cls_repr_seq_i, cls_repr_seq_j], dim=-1)
                                        hl_aux_binx_prediction_curr = self.hl_binx_pred_layer(cls_repr_seq_ij)
                                        hl_aux_binx_prediction_curr = hl_aux_binx_prediction_curr.squeeze()
                                        hl_aux_binx_predictions_tmp.append(hl_aux_binx_prediction_curr)
                                hl_aux_binx_predictions_tmp = torch.stack(hl_aux_binx_predictions_tmp)

                                if self.hl_binx_sparse_prob < 1.0:
                                    hl_binx_sparse_len = int(
                                        len(hl_aux_binx_predictions_tmp)
                                            * self.hl_binx_sparse_prob)
                                    hl_binx_sparse_pos_tmp = np.random.choice(
                                        np.arange(hl_aux_binx_predictions_tmp.size(0)),
                                                  hl_binx_sparse_len)
                                    # TODO: Temporary!!!
                                    hl_aux_binx_predictions_tmp = hl_aux_binx_predictions_tmp[
                                        hl_binx_sparse_pos_tmp]
                                    # hl_aux_binx_predictions_tmp = hl_aux_binx_predictions_tmp[:3]
                                    self.hl_binx_sparse_pos.append(hl_binx_sparse_pos_tmp)

                                hl_aux_binx_predictions.append(hl_aux_binx_predictions_tmp)

                            elif "cross_modal_dependence" == hl_aux_objective:
                                hl_aux_x_dep_predictions_tmp = []
                                for img_seq_i in range(len(cls_repr_pos_i)):
                                    # Image modality.
                                    if self.include_num_img_regional_features  is not None:
                                        img_seq_i_real = (1 + self.include_num_img_regional_features) * img_seq_i
                                        raise NotImplementedError("Not debugged yet!")
                                    else:
                                        img_seq_i_real = img_seq_i
                                    img_repr_seq_i = sequence_output_v[i][img_seq_i_real]
                                    hl_aux_x_dep_predictions_tmp.append(img_repr_seq_i)
                                hl_aux_x_dep_predictions_tmp = torch.stack(hl_aux_x_dep_predictions_tmp)
                                hl_aux_x_dep_predictions_tmp = self.cross_modal_dependence_prediction(
                                    hl_aux_x_dep_predictions_tmp)
                                hl_aux_x_dep_predictions.append(hl_aux_x_dep_predictions_tmp)

                cls_heat_map_i = self.heat_map_prediction(cls_repr_i)
                cls_heat_map.append(cls_heat_map_i)
            cls_heat_map = torch.stack(cls_heat_map)
            self.cls_heat_map = cls_heat_map
            vil_sort_prediction = cls_heat_map
        else:
            vil_sort_prediction = self.vil_sort_prediction(pooled_output)
            if self.vilbert_use_3way_logits:
                vil_sort_prediction_v = self.vil_sort_prediction(pooled_output_v)
                vil_sort_prediction_t = self.vil_sort_prediction(pooled_output_t)
                vil_sort_prediction += vil_sort_prediction_v
                vil_sort_prediction += vil_sort_prediction_t

        if self.hl_include_objectives is not None:
            if len(self.hl_include_objectives) > 0:
                for hl_aux_objective in self.hl_include_objectives:
                    aux_obj_list_idx = self.hl_include_objectives.index(hl_aux_objective)

                    if "head" == hl_aux_objective:
                        hl_aux_head_predictions = torch.stack(hl_aux_head_predictions)
                        hl_aux_predictions[aux_obj_list_idx] = hl_aux_head_predictions
                    elif ("pairwise" == hl_aux_objective
                          or "binary" == hl_aux_objective):
                        hl_aux_bin_predictions = torch.stack(hl_aux_bin_predictions)
                        hl_aux_predictions[aux_obj_list_idx] = hl_aux_bin_predictions
                    elif "binary_cross_modal" == hl_aux_objective:
                        hl_aux_binx_predictions = torch.stack(hl_aux_binx_predictions)
                        hl_aux_predictions[aux_obj_list_idx] = hl_aux_binx_predictions

                    elif "cross_modal_dependence" == hl_aux_objective:
                        hl_aux_x_dep_predictions = torch.stack(hl_aux_x_dep_predictions)
                        hl_aux_predictions[aux_obj_list_idx] = hl_aux_x_dep_predictions
            pass ####

        return (
            vil_prediction,
            vil_prediction_gqa,
            vil_logit,
            vil_binary_prediction,
            vil_tri_prediction,
            vision_prediction,
            vision_logit,
            linguistic_prediction,
            linguistic_logit,
            all_attention_mask,
            vil_sort_prediction,
            hl_aux_predictions,
            sequence_output_t,
            sequence_output_v,
        )


class SimpleClassifier(nn.Module):
    def __init__(self, in_dim, hid_dim, out_dim, dropout):
        super().__init__()
        self.logit_fc = nn.Sequential(
            nn.Linear(in_dim, hid_dim),
            GeLU(),
            BertLayerNorm(hid_dim, eps=1e-12),
            nn.Linear(hid_dim, out_dim),
            nn.Dropout(p=dropout),
        )

    def forward(self, hidden_states):
        return self.logit_fc(hidden_states)


# class SimpleClassifier(nn.Module):
#     def __init__(self, in_dim, hid_dim, out_dim, dropout):
#         super(SimpleClassifier, self).__init__()
#         layers = [
#             weight_norm(nn.Linear(in_dim, hid_dim), dim=None),
#             nn.ReLU(),
#             nn.Dropout(dropout, inplace=True),
#             weight_norm(nn.Linear(hid_dim, out_dim), dim=None)
#         ]
#         self.main = nn.Sequential(*layers)

#     def forward(self, x):
#         logits = self.main(x)
#         return logits
