def forward(self, q_emb_var, q_len, hs_emb_var, hs_len, col_emb_var, col_len, col_name_len, gt_col):
max_q_len = max(q_len)
max_hs_len = max(hs_len)
max_col_len = max(col_len)
B = len(q_len)
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_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]
# self.q_att(q_enc).transpose(1, 2): [B, dim, max_q_len]
att_val_qc = torch.bmm(col_emb.unsqueeze(1), self.q_att(q_enc).transpose(1, 2)).view(B, -1)
for idx, num in enumerate(q_len):
if num < max_q_len:
att_val_qc[idx, num:] = -100
att_prob_qc = self.softmax(att_val_qc)
q_weighted = (q_enc * att_prob_qc.unsqueeze(2)).sum(1)
# Same as the above, compute SQL history embedding weighted by column attentions
att_val_hc = torch.bmm(col_emb.unsqueeze(1), self.hs_att(hs_enc).transpose(1, 2)).view(B, -1)
for idx, num in enumerate(hs_len):
if num < max_hs_len:
att_val_hc[idx, num:] = -100
att_prob_hc = self.softmax(att_val_hc)
hs_weighted = (hs_enc * att_prob_hc.unsqueeze(2)).sum(1)
# 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
def forward(self, q_emb_var, q_len, hs_emb_var, hs_len, col_emb_var, col_len, col_name_len, gt_col):
max_q_len = max(q_len)
max_hs_len = max(hs_len)
max_col_len = max(col_len)
B = len(q_len)
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_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
att_val_qc_num = torch.bmm(col_emb.unsqueeze(1), self.q_num_att(q_enc).transpose(1, 2)).view(B, -1)
for idx, num in enumerate(q_len):
if num < max_q_len:
att_val_qc_num[idx, num:] = -100
att_prob_qc_num = self.softmax(att_val_qc_num)
q_weighted_num = (q_enc * att_prob_qc_num.unsqueeze(2)).sum(1)
# Same as the above, compute SQL history embedding weighted by column attentions
att_val_hc_num = torch.bmm(col_emb.unsqueeze(1), self.hs_num_att(hs_enc).transpose(1, 2)).view(B, -1)
for idx, num in enumerate(hs_len):
if num < max_hs_len:
att_val_hc_num[idx, num:] = -100
att_prob_hc_num = self.softmax(att_val_hc_num)
hs_weighted_num = (hs_enc * att_prob_hc_num.unsqueeze(2)).sum(1)
# agg_num_score: (B, 4)
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))
# Predict aggregators
att_val_qc = torch.bmm(col_emb.unsqueeze(1), self.q_att(q_enc).transpose(1, 2)).view(B, -1)
for idx, num in enumerate(q_len):
if num < max_q_len:
att_val_qc[idx, num:] = -100
att_prob_qc = self.softmax(att_val_qc)
q_weighted = (q_enc * att_prob_qc.unsqueeze(2)).sum(1)
# Same as the above, compute SQL history embedding weighted by column attentions
att_val_hc = torch.bmm(col_emb.unsqueeze(1), self.hs_att(hs_enc).transpose(1, 2)).view(B, -1)
for idx, num in enumerate(hs_len):
if num < max_hs_len:
att_val_hc[idx, num:] = -100
att_prob_hc = self.softmax(att_val_hc)
hs_weighted = (hs_enc * att_prob_hc.unsqueeze(2)).sum(1)
# 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
def forward(self, q_emb_var, q_len, hs_emb_var, hs_len, col_emb_var,
col_len, col_name_len):
max_q_len = max(q_len)
max_hs_len = max(hs_len)
max_col_len = max(col_len)
B = len(q_len)
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_name_encode(col_emb_var, col_name_len, col_len,
self.col_lstm)
# Predict column number: 1-3
# att_val_qc_num: (B, max_col_len, max_q_len)
att_val_qc_num = torch.bmm(col_enc,
self.q_num_att(q_enc).transpose(1, 2))
for idx, num in enumerate(col_len):
if num < max_col_len:
att_val_qc_num[idx, num:, :] = -100
for idx, num in enumerate(q_len):
if num < max_q_len:
att_val_qc_num[idx, :, num:] = -100
att_prob_qc_num = self.softmax(att_val_qc_num.view(
(-1, max_q_len))).view(B, -1, max_q_len)
# q_weighted_num: (B, hid_dim)
q_weighted_num = (q_enc.unsqueeze(1) *
att_prob_qc_num.unsqueeze(3)).sum(2).sum(1)
# Same as the above, compute SQL history embedding weighted by column attentions
# att_val_hc_num: (B, max_col_len, max_hs_len)
att_val_hc_num = torch.bmm(col_enc,
self.hs_num_att(hs_enc).transpose(1, 2))
for idx, num in enumerate(hs_len):
if num < max_hs_len:
att_val_hc_num[idx, :, num:] = -100
for idx, num in enumerate(col_len):
if num < max_col_len:
att_val_hc_num[idx, num:, :] = -100
att_prob_hc_num = self.softmax(att_val_hc_num.view(
(-1, max_hs_len))).view(B, -1, max_hs_len)
hs_weighted_num = (hs_enc.unsqueeze(1) *
att_prob_hc_num.unsqueeze(3)).sum(2).sum(1)
# self.col_num_out: (B, 3)
col_num_score = self.col_num_out(
self.col_num_out_q(q_weighted_num) +
int(self.use_hs) * self.col_num_out_hs(hs_weighted_num))
# Predict columns.
att_val_qc = torch.bmm(col_enc, self.q_att(q_enc).transpose(1, 2))
for idx, num in enumerate(q_len):
if num < max_q_len:
att_val_qc[idx, :, num:] = -100
att_prob_qc = self.softmax(att_val_qc.view(
(-1, max_q_len))).view(B, -1, max_q_len)
# q_weighted: (B, max_col_len, hid_dim)
q_weighted = (q_enc.unsqueeze(1) * att_prob_qc.unsqueeze(3)).sum(2)
# Same as the above, compute SQL history embedding weighted by column attentions
att_val_hc = torch.bmm(col_enc, self.hs_att(hs_enc).transpose(1, 2))
for idx, num in enumerate(hs_len):
if num < max_hs_len:
att_val_hc[idx, :, num:] = -100
att_prob_hc = self.softmax(att_val_hc.view(
(-1, max_hs_len))).view(B, -1, max_hs_len)
hs_weighted = (hs_enc.unsqueeze(1) * att_prob_hc.unsqueeze(3)).sum(2)
# Compute prediction scores
# self.col_out.squeeze(): (B, max_col_len)
col_score = self.col_out(
self.col_out_q(q_weighted) +
int(self.use_hs) * self.col_out_hs(hs_weighted) +
self.col_out_c(col_enc)).view(B, -1)
for idx, num in enumerate(col_len):
if num < max_col_len:
col_score[idx, num:] = -100
score = (col_num_score, col_score)
return score
def forward(self, q_emb_var, q_len, hs_emb_var, hs_len, col_emb_var,
col_len, col_name_len, gt_col):
max_q_len = max(q_len)
max_hs_len = max(hs_len)
max_col_len = max(col_len)
B = len(q_len)
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_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
att_val_qc_num = torch.bmm(col_emb.unsqueeze(1),
self.q_num_att(q_enc).transpose(1, 2)).view(
B, -1)
for idx, num in enumerate(q_len):
if num < max_q_len:
att_val_qc_num[idx, num:] = -100
att_prob_qc_num = self.softmax(att_val_qc_num)
q_weighted_num = (q_enc * att_prob_qc_num.unsqueeze(2)).sum(1)
# Same as the above, compute SQL history embedding weighted by column attentions
att_val_hc_num = torch.bmm(col_emb.unsqueeze(1),
self.hs_num_att(hs_enc).transpose(1,
2)).view(
B, -1)
for idx, num in enumerate(hs_len):
if num < max_hs_len:
att_val_hc_num[idx, num:] = -100
att_prob_hc_num = self.softmax(att_val_hc_num)
hs_weighted_num = (hs_enc * att_prob_hc_num.unsqueeze(2)).sum(1)
# 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))
# Compute attention values between selected column and question tokens.
# q_enc.transpose(1, 2): (B, hid_dim, max_q_len)
# col_emb.unsqueeze(1): (B, 1, hid_dim)
# att_val_qc: (B, max_q_len)
# print("col_emb {} q_enc {}".format(col_emb.unsqueeze(1).size(),self.q_att(q_enc).transpose(1, 2).size()))
att_val_qc = torch.bmm(col_emb.unsqueeze(1),
self.q_att(q_enc).transpose(1, 2)).view(B, -1)
# assign appended positions values -100
for idx, num in enumerate(q_len):
if num < max_q_len:
att_val_qc[idx, num:] = -100
# att_prob_qc: (B, max_q_len)
att_prob_qc = self.softmax(att_val_qc)
# q_enc: (B, max_q_len, hid_dim)
# att_prob_qc.unsqueeze(2): (B, max_q_len, 1)
# q_weighted: (B, hid_dim)
q_weighted = (q_enc * att_prob_qc.unsqueeze(2)).sum(1)
# Same as the above, compute SQL history embedding weighted by column attentions
att_val_hc = torch.bmm(col_emb.unsqueeze(1),
self.hs_att(hs_enc).transpose(1,
2)).view(B, -1)
for idx, num in enumerate(hs_len):
if num < max_hs_len:
att_val_hc[idx, num:] = -100
att_prob_hc = self.softmax(att_val_hc)
hs_weighted = (hs_enc * att_prob_hc.unsqueeze(2)).sum(1)
# 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))
score = (op_num_score, op_score)
return score