def forward(self, input_data):
        q_emb_var, q_len, hs_emb_var, hs_len, col_emb_var, col_len, col_name_len, gt_col = input_data

        B = len(q_len)
        if self.use_bert:
            q_enc = self.q_bert(q_emb_var, q_len)
        else:
            q_enc, _ = run_lstm(self.q_lstm, q_emb_var, q_len)
        hs_enc, _ = run_lstm(self.hs_lstm, hs_emb_var, hs_len)
        col_enc, _ = col_tab_name_encode(col_emb_var, col_name_len, col_len,
                                         self.col_lstm)

        # get target/predicted column's embedding
        # col_emb: (B, hid_dim)
        col_emb = []
        for b in range(B):
            col_emb.append(col_enc[b, gt_col[b]])
        col_emb = torch.stack(col_emb)
        q_weighted = plain_conditional_weighted_num(self.q_att, q_enc, q_len,
                                                    col_emb)
        hs_weighted = plain_conditional_weighted_num(self.hs_att, hs_enc,
                                                     hs_len, col_emb)
        hv_score = self.hv_out(
            self.hv_out_q(q_weighted) +
            int(self.use_hs) * self.hv_out_hs(hs_weighted) +
            self.hv_out_c(col_emb))
        SIZE_CHECK(hv_score, [B, 2])

        return hv_score
    def forward(self, input_data):
        q_emb_var, q_len, hs_emb_var, hs_len, col_emb_var, col_len, col_name_len, gt_col = input_data

        B = len(q_len)

        if self.use_bert:
            q_enc = self.q_bert(q_emb_var, q_len)
        else:
            q_enc, _ = run_lstm(self.q_lstm, q_emb_var, q_len)
        hs_enc, _ = run_lstm(self.hs_lstm, hs_emb_var, hs_len)
        col_enc, _ = col_tab_name_encode(col_emb_var, col_name_len, col_len, self.col_lstm)

        # get target/predicted column's embedding
        # col_emb: (B, hid_dim)
        col_emb = []
        for b in range(B):
            col_emb.append(col_enc[b, gt_col[b]])
        col_emb = torch.stack(col_emb) # [B, dim]
        q_weighted = plain_conditional_weighted_num(self.q_att, q_enc, q_len, col_emb)

        # Same as the above, compute SQL history embedding weighted by column attentions
        hs_weighted = plain_conditional_weighted_num(self.hs_att, hs_enc, hs_len, col_emb)
        # dat_score: (B, 4)
        dat_score = self.dat_out(self.dat_out_q(q_weighted) + int(self.use_hs)* self.dat_out_hs(hs_weighted) + self.dat_out_c(col_emb))

        return dat_score
Example #3
0
    def forward(self, input_data):
        q_emb_var, q_len, hs_emb_var, hs_len, col_emb_var, col_len, col_name_len, gt_col = input_data

        B = len(q_len)
        if self.use_bert:
            q_enc = self.q_bert(q_emb_var, q_len)
        else:
            q_enc, _ = run_lstm(self.q_lstm, q_emb_var, q_len)
        hs_enc, _ = run_lstm(self.hs_lstm, hs_emb_var, hs_len)
        col_enc, _ = col_tab_name_encode(col_emb_var, col_name_len, col_len,
                                         self.col_lstm)

        # get target/predicted column's embedding
        # col_emb: (B, hid_dim)
        col_emb = []
        for b in range(B):
            col_emb.append(col_enc[b, gt_col[b]])
        col_emb = torch.stack(col_emb)

        # Predict op number
        q_weighted_num = plain_conditional_weighted_num(
            self.q_num_att, q_enc, q_len, col_emb)

        # Same as the above, compute SQL history embedding weighted by column attentions
        hs_weighted_num = plain_conditional_weighted_num(
            self.hs_num_att, hs_enc, hs_len, col_emb)
        # op_num_score: (B, 2)
        op_num_score = self.op_num_out(
            self.op_num_out_q(q_weighted_num) +
            int(self.use_hs) * self.op_num_out_hs(hs_weighted_num) +
            self.op_num_out_c(col_emb))
        SIZE_CHECK(op_num_score, [B, 2])

        # Compute attention values between selected column and question tokens.
        q_weighted = plain_conditional_weighted_num(self.q_att, q_enc, q_len,
                                                    col_emb)

        # Same as the above, compute SQL history embedding weighted by column attentions
        hs_weighted = plain_conditional_weighted_num(self.hs_att, hs_enc,
                                                     hs_len, col_emb)

        # Compute prediction scores
        # op_score: (B, 10)
        op_score = self.op_out(
            self.op_out_q(q_weighted) +
            int(self.use_hs) * self.op_out_hs(hs_weighted) +
            self.op_out_c(col_emb))
        SIZE_CHECK(op_score, [B, 11])

        score = (op_num_score, op_score)

        return score
Example #4
0
    def forward(self, input_data):
        q_emb_var, q_len, hs_emb_var, hs_len, col_emb_var, col_len, col_name_len, gt_col = input_data

        B = len(q_len)

        if self.use_bert:
            q_enc = self.q_bert(q_emb_var, q_len)
        else:
            q_enc, _ = run_lstm(self.q_lstm, q_emb_var, q_len)
        hs_enc, _ = run_lstm(self.hs_lstm, hs_emb_var, hs_len)
        col_enc, _ = col_tab_name_encode(col_emb_var, col_name_len, col_len,
                                         self.col_lstm)

        col_emb = []
        for b in range(B):
            col_emb.append(col_enc[b, gt_col[b]])
        col_emb = torch.stack(col_emb)

        # Predict agg number
        q_weighted_num = plain_conditional_weighted_num(
            self.q_num_att, q_enc, q_len, col_emb)

        # Same as the above, compute SQL history embedding weighted by column attentions
        hs_weighted_num = plain_conditional_weighted_num(
            self.hs_num_att, hs_enc, hs_len, col_emb)
        agg_num_score = self.agg_num_out(
            self.agg_num_out_q(q_weighted_num) +
            int(self.use_hs) * self.agg_num_out_hs(hs_weighted_num) +
            self.agg_num_out_c(col_emb))
        SIZE_CHECK(agg_num_score, [B, 4])

        # Predict aggregators
        q_weighted = plain_conditional_weighted_num(self.q_att, q_enc, q_len,
                                                    col_emb)

        # Same as the above, compute SQL history embedding weighted by column attentions
        hs_weighted = plain_conditional_weighted_num(self.hs_att, hs_enc,
                                                     hs_len, col_emb)
        # agg_score: (B, 5)
        agg_score = self.agg_out(
            self.agg_out_q(q_weighted) +
            int(self.use_hs) * self.agg_out_hs(hs_weighted) +
            self.agg_out_c(col_emb))

        score = (agg_num_score, agg_score)

        return score