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, mkw_emb_var,
mkw_len):
# print("q_emb_shape:{} hs_emb_shape:{}".format(q_emb_var.size(), hs_emb_var.size()))
max_q_len = max(q_len)
max_hs_len = max(hs_len)
B = len(q_len)
# q_enc: (B, max_q_len, hid_dim)
# hs_enc: (B, max_hs_len, hid_dim)
# mkw: (B, 4, hid_dim)
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)
mkw_enc, _ = run_lstm(self.mkw_lstm, mkw_emb_var, mkw_len)
# Compute attention values between multi SQL key words and question tokens.
q_weighted = seq_conditional_weighted_num(self.q_att, q_enc, q_len,
mkw_enc)
SIZE_CHECK(q_weighted, [B, 4, self.N_h])
# Same as the above, compute SQL history embedding weighted by key words attentions
hs_weighted = seq_conditional_weighted_num(self.hs_att, hs_enc, hs_len,
mkw_enc)
# Compute prediction scores=
mulit_score = self.multi_out(
self.multi_out_q(q_weighted) +
int(self.use_hs) * self.multi_out_hs(hs_weighted) +
self.multi_out_c(mkw_enc)).view(B, 4)
return mulit_score
def forward(self, q_emb_var, q_len, hs_emb_var, hs_len):
max_q_len = max(q_len)
max_hs_len = max(hs_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)
att_np_q = np.ones((B, max_q_len))
att_val_q = torch.from_numpy(att_np_q).float()
att_val_q = Variable(att_val_q.cuda())
for idx, num in enumerate(q_len):
if num < max_q_len:
att_val_q[idx, num:] = -100
att_prob_q = self.softmax(att_val_q)
q_weighted = (q_enc * att_prob_q.unsqueeze(2)).sum(1)
# Same as the above, compute SQL history embedding weighted by column attentions
att_np_h = np.ones((B, max_hs_len))
att_val_h = torch.from_numpy(att_np_h).float()
att_val_h = Variable(att_val_h.cuda())
for idx, num in enumerate(hs_len):
if num < max_hs_len:
att_val_h[idx, num:] = -100
att_prob_h = self.softmax(att_val_h)
hs_weighted = (hs_enc * att_prob_h.unsqueeze(2)).sum(1)
# ao_score: (B, 2)
ao_score = self.ao_out(
self.ao_out_q(q_weighted) +
int(self.use_hs) * self.ao_out_hs(hs_weighted))
return ao_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)
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
def forward(self, q_emb_var, q_len, hs_emb_var, hs_len, kw_emb_var, kw_len):
max_q_len = max(q_len)
max_hs_len = max(hs_len)
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)
kw_enc, _ = run_lstm(self.kw_lstm, kw_emb_var, kw_len)
# Predict key words number: 0-3
q_weighted_num = seq_conditional_weighted_num(self.q_num_att, q_enc, q_len, kw_enc).sum(1)
# Same as the above, compute SQL history embedding weighted by key words attentions
hs_weighted_num = seq_conditional_weighted_num(self.hs_num_att, hs_enc, hs_len, kw_enc).sum(1)
# Compute prediction scores
kw_num_score = self.kw_num_out(self.kw_num_out_q(q_weighted_num) + int(self.use_hs)* self.kw_num_out_hs(hs_weighted_num))
SIZE_CHECK(kw_num_score, [B, 4])
# Predict key words: WHERE, GROUP BY, ORDER BY.
q_weighted = seq_conditional_weighted_num(self.q_att, q_enc, q_len, kw_enc)
SIZE_CHECK(q_weighted, [B, 3, self.N_h])
# Same as the above, compute SQL history embedding weighted by key words attentions
hs_weighted = seq_conditional_weighted_num(self.hs_att, hs_enc, hs_len, kw_enc)
# Compute prediction scores
kw_score = self.kw_out(self.kw_out_q(q_weighted) + int(self.use_hs)* self.kw_out_hs(hs_weighted) + self.kw_out_kw(kw_enc)).view(B,3)
score = (kw_num_score, kw_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)
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, gt_col):
max_q_len = max(q_len)
max_hs_len = max(hs_len)
max_col_len = max(col_len)
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
def forward(self, q_emb_var, q_len, hs_emb_var, hs_len, mkw_emb_var,
mkw_len):
# print("q_emb_shape:{} hs_emb_shape:{}".format(q_emb_var.size(), hs_emb_var.size()))
max_q_len = max(q_len)
max_hs_len = max(hs_len)
B = len(q_len)
# q_enc: (B, max_q_len, hid_dim)
# hs_enc: (B, max_hs_len, hid_dim)
# mkw: (B, 4, hid_dim)
q_enc, _ = run_lstm(self.q_lstm, q_emb_var, q_len)
hs_enc, _ = run_lstm(self.hs_lstm, hs_emb_var, hs_len)
mkw_enc, _ = run_lstm(self.mkw_lstm, mkw_emb_var, mkw_len)
# Compute attention values between multi SQL key words and question tokens.
# qmkw_att(q_enc).transpose(1, 2): (B, hid_dim, max_q_len)
# att_val_qmkw: (B, 4, max_q_len)
# print("mkw_enc {} q_enc {}".format(mkw_enc.size(), self.q_att(q_enc).transpose(1, 2).size()))
att_val_qmkw = torch.bmm(mkw_enc, self.q_att(q_enc).transpose(1, 2))
# assign appended positions values -100
for idx, num in enumerate(q_len):
if num < max_q_len:
att_val_qmkw[idx, :, num:] = -100
# att_prob_qmkw: (B, 4, max_q_len)
att_prob_qmkw = self.softmax(att_val_qmkw.view(
(-1, max_q_len))).view(B, -1, max_q_len)
# q_enc.unsqueeze(1): (B, 1, max_q_len, hid_dim)
# att_prob_qmkw.unsqueeze(3): (B, 4, max_q_len, 1)
# q_weighted: (B, 4, hid_dim)
q_weighted = (q_enc.unsqueeze(1) * att_prob_qmkw.unsqueeze(3)).sum(2)
# Same as the above, compute SQL history embedding weighted by key words attentions
att_val_hsmkw = torch.bmm(mkw_enc, self.hs_att(hs_enc).transpose(1, 2))
for idx, num in enumerate(hs_len):
if num < max_hs_len:
att_val_hsmkw[idx, :, num:] = -100
att_prob_hsmkw = self.softmax(att_val_hsmkw.view(
(-1, max_hs_len))).view(B, -1, max_hs_len)
hs_weighted = (hs_enc.unsqueeze(1) *
att_prob_hsmkw.unsqueeze(3)).sum(2)
# Compute prediction scores
# self.multi_out.squeeze(): (B, 4, 1) -> (B, 4)
mulit_score = self.multi_out(
self.multi_out_q(q_weighted) +
int(self.use_hs) * self.multi_out_hs(hs_weighted) +
self.multi_out_c(mkw_enc)).view(B, -1)
return mulit_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)
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
def forward(self, q_emb_var, q_len, hs_emb_var, hs_len, col_emb_var, col_len, col_name_len, col_candidates=None):
max_q_len = max(q_len)
max_hs_len = max(hs_len)
max_col_len = max(col_len)
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)
# Predict column number: 1-3
q_weighted_num = seq_conditional_weighted_num(self.q_num_att, q_enc, q_len, col_enc, col_len).sum(1)
SIZE_CHECK(q_weighted_num, [B, self.N_h])
# Same as the above, compute SQL history embedding weighted by column attentions
hs_weighted_num = seq_conditional_weighted_num(self.hs_num_att, hs_enc, hs_len, col_enc, col_len).sum(1)
SIZE_CHECK(hs_weighted_num, [B, self.N_h])
# 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.
q_weighted = seq_conditional_weighted_num(self.q_att, q_enc, q_len, col_enc)
# Same as the above, compute SQL history embedding weighted by column attentions
hs_weighted = seq_conditional_weighted_num(self.hs_att, hs_enc, hs_len, col_enc)
# 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
for col_num in range(num):
if col_candidates is not None:
if col_num not in col_candidates[idx]:
col_score[idx, col_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):
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, 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, kw_emb_var,
kw_len):
max_q_len = max(q_len)
max_hs_len = max(hs_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)
kw_enc, _ = run_lstm(self.kw_lstm, kw_emb_var, kw_len)
# Predict key words number: 0-3
att_val_qkw_num = torch.bmm(kw_enc,
self.q_num_att(q_enc).transpose(1, 2))
for idx, num in enumerate(q_len):
if num < max_q_len:
att_val_qkw_num[idx, :, num:] = -100
att_prob_qkw_num = self.softmax(att_val_qkw_num.view(
(-1, max_q_len))).view(B, -1, max_q_len)
# q_weighted: (B, hid_dim)
q_weighted_num = (q_enc.unsqueeze(1) *
att_prob_qkw_num.unsqueeze(3)).sum(2).sum(1)
# Same as the above, compute SQL history embedding weighted by key words attentions
att_val_hskw_num = torch.bmm(kw_enc,
self.hs_num_att(hs_enc).transpose(1, 2))
for idx, num in enumerate(hs_len):
if num < max_hs_len:
att_val_hskw_num[idx, :, num:] = -100
att_prob_hskw_num = self.softmax(
att_val_hskw_num.view((-1, max_hs_len))).view(B, -1, max_hs_len)
hs_weighted_num = (hs_enc.unsqueeze(1) *
att_prob_hskw_num.unsqueeze(3)).sum(2).sum(1)
# Compute prediction scores
# self.kw_num_out: (B, 4)
kw_num_score = self.kw_num_out(
self.kw_num_out_q(q_weighted_num) +
int(self.use_hs) * self.kw_num_out_hs(hs_weighted_num))
# Predict key words: WHERE, GROUP BY, ORDER BY.
att_val_qkw = torch.bmm(kw_enc, self.q_att(q_enc).transpose(1, 2))
for idx, num in enumerate(q_len):
if num < max_q_len:
att_val_qkw[idx, :, num:] = -100
att_prob_qkw = self.softmax(att_val_qkw.view(
(-1, max_q_len))).view(B, -1, max_q_len)
# q_weighted: (B, 3, hid_dim)
q_weighted = (q_enc.unsqueeze(1) * att_prob_qkw.unsqueeze(3)).sum(2)
# Same as the above, compute SQL history embedding weighted by key words attentions
att_val_hskw = torch.bmm(kw_enc, self.hs_att(hs_enc).transpose(1, 2))
for idx, num in enumerate(hs_len):
if num < max_hs_len:
att_val_hskw[idx, :, num:] = -100
att_prob_hskw = self.softmax(att_val_hskw.view(
(-1, max_hs_len))).view(B, -1, max_hs_len)
hs_weighted = (hs_enc.unsqueeze(1) * att_prob_hskw.unsqueeze(3)).sum(2)
# Compute prediction scores
# self.kw_out.squeeze(): (B, 3)
kw_score = self.kw_out(
self.kw_out_q(q_weighted) +
int(self.use_hs) * self.kw_out_hs(hs_weighted) +
self.kw_out_kw(kw_enc)).view(B, -1)
score = (kw_num_score, kw_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
def forward(self, parent_tables, foreign_keys, q_emb_var, q_len,
hs_emb_var, hs_len, col_emb_var, col_len, col_name_len,
table_emb_var, table_len, table_name_len):
max_q_len = max(q_len)
max_hs_len = max(hs_len)
max_col_len = max(col_len)
max_table_len = max(table_len)
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)
assert list(q_enc.size()) == [B, max_q_len, self.encoded_num]
hs_enc, _ = run_lstm(self.hs_lstm, hs_emb_var, hs_len)
table_tensors, col_tensors, batch_graph = self.schema_encoder(
parent_tables, foreign_keys, col_emb_var, col_name_len, col_len,
table_emb_var, table_name_len, table_len)
aggregated_schema = self.schema_aggregator(batch_graph)
SIZE_CHECK(table_tensors, [B, max_table_len, self.N_h])
SIZE_CHECK(col_tensors, [B, max_col_len, self.N_h])
q_table_weighted_num_num = seq_conditional_weighted_num(
self.q_table_num_att, q_enc, q_len, table_tensors,
table_len).sum(1)
hs_table_weighted_num_num = seq_conditional_weighted_num(
self.hs_table_num_att, hs_enc, hs_len, table_tensors,
table_len).sum(1)
q_col_weighted_num_num = seq_conditional_weighted_num(
self.q_col_num_att, q_enc, q_len, col_tensors, col_len).sum(1)
hs_col_weighted_num_num = seq_conditional_weighted_num(
self.hs_col_num_att, hs_enc, hs_len, col_tensors, col_len).sum(1)
x = self.schema_out(F.relu(aggregated_schema))
x = x + self.q_table_out(q_table_weighted_num_num)
x = x + int(self.use_hs) * self.hs_table_out(hs_table_weighted_num_num)
x = x + self.q_col_out(q_col_weighted_num_num)
x = x + int(self.use_hs) * self.hs_col_out(hs_col_weighted_num_num)
table_num_score = self.table_num_out(x)
q_table_weighted_num = seq_conditional_weighted_num(
self.q_table_att, q_enc, q_len, table_tensors, table_len).sum(1)
hs_table_weighted_num = seq_conditional_weighted_num(
self.hs_table_att, hs_enc, hs_len, table_tensors, table_len).sum(1)
q_col_weighted_num = seq_conditional_weighted_num(
self.q_col_att, q_enc, q_len, col_tensors, col_len).sum(1)
hs_col_weighted_num = seq_conditional_weighted_num(
self.hs_col_att, hs_enc, hs_len, col_tensors, col_len).sum(1)
x = self.schema_out(F.relu(aggregated_schema))
x = x + self.q_table_out(q_table_weighted_num)
x = x + int(self.use_hs) * self.hs_table_out(hs_table_weighted_num)
x = x + self.q_col_out(q_col_weighted_num)
x = x + int(self.use_hs) * self.hs_col_out(hs_col_weighted_num)
SIZE_CHECK(x, [B, self.N_h])
table_score = (self.table_att(table_tensors) * x.unsqueeze(1)).sum(2)
SIZE_CHECK(table_score, [B, max_table_len])
for idx, num in enumerate(table_len.tolist()):
if num < max_table_len:
table_score[idx, num:] = -100
return table_num_score, table_score
