def __init__(self, state_num, hidden_dim, know_vocab_size, embed_dim, embedder, lg_interpreter: LocGloInterpreter, gen_strategy, know2word_tensor, with_copy=True): super(PriorStateTracker, self).__init__() self.state_num = state_num self.hidden_dim = hidden_dim self.know_vocab_size = know_vocab_size self.embed_dim = embed_dim self.embedder = embedder self.lg_interpreter = lg_interpreter self.gen_strategy = gen_strategy self.know2word_tensor = know2word_tensor self.with_copy = with_copy self.embed_attn = Attention(embed_dim, hidden_dim) self.prior_basic_state_tracker = BasicStateTracker( self.know2word_tensor, self.state_num, self.hidden_dim, self.know_vocab_size, self.embed_dim, self.embedder, self.lg_interpreter, self.gen_strategy, self.with_copy) self.rnn_cell = nn.GRU(input_size=self.embed_dim, hidden_size=self.embed_dim, num_layers=1, batch_first=True, bidirectional=False)
class IntentionDetector(nn.Module): def __init__(self, know_vocab_size, intention_cate=4, hidden_dim=512, embed_dim=300, graph_dim=128, is_prior=True): super(IntentionDetector, self).__init__() self.know_vocab_size = know_vocab_size self.intention_cate = intention_cate self.hidden_dim = hidden_dim self.embed_dim = embed_dim self.is_prior = is_prior self.graph_dim = graph_dim if self.is_prior else 0 if not self.is_prior: self.res_attn = Attention(self.hidden_dim, self.hidden_dim) self.his_attn = Attention(self.hidden_dim, self.hidden_dim) self.state_attn = Attention(self.hidden_dim, self.embed_dim) intention_mlp_input_size = self.hidden_dim + self.hidden_dim + self.embed_dim if not self.is_prior: intention_mlp_input_size += self.hidden_dim if self.is_prior: # prior policy network, use the graph embedding intention_mlp_input_size += self.graph_dim self.intention_mlp = nn.Sequential(nn.Linear(intention_mlp_input_size, self.hidden_dim), nn.ReLU(), nn.Linear(self.hidden_dim, self.intention_cate), nn.Softmax()) def forward(self, state_emb, hidden, pv_r_u_enc, pv_r_u_len=None, r_enc=None, graph_context=None): intention_input = [] # r if not self.is_prior: r_inp, _ = self.res_attn.forward(hidden.permute(1, 0, 2), r_enc) intention_input.append(r_inp) # pv_r_u h_inp, _ = self.his_attn.forward(hidden.permute(1, 0, 2), pv_r_u_enc, mask=reverse_sequence_mask(pv_r_u_len, max_len=pv_r_u_enc.size(1))) intention_input.append(h_inp) # state s_inp, _ = self.state_attn.forward(hidden.permute(1, 0, 2), state_emb) intention_input.append(s_inp) # question_hidden intention_input.append(hidden.permute(1, 0, 2)) intention_input = torch.cat(intention_input, dim=-1).squeeze(1) # B, E + H + H [+ H] if graph_context is not None and self.is_prior: intention_input = torch.cat([intention_input, graph_context.squeeze(1)], dim=-1) # intention project intention = self.intention_mlp.forward(intention_input) # B, I return intention
def __init__(self, action_num, hidden_dim, know_vocab_size, embed_dim, embedder, lg_interpreter, gen_strategy, with_copy, know2word_tensor, know_encoder: RNNEncoder = None): super(PosteriorPolicyNetwork, self).__init__() self.action_num = action_num self.hidden_dim = hidden_dim self.know_vocab_size = know_vocab_size self.embed_dim = embed_dim self.embedder = embedder self.lg_interpreter = lg_interpreter self.gen_strategy = gen_strategy self.with_copy = with_copy self.know2word_tensor = know2word_tensor self.embed_attn = Attention(hidden_dim, embed_dim) self.hidden_attn = Attention(hidden_dim, hidden_dim) if know_encoder is not None: self.know_encoder = know_encoder else: self.know_encoder = RNNEncoder(self.embed_dim, self.hidden_dim, self.hidden_dim // 2, num_layers=1, bidirectional=True, dropout=0, embedder=self.embedder) self.embed2hidden_linear = nn.Linear(self.embed_dim, self.hidden_dim) self.intention_detector = IntentionDetector(know_vocab_size, hidden_dim=self.hidden_dim, embed_dim=self.embed_dim, is_prior=False) self.basic_policy_network = BasicPolicyNetwork( self.action_num, self.hidden_dim, self.embed_dim, self.know_vocab_size, self.embedder, self.lg_interpreter, self.know2word_tensor, self.gen_strategy, self.with_copy) self.rnn_cell = nn.GRU(input_size=self.embed_dim, hidden_size=self.embed_dim, num_layers=1, batch_first=True, bidirectional=False)
def __init__(self, action_num, hidden_dim, know_vocab_size, embed_dim, embedder, lg_interpreter, gen_strategy, with_copy, graph_db: GraphDB, know2word_tensor, intention_gumbel_softmax): super(PriorPolicyNetwork, self).__init__() self.action_num = action_num self.hidden_dim = hidden_dim self.know_vocab_size = know_vocab_size self.embed_dim = embed_dim self.embedder = embedder self.lg_interpreter = lg_interpreter self.gen_strategy = gen_strategy self.with_copy = with_copy self.graph_db = graph_db self.intention_gumbel_softmax = intention_gumbel_softmax self.know2word_tensor = know2word_tensor self.embed_attn = Attention(hidden_dim, embed_dim) self.hidden_attn = Attention(hidden_dim, hidden_dim) self.hidden_type_linear = nn.Linear(hidden_dim + 4, hidden_dim) self.embed2hidden_linear = nn.Linear(self.embed_dim, self.hidden_dim) self.intention_detector = IntentionDetector( know_vocab_size, hidden_dim=self.hidden_dim, embed_dim=self.embed_dim, graph_dim=VO.GATConfig.embed_dim, is_prior=True) self.basic_policy_network = BasicPolicyNetwork( self.action_num, self.hidden_dim, self.embed_dim, self.know_vocab_size, self.embedder, self.lg_interpreter, self.know2word_tensor, self.gen_strategy, self.with_copy) self.graph_attn = GraphAttn(VO.node_embed_dim, VO.hidden_dim) self.rnn_enc_cell = nn.GRU(input_size=self.embed_dim, hidden_size=self.embed_dim, num_layers=1, batch_first=True, bidirectional=False) self.r_gat = None gatc = VO.GATConfig self.r_gat = MyGAT(gatc.embed_dim, gatc.edge_embed_dim, gatc.flag_embed_dim, gatc.node_num, gatc.edge_num, gatc.flag_num)
def __init__(self, know2word_tensor, state_num, hidden_dim, know_vocab_size, embed_dim, embedder, lg_interpreter: LocGloInterpreter, gen_strategy="gru", with_copy=True): super(BasicStateTracker, self).__init__() self.know2word_tensor = know2word_tensor self.state_num = state_num self.hidden_dim = hidden_dim self.know_vocab_size = know_vocab_size self.embed_dim = embed_dim self.embedder = embedder self.lg_interpreter = lg_interpreter self.gen_strategy = gen_strategy self.with_copy = with_copy self.gumbel_softmax = GumbelSoftmax(normed=True) assert self.gen_strategy in ("gru", "mlp") or self.gen_strategy is None self.embed_attn = Attention(self.hidden_dim, self.embed_dim) self.hidden_attn = Attention(self.hidden_dim, self.hidden_dim) if self.with_copy: self.embed_copy_attn = Attention(self.hidden_dim, self.embed_dim) self.hidden_copy_attn = Attention(self.hidden_dim, self.hidden_dim) if self.gen_strategy == "gru": self.gru = nn.GRU(input_size=self.embed_dim + self.hidden_dim + self.embed_dim, hidden_size=self.hidden_dim, num_layers=1, dropout=0.0, batch_first=True) elif self.gen_strategy == "mlp": self.step_linear = nn.Sequential( nn.Linear(self.hidden_dim, self.state_num * self.hidden_dim), Reshape(1, [self.state_num, self.hidden_dim]), nn.Dropout(0.2), nn.Linear(self.hidden_dim, self.hidden_dim)) else: raise NotImplementedError self.hidden_projection = nn.Linear(self.hidden_dim, self.know_vocab_size) self.word_softmax = nn.Softmax(-1) self.hidden2word_projection = nn.Sequential(self.hidden_projection, self.word_softmax)
def __init__(self, know_vocab_size, intention_cate=4, hidden_dim=512, embed_dim=300, graph_dim=128, is_prior=True): super(IntentionDetector, self).__init__() self.know_vocab_size = know_vocab_size self.intention_cate = intention_cate self.hidden_dim = hidden_dim self.embed_dim = embed_dim self.is_prior = is_prior self.graph_dim = graph_dim if self.is_prior else 0 if not self.is_prior: self.res_attn = Attention(self.hidden_dim, self.hidden_dim) self.his_attn = Attention(self.hidden_dim, self.hidden_dim) self.state_attn = Attention(self.hidden_dim, self.embed_dim) intention_mlp_input_size = self.hidden_dim + self.hidden_dim + self.embed_dim if not self.is_prior: intention_mlp_input_size += self.hidden_dim if self.is_prior: # prior policy network, use the graph embedding intention_mlp_input_size += self.graph_dim self.intention_mlp = nn.Sequential(nn.Linear(intention_mlp_input_size, self.hidden_dim), nn.ReLU(), nn.Linear(self.hidden_dim, self.intention_cate), nn.Softmax())
class GraphAttn(nn.Module): def __init__(self, node_embed_dim, hidden_dim): super(GraphAttn, self).__init__() self.node_embed_dim = node_embed_dim self.hidden_dim = hidden_dim # attention read context self.attn = Attention(self.hidden_dim, self.node_embed_dim) def forward(self, node_embedding, node_efficient, head_flag_bit_matrix, h_c_t): efficient_mask = (1 - node_efficient) # B, N efficient_mask = efficient_mask | head_flag_bit_matrix node_context, _ = self.attn.forward(h_c_t, node_embedding, mask=efficient_mask.bool()) # B, 1, E return node_context
class GraphCopy(nn.Module): def __init__(self, node_embed_dim, hidden_dim): super(GraphCopy, self).__init__() self.node_embed_dim = node_embed_dim self.hidden_dim = hidden_dim self.copy_attn = Attention(self.hidden_dim, self.node_embed_dim) def forward(self, node_embedding, node_efficient, head_flag_bit_matrix, h_c_t): efficient_mask = (1 - node_efficient) # B, N efficient_mask = efficient_mask | head_flag_bit_matrix node_logits = self.copy_attn.forward(h_c_t, node_embedding, mask=efficient_mask.bool(), not_softmax=True, return_weight_only=True) # B, 1, N return node_logits
def __init__(self, node_embed_dim, hidden_dim): super(GraphAttn, self).__init__() self.node_embed_dim = node_embed_dim self.hidden_dim = hidden_dim # attention read context self.attn = Attention(self.hidden_dim, self.node_embed_dim)
def __init__(self, node_embed_dim, hidden_dim): super(GraphCopy, self).__init__() self.node_embed_dim = node_embed_dim self.hidden_dim = hidden_dim self.copy_attn = Attention(self.hidden_dim, self.node_embed_dim)
class PriorStateTracker(nn.Module): def __init__(self, state_num, hidden_dim, know_vocab_size, embed_dim, embedder, lg_interpreter: LocGloInterpreter, gen_strategy, know2word_tensor, with_copy=True): super(PriorStateTracker, self).__init__() self.state_num = state_num self.hidden_dim = hidden_dim self.know_vocab_size = know_vocab_size self.embed_dim = embed_dim self.embedder = embedder self.lg_interpreter = lg_interpreter self.gen_strategy = gen_strategy self.know2word_tensor = know2word_tensor self.with_copy = with_copy self.embed_attn = Attention(embed_dim, hidden_dim) self.prior_basic_state_tracker = BasicStateTracker( self.know2word_tensor, self.state_num, self.hidden_dim, self.know_vocab_size, self.embed_dim, self.embedder, self.lg_interpreter, self.gen_strategy, self.with_copy) self.rnn_cell = nn.GRU(input_size=self.embed_dim, hidden_size=self.embed_dim, num_layers=1, batch_first=True, bidirectional=False) def forward(self, hidden, pv_state, pv_r_u, pv_r_u_enc, gth_state=None, supervised=False): # B, S, E pv_state_emb = self.know_prob_embed(pv_state) tmp = [] s_hidden = None for i in range(pv_state_emb.size(1)): _, s_hidden = self.rnn_cell.forward(pv_state_emb[:, i:i + 1, :], s_hidden) # [1, B, H] tmp.append(s_hidden.permute(1, 0, 2)) # [B, 1, E] * S pv_state_emb = torch.cat(tmp, 1) # B, S, E # B, 1, E pv_state_emb_mean = pv_state_emb.mean(1).unsqueeze(1) # B, 1, T weight = self.embed_attn.forward( pv_state_emb_mean, pv_r_u_enc, mask=pv_r_u <= DO.PreventWord.RESERVED_MAX_INDEX, not_softmax=True, return_weight_only=True) # B, 1, T weight = torch.softmax(weight, -1) # B, 1, H ==> B, H hidden = hidden.squeeze(0) + torch.bmm(weight, pv_r_u_enc).squeeze(1) states, gumbel_states = self.prior_basic_state_tracker.forward( hidden, pv_state, pv_state_emb, pv_r_u, pv_r_u_enc, gth_state=gth_state, supervised=supervised) return states, gumbel_states def know_prob_embed(self, state_gumbel_prob): B, S, K = state_gumbel_prob.shape # K, E know_embedding = self.embedder(self.know2word_tensor) state_gumbel_embed = torch.bmm( state_gumbel_prob.reshape(B * S, 1, K), know_embedding.unsqueeze(0).expand(B * S, K, self.embed_dim)) state_gumbel_embed = state_gumbel_embed.reshape(B, S, self.embed_dim) return state_gumbel_embed
def __init__(self, embed_size, hidden_size, vocab_size, target_max_len, embedder=None, num_layers=1, dropout=0.0, attention_mode="mlp", max_pooling_k=10, attn_history_sentence=True, with_state_know=True, with_action_know=True, with_copy=True): super(RNNDecoder, self).__init__() self.embed_size = embed_size self.hidden_size = hidden_size self.vocab_size = vocab_size self.target_max_len = target_max_len self.embedder = embedder self.max_pooling_k = max_pooling_k self.num_layers = num_layers self.dropout = dropout if (self.num_layers > 1) else 0.0 self.attn_history_sentence = attn_history_sentence self.with_state_know = False self.with_action_know = False self.with_copy = with_copy and (with_state_know or with_action_know) self.rnn_input_size = self.embed_size # r_{t-1} input self.rnn_input_size += (self.hidden_size + self.embed_size + self.embed_size ) # history attn, state attn, action attn self.out_input_size = self.hidden_size if self.attn_history_sentence: # SENTENCE & WORD LEVEL ATTENTION self.history_attn_word = Attention(query_size=self.hidden_size, mode=attention_mode) if self.with_state_know: self.state_know_attn = Attention(query_size=self.hidden_size, mode=attention_mode) self.rnn_input_size += self.hidden_size if self.with_action_know: self.action_know_attn = Attention(query_size=self.hidden_size, mode=attention_mode) self.rnn_input_size += self.hidden_size if self.with_copy: # copy from knowledge self.copy_attn = Attention(query_size=self.hidden_size, mode="dot") self.sparse_copy_attn = Attention(query_size=self.hidden_size, memory_size=self.embed_size + 1, hidden_size=self.hidden_size) self.hm_linear = nn.Linear(self.hidden_size * 2, self.hidden_size) # select the most relative items self.k_max_pooling = KMaxPooling(self.max_pooling_k) # DECODER CELL self.rnn = nn.GRU(input_size=self.rnn_input_size, hidden_size=self.hidden_size, num_layers=self.num_layers, dropout=self.dropout, batch_first=True) self.init_s_attn = Attention(self.hidden_size, self.embed_size) self.init_a_attn = Attention(self.hidden_size, self.embed_size) self.s_linear = nn.Linear(self.embed_size, self.hidden_size) self.a_linear = nn.Linear(self.embed_size, self.hidden_size) if not self.with_copy: self.output_layer = nn.Sequential( nn.Dropout(p=self.dropout), nn.Linear(self.out_input_size, self.vocab_size), nn.LogSoftmax(dim=-1)) else: self.output_layer = nn.Sequential( nn.Dropout(p=self.dropout), nn.Linear(self.out_input_size, self.vocab_size)) self.softmax = nn.Softmax(dim=-1)
class RNNDecoder(nn.Module): def __init__(self, embed_size, hidden_size, vocab_size, target_max_len, embedder=None, num_layers=1, dropout=0.0, attention_mode="mlp", max_pooling_k=10, attn_history_sentence=True, with_state_know=True, with_action_know=True, with_copy=True): super(RNNDecoder, self).__init__() self.embed_size = embed_size self.hidden_size = hidden_size self.vocab_size = vocab_size self.target_max_len = target_max_len self.embedder = embedder self.max_pooling_k = max_pooling_k self.num_layers = num_layers self.dropout = dropout if (self.num_layers > 1) else 0.0 self.attn_history_sentence = attn_history_sentence self.with_state_know = False self.with_action_know = False self.with_copy = with_copy and (with_state_know or with_action_know) self.rnn_input_size = self.embed_size # r_{t-1} input self.rnn_input_size += (self.hidden_size + self.embed_size + self.embed_size ) # history attn, state attn, action attn self.out_input_size = self.hidden_size if self.attn_history_sentence: # SENTENCE & WORD LEVEL ATTENTION self.history_attn_word = Attention(query_size=self.hidden_size, mode=attention_mode) if self.with_state_know: self.state_know_attn = Attention(query_size=self.hidden_size, mode=attention_mode) self.rnn_input_size += self.hidden_size if self.with_action_know: self.action_know_attn = Attention(query_size=self.hidden_size, mode=attention_mode) self.rnn_input_size += self.hidden_size if self.with_copy: # copy from knowledge self.copy_attn = Attention(query_size=self.hidden_size, mode="dot") self.sparse_copy_attn = Attention(query_size=self.hidden_size, memory_size=self.embed_size + 1, hidden_size=self.hidden_size) self.hm_linear = nn.Linear(self.hidden_size * 2, self.hidden_size) # select the most relative items self.k_max_pooling = KMaxPooling(self.max_pooling_k) # DECODER CELL self.rnn = nn.GRU(input_size=self.rnn_input_size, hidden_size=self.hidden_size, num_layers=self.num_layers, dropout=self.dropout, batch_first=True) self.init_s_attn = Attention(self.hidden_size, self.embed_size) self.init_a_attn = Attention(self.hidden_size, self.embed_size) self.s_linear = nn.Linear(self.embed_size, self.hidden_size) self.a_linear = nn.Linear(self.embed_size, self.hidden_size) if not self.with_copy: self.output_layer = nn.Sequential( nn.Dropout(p=self.dropout), nn.Linear(self.out_input_size, self.vocab_size), nn.LogSoftmax(dim=-1)) else: self.output_layer = nn.Sequential( nn.Dropout(p=self.dropout), nn.Linear(self.out_input_size, self.vocab_size)) self.softmax = nn.Softmax(dim=-1) def decode(self, inp, hidden, hs_vectors, state_gumbel_prob, state_gumbel_embed, action_gumbel_prob, action_gumbel_embed, history_lens=None, history_word_indices=None, mask_action_prob=False, mask_state_prob=False, state_detach=False): if self.embedder is not None: inp = self.embedder(inp) # B, 1, E rnn_input = inp if self.attn_history_sentence: word_mask = None B, L_sent, H = hs_vectors.shape if history_lens is not None: # B, L_sent word_mask = reverse_sequence_mask(history_lens, max_len=L_sent).reshape( B, -1) # B, 1, H c_his_word, _ = self.history_attn_word.forward(hidden.transpose( 0, 1), hs_vectors, mask=word_mask) rnn_input = torch.cat([rnn_input, c_his_word], dim=-1) # attn state and action embed representation c_state = self.init_consult(hidden, state_gumbel_embed, self.init_s_attn) c_action = self.init_consult(hidden, action_gumbel_embed, self.init_a_attn) rnn_input = torch.cat([rnn_input, c_state, c_action], dim=-1) rnn_output, new_hidden = self.rnn(rnn_input, hidden) gene_output = self.output_layer(rnn_output) if not self.with_copy: return gene_output, new_hidden, None copy_query = new_hidden.permute(1, 0, 2) copy_logits = [] copy_word_index = [] # copy from post hs_logits = self.copy_attn.forward( copy_query, hs_vectors, not_softmax=True, return_weight_only=True) # B, 1, L_sentence copy_logits.append(hs_logits) copy_word_index.append(history_word_indices.unsqueeze(1)) s_tag = torch.zeros(list(state_gumbel_prob.shape)[:-1] + [1], device=state_gumbel_prob.device, dtype=torch.float) a_tag = torch.ones(list(action_gumbel_prob.shape)[:-1] + [1], device=action_gumbel_prob.device, dtype=torch.float) # COPY FROM PV_STATE & ACTION # B, 1, S state_copy_weight = self.sparse_copy( copy_query, torch.cat([state_gumbel_embed, s_tag], dim=-1)) # B, 1, A action_copy_weight = self.sparse_copy( copy_query, torch.cat([action_gumbel_embed, a_tag], dim=-1)) if mask_state_prob: state_copy_weight.masked_fill_( mask=torch.ones_like(state_copy_weight).bool(), value=-1e24) if mask_action_prob: action_copy_weight.masked_fill_( mask=torch.ones_like(action_copy_weight).bool(), value=-1e24) copy_logits.append(state_copy_weight) copy_logits.append(action_copy_weight) # B, V [+ (k * L_triple) [+ (k * L_triple)]] + S + A word_proba = self.softmax( torch.cat([gene_output] + copy_logits, dim=-1)).squeeze(1) # B, V gene_proba = word_proba[:, :self.vocab_size] state_gumbel_weight = word_proba[:, -( DO.state_num + DO.action_num):-DO.action_num].unsqueeze( 1) # B, 1, S action_gumbel_weight = word_proba[:, -DO.action_num:].unsqueeze( 1) # B, 1, A # B, V state_copy_prob = torch.bmm(state_gumbel_weight, state_gumbel_prob).squeeze(1) if state_detach: state_copy_prob = state_copy_prob.detach() action_copy_prob = torch.bmm(action_gumbel_weight, action_gumbel_prob).squeeze(1) l_s = vrbot_train_stage.s_copy_lambda l_a = vrbot_train_stage.a_copy_lambda if VO.train_stage == "action" else vrbot_train_stage.a_copy_lambda_mini gene_proba = gene_proba + l_s * state_copy_prob + l_a * action_copy_prob if len(copy_word_index) == 0: output = gene_proba else: if len(copy_word_index) > 1: copy_word_index = torch.cat(copy_word_index, dim=-1) else: copy_word_index = copy_word_index[0] # B, V B = word_proba.size(0) copy_proba = torch.zeros(B, self.vocab_size, device=word_proba.device) copy_proba = copy_proba.scatter_add( 1, copy_word_index.squeeze(1), word_proba[:, self.vocab_size:-(DO.state_num + DO.action_num)]) output = gene_proba + copy_proba output = torch.log(output.unsqueeze(1)) return output, new_hidden def copy(self, query, memory_word_enc, memory_word_index): B, t_num, l_trip, H = memory_word_enc.shape flatten_memory_word_enc = memory_word_enc.reshape(B, t_num * l_trip, H) flatten_selected_word_index = memory_word_index.reshape(B, -1) flatten_selected_mask = flatten_selected_word_index <= DO.PreventWord.RESERVED_MAX_INDEX match_logits = self.copy_attn.forward(query, flatten_memory_word_enc, not_softmax=True)[1].squeeze(1) match_logits = match_logits.masked_fill(flatten_selected_mask, -1e24) return match_logits, flatten_selected_word_index def sparse_copy(self, query, word_enc): # B, 1, S copy_weight = self.sparse_copy_attn.forward(query, word_enc, not_softmax=True, return_weight_only=True) return copy_weight @staticmethod def init_consult(hidden, state_emb, attn, state_linear=None): state_attn_emb, _ = attn.forward(hidden.permute(1, 0, 2), state_emb) # B, 1, E if state_linear is not None: state_attn_emb = state_linear(state_attn_emb) # B, 1, H return state_attn_emb.permute(1, 0, 2) # 1, B, H else: return state_attn_emb def forward(self, hidden, inputs, hs_vectors, state_gumbel_prob, state_gumbel_embed, action_gumbel_prob, action_gumbel_embed, history_lens=None, history_word_indices=None, mask_action_prob=False, mask_state_prob=False, state_detach=False): if self.training: assert inputs is not None, "In training stage, inputs should not be None" batch_size, s_max_len = inputs.shape else: batch_size = hidden.size(1) rnn_input = hidden.new_ones( batch_size, 1, dtype=torch.long) * DO.PreventWord.SOS_ID # SOS # consult the state and action hidden = self.init_consult_state_action(action_gumbel_embed, hidden, mask_action_prob, state_gumbel_embed) if self.training: valid_lengths = s_max_len - ( inputs == DO.PreventWord.EOS_ID).long().cumsum(1).sum(1) valid_lengths, indices = valid_lengths.sort(descending=True) lengths_tag = torch.arange( 0, s_max_len, device=TrainOption.device).unsqueeze(0).expand(batch_size, -1) # max_len batch_num_valid = (valid_lengths.unsqueeze(1).expand( -1, s_max_len) > lengths_tag).long().sum(0) batch_num_valid = batch_num_valid.cpu().numpy().tolist() # B, max_len, vocab_size output_placeholder = torch.zeros(batch_size, self.target_max_len, self.vocab_size, dtype=torch.float, device=TrainOption.device) # input Tensor index_select hidden = index_select_if_not_none(hidden, indices, 1) inputs = index_select_if_not_none(inputs, indices, 0) hs_vectors = index_select_if_not_none(hs_vectors, indices, 0) state_gumbel_embed = index_select_if_not_none( state_gumbel_embed, indices, 0) state_gumbel_prob = index_select_if_not_none( state_gumbel_prob, indices, 0) action_gumbel_embed = index_select_if_not_none( action_gumbel_embed, indices, 0) action_gumbel_prob = index_select_if_not_none( action_gumbel_prob, indices, 0) history_word_indices = index_select_if_not_none( history_word_indices, indices, 0) history_lens = index_select_if_not_none(history_lens, indices, 0) for dec_step, vb in enumerate(batch_num_valid): if vb <= 0: break if inputs is not None and self.training: # teacher forcing in training stage # B, 1 rnn_input = inputs[:, dec_step].unsqueeze(1) rnn_dec = self.decode( rnn_input[:vb, :], hidden[:, :vb, :], hs_vectors[:vb, ...] if hs_vectors is not None else None, state_gumbel_prob[:vb, ...] if state_gumbel_prob is not None else None, state_gumbel_embed[:vb, ...] if state_gumbel_embed is not None else None, action_gumbel_prob[:vb, ...] if action_gumbel_prob is not None else None, action_gumbel_embed[:vb, ...] if action_gumbel_embed is not None else None, history_lens[:vb] if history_lens is not None else None, history_word_indices[:vb] if history_word_indices is not None else None, mask_action_prob, mask_state_prob, state_detach) word_output, hidden = rnn_dec # B, max_len, vocab_size output_placeholder[:vb, dec_step:dec_step + 1, :] = word_output if not (inputs is not None and self.training): # inference in testing or valid stage rnn_input = word_output.argmax(dim=-1) _, rev_indices = indices.sort() output_placeholder = output_placeholder.index_select( 0, rev_indices) return output_placeholder else: global_indices = torch.arange(0, batch_size, dtype=torch.long, device=TrainOption.device) output_placeholder = torch.zeros(batch_size, self.target_max_len, self.vocab_size, dtype=torch.float, device=TrainOption.device) # decode for i in range(self.target_max_len): word_output, hidden = self.decode(rnn_input, hidden, hs_vectors, history_lens) # B, 1 next_step_input = word_output.argmax(dim=-1) # B, continue_tag = (next_step_input.squeeze(1) != DO.PreventWord.EOS_ID).long() _, local_indices = continue_tag.sort(descending=True) output_placeholder[global_indices, i:i + 1, :] = word_output b_case = continue_tag.sum().item() if b_case <= 0: break # B', local_indices = local_indices[:b_case] global_indices = global_indices.index_select(0, local_indices) rnn_input = next_step_input # input next time rnn_input = rnn_input.index_select(0, local_indices) hidden = hidden.index_select(1, local_indices) # history_info & memory hs_vectors = index_select_if_not_none(hs_vectors, local_indices, 0) state_know_key = index_select_if_not_none( state_know_key, local_indices, 0) state_know_value = index_select_if_not_none( state_know_value, local_indices, 0) state_know_word_enc = index_select_if_not_none( state_know_word_enc, local_indices, 0) state_know_word = index_select_if_not_none( state_know_word, local_indices, 0) state_know_len = index_select_if_not_none( state_know_len, local_indices, 0) action_know_key = index_select_if_not_none( action_know_key, local_indices, 0) action_know_value = index_select_if_not_none( action_know_value, local_indices, 0) action_know_word_enc = index_select_if_not_none( action_know_word_enc, local_indices, 0) action_know_word = index_select_if_not_none( action_know_word, local_indices, 0) action_know_len = index_select_if_not_none( action_know_len, local_indices, 0) history_lens = index_select_if_not_none( history_lens, local_indices, 0) return output_placeholder def init_consult_state_action(self, action_gumbel_embed, hidden, mask_action_prob, state_gumbel_embed): if mask_action_prob: s_emb = self.init_consult(hidden, state_gumbel_embed, self.init_s_attn, self.s_linear) hidden = hidden + s_emb else: s_emb = self.init_consult(hidden, state_gumbel_embed, self.init_s_attn, self.s_linear) a_emb = self.init_consult(hidden, action_gumbel_embed, self.init_a_attn, self.a_linear) hidden = hidden + (s_emb + a_emb) / 2.0 return hidden
def __init__(self, action_num, hidden_dim, embed_dim, know_vocab_size, embedder, kw_interpreter: LocGloInterpreter, know2word_tensor, gen_strategy=None, with_copy=True): super(BasicPolicyNetwork, self).__init__() self.action_num = action_num self.hidden_dim = hidden_dim self.embed_dim = embed_dim self.know_vocab_size = know_vocab_size self.embedder = embedder self.kw_interpreter = kw_interpreter self.gen_strategy = gen_strategy self.with_copy = with_copy self.know2word_tensor = know2word_tensor self.gumbel_softmax = GumbelSoftmax(normed=True) self.graph_copy = GraphCopy(VO.node_embed_dim, VO.hidden_dim) assert self.gen_strategy in ("gru", "mlp") if self.with_copy: self.embed_copy_attn = Attention(self.hidden_dim, self.hidden_dim) self.hidden_copy_attn = Attention(self.hidden_dim, self.hidden_dim) if self.gen_strategy == "gru": self.embed_attn = Attention(self.hidden_dim, self.embed_dim) self.hidden_attn = Attention(self.hidden_dim, self.hidden_dim) self.embed2hidden_linear = nn.Linear(self.embed_dim, self.hidden_dim) self.gru = nn.GRU(input_size=self.embed_dim + self.embed_dim + self.hidden_dim, hidden_size=self.hidden_dim, num_layers=1, dropout=0.0, batch_first=True) elif self.gen_strategy == "mlp": self.eh_linear = nn.Sequential( nn.Linear(self.hidden_dim + self.embed_dim, self.hidden_dim), nn.ReLU(), nn.Linear(self.hidden_dim, self.hidden_dim)) self.trans_attn = Attention( self.hidden_dim, self.hidden_dim) # for interaction between state self.action_attn = Attention( self.hidden_dim, self.hidden_dim) # for attentively read state self.pred_proj = nn.Sequential( nn.Linear(self.hidden_dim + self.hidden_dim + self.embed_dim, self.hidden_dim * self.action_num), Reshape(1, [self.action_num, self.hidden_dim])) else: raise NotImplementedError self.gen_linear = nn.Linear(self.hidden_dim, self.know_vocab_size) self.hidden_projection = nn.Linear(self.hidden_dim, self.know_vocab_size) self.word_softmax = nn.Softmax(-1) self.hidden2word_projection = nn.Sequential(self.hidden_projection, self.word_softmax)
class BasicPolicyNetwork(nn.Module): def __init__(self, action_num, hidden_dim, embed_dim, know_vocab_size, embedder, kw_interpreter: LocGloInterpreter, know2word_tensor, gen_strategy=None, with_copy=True): super(BasicPolicyNetwork, self).__init__() self.action_num = action_num self.hidden_dim = hidden_dim self.embed_dim = embed_dim self.know_vocab_size = know_vocab_size self.embedder = embedder self.kw_interpreter = kw_interpreter self.gen_strategy = gen_strategy self.with_copy = with_copy self.know2word_tensor = know2word_tensor self.gumbel_softmax = GumbelSoftmax(normed=True) self.graph_copy = GraphCopy(VO.node_embed_dim, VO.hidden_dim) assert self.gen_strategy in ("gru", "mlp") if self.with_copy: self.embed_copy_attn = Attention(self.hidden_dim, self.hidden_dim) self.hidden_copy_attn = Attention(self.hidden_dim, self.hidden_dim) if self.gen_strategy == "gru": self.embed_attn = Attention(self.hidden_dim, self.embed_dim) self.hidden_attn = Attention(self.hidden_dim, self.hidden_dim) self.embed2hidden_linear = nn.Linear(self.embed_dim, self.hidden_dim) self.gru = nn.GRU(input_size=self.embed_dim + self.embed_dim + self.hidden_dim, hidden_size=self.hidden_dim, num_layers=1, dropout=0.0, batch_first=True) elif self.gen_strategy == "mlp": self.eh_linear = nn.Sequential( nn.Linear(self.hidden_dim + self.embed_dim, self.hidden_dim), nn.ReLU(), nn.Linear(self.hidden_dim, self.hidden_dim)) self.trans_attn = Attention( self.hidden_dim, self.hidden_dim) # for interaction between state self.action_attn = Attention( self.hidden_dim, self.hidden_dim) # for attentively read state self.pred_proj = nn.Sequential( nn.Linear(self.hidden_dim + self.hidden_dim + self.embed_dim, self.hidden_dim * self.action_num), Reshape(1, [self.action_num, self.hidden_dim])) else: raise NotImplementedError self.gen_linear = nn.Linear(self.hidden_dim, self.know_vocab_size) self.hidden_projection = nn.Linear(self.hidden_dim, self.know_vocab_size) self.word_softmax = nn.Softmax(-1) self.hidden2word_projection = nn.Sequential(self.hidden_projection, self.word_softmax) def forward(self, hidden, state_emb, state, pv_r_u_enc, pv_r_u_len, r=None, r_enc=None, r_mask=None, mask_gen=False, gth_action=None, supervised=False, node_embedding=None, head_nodes=None, node_efficient=None, head_flag_bit=None): if self.gen_strategy == "gru": return self.forward_gru( hidden, state_emb, pv_r_u_enc, pv_r_u_len, r=r, r_enc=r_enc, r_mask=r_mask, gth_action=gth_action, mask_gen=mask_gen, supervised=supervised, node_embedding=node_embedding, # B, N, E head_nodes=head_nodes, # B, N node_efficient=node_efficient, # B, N head_flag_bit=head_flag_bit) elif self.gen_strategy == "mlp": return self.forward_mlp(hidden, state_emb, r=r, r_enc=r_enc, r_mask=r_mask, mask_gen=mask_gen), None else: raise NotImplementedError def forward_gru( self, hidden, state_emb, pv_r_u_enc, pv_r_u_len, r=None, r_enc=None, r_mask=None, gth_action=None, mask_gen=False, supervised=False, node_embedding=None, # B, N, E head_nodes=None, # B, N node_efficient=None, # B, N head_flag_bit=None): batch_size = hidden.size(0) # B, 1, E state_context, _ = self.embed_attn.forward(hidden.unsqueeze(1), state_emb) hidden = hidden + self.embed2hidden_linear(state_context).squeeze(1) hidden = hidden.unsqueeze(0) # 1, B, H # init input step_input = torch.zeros(batch_size, 1, self.know_vocab_size, dtype=torch.float, device=TO.device) step_input[:, :, 0] = 1.0 # B, 1, K step_input_embed = self.know_prob_embed(step_input) # B, 1, E actions = [] gumbel_actions = [] for i in range(self.action_num): # B, 1, E + H state_context, _ = self.embed_attn.forward(hidden.permute(1, 0, 2), state_emb) pv_r_u_mask = reverse_sequence_mask(pv_r_u_len, pv_r_u_enc.size(1)) post_context, _ = self.hidden_attn.forward(hidden.permute(1, 0, 2), pv_r_u_enc, mask=pv_r_u_mask) pv_s_input = torch.cat( [step_input_embed, state_context, post_context], dim=-1) next_action_hidden, hidden = self.gru.forward(pv_s_input, hidden) probs = self.action_pred(batch_size, next_action_hidden, r=r, r_enc=r_enc, r_mask=r_mask, mask_gen=mask_gen, node_embedding=node_embedding, head_nodes=head_nodes, node_efficient=node_efficient, head_flag_bit=head_flag_bit) actions.append(probs) if self.training and TO.auto_regressive and ( gth_action is not None) and supervised and (not TO.no_action_super): gth_step_input = gth_action[:, i:i + 1] gth_step_input = one_hot_scatter(gth_step_input, self.know_vocab_size, dtype=torch.float) step_input_embed = self.know_prob_embed(gth_step_input) else: if TO.auto_regressive: if self.training: gumbel_probs = self.action_gumbel_softmax_sampling( probs) else: max_indices = probs.argmax(-1) gumbel_probs = one_hot_scatter(max_indices, probs.size(2), dtype=torch.float) step_input_embed = self.know_prob_embed(gumbel_probs) gumbel_actions.append(gumbel_probs) else: step_input_embed = self.know_prob_embed(probs) actions = torch.cat(actions, dim=1) if len(gumbel_actions) == 0: return actions, None gumbel_actions = torch.cat(gumbel_actions, dim=1) return actions, gumbel_actions def forward_mlp(self, hidden, state_emb, r=None, r_enc=None, r_mask=None, mask_gen=False): batch_size, state_num = state_emb.size(0), state_emb.size(1) # B, S, H expanded_hidden = hidden.unsqueeze(0).permute(1, 0, 2).expand( batch_size, state_num, self.hidden_dim) # B, S, H deep_input = self.eh_linear( torch.cat([expanded_hidden, state_emb], dim=2)) # B, S, H deep_inner, _ = self.trans_attn.forward(deep_input, deep_input) # B, 1, H deep_output, _ = self.action_attn.forward(hidden.unsqueeze(1), deep_inner) deep_output = deep_output.squeeze(1) # B, H + H + E proj_input = torch.cat([deep_output, hidden, state_emb.sum(1)], dim=-1) # B, A, H (for generation or for copy from ) pred_action_hidden = self.pred_proj.forward(proj_input) return self.action_pred(batch_size, pred_action_hidden, r, r_enc, r_mask, mask_gen=mask_gen) def action_pred( self, batch_size, pred_action_hidden, r=None, r_enc=None, r_mask=None, mask_gen=False, node_embedding=None, # B, N, E head_nodes=None, # B, N node_efficient=None, # B, N head_flag_bit=None): # B, N logits = [] indexs = None action_num = pred_action_hidden.size(1) action_gen_logits = self.gen_linear(pred_action_hidden) logits.append(action_gen_logits) if r is not None: r_know_index = self.kw_interpreter.glo2loc(r) r_inner_mask = (r_know_index == 0) if r_mask is not None: r_mask = r_inner_mask | r_mask else: r_mask = r_inner_mask r_logits = self.hidden_copy_attn.forward(pred_action_hidden, r_enc, mask=r_mask, not_softmax=True, return_weight_only=True) logits.append(r_logits) indexs = r_know_index # B, Tr if node_embedding is not None: node_copy_logits = self.graph_copy.forward(node_embedding, node_efficient, head_flag_bit, pred_action_hidden) logits.append(node_copy_logits) indexs = head_nodes # B, N if (r is None) or (not VO.ppn_dq): if len(logits) >= 1: logits = torch.cat(logits, -1) else: raise RuntimeError probs = self.word_softmax(logits) if indexs is None: return probs if not mask_gen: gen_probs = probs[:, :, :self.know_vocab_size] copy_probs = probs[:, :, self.know_vocab_size:] else: gen_probs = 0.0 copy_probs = probs else: g_lambda = 0.05 gen_probs = self.word_softmax(logits[0]) * (1 - g_lambda) copy_probs = self.word_softmax(logits[1]) * g_lambda copy_probs_placeholder = torch.zeros(batch_size, action_num, self.know_vocab_size, device=TO.device) if indexs is not None: expand_indexs = indexs.unsqueeze(1).expand(-1, action_num, -1) copy_probs = copy_probs_placeholder.scatter_add( 2, expand_indexs, copy_probs) # B, A, K action_probs = gen_probs + copy_probs # B, A, K return action_probs def action_gumbel_softmax_sampling(self, probs): gumbel_probs = self.gumbel_softmax.forward(probs, vrbot_train_stage.a_tau) return gumbel_probs def know_prob_embed(self, action_prob): B, A, K = action_prob.shape # K, E know_embedding = self.embedder(self.know2word_tensor) action_embed = torch.bmm( action_prob.reshape(B * A, 1, K), know_embedding.unsqueeze(0).expand(B * A, K, self.embed_dim)) action_embed = action_embed.reshape(B, A, self.embed_dim) return action_embed
class BasicStateTracker(nn.Module): def __init__(self, know2word_tensor, state_num, hidden_dim, know_vocab_size, embed_dim, embedder, lg_interpreter: LocGloInterpreter, gen_strategy="gru", with_copy=True): super(BasicStateTracker, self).__init__() self.know2word_tensor = know2word_tensor self.state_num = state_num self.hidden_dim = hidden_dim self.know_vocab_size = know_vocab_size self.embed_dim = embed_dim self.embedder = embedder self.lg_interpreter = lg_interpreter self.gen_strategy = gen_strategy self.with_copy = with_copy self.gumbel_softmax = GumbelSoftmax(normed=True) assert self.gen_strategy in ("gru", "mlp") or self.gen_strategy is None self.embed_attn = Attention(self.hidden_dim, self.embed_dim) self.hidden_attn = Attention(self.hidden_dim, self.hidden_dim) if self.with_copy: self.embed_copy_attn = Attention(self.hidden_dim, self.embed_dim) self.hidden_copy_attn = Attention(self.hidden_dim, self.hidden_dim) if self.gen_strategy == "gru": self.gru = nn.GRU(input_size=self.embed_dim + self.hidden_dim + self.embed_dim, hidden_size=self.hidden_dim, num_layers=1, dropout=0.0, batch_first=True) elif self.gen_strategy == "mlp": self.step_linear = nn.Sequential( nn.Linear(self.hidden_dim, self.state_num * self.hidden_dim), Reshape(1, [self.state_num, self.hidden_dim]), nn.Dropout(0.2), nn.Linear(self.hidden_dim, self.hidden_dim)) else: raise NotImplementedError self.hidden_projection = nn.Linear(self.hidden_dim, self.know_vocab_size) self.word_softmax = nn.Softmax(-1) self.hidden2word_projection = nn.Sequential(self.hidden_projection, self.word_softmax) def state_gumbel_softmax_sampling(self, probs): gumbel_probs = self.gumbel_softmax.forward(probs, vrbot_train_stage.s_tau) return gumbel_probs def forward(self, hidden, pv_state, pv_state_emb, pv_r_u, pv_r_u_enc, gth_state=None, supervised=False): batch_size = pv_state_emb.size(0) states = [] gumbel_states = [] multi_hiddens = None step_input_embed = None if self.gen_strategy == "gru": hidden = hidden.unsqueeze(0) # 1, B, H step_input = torch.zeros(batch_size, 1, self.know_vocab_size, dtype=torch.float, device=TO.device) step_input[:, :, 0] = 1.0 # B, 1, K step_input_embed = self.know_prob_embed(step_input) # B, 1, E elif self.gen_strategy == "mlp": multi_hiddens = self.step_linear(hidden) # B, S, H else: raise NotImplementedError for i in range(self.state_num): if self.gen_strategy == "gru": pv_s_context, _ = self.embed_attn.forward( hidden.permute(1, 0, 2), pv_state_emb) pv_r_u_context, _ = self.hidden_attn.forward( hidden.permute(1, 0, 2), pv_r_u_enc) pv_s_input = torch.cat( [pv_s_context, pv_r_u_context, step_input_embed], dim=-1) # B, 1, E + H + E next_state_hidden, hidden = self.gru.forward( pv_s_input, hidden) # B, 1, H | 1, B, H elif self.gen_strategy == "mlp": next_state_hidden = multi_hiddens[:, i:i + 1, :] # B, 1, H else: raise NotImplementedError next_state = self.hidden_projection(next_state_hidden) logits = [next_state] indexs = [] if self.with_copy: pv_state_weight = self.embed_copy_attn.forward( next_state_hidden, pv_state_emb, mask=None, not_softmax=True, return_weight_only=True) logits.append(pv_state_weight) pv_r_u_know_index = self.lg_interpreter.glo2loc(pv_r_u) pv_r_u_mask = (pv_r_u_know_index == 0) pv_r_u_weight = self.hidden_copy_attn.forward( next_state_hidden, pv_r_u_enc, mask=pv_r_u_mask, not_softmax=True, return_weight_only=True) logits.append(pv_r_u_weight) indexs.append(pv_r_u_know_index) logits = torch.cat(logits, -1) indexs = torch.cat(indexs, -1).unsqueeze(1) probs = self.word_softmax(logits) if self.with_copy: gen_probs = probs[:, :, :self.know_vocab_size] pv_state_copy_probs = probs[:, :, self.know_vocab_size:self. know_vocab_size + DO.state_num] pv_state_copy_probs = torch.bmm(pv_state_copy_probs, pv_state) copy_probs = probs[:, :, self.know_vocab_size + DO.state_num:] copy_probs_placeholder = torch.zeros(batch_size, 1, self.know_vocab_size, device=TO.device) copy_probs = copy_probs_placeholder.scatter_add( 2, indexs, copy_probs) probs = gen_probs + pv_state_copy_probs + copy_probs states.append(probs) if self.training and TO.auto_regressive and ( gth_state is not None) and supervised: gth_step_input = gth_state[:, i:i + 1] # B, 1 gth_step_input = one_hot_scatter(gth_step_input, self.know_vocab_size, dtype=torch.float) step_input_embed = self.know_prob_embed(gth_step_input) else: if TO.auto_regressive: if self.training: gumbel_probs = self.state_gumbel_softmax_sampling( probs) # gumbel-softmax else: max_indices = probs.argmax(-1) # B, S gumbel_probs = one_hot_scatter(max_indices, probs.size(2), dtype=torch.float) step_input_embed = self.know_prob_embed(gumbel_probs) gumbel_states.append(gumbel_probs) else: step_input_embed = self.know_prob_embed(probs) # B, 1, E states = torch.cat(states, dim=1) if len(gumbel_states) == 0: return states, None gumbel_states = torch.cat(gumbel_states, dim=1) return states, gumbel_states def know_prob_embed(self, state_prob): B, S, K = state_prob.shape know_embedding = self.embedder(self.know2word_tensor) state_embed = torch.bmm( state_prob.reshape(B * S, 1, K), know_embedding.unsqueeze(0).expand(B * S, K, self.embed_dim)) state_embed = state_embed.reshape(B, S, self.embed_dim) return state_embed