semi_sup_net.py

import torch

from torch import nn
import torch.nn.functional as F
from torch.autograd import Variable

import numpy as np
import math
from config import global_config as cfg
import copy, random, time, logging


def cuda_(var):
    return var.cuda() if cfg.cuda else var


def toss_(p):
    return random.randint(0, 99) <= p


def nan(v):
    return np.isnan(np.sum(v.data.cpu().numpy()))


def get_sparse_input(x_input):
    """
    get a sparse matrix of x_input: [T,B,V] where x_sparse[i][j][k]=1, and others = 1e-8
    :param x_input: *Tensor* of [T,B]
    :return: *Tensor* in shape [B,T,V]
    """
    # indexes that will make no effect in copying
    sw = time.time()
    print('sparse input start: %s' % sw)
    ignore_index = [0]
    result = torch.normal(mean=0, std=torch.zeros(x_input.size(0), x_input.size(1), cfg.vocab_size))
    for t in range(x_input.size(0)):
        for b in range(x_input.size(1)):
            if x_input[t][b] not in ignore_index:
                result[t][b][x_input[t][b]] = 1.0
    print('sparse input end %s' % time.time())
    return result.transpose(0, 1)


def get_sparse_input_efficient(x_input_np):
    ignore_index = [0]
    result = np.zeros((x_input_np.shape[0], x_input_np.shape[1], cfg.vocab_size), dtype=np.float32)
    result.fill(1e-10)
    for t in range(x_input_np.shape[0]):
        for b in range(x_input_np.shape[1]):
            if x_input_np[t][b] not in ignore_index:
                result[t,b,x_input_np[t][b]] = 1.0
    result_np = result.transpose((1, 0, 2))
    result = torch.from_numpy(result_np).float()
    return result


def shift(pz_proba):
    first_input = np.zeros((pz_proba.size(1), pz_proba.size(2)))
    first_input.fill(1e-12)
    first_input = cuda_(Variable(torch.from_numpy(first_input)).float())
    pz_proba = list(pz_proba)[:-1]
    pz_proba.insert(0, first_input)
    pz_proba = torch.stack(pz_proba, 0)
    return pz_proba.contiguous()


class Encoder(nn.Module):
    def __init__(self, input_size, embed_size, hidden_size, n_layers, dropout):
        super(Encoder, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.embed_size = embed_size
        self.n_layers = n_layers
        self.dropout = dropout
        self.embedding = nn.Embedding(input_size, embed_size)
        self.gru = nn.GRU(embed_size, hidden_size, n_layers, dropout=self.dropout, bidirectional=True)

    def forward(self, input_seqs, hidden=None):
        embedded = self.embedding(input_seqs)
        outputs, hidden = self.gru(embedded, hidden)
        outputs = outputs[:, :, :self.hidden_size] + outputs[:, :, self.hidden_size:]  # Sum bidirectional outputs
        return outputs, hidden


class DynamicEncoder(nn.Module):
    def __init__(self, input_size, embed_size, hidden_size, n_layers, dropout):
        super().__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.embed_size = embed_size
        self.n_layers = n_layers
        self.dropout = dropout
        self.embedding = nn.Embedding(input_size, embed_size)
        self.gru = nn.GRU(embed_size, hidden_size, n_layers, dropout=self.dropout, bidirectional=True)

    def forward(self, input_seqs, input_lens, hidden=None):
        """
        forward procedure. No need for inputs to be sorted
        :param input_seqs: Variable of [T,B]
        :param hidden:
        :param input_lens: *numpy array* of len for each input sequence
        :return:
        """
        batch_size = input_seqs.size(1)
        embedded = self.embedding(input_seqs)
        embedded = embedded.transpose(0, 1)  # [B,T,E]
        sort_idx = np.argsort(-input_lens)
        unsort_idx = cuda_(torch.LongTensor(np.argsort(sort_idx)))
        input_lens = input_lens[sort_idx]
        sort_idx = cuda_(torch.LongTensor(sort_idx))
        embedded = embedded[sort_idx].transpose(0, 1)  # [T,B,E]
        packed = torch.nn.utils.rnn.pack_padded_sequence(embedded, input_lens)
        outputs, hidden = self.gru(packed, hidden)

        outputs, _ = torch.nn.utils.rnn.pad_packed_sequence(outputs)
        outputs = outputs[:,:,:self.hidden_size] + outputs[:,:,self.hidden_size:]
        outputs = outputs.transpose(0, 1)[unsort_idx].transpose(0, 1).contiguous()
        hidden = hidden.transpose(0, 1)[unsort_idx].transpose(0, 1).contiguous()
        return outputs, hidden


class Attn(nn.Module):
    def __init__(self, hidden_size):
        super(Attn, self).__init__()
        self.hidden_size = hidden_size
        self.attn = nn.Linear(self.hidden_size * 2, hidden_size)
        self.v = nn.Linear(self.hidden_size, 1)

    def forward(self, hidden, encoder_outputs, normalize=True):
        encoder_outputs = encoder_outputs.transpose(0, 1)  # [B,T,H]
        attn_energies = self.score(hidden, encoder_outputs)
        normalized_energy = F.softmax(attn_energies, dim=2)  # [B,1,T]
        context = torch.bmm(normalized_energy, encoder_outputs)  # [B,1,H]
        return context.transpose(0, 1)  # [1,B,H]

    def score(self, hidden, encoder_outputs):
        max_len = encoder_outputs.size(1)
        H = hidden.repeat(max_len, 1, 1).transpose(0, 1)
        energy = self.attn(torch.cat([H, encoder_outputs], 2))  # [B,T,2H]->[B,T,H]
        energy = self.v(F.tanh(energy)).transpose(1,2) # [B,1,T]
        return energy


class MultiTurnInferenceDecoder_Z(nn.Module):
    """
    Inference network: copying version of Q_phi(z_t|s_t,m_t) <- Q_phi(z_ti|s_t,m_t,z_t[1..i-1])
    """

    def __init__(self, embed_size, hidden_size, vocab_size, dropout_rate):
        super().__init__()
        self.gru = nn.GRU(embed_size, hidden_size, dropout=dropout_rate)
        self.w1 = nn.Linear(hidden_size, vocab_size)
        self.mu = nn.Linear(vocab_size, embed_size, bias=False)
        self.log_sigma = nn.Linear(vocab_size, embed_size)
        self.dropout_rate = dropout_rate
        self.vocab_size = vocab_size
        self.proj_copy1 = nn.Linear(hidden_size, hidden_size)
        self.proj_copy2 = nn.Linear(hidden_size, hidden_size)
        self.proj_copy3 = nn.Linear(hidden_size, hidden_size)
        self.dropout = nn.Dropout(self.dropout_rate)

    def forward(self, u_input, u_enc_out, pv_pz_proba, pv_z_dec_out, m_input, m_enc_out, embed_z, last_hidden,
                rand_eps, u_input_np, m_input_np):
        """
        Similar to base class method
        :param m_input:
        :param u_input:
        :param u_enc_out:
        :param m_enc_out:
        :param embed_z:
        :param last_hidden:
        :param rand_eps:
        :return:
        """
        sparse_u_input = Variable(get_sparse_input_efficient(u_input_np), requires_grad=False)  # [B,T,V]
        sparse_m_input = Variable(get_sparse_input_efficient(m_input_np), requires_grad=False)  # [B,T,V]

        # if cfg.cuda: sparse_m_input = sparse_m_input.cuda()
        # if cfg.cuda: sparse_u_input = sparse_u_input.cuda()

        embed_z = self.dropout(embed_z)
        gru_out, last_hidden = self.gru(embed_z, last_hidden)
        gen_score = self.w1(gru_out).squeeze(0) # [B,V]
        u_copy_score = F.tanh(self.proj_copy1(u_enc_out.transpose(0, 1)))  # [B,T,H]
        m_copy_score = F.tanh(self.proj_copy2(m_enc_out.transpose(0, 1)))
        if not cfg.force_stable:
            # unstable version of copynet for small dataset
            u_copy_score = torch.exp(torch.matmul(u_copy_score, gru_out.squeeze(0).unsqueeze(2)).squeeze(2))  # [B,T]
            m_copy_score = torch.exp(torch.matmul(m_copy_score, gru_out.squeeze(0).unsqueeze(2)).squeeze(2))  # [B,T]
            u_copy_score, m_copy_score = u_copy_score.cpu(), m_copy_score.cpu()
            u_copy_score = torch.log(torch.bmm(u_copy_score.unsqueeze(1), sparse_u_input)).squeeze(1)  # [B,V]
            m_copy_score = torch.log(torch.bmm(m_copy_score.unsqueeze(1), sparse_m_input)).squeeze(1)  # [B,V]
        else:
            # stable version of copynet
            u_copy_score = torch.matmul(u_copy_score, gru_out.squeeze(0).unsqueeze(2)).squeeze(2)
            m_copy_score = torch.matmul(m_copy_score, gru_out.squeeze(0).unsqueeze(2)).squeeze(2)
            u_copy_score, m_copy_score = u_copy_score.cpu(), m_copy_score.cpu()
            u_copy_score_max, m_copy_score_max = torch.max(u_copy_score, dim=1, keepdim=True)[0], \
                                                 torch.max(m_copy_score, dim=1, keepdim=True)[0]
            u_copy_score = torch.exp(u_copy_score - u_copy_score_max)  # [B,T]
            m_copy_score = torch.exp(m_copy_score - m_copy_score_max)  # [B,T]
            # u_copy_score, m_copy_score = u_copy_score.cpu(), m_copy_score.cpu()
            u_copy_score = torch.log(torch.bmm(u_copy_score.unsqueeze(1), sparse_u_input)).squeeze(
                1) + u_copy_score_max  # [B,V]
            m_copy_score = torch.log(torch.bmm(m_copy_score.unsqueeze(1), sparse_m_input)).squeeze(
                1) + m_copy_score_max  # [B,V]
        u_copy_score, m_copy_score = cuda_(u_copy_score), cuda_(m_copy_score)
        if pv_pz_proba is not None:
            pv_pz_proba = shift(pv_pz_proba)
            pv_z_copy_score = F.tanh(self.proj_copy3(pv_z_dec_out.transpose(0, 1)))  # [B,T,H]
            if cfg.force_stable:
                pv_z_copy_score = torch.exp(
                    torch.matmul(pv_z_copy_score, gru_out.squeeze(0).unsqueeze(2)).squeeze(2))  # [B,T]
                pv_z_copy_score = torch.log(
                    torch.bmm(pv_z_copy_score.unsqueeze(1), pv_pz_proba.transpose(0, 1))).squeeze(
                    1)  # [B,V]
            else:
                pv_z_copy_score = torch.matmul(pv_z_copy_score, gru_out.squeeze(0).unsqueeze(2)).squeeze(2)
                pv_z_copy_score_max = torch.max(pv_z_copy_score, dim=1, keepdim=True)[0]
                pv_z_copy_score = torch.exp(pv_z_copy_score - pv_z_copy_score_max)
                pv_z_copy_score = torch.log(
                    torch.bmm(pv_z_copy_score.unsqueeze(1), pv_pz_proba.transpose(0, 1))).squeeze(
                    1) + pv_z_copy_score_max  # [B,V]
            scores = F.softmax(torch.cat([gen_score, u_copy_score, m_copy_score, pv_z_copy_score], dim=1), dim=1)
            gen_score, u_copy_score, m_copy_score, pv_z_copy_score = tuple(
                torch.split(scores, gen_score.size(1), dim=1))
            proba = gen_score + u_copy_score + m_copy_score + pv_z_copy_score
        else:
            scores = F.softmax(torch.cat([gen_score, u_copy_score, m_copy_score], dim=1), dim=1)
            gen_score, u_copy_score, m_copy_score = tuple(
                torch.split(scores, gen_score.size(1), dim=1))
            proba = gen_score + u_copy_score + m_copy_score
        appr_emb = self.mu(proba).unsqueeze(0)
        # log_sigma_ae = self.log_sigma(proba)
        # sigma_ae = torch.exp(log_sigma_ae)
        # sampled_ae = appr_emb + torch.mul(sigma_ae, rand_eps)
        return appr_emb, gru_out, last_hidden, proba, appr_emb, None


class MultiTurnPriorDecoder_Z(nn.Module):
    def __init__(self, embed_size, hidden_size, vocab_size, dropout_rate):
        super().__init__()
        self.gru = nn.GRU(embed_size, hidden_size, dropout=dropout_rate)
        self.w1 = nn.Linear(hidden_size, vocab_size)
        self.proj_copy1 = nn.Linear(hidden_size, hidden_size)
        self.proj_copy2 = nn.Linear(hidden_size, hidden_size)
        self.mu = nn.Linear(vocab_size, embed_size, bias=False)
        self.log_sigma = nn.Linear(vocab_size, embed_size)
        self.dropout_rate = dropout_rate
        self.dropout = nn.Dropout(dropout_rate)

    def forward(self, u_input, u_enc_out, pv_pz_proba, pv_z_dec_out, embed_z, last_hidden, rand_eps, u_input_np,
                m_input_np):
        sparse_u_input = Variable(get_sparse_input_efficient(u_input_np), requires_grad=False)
        embed_z = self.dropout(embed_z)
        gru_out, last_hidden = self.gru(embed_z, last_hidden)
        gen_score = self.w1(gru_out).squeeze(0)
        u_copy_score = F.tanh(self.proj_copy1(u_enc_out.transpose(0, 1)))  # [B,T,H]
        if not cfg.force_stable:
            u_copy_score = torch.exp(torch.matmul(u_copy_score, gru_out.squeeze(0).unsqueeze(2)).squeeze(2))  # [B,T]
            u_copy_score = u_copy_score.cpu()
            u_copy_score = torch.log(torch.bmm(u_copy_score.unsqueeze(1), sparse_u_input)).squeeze(1)  # [B,V]
        else:
            # stable version of copynet
            u_copy_score = torch.matmul(u_copy_score, gru_out.squeeze(0).unsqueeze(2)).squeeze(2)
            u_copy_score = u_copy_score.cpu()
            u_copy_score_max = torch.max(u_copy_score, dim=1, keepdim=True)[0]
            u_copy_score = torch.exp(u_copy_score - u_copy_score_max)  # [B,T]
            u_copy_score = torch.log(torch.bmm(u_copy_score.unsqueeze(1), sparse_u_input)).squeeze(
                1) + u_copy_score_max  # [B,V]
        u_copy_score = cuda_(u_copy_score)
        if pv_pz_proba is not None:
            pv_pz_proba = shift(pv_pz_proba)
            pv_z_copy_score = F.tanh(self.proj_copy2(pv_z_dec_out.transpose(0, 1)))  # [B,T,H]
            if cfg.force_stable:
                pv_z_copy_score = torch.exp(
                    torch.matmul(pv_z_copy_score, gru_out.squeeze(0).unsqueeze(2)).squeeze(2))  # [B,T]
                pv_z_copy_score = torch.log(
                    torch.bmm(pv_z_copy_score.unsqueeze(1), pv_pz_proba.transpose(0, 1))).squeeze(
                    1)  # [B,V]
            else:
                pv_z_copy_score = torch.matmul(pv_z_copy_score, gru_out.squeeze(0).unsqueeze(2)).squeeze(2)
                pv_z_copy_score_max = torch.max(pv_z_copy_score, dim=1, keepdim=True)[0]
                pv_z_copy_score = torch.exp(pv_z_copy_score - pv_z_copy_score_max)
                pv_z_copy_score = torch.log(
                    torch.bmm(pv_z_copy_score.unsqueeze(1), pv_pz_proba.transpose(0, 1))).squeeze(
                    1) + pv_z_copy_score_max  # [B,V]
            scores = F.softmax(torch.cat([gen_score, u_copy_score, pv_z_copy_score], dim=1), dim=1)
            gen_score, u_copy_score, pv_z_copy_score = tuple(torch.split(scores, gen_score.size(1), dim=1))
            proba = gen_score + u_copy_score + pv_z_copy_score  # [B,V]
        else:
            scores = F.softmax(torch.cat([gen_score, u_copy_score], dim=1), dim=1)
            gen_score, u_copy_score = tuple(torch.split(scores, gen_score.size(1), dim=1))
            proba = gen_score + u_copy_score  # [B,V]
        appr_emb = self.mu(proba).unsqueeze(0)
        return appr_emb, gru_out, last_hidden, proba, appr_emb, None


class ResponseDecoder(nn.Module):
    """
    Response decoder: P_theta(m_t|s_t, z_t) <- P_theta(m_ti|s_t, z_t, m_t[1..i-1])
    This is a deterministic decoder.
    """

    def __init__(self, embed_size, hidden_size, vocab_size, degree_size, dropout_rate):
        super().__init__()
        self.emb = nn.Embedding(vocab_size, embed_size)
        self.attn_z = Attn(hidden_size)
        self.attn_u = Attn(hidden_size)
        self.w4 = nn.Linear(hidden_size, hidden_size)
        self.gate_z = nn.Linear(hidden_size, hidden_size)
        self.w5 = nn.Linear(hidden_size, hidden_size)
        self.gru = nn.GRU(embed_size + hidden_size + degree_size, hidden_size, dropout=dropout_rate)
        self.proj = nn.Linear(hidden_size * 3, vocab_size)
        self.proj_copy1 = nn.Linear(hidden_size, hidden_size)
        self.dropout_rate = dropout_rate
        self.dropout = nn.Dropout(dropout_rate)

    def forward(self, z_enc_out, pz_proba, u_enc_out, m_t_input, degree_input, last_hidden):
        """
        decode the response: P(m|u,z)
        :param degree_input: [B,D]
        :param pz_proba: [Tz,B,V], output of the prior decoder
        :param z_enc_out: [Tz,B,H]
        :param u_enc_out: [T,B,H]
        :param m_t_input: [1,B]
        :param last_hidden:
        :return: proba: [1,B,V]
        """
        m_embed = self.emb(m_t_input)
        pz_proba = shift(pz_proba)

        z_context = self.attn_z(last_hidden, z_enc_out)
        u_context = self.attn_u(last_hidden, u_enc_out)
        d_control = z_context + torch.mul(F.sigmoid(self.gate_z(z_context)), u_context)
        embed = torch.cat([d_control, m_embed, degree_input.unsqueeze(0)], dim=2)
        embed = self.dropout(embed)
        gru_out, last_hidden = self.gru(embed, last_hidden)
        gen_score = self.proj(torch.cat([z_context, u_context, gru_out], 2)).squeeze(0)
        z_copy_score = F.tanh(self.proj_copy1(z_enc_out.transpose(0, 1)))  # [B,T,H]
        if not cfg.force_stable:
            z_copy_score = torch.exp(torch.matmul(z_copy_score, gru_out.squeeze(0).unsqueeze(2)).squeeze(2))  # [B,T]
            z_copy_score = torch.log(torch.bmm(z_copy_score.unsqueeze(1), pz_proba.transpose(0, 1))).squeeze(1)  # [B,V]
        else:
            z_copy_score = torch.matmul(z_copy_score, gru_out.squeeze(0).unsqueeze(2)).squeeze(2)
            z_copy_score_max = torch.max(z_copy_score, dim=1, keepdim=True)[0]
            z_copy_score = torch.exp(z_copy_score - z_copy_score_max)
            z_copy_score = torch.log(torch.bmm(z_copy_score.unsqueeze(1), pz_proba.transpose(0, 1)))
            z_copy_score = z_copy_score.squeeze(1) + z_copy_score_max
        scores = F.softmax(torch.cat([gen_score, z_copy_score], dim=1), dim=1)
        gen_score, z_copy_score = tuple(torch.split(scores, gen_score.size(1), dim=1))
        proba = gen_score + z_copy_score  # [B,V]
        return proba, last_hidden, gru_out


class MultinomialKLDivergenceLoss(nn.Module):
    def __init__(self, special_tokens=[]):
        super().__init__()
        self.special_tokens = special_tokens

    def forward(self, p_proba, q_proba): # [B, T, V]
        mask = torch.ones(p_proba.size(0), p_proba.size(1))
        cnt = 0
        for i in range(q_proba.size(0)):
            flg = False
            for j in range(q_proba.size(1)):
                topv, topi = torch.max(q_proba[i,j], -1)
                if flg:
                    mask[i,j] = 0
                else:
                    mask[i,j] = 1
                    cnt += 1
                if topi.item() in self.special_tokens:
                    flg = True
        mask = cuda_(Variable(mask))
        loss = q_proba * (torch.log(q_proba) - torch.log(p_proba))
        masked_loss = torch.sum(mask.unsqueeze(-1) * loss)
        return masked_loss / (cnt + 1e-10)


class SemiSupervisedSEDST(nn.Module):
    def __init__(self, embed_size, hidden_size, vocab_size, degree_size, layer_num, dropout_rate, z_length, alpha,
                 max_ts, beam_search=False, teacher_force=100, **kwargs):
        super().__init__()
        self.u_encoder = DynamicEncoder(vocab_size, embed_size, hidden_size, layer_num, dropout_rate)
        self.m_encoder = DynamicEncoder(vocab_size, embed_size, hidden_size, layer_num, dropout_rate)
        self.m_decoder = ResponseDecoder(embed_size, hidden_size, vocab_size, degree_size, dropout_rate)
        self.qz_decoder = MultiTurnInferenceDecoder_Z(embed_size, hidden_size, vocab_size, dropout_rate)  # posterior
        self.pz_decoder = MultiTurnPriorDecoder_Z(embed_size, hidden_size, vocab_size, dropout_rate)  # prior

        self.embed_size = embed_size
        self.vocab = kwargs['vocab']

        self.pr_loss = nn.NLLLoss(ignore_index=0)
        self.q_loss = nn.NLLLoss(ignore_index=0)
        self.dec_loss = nn.NLLLoss(ignore_index=0)
        self.kl_loss = MultinomialKLDivergenceLoss(special_tokens=[self.vocab.encode(x) for x in ['EOS_Z1','EOS_Z2',
                                                                                                  '</s>', '<pad>']])

        self.z_length = z_length
        self.alpha = alpha
        self.max_ts = max_ts
        self.beam_search = beam_search
        self.teacher_force = teacher_force

        if self.beam_search:
            self.beam_size = kwargs['beam_size']
            self.eos_token_idx = kwargs['eos_token_idx']

    def forward(self, u_input, u_input_np, m_input, m_input_np, z_input, u_len, m_len, turn_states, z_supervised,
                p_input, p_input_np, p_len,
                degree_input, mode):
        if mode == 'train' or mode == 'valid':
            if not z_supervised:
                z_input = None
            pz_proba, qz_proba, pm_dec_proba, pz_mu, pz_log_sigma, qz_mu, qz_log_sigma, turn_states = \
                self.forward_turn(u_input, u_len, m_input=m_input, m_len=m_len, z_input=z_input, is_train=True,
                                  turn_states=turn_states, degree_input=degree_input, u_input_np=u_input_np,
                                  m_input_np=m_input_np,p_input=p_input,p_input_np=p_input_np,p_len=p_len)
            if z_supervised:
                loss, pr_loss, m_loss, q_loss = self.supervised_loss(torch.log(pz_proba), torch.log(qz_proba),
                                                                     torch.log(pm_dec_proba), z_input, m_input)
                return loss, pr_loss, m_loss, q_loss, turn_states
            else:
                loss, m_loss, kl_div_loss = self.unsupervised_loss(qz_mu, qz_log_sigma, pz_mu, pz_log_sigma,
                                                                   torch.log(pm_dec_proba), m_input, pz_proba, qz_proba)
            return loss, m_loss, kl_div_loss, turn_states
        elif mode == 'test':
            m_output_index, pz_index, turn_states = self.forward_turn(u_input, u_len=u_len, is_train=False,
                                                                      turn_states=turn_states,
                                                                      degree_input=degree_input,
                                                                      u_input_np=u_input_np, m_input_np=m_input_np,
                                                                      p_input=p_input, p_input_np=p_input_np,
                                                                      p_len=p_len
                                                                      )
            return m_output_index, pz_index, turn_states


    def forward_turn(self, u_input, u_len, turn_states, is_train, degree_input, u_input_np, m_input_np=None,
                     m_input=None, m_len=None, z_input=None,
                     p_input=None, p_input_np=None, p_len=None,test_type='pr'):
        """
        compute required outputs for a single dialogue turn. Turn state{Dict} will be updated in each call.
        :param u_input_np:
        :param m_input_np:
        :param u_len:
        :param turn_states:
        :param is_train:
        :param u_input: [T,B]
        :param m_input: [T,B]
        :param z_input: [T,B]
        :return:
        """
        pv_pz_proba = turn_states.get('pv_pz_proba', None)
        pv_z_outs = turn_states.get('pv_z_dec_outs', None)
        pv_qz_proba = turn_states.get('pv_qz_proba', None)
        pv_qz_outs = turn_states.get('pv_qz_dec_outs', None)
        batch_size = u_input.size(1)
        u_enc_out, u_enc_hidden = self.u_encoder(u_input, u_len)
        last_hidden = u_enc_hidden[:-1]
        # initial approximate embedding: SOS token initialized with all zero
        # Pi(z|u)
        pz_ae = cuda_(Variable(torch.zeros(1, batch_size, self.embed_size)))
        pz_proba, pz_mu, pz_log_sigma = [], [], []
        pz_dec_outs = []
        z_length = z_input.size(0) if z_input is not None else self.z_length
        for t in range(z_length):
            if cfg.sampling:
                rand_eps = Variable(torch.normal(means=torch.zeros(1, batch_size, cfg.embedding_size), std=1))
            else:
                rand_eps = Variable(torch.zeros(1, batch_size, cfg.embedding_size))
            if cfg.cuda: rand_eps = rand_eps.cuda()

            pz_ae, last_hidden, pz_dec_out, proba, appr_emb, log_sigma_ae = \
                self.pz_decoder(u_input=u_input, u_enc_out=u_enc_out, pv_pz_proba=pv_pz_proba, pv_z_dec_out=pv_z_outs,
                                embed_z=pz_ae, last_hidden=last_hidden, rand_eps=rand_eps, u_input_np=u_input_np,
                                m_input_np=m_input_np)
            pz_proba.append(proba)
            pz_mu.append(appr_emb)
            pz_log_sigma.append(log_sigma_ae)
            pz_dec_outs.append(pz_dec_out)
        pz_dec_outs = torch.cat(pz_dec_outs, dim=0)  # [Tz,B,H]
        pz_proba, pz_mu = torch.stack(pz_proba, dim=0), torch.stack(pz_mu, dim=0)
        # P(m|z,u)
        m_tm1 = cuda_(Variable(torch.ones(1, batch_size).long()))  # GO token
        pm_dec_proba, m_dec_outs = [],[]

        turn_states['pv_z_dec_outs'], turn_states['pv_pz_proba'] = pz_dec_outs, pz_proba

        if is_train or test_type=='post':
            m_length = m_input.size(0)  # Tm
            for t in range(m_length):
                teacher_forcing = toss_(self.teacher_force)
                proba, last_hidden, dec_out = self.m_decoder(pz_dec_outs, pz_proba, u_enc_out, m_tm1, degree_input, last_hidden)
                if teacher_forcing:
                    m_tm1 = m_input[t].view(1, -1)
                else:
                    _, m_tm1 = torch.topk(proba, 1)
                    m_tm1 = m_tm1.view(1, -1)
                pm_dec_proba.append(proba)
                m_dec_outs.append(dec_out)

            pm_dec_proba = torch.stack(pm_dec_proba, dim=0)  # [T,B,V]

            # Q(z|u,m)
            u_enc_out, u_enc_hidden = self.m_encoder(u_input, u_len)
            m_enc_out, m_enc_hidden = self.m_encoder(m_input, m_len)

            last_hidden = u_enc_hidden[:-1]

            qz_ae = cuda_(Variable(torch.zeros(1, batch_size, self.embed_size)))
            qz_proba, qz_mu, qz_log_sigma, qz_dec_outs = [], [], [], []
            for t in range(z_length):
                if cfg.sampling:
                    rand_eps = self.alpha * Variable(torch.normal(means=torch.zeros(1, batch_size, cfg.embedding_size), std=1))
                else:
                    rand_eps = Variable(torch.zeros(1, batch_size, cfg.embedding_size))
                if cfg.cuda: rand_eps = rand_eps.cuda()
                qz_ae, gru_out, last_hidden, proba, appr_emb, log_sigma_ae = \
                    self.qz_decoder(u_input=u_input, u_enc_out=u_enc_out, pv_pz_proba=pv_qz_proba,
                                 pv_z_dec_out=pv_qz_outs,
                                 m_input=m_input, m_enc_out=m_enc_out, u_input_np=u_input_np, m_input_np=m_input_np,
                                 embed_z=qz_ae, last_hidden=last_hidden, rand_eps=rand_eps)
                qz_proba.append(proba)
                qz_mu.append(appr_emb)
                qz_log_sigma.append(log_sigma_ae)
                qz_dec_outs.append(gru_out)
            qz_proba, qz_mu = torch.stack(qz_proba, dim=0), torch.stack(qz_mu, dim=0)
            qz_dec_outs = torch.cat(qz_dec_outs, dim=0)
            turn_states['pv_qz_dec_outs'], turn_states['pv_qz_proba'] = qz_dec_outs, qz_proba

            if is_train:
                return pz_proba, qz_proba, pm_dec_proba, pz_mu, pz_log_sigma, qz_mu, qz_log_sigma, turn_states
            else:
                qz_index = self.pz_max_sampling(qz_proba)
                return None, qz_index, turn_states
        else:
            if not self.beam_search:
                m_output_index = self.greedy_decode(pz_dec_outs, pz_proba, u_enc_out, m_tm1, last_hidden, degree_input)
            else:
                m_output_index = self.beam_search_decode(pz_dec_outs, pz_proba, u_enc_out, m_tm1, last_hidden,
                                                         degree_input,
                                                         self.eos_token_idx)
            pz_index = self.pz_max_sampling(pz_proba)
            return m_output_index, pz_index, turn_states

    def greedy_decode(self, pz_dec_outs, pz_proba, u_enc_out, m_tm1, last_hidden, degree_input):
        """
        greedy decoding of the response
        :param pz_dec_outs:
        :param u_enc_out:
        :param m_tm1:
        :param last_hidden:
        :return: nested-list
        """
        decoded = []
        for t in range(self.max_ts):
            proba, last_hidden, _ = self.m_decoder(pz_dec_outs, pz_proba, u_enc_out, m_tm1, degree_input, last_hidden)
            mt_proba, mt_index = torch.topk(proba, 1)  # [B,1]
            mt_index = mt_index.data.view(-1)
            decoded.append(mt_index)
            m_tm1 = cuda_(Variable(mt_index).view(1, -1))
        decoded = torch.stack(decoded, dim=0).transpose(0, 1)
        decoded = list(decoded)
        return [list(_) for _ in decoded]

    def pz_max_sampling(self, pz_proba):
        """
        Max-sampling procedure of pz during testing.
        :param pz_proba: # [Tz, B, Vz]
        :return: nested-list: B * [T]
        """
        pz_proba = pz_proba.data
        z_proba, z_token = torch.topk(pz_proba, 1, dim=2)  # [Tz, B, 1]
        z_token = list(z_token.squeeze(2).transpose(0, 1))
        return [list(_) for _ in z_token]

    def pz_selective_sampling(self, pz_proba):
        """
        Selective sampling of pz
        """
        if cfg.spv_proportion == 0:
            return self.pz_max_sampling(pz_proba)
        pz_proba = pz_proba.data
        z_proba, z_token = torch.topk(pz_proba, pz_proba.size(0), dim=2)
        z_token = z_token.transpose(0, 1)  # [B,Tz,top_Tz]
        all_sampled_z = []
        for b in range(z_token.size(0)):
            sampled_z = []
            for t in range(z_token.size(1)):
                for i in range(z_token.size(2)):
                    if z_token[b][t][i] not in sampled_z:
                        sampled_z.append(z_token[b][t][i])
                        break
            all_sampled_z.append(sampled_z)
        return all_sampled_z

    def beam_search_decode_single(self, pz_dec_outs, pz_proba, u_enc_out, m_tm1, last_hidden, degree_input,
                                  eos_token_id):
        """
        Single beam search decoding. Batch size have to be 1.
        :param eos_token_id:
        :param degree_input:
        :param last_hidden:
        :param m_tm1:
        :param pz_dec_outs: [T,1,H]
        :param pz_proba: [T,1,V]
        :param u_enc_out: [T,1,H]
        :return:
        """
        eos_token_id = self.vocab.encode(cfg.eos_m_token)
        batch_size = pz_dec_outs.size(1)
        if batch_size != 1:
            raise ValueError('"Beam search single" requires batch size to be 1')

        class BeamState:
            def __init__(self, score, last_hidden, decoded, length):
                """
                Beam state in beam decoding
                :param score: sum of log-probabilities
                :param last_hidden: last hidden
                :param decoded: list of *Variable[1*1]* of all decoded words
                :param length: current decoded sentence length
                """
                self.score = score
                self.last_hidden = last_hidden
                self.decoded = decoded
                self.length = length

            def update_clone(self, score_incre, last_hidden, decoded_t):
                decoded = copy.copy(self.decoded)
                decoded.append(decoded_t)
                clone = BeamState(self.score + score_incre, last_hidden, decoded, self.length + 1)
                return clone

        def beam_result_valid(decoded_t):
            pz_max_samples = self.pz_selective_sampling(pz_proba)
            requested, start = [], False
            t = 0
            while t < len(pz_max_samples[0]) and pz_max_samples[0][t] != self.vocab.encode('EOS_Z1'):
                t += 1
            t += 1
            while t < len(pz_max_samples[0]) and pz_max_samples[0][t] != self.vocab.encode('EOS_Z2'):
                requested.append(self.vocab.decode(pz_max_samples[0][t]))
                t += 1
            decoded_t = [_.view(-1).data[0] for _ in decoded_t]
            decoded_sentence = self.vocab.sentence_decode(decoded_t, cfg.eos_m_token)

            requested = set(requested).intersection(['address', 'food', 'pricerange', 'phone', 'postcode'])
            # return True
            for rq in requested:
                if '%s SLOT' % rq not in decoded_sentence:
                    #print('Fail %s' % decoded_sentence)
                    return False
            #print('Success %s' % decoded_sentence)
            return True

        def score_bonus(state, decoded):
            """
            bonus scheme: bonus per token, or per new decoded slot.
            :param state:
            :return:
            """
            bonus = cfg.beam_len_bonus
            decoded = self.vocab.decode(decoded)
            decoded_t = [_.view(-1).data[0] for _ in state.decoded]
            decoded_sentence = self.vocab.sentence_decode(decoded_t, cfg.eos_m_token)
            decoded_sentence = decoded_sentence.split()
            if len(decoded_sentence) >= 1 and decoded_sentence[-1] == decoded: # repeated words
                bonus -= 10000
            if decoded == '**unknown**':
                bonus -= 3.0
            return bonus

        def soft_score_incre(score, turn):
            return score

        finished, failed = [], []
        states = []  # sorted by score decreasingly
        dead_k = 0
        states.append(BeamState(0, last_hidden, [m_tm1], 0))
        for t in range(self.max_ts):
            new_states = []
            k = 0
            while k < len(states) and k < self.beam_size - dead_k:
                state = states[k]
                last_hidden, m_tm1 = state.last_hidden, state.decoded[-1]
                proba, last_hidden = self.m_decoder(pz_dec_outs, pz_proba, u_enc_out, m_tm1, degree_input, last_hidden)
                proba = torch.log(proba)
                mt_proba, mt_index = torch.topk(proba, self.beam_size - dead_k)  # [1,K]
                for new_k in range(self.beam_size - dead_k):
                    score_incre = soft_score_incre(mt_proba[0][new_k].data[0], t) + score_bonus(state, mt_index[0][new_k].data[0])
                    if len(new_states) >= self.beam_size - dead_k and state.score + score_incre < new_states[-1].score:
                        break
                    decoded_t = mt_index[0][new_k]
                    if self.vocab.decode(decoded_t.data[0]) == cfg.eos_m_token:
                        if beam_result_valid(state.decoded):
                            finished.append(state)
                            dead_k += 1
                        else:
                            failed.append(state)
                    else:
                        decoded_t = decoded_t.view(1, -1)
                        new_state = state.update_clone(score_incre, last_hidden, decoded_t)
                        new_states.append(new_state)
                        #beam_result_valid(new_state.decoded)
                        #print(self.vocab.decode(decoded_t.view(-1).data[0]), t, new_k)
                k += 1
            if self.beam_size - dead_k < 0:
                break
            new_states = new_states[:self.beam_size - dead_k]
            new_states.sort(key=lambda x: -x.score)
            states = new_states

            if t == self.max_ts - 1 and not finished:
                finished = failed
                if not finished:
                    finished.append(states[0])

        finished.sort(key=lambda x: -x.score)
        decoded_t = finished[0].decoded
        decoded_t = [_.view(-1).data[0] for _ in decoded_t]
        decoded_sentence = self.vocab.sentence_decode(decoded_t, cfg.eos_m_token)
        print(decoded_sentence) 
        generated = torch.cat(finished[0].decoded, dim=1).data  # [B=1, T]
        return generated

    def beam_search_decode(self, pz_dec_outs, pz_proba, u_enc_out, m_tm1, last_hidden, degree_input, eos_token_id):
        vars = torch.split(pz_dec_outs, 1, dim=1), torch.split(pz_proba, 1, dim=1), torch.split(u_enc_out, 1,
                                                                                                dim=1), torch.split(
            m_tm1, 1, dim=1), torch.split(last_hidden, 1, dim=1), torch.split(degree_input, 1, dim=0)
        decoded = []
        for pz_dec_out_s, pz_proba_s, u_enc_out_s, m_tm1_s, last_hidden_s, degree_input_s in zip(*vars):
            decoded_s = self.beam_search_decode_single(pz_dec_out_s, pz_proba_s, u_enc_out_s, m_tm1_s, last_hidden_s,
                                                       degree_input_s, eos_token_id)
            decoded.append(decoded_s)
        return [list(_.view(-1)) for _ in decoded]

    def supervised_loss(self, pz_proba, qz_proba, pm_dec_proba, z_input, m_input):
        pr_loss = self.pr_loss(pz_proba.view(-1, pz_proba.size(2)), z_input.view(-1))
        m_loss = self.dec_loss(pm_dec_proba.view(-1, pm_dec_proba.size(2)), m_input.view(-1))
        q_loss = self.q_loss(qz_proba.view(-1, pz_proba.size(2)), z_input.view(-1))

        pr_loss, m_loss, q_loss = pr_loss, m_loss, q_loss

        if cfg.pretrain:
            loss = q_loss
        else:
            loss = pr_loss + m_loss + q_loss
        return loss, pr_loss, m_loss, q_loss

    def unsupervised_loss(self, mu_q, log_sigma_q, mu_p, log_sigma_p, pm_dec_proba, m_input, pz_proba, qz_proba):
        m_loss = self.dec_loss(pm_dec_proba.view(-1, pm_dec_proba.size(2)), m_input.view(-1))
        kl_div_loss = self.kl_loss(pz_proba, qz_proba.detach())

        m_loss, kl_div_loss = m_loss, kl_div_loss * self.alpha

        loss = m_loss + kl_div_loss
        return loss, m_loss, kl_div_loss

    def self_adjust(self, epoch):
        pass