def forward(self, input_data):
q_emb_var, q_len, hs_emb_var, hs_len, mkw_emb_var, mkw_len = input_data
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, input_data):
q_emb_var, q_len, hs_emb_var, hs_len, kw_emb_var, kw_len = 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)
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, input_data):
q_emb_var, q_len, hs_emb_var, hs_len, col_emb_var, col_len, col_name_len = input_data[
0:7]
col_candidates = input_data[7] if len(input_data) == 7 else None
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