text: from itertools import chain
  import numpy as np
  import math
  import argparse
  import torch
  import torch.nn as nn
  import torch.nn.init as init
  import torch.nn.functional as F
  import torch.optim as optim
  from transformers import BertModel, LayoutLMModel, AutoModel, LayoutLMConfig, AutoConfig
  
  from retrieval_based.data_provider.utils import load_word_embedding
  from retrieval_based.network.layout import PretrainedLayoutEmbedding
  from retrieval_based.network.attention import TransformerBlock, Identity_TransformerBlock
  from retrieval_based.network.skim_attention import *
  from retrieval_based.network import *
  
  from config_file import *
  
  
  class LayoutDeepMatch(nn.Module):
      def __init__(self, args):
          self.args = args
          super(LayoutDeepMatch, self).__init__()
          if args.use_bert:
              self.embeddings = AutoModel.from_pretrained(args.pretrained_model)
              self.ffn = nn.Linear(in_features=768, out_features=args.emb_size)
              print('Using %s' % self.embeddings.__class__.__name__)
          else:
              word_embeddings, word2id, id2word = load_word_embedding(glove_path)
              n_vocab, n_dim = word_embeddings.shape
              self.embeddings = nn.Embedding(num_embeddings=n_vocab, embedding_dim=n_dim)
              print('Using Glove embedding.')
  
          self.layout_embedding = PretrainedLayoutEmbedding(load_weight=False)
          self.layout_embedding.requires_grad_(True)
          self.layout_ffn = nn.Linear(in_features=768, out_features=args.emb_size)
          self.layoutlm_layernorm = nn.LayerNorm(normalized_shape=(args.emb_size))
  
          self.turn_layout_W = nn.Parameter(torch.Tensor(DOC_LEN + 2, MAX_TURN))
          self.turn_layernorm = nn.LayerNorm(normalized_shape=(args.emb_size))
          self.n_layer = args.n_layer
  
          self.transformer_u_self = nn.ModuleList([Identity_TransformerBlock()] + [TransformerBlock(input_size=args.emb_size, is_layer_norm=True)] * (args.n_layer - 1))
          self.transformer_r_self = nn.ModuleList([Identity_TransformerBlock()] + [TransformerBlock(input_size=args.emb_size, is_layer_norm=True)] * (args.n_layer - 1))
          self.transformer_p_self = nn.ModuleList([Identity_TransformerBlock()] + [TransformerBlock(input_size=args.emb_size, is_layer_norm=True)] * (args.n_layer - 1))
  
          self.transformer_urr_cross = nn.ModuleList([TransformerBlock(input_size=args.emb_size, is_layer_norm=True)] * args.n_layer)
          self.transformer_ruu_cross = nn.ModuleList([TransformerBlock(input_size=args.emb_size, is_layer_norm=True)] * args.n_layer)
          self.transformer_upp_cross = nn.ModuleList([TransformerBlock(input_size=args.emb_size, is_layer_norm=True)] * args.n_layer)
          self.transformer_puu_cross = nn.ModuleList([TransformerBlock(input_size=args.emb_size, is_layer_norm=True)] * args.n_layer)
          self.transformer_prr_cross = nn.ModuleList([TransformerBlock(input_size=args.emb_size, is_layer_norm=True)] * args.n_layer)
          self.transformer_rpp_cross = nn.ModuleList([TransformerBlock(input_size=args.emb_size, is_layer_norm=True)] * args.n_layer)
  
          self.cnn_3d_1 = nn.Conv3d(in_channels=args.n_layer * 2, out_channels=32, kernel_size=(3, 3, 3), padding=(1, 1, 1))
          self.maxpooling3d_1 = nn.MaxPool3d(kernel_size=(3, 3, 3), stride=(3, 3, 3), padding=(1, 1, 1))
          self.cnn_3d_2 = nn.Conv3d(in_channels=32, out_channels=16, kernel_size=(3, 3, 3), padding=(1, 1, 1))
          self.maxpooling3d_2 = nn.MaxPool3d(kernel_size=(3, 3, 3), stride=(3, 3, 3), padding=(1, 1, 1))
          self.cnn3d_doc1 = nn.Conv3d(in_channels=args.n_layer * 2, out_channels=32, kernel_size=(3, 3, 3), padding=(1, 1, 1))
          self.cnn3d_doc2 = nn.Conv3d(in_channels=32, out_channels=16, kernel_size=(3, 3, 3), padding=(1, 1, 1))
          in_d1, in_d2, in_d3 = 1, int((int((DOC_LEN + 2 + 2) / 3) + 2) / 3), int((int((MAX_SEQ_LEN + 2 + 2) / 3) + 2) / 3)
          self.doc_linear = nn.Linear(in_features=16 * in_d1 * in_d2 * in_d3, out_features=288, bias=True)
  
          self.classifier = nn.Linear(in_features=288 * 2, out_features=1)
  
          self.dropout = nn.Dropout(0.2)
  
          self.debug_data = {'bU': torch.randint(low=0, high=100, size=(16, 10, 25)),
                             'bR': torch.randint(low=0, high=100, size=(16, 25)),
                             'bP': torch.randint(low=0, high=100, size=(16, 512)),
                             'bl': None,
                             'b_turn': torch.randint(low=1, high=MAX_TURN, size=(16, 1))}
  
          self.init_weight()
  
      def init_weight(self):
          init.uniform_(self.turn_layout_W)
          # init.xavier_normal_(self.cnn_3d_1.weight, 0.01)
          # init.xavier_normal_(self.cnn_3d_2.weight, 0.01)
          init.uniform_(self.cnn_3d_1.weight, -0.01, 0.01)        # ref: paddle-DAM
          init.constant_(self.cnn_3d_1.bias, 0.0)                 # ref: paddle-DAM
          init.uniform_(self.cnn_3d_2.weight, -0.01, 0.01)        # ref: paddle-DAM
          init.constant_(self.cnn_3d_2.bias, 0.0)                 # ref: paddle-DAM
  
          init.uniform_(self.cnn3d_doc1.weight, -0.01, 0.01)
          init.constant_(self.cnn3d_doc1.bias, 0.0)
          init.uniform_(self.cnn3d_doc2.weight, -0.01, 0.01)
          init.constant_(self.cnn3d_doc2.bias, 0.0)
  
          init.kaiming_normal_(self.doc_linear.weight)
          init.kaiming_normal_(self.classifier.weight)
          init.constant_(self.classifier.bias, 0.0)               # ref: paddle-DAM
  
      def Encoding_Layer(self, bU, bR, bP, bl):
          if self.args.use_bert:
              bU_embedding, bR_embedding, bP_embedding = [self.ffn(self.embeddings(b, return_dict=True)['last_hidden_state']) for b in [bU, bR, bP]]
          else:
              bU_embedding, bR_embedding, bP_embedding = [self.embeddings(b) for b in [bU, bR, bP]]
          bl_embedding = self.layout_ffn(self.layout_embedding(bl, operation=''))
  
          # bU_embedding = self.dropout(bU_embedding)
          bR_embedding = self.dropout(bR_embedding)
          # bP_embedding = self.dropout(bP_embedding)
          # bl_embedding = self.dropout(bl_embedding)
  
          return bU_embedding, bR_embedding, bP_embedding, bl_embedding
  
      def Representation_Layer(self, bU_embedding_self, bR_embedding_self):
          batch_size, max_utterances, max_u_words, emb_size = bU_embedding_self.shape
          batch_size, max_r_words, emb_size = bR_embedding_self.shape
  
          bU_self = []
          bR_self = []
          bU_cross, bR_cross = [], []
          for i in range(self.n_layer):
              bU_embedding_self = bU_embedding_self.view(-1, max_u_words, emb_size)
              bU_embedding_self = self.transformer_u_self[i](bU_embedding_self, bU_embedding_self, bU_embedding_self)
              bU_embedding_self = bU_embedding_self.view(batch_size, max_utterances, max_u_words, emb_size)
              bU_self.append(bU_embedding_self)
  
              bR_embedding_self = self.transformer_r_self[i](bR_embedding_self, bR_embedding_self, bR_embedding_self)
              bR_self.append(bR_embedding_self)
  
              bR_embedding_repeated = bR_embedding_self.view(batch_size, 1, max_r_words, emb_size).repeat(1, max_utterances, 1, 1)
  
              bU_embedding_self = bU_embedding_self.view(-1, max_u_words, emb_size)
              bR_embedding_self = bR_embedding_self.view(-1, max_r_words, emb_size)
              bR_embedding_repeated = bR_embedding_repeated.view(-1, max_r_words, emb_size)
  
              bU_embedding_cross = self.transformer_urr_cross[i](bU_embedding_self, bR_embedding_repeated, bR_embedding_repeated)
              bR_embedding_cross = self.transformer_ruu_cross[i](bR_embedding_repeated, bU_embedding_self, bU_embedding_self)
              bU_embedding_cross = bU_embedding_cross.view(batch_size, max_utterances, max_u_words, -1)
              bR_embedding_cross = bR_embedding_cross.view(batch_size, max_utterances, max_r_words, -1)
              bU_cross.append(bU_embedding_cross)
              bR_cross.append(bR_embedding_cross)
  
          bU_self = torch.stack(bU_self, dim=1)
          bR_self = torch.stack(bR_self, dim=1)
          bU_cross = torch.stack(bU_cross, dim=1)
          bR_cross = torch.stack(bR_cross, dim=1)
          return bU_self, bR_self, bU_cross, bR_cross
  
      def Document_Interaction_Layer(self, bR_self, bP_embedding_self):
          batch_size, n_doc, max_p_words, emb_size = bP_embedding_self.shape
          batch_size, n_layer, max_r_words, emb_size = bR_self.shape
  
          bP_self, bR_cross_d, bP_cross_d = [], [], []
          for i in range(self.n_layer):
              bP_embedding_self = bP_embedding_self.view(-1, max_p_words, emb_size)
              bP_embedding_self = self.transformer_p_self[i](bP_embedding_self, bP_embedding_self, bP_embedding_self)
              bP_embedding_self = bP_embedding_self.view(batch_size, n_doc, max_p_words, emb_size)
              bP_self.append(bP_embedding_self)
  
              bP_embedding_self = bP_embedding_self.view(-1, max_p_words, emb_size)
  
              bR_embedding_cross_d = self.transformer_rpp_cross[i](bR_self[:, i, :, :], bP_embedding_self, bP_embedding_self)
              bP_embedding_cross_d = self.transformer_prr_cross[i](bP_embedding_self, bR_self[:, i, :, :], bR_self[:, i, :, :])
              bR_embedding_cross_d = bR_embedding_cross_d.view(batch_size, n_doc, max_r_words, -1)
              bP_embedding_cross_d = bP_embedding_cross_d.view(batch_size, n_doc, max_p_words, -1)
              bR_cross_d.append(bR_embedding_cross_d)
              bP_cross_d.append(bP_embedding_cross_d)
  
          bP_self = torch.stack(bP_self, dim=1)
          bR_cross_d = torch.stack(bR_cross_d, dim=1)
          bP_cross_d = torch.stack(bP_cross_d, dim=1)
          return bP_self, bR_cross_d, bP_cross_d
  
      def Matching_Layer(self, bU, bR):
          # batch_size, max_utterances, n_layer, max_u_words, hidden_size = bU.shape
          # batch_size, max_utterances, n_layer, max_r_words, hidden_size = bR.shape
          M = torch.einsum('bmnuh,bmnrh->bmnur', bU, bR)
          return M
  
      def Aggregation_Layer(self, M):
          Z = F.elu(self.cnn_3d_1(M))
          Z = self.maxpooling3d_1(Z)
          Z = F.elu(self.cnn_3d_2(Z))
          Z = self.maxpooling3d_2(Z)
          Z = Z.view(M.shape[0], -1)  # (bsz, *)
          return Z
  
      def Document_Aggregation_Layer(self, M):
          Z = F.elu(self.cnn3d_doc1(M))
          Z = self.maxpooling3d_1(Z)
          Z = F.elu(self.cnn3d_doc2(Z))
          Z = self.maxpooling3d_2(Z)
          Z = Z.view(M.shape[0], -1)
          return Z
  
  
      def forward(self, bU, bR, bP, bl, b_turn, is_test=False):
          """
          :param bU: batch utterance, size: (batch_size, max_utterances, max_u_words)
          :param bR: batch responses, size: (batch_size, max_r_words)
          :param bP: batch profiles, size: (batch_size, max_num_profiels, max_p_words)
          :return: scores, size: (batch_size, )
          """
          bP = bP.view(bP.shape[0], 1, bP.shape[1])
  
          su1, su2, su3 = bU.shape
          sr1, sr2 = bR.shape
          sp1, sp2, sp3 = bP.shape
          sl1, sl2, sl3 = bl.shape
  
          # Encoding Layer
          bU = bU.view(-1, su3)
          bR = bR.view(-1, sr2)
          bP = bP.view(-1, sp3)
  
          ## U, R, D
          bU_embedding, bR_embedding, bP_embedding, bl_embedding = self.Encoding_Layer(bU, bR, bP, bl)
  
          bU_embedding = bU_embedding.view(su1, su2, su3, -1)     # (batch_size=16, max_utterances=10, max_u_words=25, emb_size=300)
          bR_embedding = bR_embedding.view(sr1, sr2, -1)          # (batch_size=16, max_r_words=25, emb_size=300)
          bP_embedding = bP_embedding.view(sp1, sp2, sp3, -1)     # (batch_size=16, max_num_profiels=1, max_p_words=512, emb_size=300)
          bl_embedding = bl_embedding.view(sl1, 1, sl2, -1)
  
          bP_embedding = bP_embedding + bl_embedding
          bP_embedding = self.dropout(self.layoutlm_layernorm(bP_embedding))
  
          bt_embedding = torch.einsum('bnde,dt->bnte', bl_embedding, torch.softmax(self.turn_layout_W, dim=0))
          bt_embedding = torch.mean(bt_embedding, dim=1).view(sl1, su2, 1, -1).repeat(1, 1, su3, 1)
          bU_embedding = bU_embedding + bt_embedding
          bU_embedding = self.dropout(self.layoutlm_layernorm(bU_embedding))
  
          bU_self, bR_self, bU_cross, bR_cross = self.Representation_Layer(bU_embedding, bR_embedding)
          bP_self_d, bR_cross_d, bP_cross_d = self.Document_Interaction_Layer(bR_self, bP_embedding)
          bR_self_d = bR_self.view(sr1, self.n_layer, 1, sr2, -1).repeat(1, 1, sp2, 1, 1)
          bR_self = bR_self.view(sr1, self.n_layer, 1, sr2, -1).repeat(1, 1, su2, 1, 1)
          # print('self & cross: ', bU_self.shape, bR_self.shape, bU_cross.shape, bR_cross.shape)
          # print('self & cross: ', bP_self_d.shape, bR_cross_d.shape, bP_cross_d.shape)
  
          M_self = self.Matching_Layer(bU_self, bR_self)
          M_cross = self.Matching_Layer(bU_cross, bR_cross)
          M_self_d = self.Matching_Layer(bP_self_d, bR_self_d)
          M_cross_d = self.Matching_Layer(bP_cross_d, bR_cross_d)
          # print(torch.mean(M_self), torch.mean(M_cross))
          # print(M_self_d.shape, M_cross_d.shape)
  
          M = torch.cat([M_self, M_cross], dim=1)
          M_d = torch.cat([M_self_d, M_cross_d], dim=1)   # (16, 10, 1, 512, 300)
          # print(M.shape)
  
          m = self.Aggregation_Layer(M)
          # m = self.Aggregation_Layer_181(M, b_turn)
          m_d = self.Document_Aggregation_Layer(M_d)
          # print(m_d.shape)
          m_d = self.doc_linear(m_d)
  
          # print(m.shape, m_d.shape)
  
          m = torch.cat([m, m_d], dim=-1)
  
          logits = self.classifier(m)
          # print(logits)
          logits = torch.sigmoid(logits)
          # print(logits)
          logits = logits.squeeze(-1)
          return logits
  
  
  if __name__ == '__main__':
      parser = argparse.ArgumentParser()
      parser.add_argument("--seed", default=0, type=int, help="Random seed.")
      parser.add_argument("--level", default='word', type=str, help="The dataset used for training and test.")
      parser.add_argument("--gru_hidden", default=300, type=int, help="The hidden size of GRU in layer 1")
      parser.add_argument("--emb_size", default=300, type=int, help="The embedding size")
      parser.add_argument("--weight_decay", default=0, type=float, help="weight decay coefficient")
      parser.add_argument("--gamma", default=0.3, type=float, help="Threshold.")
      parser.add_argument("--decay", default=0.9, type=float, help="Decay rate.")
      parser.add_argument("--pretrained_model",
                          default='bert-base-uncased',
                          type=str,
                          help='Follow what includes in Hugging Face Models. Options: [bert-base-uncased, allenai/longformer-base-4096]')
      parser.add_argument("--use_bert", default=False, type=bool, help='Use bert to embed or glove embedding?')
      parser.add_argument("--n_layer", default=5, type=int, help='The number of transformer layers for DAM.')
      args = parser.parse_args()
  
      net = LayoutDGMN(args)
      bU, bR, bP, bl, b_turn = net.debug_data['bU'], net.debug_data['bR'], net.debug_data['bP'], net.debug_data['bl'], net.debug_data['b_turn']
      print(b_turn.view(-1))
      g = net(bU, bR, bP, bl, b_turn)
      print(g)
  