<reponame>Coldog2333/DGDS
import torch
import torch.nn as nn
import torch.nn.init as init
import torch.nn.functional as F
from transformers import AutoModel

from retrieval_based.data_provider.utils import load_word_embedding
from config_file import *


class VanillaSMN(nn.Module):
    def __init__(self, args):
        self.args = args
        super(VanillaSMN, 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.gru_sentence = nn.GRU(input_size=args.emb_size, hidden_size=args.gru_hidden, batch_first=True)

        self.A = nn.Parameter(torch.Tensor(args.gru_hidden, args.gru_hidden))

        self.cnn_2d = nn.Conv2d(in_channels=2, out_channels=8, kernel_size=(3, 3))
        self.maxpooling = nn.MaxPool2d(kernel_size=(3, 3), stride=(3, 3))
        in_d1, in_d2 = int((MAX_SEQ_LEN - 2) / 3), int((MAX_SEQ_LEN - 2) / 3)
        self.context_linear = nn.Linear(in_features=8 * in_d1 * in_d2, out_features=50)
        self.gru_agg = nn.GRU(input_size=50, hidden_size=args.gru_hidden, batch_first=True)

        self.classifier = nn.Linear(in_features=args.gru_hidden, 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.xavier_normal_(self.A)                         # ref: YuWU-SCN
        init.kaiming_normal_(self.cnn_2d.weight)            # ref: YuWU-SCN
        init.xavier_uniform_(self.context_linear.weight)
        init.kaiming_normal_(self.classifier.weight)

        # ref: YuWU-SCN
        for weights in [self.gru_sentence.weight_hh_l0, self.gru_sentence.weight_ih_l0]:
            init.orthogonal_(weights)
        for weights in [self.gru_agg.weight_hh_l0, self.gru_agg.weight_ih_l0]:
            init.orthogonal_(weights)

    def Encoding_Layer(self, bU, bR):
        if self.args.use_bert:
            bU_embedding, bR_embedding = [self.ffn(self.embeddings(b, return_dict=True)['last_hidden_state']) for b in [bU, bR]]
        else:
            bU_embedding, bR_embedding = [self.embeddings(embedding) for embedding in [bU, bR]]
        bU_embedding, bR_embedding = [self.dropout(embedding) for embedding in [bU_embedding, bR_embedding]]
        return bU_embedding, bR_embedding

    def Matching_Layer(self, bR_embedding, bU_embedding, bR_gru_embedding, bU_gru_embedding):
        M1 = torch.einsum('brd,btud->btru', bR_embedding, bU_embedding)
        M2 = torch.einsum('brd,dd,btud->btru', bR_gru_embedding, self.A, bU_gru_embedding)
        M = torch.stack([M1, M2], dim=2)
        return M

    def Aggregation_Layer(self, M, m_type):
        batch_size, max_utterances, n_view, max_u_words, max_r_words = M.shape
        M = M.view(-1, 2, max_u_words, max_r_words)

        # print(M.shape)
        Z = F.relu(self.cnn_2d(M))
        # print(Z.shape)
        Z = self.maxpooling(Z)
        # print(Z.shape)
        Z = Z.view(batch_size, max_utterances, -1)  # (bsz*max_utterances, *)

        # print(Z.shape)

        if 1 == m_type:
            matching_vector = torch.tanh(self.context_linear(Z))
            _, matching_vector = self.gru_agg(matching_vector)
            matching_vector = matching_vector.squeeze()
        else:
            matching_vector = Z.view(batch_size, -1)
            matching_vector = torch.tanh(self.doc_linear(matching_vector.squeeze()))

        return matching_vector

    def forward(self, inputs):
        """
        :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, )
        """
        bU = inputs['context']
        bR = inputs['response']
        bP = inputs['doc_text']
        bl = inputs['layout']

        bP = bP.view(bP.shape[0], 1, bP.shape[1])

        su1, su2, su3 = bU.shape
        sr1, sr2 = bR.shape
        sp1, sp2, sp3 = bP.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 = self.Encoding_Layer(bU, bR)

        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)

        # print(bU_embedding.shape, bR_embedding.shape, bP_embedding.shape)

        bU_embedding = bU_embedding.view(-1, su3, self.args.emb_size)     # (batch_size=16, max_utterances=10, max_u_words=25, emb_size=300)
        bR_embedding = bR_embedding.view(-1, sr2, self.args.emb_size)     # (batch_size=16, max_r_words=25, emb_size=300)

        bU_gru_embedding, _ = self.gru_sentence(bU_embedding)
        bR_gru_embedding, _ = self.gru_sentence(bR_embedding)

        # Matching Layer
        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)
        bU_gru_embedding = bU_gru_embedding.view(su1, su2, su3, -1)     # (batch_size=16, max_utterances=10, max_u_words=25, emb_size=300)
        bR_gru_embedding = bR_gru_embedding.view(sr1, sr2, -1)          # (batch_size=16, max_r_words=25, emb_size=300)

        M1 = self.Matching_Layer(bR_embedding, bU_embedding, bR_gru_embedding, bU_gru_embedding)
        m = self.Aggregation_Layer(M1, 1)
        logits = self.classifier(m)
        logits = torch.sigmoid(logits)
        logits = logits.squeeze(-1)
        return logits


class DSMN(nn.Module):
    def __init__(self, args):
        self.args = args
        super(DSMN, 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.gru_sentence = nn.GRU(input_size=args.emb_size, hidden_size=args.gru_hidden, batch_first=True)

        self.A = nn.Parameter(torch.Tensor(args.gru_hidden, args.gru_hidden))

        self.cnn_2d = nn.Conv2d(in_channels=2, out_channels=8, kernel_size=(3, 3))
        self.maxpooling = nn.MaxPool2d(kernel_size=(3, 3), stride=(3, 3))
        in_d1, in_d2 = int((MAX_SEQ_LEN - 2) / 3), int((MAX_SEQ_LEN - 2) / 3)
        self.gru_agg = nn.GRU(input_size=8 * in_d1 * in_d2, hidden_size=args.gru_hidden, batch_first=True)

        in_d1, in_d2 = int((MAX_SEQ_LEN - 2) / 3), int(DOC_LEN / 3)
        self.doc_linear = nn.Linear(in_features=8 * in_d1 * in_d2, out_features=args.gru_hidden)
        self.classifier = nn.Linear(in_features=args.gru_hidden * 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.xavier_normal_(self.A)                         # ref: YuWU-SCN
        init.kaiming_normal_(self.cnn_2d.weight)            # ref: YuWU-SCN
        init.xavier_uniform_(self.doc_linear.weight)        # if we use xavier_normal_, the model cannot converge well with lr=1e-4
        init.kaiming_normal_(self.classifier.weight)

        # ref: YuWU-SCN
        for weights in [self.gru_sentence.weight_hh_l0, self.gru_sentence.weight_ih_l0]:
            init.orthogonal_(weights)
        for weights in [self.gru_agg.weight_hh_l0, self.gru_agg.weight_ih_l0]:
            init.orthogonal_(weights)

    def Encoding_Layer(self, bU, bR, bP):
        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(embedding) for embedding in [bU, bR, bP]]
        bU_embedding, bR_embedding, bP_embedding = [self.dropout(embedding) for embedding in [bU_embedding, bR_embedding, bP_embedding]]
        return bU_embedding, bR_embedding, bP_embedding

    def Matching_Layer(self, bR_embedding, bU_embedding, bR_gru_embedding, bU_gru_embedding):
        M1 = torch.einsum('brd,btud->btru', bR_embedding, bU_embedding)
        M2 = torch.einsum('brd,dd,btud->btru', bR_gru_embedding, self.A, bU_gru_embedding)
        M = torch.stack([M1, M2], dim=2)
        return M

    def Aggregation_Layer(self, M, m_type):
        batch_size, max_utterances, n_view, max_u_words, max_r_words = M.shape
        M = M.view(-1, 2, max_u_words, max_r_words)

        # print(M.shape)
        Z = F.relu(self.cnn_2d(M))
        # print(Z.shape)
        Z = self.maxpooling(Z)
        # print(Z.shape)
        Z = Z.view(batch_size, max_utterances, -1)  # (bsz*max_utterances, *)
        # print(Z.shape)

        if 1 == m_type:
            _, matching_vector = self.gru_agg(Z)
            matching_vector = matching_vector.squeeze()
        else:
            matching_vector = Z.view(batch_size, -1)
            matching_vector = torch.tanh(self.doc_linear(matching_vector.squeeze()))

        return matching_vector

    def forward(self,inputs):
        """
        :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, )
        """
        bU = inputs['context']
        bR = inputs['response']
        bP = inputs['doc_text']
        bl = inputs['layout']


        bP = bP.view(bP.shape[0], 1, bP.shape[1])

        su1, su2, su3 = bU.shape
        sr1, sr2 = bR.shape
        sp1, sp2, sp3 = bP.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 = self.Encoding_Layer(bU, bR, bP)

        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)

        # print(bU_embedding.shape, bR_embedding.shape, bP_embedding.shape)

        bU_embedding = bU_embedding.view(-1, su3, self.args.emb_size)     # (batch_size=16, max_utterances=10, max_u_words=25, emb_size=300)
        bR_embedding = bR_embedding.view(-1, sr2, self.args.emb_size)     # (batch_size=16, max_r_words=25, emb_size=300)
        bP_embedding = bP_embedding.view(-1, sp3, self.args.emb_size)     # (batch_size=16, max_num_profiels=1, max_p_words=512, emb_size=300)

        bU_gru_embedding, _ = self.gru_sentence(bU_embedding)
        bR_gru_embedding, _ = self.gru_sentence(bR_embedding)
        bP_gru_embedding, _ = self.gru_sentence(bP_embedding)

        # Matching Layer
        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)
        bU_gru_embedding = bU_gru_embedding.view(su1, su2, su3, -1)     # (batch_size=16, max_utterances=10, max_u_words=25, emb_size=300)
        bR_gru_embedding = bR_gru_embedding.view(sr1, sr2, -1)          # (batch_size=16, max_r_words=25, emb_size=300)
        bP_gru_embedding = bP_gru_embedding.view(sp1, sp2, sp3, -1)     # (batch_size=16, max_num_profiels=1, max_p_words=512, emb_size=300)

        M1 = self.Matching_Layer(bR_embedding, bU_embedding, bR_gru_embedding, bU_gru_embedding)
        M2 = self.Matching_Layer(bR_embedding, bP_embedding, bR_gru_embedding, bP_gru_embedding)
        # print(M1.shape, M2.shape)
        # print(M.shape)                                                  # (batch_size=16, max_utterances=10, n_view=2, max_u_words=25, max_r_words=25)

        m1 = self.Aggregation_Layer(M1, 1)
        m2 = self.Aggregation_Layer(M2, 2)

        m = torch.cat([m1, m2], dim=-1)
        logits = self.classifier(m)
        logits = torch.sigmoid(logits)
        logits = logits.squeeze(-1)
        return logits

import torch
import torch.nn as nn
import torch.nn.init as init
import torch.nn.functional as F
from transformers import AutoModel

from retrieval_based.data_provider.utils import load_word_embedding
from config_file import *


class VanillaSMN(nn.Module):
    def __init__(self, args):
        self.args = args
        super(VanillaSMN, 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.gru_sentence = nn.GRU(input_size=args.emb_size, hidden_size=args.gru_hidden, batch_first=True)

        self.A = nn.Parameter(torch.Tensor(args.gru_hidden, args.gru_hidden))

        self.cnn_2d = nn.Conv2d(in_channels=2, out_channels=8, kernel_size=(3, 3))
        self.maxpooling = nn.MaxPool2d(kernel_size=(3, 3), stride=(3, 3))
        in_d1, in_d2 = int((MAX_SEQ_LEN - 2) / 3), int((MAX_SEQ_LEN - 2) / 3)
        self.context_linear = nn.Linear(in_features=8 * in_d1 * in_d2, out_features=50)
        self.gru_agg = nn.GRU(input_size=50, hidden_size=args.gru_hidden, batch_first=True)

        self.classifier = nn.Linear(in_features=args.gru_hidden, 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.xavier_normal_(self.A)                         # ref: YuWU-SCN
        init.kaiming_normal_(self.cnn_2d.weight)            # ref: YuWU-SCN
        init.xavier_uniform_(self.context_linear.weight)
        init.kaiming_normal_(self.classifier.weight)

        # ref: YuWU-SCN
        for weights in [self.gru_sentence.weight_hh_l0, self.gru_sentence.weight_ih_l0]:
            init.orthogonal_(weights)
        for weights in [self.gru_agg.weight_hh_l0, self.gru_agg.weight_ih_l0]:
            init.orthogonal_(weights)

    def Encoding_Layer(self, bU, bR):
        if self.args.use_bert:
            bU_embedding, bR_embedding = [self.ffn(self.embeddings(b, return_dict=True)['last_hidden_state']) for b in [bU, bR]]
        else:
            bU_embedding, bR_embedding = [self.embeddings(embedding) for embedding in [bU, bR]]
        bU_embedding, bR_embedding = [self.dropout(embedding) for embedding in [bU_embedding, bR_embedding]]
        return bU_embedding, bR_embedding

    def Matching_Layer(self, bR_embedding, bU_embedding, bR_gru_embedding, bU_gru_embedding):
        M1 = torch.einsum('brd,btud->btru', bR_embedding, bU_embedding)
        M2 = torch.einsum('brd,dd,btud->btru', bR_gru_embedding, self.A, bU_gru_embedding)
        M = torch.stack([M1, M2], dim=2)
        return M

    def Aggregation_Layer(self, M, m_type):
        batch_size, max_utterances, n_view, max_u_words, max_r_words = M.shape
        M = M.view(-1, 2, max_u_words, max_r_words)

        # print(M.shape)
        Z = F.relu(self.cnn_2d(M))
        # print(Z.shape)
        Z = self.maxpooling(Z)
        # print(Z.shape)
        Z = Z.view(batch_size, max_utterances, -1)  # (bsz*max_utterances, *)

        # print(Z.shape)

        if 1 == m_type:
            matching_vector = torch.tanh(self.context_linear(Z))
            _, matching_vector = self.gru_agg(matching_vector)
            matching_vector = matching_vector.squeeze()
        else:
            matching_vector = Z.view(batch_size, -1)
            matching_vector = torch.tanh(self.doc_linear(matching_vector.squeeze()))

        return matching_vector

    def forward(self, inputs):
        """
        :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, )
        """
        bU = inputs['context']
        bR = inputs['response']
        bP = inputs['doc_text']
        bl = inputs['layout']

        bP = bP.view(bP.shape[0], 1, bP.shape[1])

        su1, su2, su3 = bU.shape
        sr1, sr2 = bR.shape
        sp1, sp2, sp3 = bP.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 = self.Encoding_Layer(bU, bR)

        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)

        # print(bU_embedding.shape, bR_embedding.shape, bP_embedding.shape)

        bU_embedding = bU_embedding.view(-1, su3, self.args.emb_size)     # (batch_size=16, max_utterances=10, max_u_words=25, emb_size=300)
        bR_embedding = bR_embedding.view(-1, sr2, self.args.emb_size)     # (batch_size=16, max_r_words=25, emb_size=300)

        bU_gru_embedding, _ = self.gru_sentence(bU_embedding)
        bR_gru_embedding, _ = self.gru_sentence(bR_embedding)

        # Matching Layer
        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)
        bU_gru_embedding = bU_gru_embedding.view(su1, su2, su3, -1)     # (batch_size=16, max_utterances=10, max_u_words=25, emb_size=300)
        bR_gru_embedding = bR_gru_embedding.view(sr1, sr2, -1)          # (batch_size=16, max_r_words=25, emb_size=300)

        M1 = self.Matching_Layer(bR_embedding, bU_embedding, bR_gru_embedding, bU_gru_embedding)
        m = self.Aggregation_Layer(M1, 1)
        logits = self.classifier(m)
        logits = torch.sigmoid(logits)
        logits = logits.squeeze(-1)
        return logits


class DSMN(nn.Module):
    def __init__(self, args):
        self.args = args
        super(DSMN, 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.gru_sentence = nn.GRU(input_size=args.emb_size, hidden_size=args.gru_hidden, batch_first=True)

        self.A = nn.Parameter(torch.Tensor(args.gru_hidden, args.gru_hidden))

        self.cnn_2d = nn.Conv2d(in_channels=2, out_channels=8, kernel_size=(3, 3))
        self.maxpooling = nn.MaxPool2d(kernel_size=(3, 3), stride=(3, 3))
        in_d1, in_d2 = int((MAX_SEQ_LEN - 2) / 3), int((MAX_SEQ_LEN - 2) / 3)
        self.gru_agg = nn.GRU(input_size=8 * in_d1 * in_d2, hidden_size=args.gru_hidden, batch_first=True)

        in_d1, in_d2 = int((MAX_SEQ_LEN - 2) / 3), int(DOC_LEN / 3)
        self.doc_linear = nn.Linear(in_features=8 * in_d1 * in_d2, out_features=args.gru_hidden)
        self.classifier = nn.Linear(in_features=args.gru_hidden * 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.xavier_normal_(self.A)                         # ref: YuWU-SCN
        init.kaiming_normal_(self.cnn_2d.weight)            # ref: YuWU-SCN
        init.xavier_uniform_(self.doc_linear.weight)        # if we use xavier_normal_, the model cannot converge well with lr=1e-4
        init.kaiming_normal_(self.classifier.weight)

        # ref: YuWU-SCN
        for weights in [self.gru_sentence.weight_hh_l0, self.gru_sentence.weight_ih_l0]:
            init.orthogonal_(weights)
        for weights in [self.gru_agg.weight_hh_l0, self.gru_agg.weight_ih_l0]:
            init.orthogonal_(weights)

    def Encoding_Layer(self, bU, bR, bP):
        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(embedding) for embedding in [bU, bR, bP]]
        bU_embedding, bR_embedding, bP_embedding = [self.dropout(embedding) for embedding in [bU_embedding, bR_embedding, bP_embedding]]
        return bU_embedding, bR_embedding, bP_embedding

    def Matching_Layer(self, bR_embedding, bU_embedding, bR_gru_embedding, bU_gru_embedding):
        M1 = torch.einsum('brd,btud->btru', bR_embedding, bU_embedding)
        M2 = torch.einsum('brd,dd,btud->btru', bR_gru_embedding, self.A, bU_gru_embedding)
        M = torch.stack([M1, M2], dim=2)
        return M

    def Aggregation_Layer(self, M, m_type):
        batch_size, max_utterances, n_view, max_u_words, max_r_words = M.shape
        M = M.view(-1, 2, max_u_words, max_r_words)

        # print(M.shape)
        Z = F.relu(self.cnn_2d(M))
        # print(Z.shape)
        Z = self.maxpooling(Z)
        # print(Z.shape)
        Z = Z.view(batch_size, max_utterances, -1)  # (bsz*max_utterances, *)
        # print(Z.shape)

        if 1 == m_type:
            _, matching_vector = self.gru_agg(Z)
            matching_vector = matching_vector.squeeze()
        else:
            matching_vector = Z.view(batch_size, -1)
            matching_vector = torch.tanh(self.doc_linear(matching_vector.squeeze()))

        return matching_vector

    def forward(self,inputs):
        """
        :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, )
        """
        bU = inputs['context']
        bR = inputs['response']
        bP = inputs['doc_text']
        bl = inputs['layout']


        bP = bP.view(bP.shape[0], 1, bP.shape[1])

        su1, su2, su3 = bU.shape
        sr1, sr2 = bR.shape
        sp1, sp2, sp3 = bP.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 = self.Encoding_Layer(bU, bR, bP)

        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)

        # print(bU_embedding.shape, bR_embedding.shape, bP_embedding.shape)

        bU_embedding = bU_embedding.view(-1, su3, self.args.emb_size)     # (batch_size=16, max_utterances=10, max_u_words=25, emb_size=300)
        bR_embedding = bR_embedding.view(-1, sr2, self.args.emb_size)     # (batch_size=16, max_r_words=25, emb_size=300)
        bP_embedding = bP_embedding.view(-1, sp3, self.args.emb_size)     # (batch_size=16, max_num_profiels=1, max_p_words=512, emb_size=300)

        bU_gru_embedding, _ = self.gru_sentence(bU_embedding)
        bR_gru_embedding, _ = self.gru_sentence(bR_embedding)
        bP_gru_embedding, _ = self.gru_sentence(bP_embedding)

        # Matching Layer
        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)
        bU_gru_embedding = bU_gru_embedding.view(su1, su2, su3, -1)     # (batch_size=16, max_utterances=10, max_u_words=25, emb_size=300)
        bR_gru_embedding = bR_gru_embedding.view(sr1, sr2, -1)          # (batch_size=16, max_r_words=25, emb_size=300)
        bP_gru_embedding = bP_gru_embedding.view(sp1, sp2, sp3, -1)     # (batch_size=16, max_num_profiels=1, max_p_words=512, emb_size=300)

        M1 = self.Matching_Layer(bR_embedding, bU_embedding, bR_gru_embedding, bU_gru_embedding)
        M2 = self.Matching_Layer(bR_embedding, bP_embedding, bR_gru_embedding, bP_gru_embedding)
        # print(M1.shape, M2.shape)
        # print(M.shape)                                                  # (batch_size=16, max_utterances=10, n_view=2, max_u_words=25, max_r_words=25)

        m1 = self.Aggregation_Layer(M1, 1)
        m2 = self.Aggregation_Layer(M2, 2)

        m = torch.cat([m1, m2], dim=-1)
        logits = self.classifier(m)
        logits = torch.sigmoid(logits)
        logits = logits.squeeze(-1)
        return logits
# ref: YuWU-SCN # ref: YuWU-SCN # ref: YuWU-SCN # print(M.shape) # print(Z.shape) # print(Z.shape) # (bsz*max_utterances, *) # print(Z.shape) :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, ) # Encoding Layer ## U, R, D # (batch_size=16, max_utterances=10, max_u_words=25, emb_size=300) # (batch_size=16, max_r_words=25, emb_size=300) # print(bU_embedding.shape, bR_embedding.shape, bP_embedding.shape) # (batch_size=16, max_utterances=10, max_u_words=25, emb_size=300) # (batch_size=16, max_r_words=25, emb_size=300) # Matching Layer # (batch_size=16, max_utterances=10, max_u_words=25, emb_size=300) # (batch_size=16, max_r_words=25, emb_size=300) # (batch_size=16, max_utterances=10, max_u_words=25, emb_size=300) # (batch_size=16, max_r_words=25, emb_size=300) # ref: YuWU-SCN # ref: YuWU-SCN # if we use xavier_normal_, the model cannot converge well with lr=1e-4 # ref: YuWU-SCN # print(M.shape) # print(Z.shape) # print(Z.shape) # (bsz*max_utterances, *) # print(Z.shape) :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, ) # Encoding Layer ## U, R, D # (batch_size=16, max_utterances=10, max_u_words=25, emb_size=300) # (batch_size=16, max_r_words=25, emb_size=300) # (batch_size=16, max_num_profiels=1, max_p_words=512, emb_size=300) # print(bU_embedding.shape, bR_embedding.shape, bP_embedding.shape) # (batch_size=16, max_utterances=10, max_u_words=25, emb_size=300) # (batch_size=16, max_r_words=25, emb_size=300) # (batch_size=16, max_num_profiels=1, max_p_words=512, emb_size=300) # Matching Layer # (batch_size=16, max_utterances=10, max_u_words=25, emb_size=300) # (batch_size=16, max_r_words=25, emb_size=300) # (batch_size=16, max_num_profiels=1, max_p_words=512, emb_size=300) # (batch_size=16, max_utterances=10, max_u_words=25, emb_size=300) # (batch_size=16, max_r_words=25, emb_size=300) # (batch_size=16, max_num_profiels=1, max_p_words=512, emb_size=300) # print(M1.shape, M2.shape) # print(M.shape)                                                  # (batch_size=16, max_utterances=10, n_view=2, max_u_words=25, max_r_words=25)