class SQLNet(nn.Module):
def __init__(self,
word_emb,
N_word,
N_h=120,
N_depth=2,
gpu=False,
trainable_emb=False,
db_content=0):
super(SQLNet, self).__init__()
self.trainable_emb = trainable_emb
self.db_content = db_content
self.gpu = gpu
self.N_h = N_h
self.N_depth = N_depth
self.max_col_num = 45
self.max_tok_num = 200
self.SQL_TOK = [
'<UNK>', '<END>', 'WHERE', 'AND', 'EQL', 'GT', 'LT', '<BEG>'
]
self.COND_OPS = ['EQL', 'GT', 'LT']
# the model actually doesn't use type embedding when db_content == 1
if db_content == 0:
is_train = True
else:
is_train = False
self.agg_type_embed_layer = WordEmbedding(word_emb,
N_word,
gpu,
self.SQL_TOK,
trainable=is_train)
self.sel_type_embed_layer = WordEmbedding(word_emb,
N_word,
gpu,
self.SQL_TOK,
trainable=is_train)
self.cond_type_embed_layer = WordEmbedding(word_emb,
N_word,
gpu,
self.SQL_TOK,
trainable=is_train)
self.embed_layer = WordEmbedding(word_emb,
N_word,
gpu,
self.SQL_TOK,
trainable=trainable_emb)
# Predict aggregator
self.agg_pred = AggPredictor(N_word, N_h, N_depth)
# Predict select column + condition number and columns
self.selcond_pred = SelCondPredictor(N_word, N_h, N_depth, gpu,
db_content)
# Predict condition operators and string values
self.op_str_pred = CondOpStrPredictor(N_word, N_h, N_depth,
self.max_col_num,
self.max_tok_num, gpu,
db_content)
self.CE = nn.CrossEntropyLoss()
self.softmax = nn.Softmax()
self.log_softmax = nn.LogSoftmax()
self.bce_logit = nn.BCEWithLogitsLoss()
if gpu:
self.cuda()
def get_str_index(self, all_toks, this_str):
cur_seq = []
tok_gt_1 = [t for t in all_toks if len(t) > 1]
if this_str in all_toks:
all_str = [['<BEG>'], this_str, ['<END>']]
cur_seq = [
all_toks.index(s) if s in all_toks else 0 for s in all_str
]
elif len(tok_gt_1) > 0:
flag = False
for tgt in tok_gt_1:
if set(tgt).issubset(this_str):
not_tgt = [x for x in this_str if x not in tgt]
if len(not_tgt) > 0:
not_tgt = [[x] for x in not_tgt]
all_str = [tgt] + not_tgt
else:
all_str = [tgt]
beg_ind = all_toks.index(
['<BEG>']) if ['<BEG>'] in all_toks else 0
end_ind = all_toks.index(
['<END>']) if ['<END>'] in all_toks else 0
cur_seq = sorted([
all_toks.index(s) if s in all_toks else 0
for s in all_str
])
cur_seq = [beg_ind] + cur_seq + [end_ind]
elif set(this_str).issubset(tgt):
all_str = [['<BEG>'], tgt, ['<END>']]
cur_seq = [
all_toks.index(s) if s in all_toks else 0
for s in all_str
]
if len(cur_seq) > 0:
flag = True
break
if not flag:
all_str = [['<BEG>']] + [[x] for x in this_str] + [['<END>']]
cur_seq = [
all_toks.index(s) if s in all_toks else 0 for s in all_str
]
else:
all_str = [['<BEG>']] + [[x] for x in this_str] + [['<END>']]
cur_seq = [
all_toks.index(s) if s in all_toks else 0 for s in all_str
]
return cur_seq
def generate_gt_where_seq(self, q, col, query):
"""
cur_seq is the indexes (in question toks) of string value in each where cond
"""
ret_seq = []
for cur_q, cur_col, cur_query in zip(q, col, query):
cur_values = []
st = cur_query.index(u'WHERE') + 1 if \
u'WHERE' in cur_query else len(cur_query)
all_toks = [['<BEG>']] + cur_q + [['<END>']]
while st < len(cur_query):
ed = len(cur_query) if 'AND' not in cur_query[st:] \
else cur_query[st:].index('AND') + st
if 'EQL' in cur_query[st:ed]:
op = cur_query[st:ed].index('EQL') + st
elif 'GT' in cur_query[st:ed]:
op = cur_query[st:ed].index('GT') + st
elif 'LT' in cur_query[st:ed]:
op = cur_query[st:ed].index('LT') + st
else:
raise RuntimeError("No operator in it!")
this_str = cur_query[op + 1:ed]
cur_seq = self.get_str_index(all_toks, this_str)
cur_values.append(cur_seq)
st = ed + 1
ret_seq.append(cur_values)
return ret_seq
def forward(self,
q,
col,
col_num,
q_type,
col_type,
pred_entry,
gt_where=None,
gt_cond=None,
gt_sel=None):
B = len(q)
pred_agg, pred_sel, pred_cond = pred_entry
agg_score = None
sel_cond_score = None
cond_op_str_score = None
# Predict aggregator
if self.trainable_emb:
if pred_agg:
x_emb_var, x_len = self.agg_embed_layer.gen_x_batch(q, col)
col_inp_var, col_name_len, col_len = \
self.agg_embed_layer.gen_col_batch(col)
max_x_len = max(x_len)
agg_score = self.agg_pred(x_emb_var,
x_len,
col_inp_var,
col_name_len,
col_len,
col_num,
gt_sel=gt_sel)
if pred_sel:
x_emb_var, x_len = self.sel_embed_layer.gen_x_batch(q, col)
col_inp_var, col_name_len, col_len = \
self.sel_embed_layer.gen_col_batch(col)
max_x_len = max(x_len)
sel_score = self.selcond_pred(x_emb_var, x_len, col_inp_var,
col_name_len, col_len, col_num)
if pred_cond:
x_emb_var, x_len = self.cond_embed_layer.gen_x_batch(q, col)
col_inp_var, col_name_len, col_len = \
self.cond_embed_layer.gen_col_batch(col)
max_x_len = max(x_len)
cond_score = self.cond_pred(x_emb_var, x_len, col_inp_var,
col_name_len, col_len, col_num,
gt_where, gt_cond)
elif self.db_content == 0:
x_emb_var, x_len = self.embed_layer.gen_x_batch(q,
col,
is_list=True,
is_q=True)
col_inp_var, col_len = self.embed_layer.gen_x_batch(col,
col,
is_list=True)
agg_emb_var = self.embed_layer.gen_agg_batch(q)
max_x_len = max(x_len)
if pred_agg:
# x_type_agg_emb_var, _ = self.agg_type_embed_layer.gen_xc_type_batch(q_type, is_list=True)
agg_score = self.agg_pred(x_emb_var, x_len, agg_emb_var,
col_inp_var, col_len)
if pred_sel:
x_type_sel_emb_var, _ = self.sel_type_embed_layer.gen_xc_type_batch(
q_type, is_list=True)
sel_cond_score = self.selcond_pred(x_emb_var, x_len,
col_inp_var, col_len,
x_type_sel_emb_var, gt_sel)
if pred_cond:
x_type_cond_emb_var, _ = self.cond_type_embed_layer.gen_xc_type_batch(
q_type, is_list=True)
cond_op_str_score = self.op_str_pred(x_emb_var, x_len,
col_inp_var, col_len,
x_type_cond_emb_var,
gt_where, gt_cond,
sel_cond_score)
else:
x_emb_var, x_len = self.embed_layer.gen_x_batch(q,
col,
is_list=True,
is_q=True)
col_inp_var, col_len = self.embed_layer.gen_x_batch(col,
col,
is_list=True)
x_type_emb_var, x_type_len = self.embed_layer.gen_x_batch(
q_type, col, is_list=True, is_q=True)
col_type_inp_var, col_type_len = self.embed_layer.gen_x_batch(
col_type, col_type, is_list=True)
agg_emb_var = self.embed_layer.gen_agg_batch(q)
max_x_len = max(x_len)
if pred_agg:
agg_score = self.agg_pred(x_emb_var, x_len, agg_emb_var,
col_inp_var, col_len)
if pred_sel:
sel_cond_score = self.selcond_pred(x_emb_var, x_len,
col_inp_var, col_len,
x_type_emb_var, gt_sel)
if pred_cond:
cond_op_str_score = self.op_str_pred(x_emb_var, x_len,
col_inp_var, col_len,
x_type_emb_var, gt_where,
gt_cond, sel_cond_score)
return (agg_score, sel_cond_score, cond_op_str_score)
def loss(self, score, truth_num, pred_entry, gt_where):
pred_agg, pred_sel, pred_cond = pred_entry
agg_score, sel_cond_score, cond_op_str_score = score
cond_num_score, sel_score, cond_col_score = sel_cond_score
cond_op_score, cond_str_score = cond_op_str_score
loss = 0
if pred_agg:
agg_truth = map(lambda x: x[0], truth_num)
data = torch.from_numpy(np.array(agg_truth))
if self.gpu:
agg_truth_var = Variable(data.cuda())
else:
agg_truth_var = Variable(data)
loss += self.CE(agg_score, agg_truth_var)
if pred_sel:
sel_truth = map(lambda x: x[1], truth_num)
data = torch.from_numpy(np.array(sel_truth))
if self.gpu:
sel_truth_var = Variable(data.cuda())
else:
sel_truth_var = Variable(data)
loss += self.CE(sel_score, sel_truth_var)
if pred_cond:
B = len(truth_num)
# Evaluate the number of conditions
cond_num_truth = map(lambda x: x[2], truth_num)
data = torch.from_numpy(np.array(cond_num_truth))
if self.gpu:
cond_num_truth_var = Variable(data.cuda())
else:
cond_num_truth_var = Variable(data)
loss += self.CE(cond_num_score, cond_num_truth_var)
# Evaluate the columns of conditions
T = len(cond_col_score[0])
truth_prob = np.zeros((B, T), dtype=np.float32)
for b in range(B):
if len(truth_num[b][3]) > 0:
truth_prob[b][list(truth_num[b][3])] = 1
data = torch.from_numpy(truth_prob)
if self.gpu:
cond_col_truth_var = Variable(data.cuda())
else:
cond_col_truth_var = Variable(data)
sigm = nn.Sigmoid()
cond_col_prob = sigm(cond_col_score)
bce_loss = -torch.mean(3 * (cond_col_truth_var * \
torch.log(cond_col_prob + 1e-10)) + \
(1 - cond_col_truth_var) * torch.log(1 - cond_col_prob + 1e-10))
loss += bce_loss
# Evaluate the operator of conditions
for b in range(len(truth_num)):
if len(truth_num[b][4]) == 0:
continue
data = torch.from_numpy(np.array(truth_num[b][4]))
if self.gpu:
cond_op_truth_var = Variable(data.cuda())
else:
cond_op_truth_var = Variable(data)
cond_op_pred = cond_op_score[b, :len(truth_num[b][4])]
loss += (self.CE(cond_op_pred, cond_op_truth_var) \
/ len(truth_num))
# Evaluate the strings of conditions
for b in range(len(gt_where)):
for idx in range(len(gt_where[b])):
cond_str_truth = gt_where[b][idx]
if len(cond_str_truth) == 1:
continue
data = torch.from_numpy(np.array(cond_str_truth[1:]))
if self.gpu:
cond_str_truth_var = Variable(data.cuda())
else:
cond_str_truth_var = Variable(data)
str_end = len(cond_str_truth) - 1
cond_str_pred = cond_str_score[b, idx, :str_end]
loss += (self.CE(cond_str_pred, cond_str_truth_var) \
/ (len(gt_where) * len(gt_where[b])))
return loss
def check_acc(self,
vis_info,
pred_queries,
gt_queries,
pred_entry,
error_print=False):
def pretty_print(vis_data, pred_query, gt_query):
print "\n----------detailed error prints-----------"
try:
print 'question: ', vis_data[0]
print 'question_tok: ', vis_data[3]
print 'headers: (%s)' % (' || '.join(vis_data[1]))
print 'query:', vis_data[2]
print "target query: ", gt_query
print "pred query: ", pred_query
except:
print "\n------skipping print: decoding problem ----------------------"
def gen_cond_str(conds, header):
if len(conds) == 0:
return 'None'
cond_str = []
for cond in conds:
cond_str.append(header[cond[0]] + ' ' +
self.COND_OPS[cond[1]] + ' ' +
unicode(cond[2]).lower())
return 'WHERE ' + ' AND '.join(cond_str)
pred_agg, pred_sel, pred_cond = pred_entry
B = len(gt_queries)
tot_err = agg_err = sel_err = cond_err = 0.0
cond_num_err = cond_col_err = cond_op_err = cond_val_err = 0.0
agg_ops = ['None', 'MAX', 'MIN', 'COUNT', 'SUM', 'AVG']
for b, (pred_qry, gt_qry,
vis_data) in enumerate(zip(pred_queries, gt_queries,
vis_info)):
good = True
if pred_agg:
agg_pred = pred_qry['agg']
agg_gt = gt_qry['agg']
if agg_pred != agg_gt:
agg_err += 1
good = False
if pred_sel:
sel_pred = pred_qry['sel']
sel_gt = gt_qry['sel']
if sel_pred != sel_gt:
sel_err += 1
good = False
if pred_cond:
cond_pred = pred_qry['conds']
cond_gt = gt_qry['conds']
flag = True
if len(cond_pred) != len(cond_gt):
flag = False
cond_num_err += 1
if flag and set(x[0] for x in cond_pred) != \
set(x[0] for x in cond_gt):
flag = False
cond_col_err += 1
for idx in range(len(cond_pred)):
if not flag:
break
gt_idx = tuple(x[0]
for x in cond_gt).index(cond_pred[idx][0])
if flag and cond_gt[gt_idx][1] != cond_pred[idx][1]:
flag = False
cond_op_err += 1
for idx in range(len(cond_pred)):
if not flag:
break
gt_idx = tuple(x[0]
for x in cond_gt).index(cond_pred[idx][0])
if flag and unicode(cond_gt[gt_idx][2]).lower() != \
unicode(cond_pred[idx][2]).lower():
flag = False
cond_val_err += 1
if not flag:
cond_err += 1
good = False
if not good:
if error_print:
pretty_print(vis_data, pred_qry, gt_qry)
tot_err += 1
return np.array((agg_err, sel_err, cond_err, cond_num_err,
cond_col_err, cond_op_err, cond_val_err)), tot_err
def gen_query(self,
score,
q,
col,
raw_q,
raw_col,
pred_entry,
verbose=False):
def merge_tokens(tok_list, raw_tok_str):
"""
tok_list: list of string words in current cond
raw_tok_str: list of words in question
"""
tok_str = raw_tok_str.lower()
alphabet = 'abcdefghijklmnopqrstuvwxyz0123456789$('
special = {
'-LRB-': '(',
'-RRB-': ')',
'-LSB-': '[',
'-RSB-': ']',
'``': '"',
'\'\'': '"',
'--': u'\u2013'
}
ret = ''
double_quote_appear = 0
tok_list = [x for gx in tok_list for x in gx]
for raw_tok in tok_list:
if not raw_tok:
continue
tok = special.get(raw_tok, raw_tok)
if tok == '"':
double_quote_appear = 1 - double_quote_appear
if len(ret) == 0:
pass
elif len(ret) > 0 and ret + ' ' + tok in tok_str:
ret = ret + ' '
elif len(ret) > 0 and ret + tok in tok_str:
pass
elif tok == '"':
if double_quote_appear:
ret = ret + ' '
elif tok[0] not in alphabet:
pass
elif (ret[-1] not in ['(', '/', u'\u2013', '#', '$', '&']) \
and (ret[-1] != '"' or not double_quote_appear):
ret = ret + ' '
ret = ret + tok
return ret.strip()
pred_agg, pred_sel, pred_cond = pred_entry
agg_score, sel_cond_score, cond_op_str_score = score
cond_num_score, sel_score, cond_col_score = [
x.data.cpu().numpy() for x in sel_cond_score
]
cond_op_score, cond_str_score = [
x.data.cpu().numpy() for x in cond_op_str_score
]
ret_queries = []
if pred_agg:
B = len(agg_score)
elif pred_sel:
B = len(sel_score)
elif pred_cond:
B = len(cond_num_score)
for b in range(B):
cur_query = {}
if pred_agg:
cur_query['agg'] = np.argmax(agg_score[b].data.cpu().numpy())
if pred_sel:
cur_query['sel'] = np.argmax(sel_score[b])
if pred_cond:
cur_query['conds'] = []
cond_num = np.argmax(cond_num_score[b])
all_toks = [['<BEG>']] + q[b] + [['<END>']]
max_idxes = np.argsort(-cond_col_score[b])[:cond_num]
for idx in range(cond_num):
cur_cond = []
cur_cond.append(max_idxes[idx])
cur_cond.append(np.argmax(cond_op_score[b][idx]))
cur_cond_str_toks = []
for str_score in cond_str_score[b][idx]:
str_tok = np.argmax(str_score[:len(all_toks)])
str_val = all_toks[str_tok]
if str_val == ['<END>']:
break
# add string word/grouped words to current cond str tokens ["w1", "w2" ...]
cur_cond_str_toks.append(str_val)
cur_cond.append(merge_tokens(cur_cond_str_toks, raw_q[b]))
cur_query['conds'].append(cur_cond)
ret_queries.append(cur_query)
return ret_queries
class SQLNet(nn.Module):
def __init__(self,
word_emb,
N_word,
N_h=120,
N_depth=2,
use_ca=True,
gpu=True,
trainable_emb=False,
db_content=0):
super(SQLNet, self).__init__()
self.trainable_emb = trainable_emb
self.db_content = db_content
self.use_ca = use_ca
self.gpu = gpu
self.N_h = N_h
self.N_depth = N_depth
self.max_col_num = 45
self.max_tok_num = 200
self.SQL_TOK = [
'<UNK>', '<END>', 'WHERE', 'AND', 'OR', '==', '>', '<', '!=',
'<BEG>'
]
self.COND_OPS = ['>', '<', '==', '!=']
#the model actually doesn't use type embedding when db_content == 1
if db_content == 0:
is_train = True
else:
is_train = False
# self.sel_num_type_embed_layer = WordEmbedding(word_emb, N_word, gpu,
# self.SQL_TOK, trainable=is_train)
self.agg_type_embed_layer = WordEmbedding(word_emb,
N_word,
gpu,
self.SQL_TOK,
trainable=is_train)
self.sel_type_embed_layer = WordEmbedding(word_emb,
N_word,
gpu,
self.SQL_TOK,
trainable=is_train)
self.cond_type_embed_layer = WordEmbedding(word_emb,
N_word,
gpu,
self.SQL_TOK,
trainable=is_train)
self.where_rela_type_embed_layer = WordEmbedding(word_emb,
N_word,
gpu,
self.SQL_TOK,
trainable=is_train)
self.embed_layer = WordEmbedding(word_emb,
N_word,
gpu,
self.SQL_TOK,
trainable=trainable_emb)
# # Predict selected column number
# self.sel_num = SelNumPredictor(N_word, N_h, N_depth)
#
# # Predict which columns are selected
# self.sel_pred = SelPredictor(N_word, N_h, N_depth, self.max_tok_num, use_ca=use_ca)
#Predict aggregator
self.agg_pred = AggPredictor(N_word, N_h, N_depth)
# # Predict number of conditions, condition columns, condition operations and condition values
# self.cond_pred = SQLNetCondPredictor(N_word, N_h, N_depth, self.max_col_num, self.max_tok_num,use_ca, gpu, db_content)
#Predict selected column number + select column + condition number and columns
self.selcond_pred = SelCondPredictor(N_word, N_h, N_depth, gpu,
db_content)
#Predict condition operators and string values
self.op_str_pred = CondOpStrPredictor(N_word, N_h, N_depth,
self.max_col_num,
self.max_tok_num, gpu,
db_content)
# Predict conditions' relation
self.where_rela_pred = WhereRelationPredictor(N_word,
N_h,
N_depth,
use_ca=use_ca)
self.CE = nn.CrossEntropyLoss()
self.softmax = nn.Softmax()
self.log_softmax = nn.LogSoftmax()
self.bce_logit = nn.BCEWithLogitsLoss()
if gpu:
self.cuda()
def get_str_index(self, all_toks, this_str):
cur_seq = []
tok_gt_1 = [t for t in all_toks if len(t) > 1]
if this_str in all_toks:
all_str = [['<BEG>'], this_str, ['<END>']]
cur_seq = [
all_toks.index(s) if s in all_toks else 0 for s in all_str
]
elif len(tok_gt_1) > 0:
flag = False
for tgt in tok_gt_1:
if set(tgt).issubset(this_str):
not_tgt = [x for x in this_str if x not in tgt]
if len(not_tgt) > 0:
not_tgt = [[x] for x in not_tgt]
all_str = [tgt] + not_tgt
else:
all_str = [tgt]
beg_ind = all_toks.index(
['<BEG>']) if ['<BEG>'] in all_toks else 0
end_ind = all_toks.index(
['<END>']) if ['<END>'] in all_toks else 0
cur_seq = sorted([
all_toks.index(s) if s in all_toks else 0
for s in all_str
])
cur_seq = [beg_ind] + cur_seq + [end_ind]
elif set(this_str).issubset(tgt):
all_str = [['<BEG>'], tgt, ['<END>']]
cur_seq = [
all_toks.index(s) if s in all_toks else 0
for s in all_str
]
if len(cur_seq) > 0:
flag = True
break
if not flag:
all_str = [['<BEG>']] + [[x] for x in this_str] + [['<END>']]
cur_seq = [
all_toks.index(s) if s in all_toks else 0 for s in all_str
]
else:
all_str = [['<BEG>']] + [[x] for x in this_str] + [['<END>']]
cur_seq = [
all_toks.index(s) if s in all_toks else 0 for s in all_str
]
return cur_seq
def generate_gt_where_seq_test(self, q, gt_cond_seq):
ret_seq = []
for cur_q, ans in zip(q, gt_cond_seq):
q_toks = []
q_toks_cnt = []
cur_q_join = []
cnt = 0
for toks in cur_q:
cur_q_join.append(u"".join(toks))
cnt1 = 0
for tok in toks:
q_toks.append(tok)
cnt1 = cnt1 + len(tok)
cnt = cnt + cnt1
q_toks_cnt.append(cnt)
#for tok in q_toks:
# print("{}".format(tok.encode('utf-8')))
#print("q_toks:{}".format(len(q_toks)))
temp_q = u"".join(cur_q_join)
#print("temp_q:{}".format(temp_q.encode('utf-8')))
#cur_q = [u'<BEG>'] + cur_q + [u'<END>']
#cur_q = [u'<BEG>'] + cur_q_join + [u'<END>']
#print("cur_q:{}".format(cur_q.encode('utf-8')))
record = []
record_cond = []
for cond in ans:
#print("cond[2]:{}".format(cond[2].encode('utf-8')))
if cond[2] not in temp_q:
record.append((False, cond[2]))
else:
record.append((True, cond[2]))
for idx, item in enumerate(record):
temp_ret_seq = []
if item[0]:
temp_ret_seq.append(0)
start_idx = -1
end_idx = -1
start_idx_org = temp_q.index(item[1]) + 1
end_idx_org = start_idx_org + len(item[1]) - 1
for idx, cnt in enumerate(q_toks_cnt):
if start_idx_org <= cnt:
start_idx = idx
break
for idx, cnt in enumerate(q_toks_cnt):
if end_idx_org <= cnt:
end_idx = idx + 1
break
if end_idx == -1:
end_idx = len(q_toks_cnt) + 1
#temp_ret_seq.extend(list(range(temp_q.index(item[1])+1,temp_q.index(item[1])+len(item[1])+1)))
#print("start_idx:{} end_idx:{}".format(start_idx, end_idx))
temp_ret_seq.extend(list(range(start_idx + 1,
end_idx + 1)))
temp_ret_seq.append(len(q_toks_cnt) + 1)
else:
temp_ret_seq.extend([0, len(q_toks_cnt) + 1])
#print("temp_ret_sql:{}".format(temp_ret_seq))
record_cond.append(temp_ret_seq)
ret_seq.append(record_cond)
return ret_seq
def generate_gt_where_seq(self, q, col):
ret_seq = []
for cur_q, cur_query in zip(q, col):
cur_values = []
st = cur_query.index(u'WHERE') + 1 if \
u'WHERE' in cur_query else len(cur_query)
all_toks = [['<BEG>']] + cur_q + [['<END>']]
while st < len(cur_query):
ed = len(cur_query) if 'AND' not in cur_query[st:] \
else cur_query[st:].index('AND') + st
if '==' in cur_query[st:ed]:
op = cur_query[st:ed].index('==') + st
elif '>' in cur_query[st:ed]:
op = cur_query[st:ed].index('>') + st
elif '<' in cur_query[st:ed]:
op = cur_query[st:ed].index('<') + st
elif '>=' in cur_query[st:ed]:
op = cur_query[st:ed].index('>=') + st
elif '<=' in cur_query[st:ed]:
op = cur_query[st:ed].index('<=') + st
elif '!=' in cur_query[st:ed]:
op = cur_query[st:ed].index('!=') + st
else:
raise RuntimeError("No operator in it!")
this_str = cur_query[op + 1:ed]
cur_seq = self.get_str_index(all_toks, this_str)
cur_values.append(cur_seq)
st = ed + 1
ret_seq.append(cur_values)
# ret_seq = []
# for cur_q, ans in zip(q, gt_cond_seq):
# temp_q = u"".join(cur_q)
# cur_q = [u'<BEG>'] + cur_q + [u'<END>']
# record = []
# record_cond = []
# for cond in ans:
# if cond[2] not in temp_q:
# record.append((False, cond[2]))
# else:
# record.append((True, cond[2]))
# for idx, item in enumerate(record):
# temp_ret_seq = []
# if item[0]:
# temp_ret_seq.append(0)
# temp_ret_seq.extend(list(range(temp_q.index(item[1])+1,temp_q.index(item[1])+len(item[1])+1)))
# temp_ret_seq.append(len(cur_q)-1)
# else:
# temp_ret_seq.append([0,len(cur_q)-1])
# record_cond.append(temp_ret_seq)
# ret_seq.append(record_cond)
return ret_seq
def forward(self,
q,
col,
col_num,
q_type,
col_type,
gt_where=None,
gt_cond=None,
gt_sel=None,
gt_sel_num=None):
B = len(q)
# pred_sel_num, pred_agg, pred_sel, pred_cond, pred_where_rela = pred_entry
pred_agg = True
pred_sel = True
pred_cond = True
pred_sel_num = True
pred_where_rela = True
agg_score = None
sel_cond_score = None
cond_op_str_score = None
if self.trainable_emb:
if pred_agg:
x_emb_var, x_len = self.agg_embed_layer.gen_x_batch(q, col)
col_inp_var, col_name_len, col_len = \
self.agg_embed_layer.gen_col_batch(col)
max_x_len = max(x_len)
agg_score = self.agg_pred(x_emb_var,
x_len,
col_inp_var,
col_name_len,
col_len,
col_num,
gt_sel=gt_sel)
if pred_sel:
x_emb_var, x_len = self.sel_embed_layer.gen_x_batch(q, col)
col_inp_var, col_name_len, col_len = \
self.sel_embed_layer.gen_col_batch(col)
max_x_len = max(x_len)
sel_score = self.selcond_pred(x_emb_var, x_len, col_inp_var,
col_name_len, col_len, col_num)
if pred_cond:
x_emb_var, x_len = self.cond_embed_layer.gen_x_batch(q, col)
col_inp_var, col_name_len, col_len = \
self.cond_embed_layer.gen_col_batch(col)
max_x_len = max(x_len)
cond_score = self.cond_pred(x_emb_var, x_len, col_inp_var,
col_name_len, col_len, col_num,
gt_where, gt_cond)
elif self.db_content == 0:
x_emb_var, x_len = self.embed_layer.gen_x_batch(q,
col,
is_list=True,
is_q=True)
#col_inp_var, col_len = self.embed_layer.gen_x_batch(col, col, is_list=True)
col_inp_var, col_name_len, col_len = self.embed_layer.gen_col_batch(
col)
agg_emb_var = self.embed_layer.gen_agg_batch(q)
max_x_len = max(x_len)
if pred_sel_num and pred_agg and pred_sel:
sel_num_score = self.sel_num(x_emb_var, x_len, col_inp_var,
col_name_len, col_len, col_num)
if gt_sel_num:
pre_sel_num = gt_sel_num
else:
pr_sel_num = np.argmax(sel_num_score.data.cpu().numpy(),
axis=1)
x_type_sel_emb_var, _ = self.sel_type_embed_layer.gen_xc_type_batch(
q_type, is_list=True)
sel_cond_score = self.selcond_pred(x_emb_var, x_len, col_inp_var,
col_len, x_type_sel_emb_var,
gt_sel)
if gt_sel:
pr_sel = gt_sel
else:
num = np.argmax(sel_num_score.data.cpu().numpy(), axis=1)
sel = sel_cond_score.data.cpu().numpy()
pr_sel = [
list(np.argsort(-sel[b])[:num[b]]) for b in range(len(num))
]
agg_score = self.agg_pred(x_emb_var,
x_len,
agg_emb_var,
col_inp_var,
col_len,
gt_sel=pr_sel,
gt_sel_num=pr_sel_num)
# if pred_agg:
# #x_type_agg_emb_var, _ = self.agg_type_embed_layer.gen_xc_type_batch(q_type, is_list=True)
# agg_score = self.agg_pred(x_emb_var, x_len, agg_emb_var, col_inp_var, col_len)
#
# if pred_sel:
# x_type_sel_emb_var, _ = self.sel_type_embed_layer.gen_xc_type_batch(q_type, is_list=True)
# sel_cond_score = self.selcond_pred(x_emb_var, x_len, col_inp_var, col_len, x_type_sel_emb_var,
# gt_sel)
if pred_cond:
x_type_cond_emb_var, _ = self.cond_type_embed_layer.gen_xc_type_batch(
q_type, is_list=True)
cond_op_str_score = self.op_str_pred(x_emb_var, x_len,
col_inp_var, col_len,
x_type_cond_emb_var,
gt_where, gt_cond,
sel_cond_score)
if pred_where_rela:
where_rela_score = self.where_rela_pred(
x_emb_var, x_len, col_inp_var, col_name_len, col_len,
col_num)
else:
x_emb_var, x_len = self.embed_layer.gen_x_batch(q,
col,
is_list=True,
is_q=True)
col_inp_var, col_name_len, col_len = self.embed_layer.gen_col_batch(
col)
x_type_emb_var, x_type_len = self.embed_layer.gen_x_batch(
q_type, col, is_list=True, is_q=True)
#x_type_cond_emb_var, _ = self.cond_type_embed_layer.gen_xc_type_batch(q_type, is_list=True)
#col_type_inp_var, col_type_len = self.embed_layer.gen_x_batch(col_type, col_type, is_list=True)
#print("x_var_shape:{}, x_type_var_shape:{}".format(x_emb_var, x_type_emb_var))
sel_cond_score = self.selcond_pred(x_emb_var, x_len, col_inp_var,
col_name_len, col_len,
x_type_emb_var, gt_sel)
agg_score = self.agg_pred(x_emb_var, x_len, col_inp_var, col_len,
col_name_len, x_type_emb_var, gt_sel,
sel_cond_score)
cond_op_str_score = self.op_str_pred(x_emb_var, x_len, col_inp_var,
col_len, col_name_len,
x_type_emb_var, gt_where,
gt_cond, sel_cond_score)
where_rela_score = self.where_rela_pred(x_emb_var, x_len,
col_inp_var, col_name_len,
col_len, col_num,
x_type_emb_var)
# sel_num_score = self.sel_num(x_emb_var, x_len, col_inp_var, col_name_len, col_len, col_num, x_type_emb_var)
#
# if gt_sel_num:
# pr_sel_num = gt_sel_num
# else:
# pr_sel_num = np.argmax(sel_num_score.data.cpu().numpy(), axis=1)
# sel_score = self.sel_pred(x_emb_var, x_len, col_inp_var, col_name_len, col_len, col_num, x_type_emb_var)
#
# if gt_sel:
# pr_sel = gt_sel
# else:
# num = np.argmax(sel_num_score.data.cpu().numpy(), axis=1) +1
# sel = sel_score.data.cpu().numpy()
# pr_sel = [list(np.argsort(-sel[b])[:num[b]]) for b in range(len(num))]
# agg_score = self.agg_pred(x_emb_var, x_len, col_inp_var, col_name_len, col_len, col_num, x_type_emb_var, gt_sel=pr_sel,
# gt_sel_num=pr_sel_num)
# cond_score = self.cond_pred(x_emb_var, x_len, col_inp_var, col_name_len, col_len, col_num, x_type_emb_var, gt_where, gt_cond)
# where_rela_score = self.where_rela_pred(x_emb_var, x_len, col_inp_var, col_name_len, col_len, col_num, x_type_emb_var)
return (sel_cond_score, agg_score, cond_op_str_score, where_rela_score)
def loss(self, score, truth_num, gt_where):
sel_cond_score, agg_score, cond_op_str_score, where_rela_score = score
sel_num_score, cond_num_score, sel_score, cond_col_score = sel_cond_score
cond_op_score, cond_str_score = cond_op_str_score
B = len(truth_num)
loss = 0
# Evaluate select number
sel_num_truth = map(lambda x: x[0], truth_num)
sel_num_truth = torch.from_numpy(np.array(sel_num_truth))
if self.gpu:
sel_num_truth = Variable(sel_num_truth.cuda())
else:
sel_num_truth = Variable(sel_num_truth)
loss += self.CE(sel_num_score, sel_num_truth)
# Evaluate select column
T = len(sel_score[0])
truth_prob = np.zeros((B, T), dtype=np.float32)
for b in range(B):
truth_prob[b][list(truth_num[b][1])] = 1
data = torch.from_numpy(truth_prob)
if self.gpu:
sel_col_truth_var = Variable(data.cuda())
else:
sel_col_truth_var = Variable(data)
sigm = nn.Sigmoid()
sel_col_prob = sigm(sel_score)
bce_loss = -torch.mean(
3 * (sel_col_truth_var * torch.log(sel_col_prob + 1e-10)) +
(1 - sel_col_truth_var) * torch.log(1 - sel_col_prob + 1e-10))
loss += bce_loss
# Evaluate select aggregation
for b in range(len(truth_num)):
data = torch.from_numpy(np.array(truth_num[b][2]))
if self.gpu:
sel_agg_truth_var = Variable(data.cuda())
else:
sel_agg_truth_var = Variable(data)
sel_agg_pred = agg_score[b, :len(truth_num[b][1])]
loss += (self.CE(sel_agg_pred, sel_agg_truth_var)) / len(truth_num)
# Evaluate the number of conditions
cond_num_truth = map(lambda x: x[3], truth_num)
data = torch.from_numpy(np.array(cond_num_truth))
if self.gpu:
try:
cond_num_truth_var = Variable(data.cuda())
except:
print "cond_num_truth_var error"
print data
exit(0)
else:
cond_num_truth_var = Variable(data)
loss += self.CE(cond_num_score, cond_num_truth_var)
# Evaluate the columns of conditions
T = len(cond_col_score[0])
truth_prob = np.zeros((B, T), dtype=np.float32)
for b in range(B):
if len(truth_num[b][4]) > 0:
truth_prob[b][list(truth_num[b][4])] = 1
data = torch.from_numpy(truth_prob)
if self.gpu:
cond_col_truth_var = Variable(data.cuda())
else:
cond_col_truth_var = Variable(data)
sigm = nn.Sigmoid()
cond_col_prob = sigm(cond_col_score)
bce_loss = -torch.mean(
3 * (cond_col_truth_var * torch.log(cond_col_prob + 1e-10)) +
(1 - cond_col_truth_var) * torch.log(1 - cond_col_prob + 1e-10))
loss += bce_loss
# Evaluate the operator of conditions
for b in range(len(truth_num)):
if len(truth_num[b][5]) == 0:
continue
data = torch.from_numpy(np.array(truth_num[b][5]))
if self.gpu:
cond_op_truth_var = Variable(data.cuda())
else:
cond_op_truth_var = Variable(data)
cond_op_pred = cond_op_score[b, :len(truth_num[b][5])]
try:
loss += (self.CE(cond_op_pred, cond_op_truth_var) /
len(truth_num))
except:
print cond_op_pred
print cond_op_truth_var
exit(0)
# Evaluate the strings of conditions
for b in range(len(gt_where)):
#print(gt_where[b])
for idx in range(len(gt_where[b])):
cond_str_truth = gt_where[b][idx]
#print("{}{}".format(cond_str_truth, len(cond_str_truth)))
if len(cond_str_truth) == 2:
continue
#print("cond_str_truth:{}".format(cond_str_truth[1:]))
#for tok in cond_str_truth[1:]:
#print("cond_str_tr_tok:{}".format(tok))
#print(' ')
data = torch.from_numpy(np.array(cond_str_truth[1:]))
#print("data:{}{}".format(data, data.shape))
if self.gpu:
cond_str_truth_var = Variable(data.cuda())
else:
cond_str_truth_var = Variable(data)
str_end = len(cond_str_truth) - 1
#print("cond_str_score:{} str_end:{}".format(cond_str_score.shape, str_end))
cond_str_pred = cond_str_score[b, idx, :str_end]
#print ("cond_str_score:{}cond_str_pred:{}".format(cond_str_score.shape, cond_str_pred.shape))
#print("cond_str_pred:{}".format(cond_str_pred))
loss += (self.CE(cond_str_pred, cond_str_truth_var) \
/ (len(gt_where) * len(gt_where[b])))
# Evaluate condition relationship, and / or
where_rela_truth = map(lambda x: x[6], truth_num)
data = torch.from_numpy(np.array(where_rela_truth))
if self.gpu:
try:
where_rela_truth = Variable(data.cuda())
except:
print "where_rela_truth error"
print data
exit(0)
else:
where_rela_truth = Variable(data)
loss += self.CE(where_rela_score, where_rela_truth)
return loss
# def loss(self, score, truth_num, pred_entry, gt_where): #edited by qwy
# pred_agg, pred_sel, pred_cond = pred_entry
# sel_num_score, agg_score, sel_cond_score, cond_op_str_score, where_rela_score = score
#
# cond_num_score, sel_score, cond_col_score = sel_cond_score
# cond_op_score, cond_str_score = cond_op_str_score
#
# loss = 0
#
# sel_num_truth = map(lambda x:x[0], truth_num)
# sel_num_truth = torch.from_numpy(np.array(sel_num_truth))
# if self.gpu:
# sel_num_truth = Variable(sel_num_truth.cuda())
# else:
# sel_num_truth = Variable(sel_num_truth)
# loss += self.CE(sel_num_score, sel_num_truth)
#
# if pred_sel:
# sel_truth = map(lambda x:x[1], truth_num)
# data = torch.from_numpy(np.array(sel_truth))
# if self.gpu:
# sel_truth_var = Variable(data.cuda())
# else:
# sel_truth_var = Variable(data)
#
# loss += self.CE(sel_score, sel_truth_var)
#
# if pred_agg:
# agg_truth = map(lambda x:x[2], truth_num)
# data = torch.from_numpy(np.array(agg_truth))
# if self.gpu:
# agg_truth_var = Variable(data.cuda())
# else:
# agg_truth_var = Variable(data)
#
# loss += self.CE(agg_score, agg_truth_var)
#
# if pred_cond:
# B = len(truth_num)
# #Evaluate the number of conditions
# cond_num_truth = map(lambda x:x[3], truth_num)
# data = torch.from_numpy(np.array(cond_num_truth))
# if self.gpu:
# cond_num_truth_var = Variable(data.cuda())
# else:
# cond_num_truth_var = Variable(data)
# loss += self.CE(cond_num_score, cond_num_truth_var)
#
# #Evaluate the columns of conditions
# T = len(cond_col_score[0])
# truth_prob = np.zeros((B, T), dtype=np.float32)
# for b in range(B):
# if len(truth_num[b][4]) > 0:
# truth_prob[b][list(truth_num[b][4])] = 1
# data = torch.from_numpy(truth_prob)
# if self.gpu:
# cond_col_truth_var = Variable(data.cuda())
# else:
# cond_col_truth_var = Variable(data)
#
# sigm = nn.Sigmoid()
# cond_col_prob = sigm(cond_col_score)
# bce_loss = -torch.mean( 3*(cond_col_truth_var * \
# torch.log(cond_col_prob+1e-10)) + \
# (1-cond_col_truth_var) * torch.log(1-cond_col_prob+1e-10) )
# loss += bce_loss
#
# #Evaluate the operator of conditions
# for b in range(len(truth_num)):
# if len(truth_num[b][5]) == 0:
# continue
# data = torch.from_numpy(np.array(truth_num[b][5]))
# if self.gpu:
# cond_op_truth_var = Variable(data.cuda())
# else:
# cond_op_truth_var = Variable(data)
# cond_op_pred = cond_op_score[b, :len(truth_num[b][5])]
# loss += (self.CE(cond_op_pred, cond_op_truth_var) \
# / len(truth_num))
#
# #Evaluate the strings of conditions
# for b in range(len(gt_where)):
# for idx in range(len(gt_where[b])):
# cond_str_truth = gt_where[b][idx]
# if len(cond_str_truth) == 1:
# continue
# data = torch.from_numpy(np.array(cond_str_truth[1:]))
# if self.gpu:
# cond_str_truth_var = Variable(data.cuda())
# else:
# cond_str_truth_var = Variable(data)
# str_end = len(cond_str_truth)-1
# cond_str_pred = cond_str_score[b, idx, :str_end]
# loss += (self.CE(cond_str_pred, cond_str_truth_var) \
# / (len(gt_where) * len(gt_where[b])))
#
# where_rela_truth = map(lambda x: x[6], truth_num)
# data = torch.from_numpy(np.array(where_rela_truth))
# if self.gpu:
# try:
# where_rela_truth = Variable(data.cuda())
# except:
# print "where_rela_truth error"
# print data
# exit(0)
# else:
# where_rela_truth = Variable(data)
# loss += self.CE(where_rela_score, where_rela_truth)
#
# return loss
def check_acc(self, vis_info, pred_queries, gt_queries):
def pretty_print(vis_data, pred_query, gt_query):
print "\n----------detailed error prints-----------"
try:
print 'question: ', vis_data[0]
print 'question_tok: ', vis_data[3]
print 'headers: (%s)' % (' || '.join(vis_data[1]))
print 'query:', vis_data[2]
print "target query: ", gt_query
print "pred query: ", pred_query
except:
print "\n------skipping print: decoding problem ----------------------"
def gen_cond_str(conds, header):
if len(conds) == 0:
return 'None'
cond_str = []
for cond in conds:
cond_str.append(header[cond[0]] + ' ' +
self.COND_OPS[cond[1]] + ' ' +
unicode(cond[2]).lower())
return 'WHERE ' + ' AND '.join(cond_str)
"""
B = len(gt_queries)
tot_err = agg_err = sel_err = cond_err = 0.0
cond_num_err = cond_col_err = cond_op_err = cond_val_err = 0.0
agg_ops = ['None', 'MAX', 'MIN', 'COUNT', 'SUM', 'AVG']
for b, (pred_qry, gt_qry, vis_data) in enumerate(zip(pred_queries, gt_queries, vis_info)):
good = True
if pred_agg:
agg_pred = pred_qry['agg']
agg_gt = gt_qry['agg']
if agg_pred != agg_gt:
agg_err += 1
good = False
if pred_sel:
sel_pred = pred_qry['sel']
sel_gt = gt_qry['sel']
if sel_pred != sel_gt:
sel_err += 1
good = False
if pred_cond:
cond_pred = pred_qry['conds']
cond_gt = gt_qry['conds']
flag = True
if len(cond_pred) != len(cond_gt):
flag = False
cond_num_err += 1
if flag and set(x[0] for x in cond_pred) != \
set(x[0] for x in cond_gt):
flag = False
cond_col_err += 1
for idx in range(len(cond_pred)):
if not flag:
break
gt_idx = tuple(x[0] for x in cond_gt).index(cond_pred[idx][0])
if flag and cond_gt[gt_idx][1] != cond_pred[idx][1]:
flag = False
cond_op_err += 1
for idx in range(len(cond_pred)):
if not flag:
break
gt_idx = tuple(
x[0] for x in cond_gt).index(cond_pred[idx][0])
if flag and unicode(cond_gt[gt_idx][2]).lower() != \
unicode(cond_pred[idx][2]).lower():
flag = False
cond_val_err += 1
if not flag:
cond_err += 1
good = False
if not good:
if error_print:
pretty_print(vis_data, pred_qry, gt_qry)
tot_err += 1
return np.array((agg_err, sel_err, cond_err, cond_num_err, cond_col_err, cond_op_err, cond_val_err)), tot_err
"""
tot_err = sel_num_err = agg_err = sel_err = 0.0
cond_num_err = cond_col_err = cond_op_err = cond_val_err = cond_rela_err = 0.0
for b, (pred_qry, gt_qry) in enumerate(zip(pred_queries, gt_queries)):
good = True
sel_pred, agg_pred, where_rela_pred = pred_qry['sel'], pred_qry[
'agg'], pred_qry['cond_conn_op']
sel_gt, agg_gt, where_rela_gt = gt_qry['sel'], gt_qry[
'agg'], gt_qry['cond_conn_op']
if where_rela_gt != where_rela_pred:
good = False
cond_rela_err += 1
if len(sel_pred) != len(sel_gt):
good = False
sel_num_err += 1
pred_sel_dict = {
k: v
for k, v in zip(list(sel_pred), list(agg_pred))
}
gt_sel_dict = {k: v for k, v in zip(sel_gt, agg_gt)}
if set(sel_pred) != set(sel_gt):
good = False
sel_err += 1
agg_pred = [pred_sel_dict[x] for x in sorted(pred_sel_dict.keys())]
agg_gt = [gt_sel_dict[x] for x in sorted(gt_sel_dict.keys())]
if agg_pred != agg_gt:
good = False
agg_err += 1
cond_pred = pred_qry['conds']
cond_gt = gt_qry['conds']
if len(cond_pred) != len(cond_gt):
good = False
cond_num_err += 1
else:
cond_op_pred, cond_op_gt = {}, {}
cond_val_pred, cond_val_gt = {}, {}
for p, g in zip(cond_pred, cond_gt):
cond_op_pred[p[0]] = p[1]
cond_val_pred[p[0]] = p[2]
cond_op_gt[g[0]] = g[1]
cond_val_gt[g[0]] = g[2]
if set(cond_op_pred.keys()) != set(cond_op_gt.keys()):
cond_col_err += 1
good = False
where_op_pred = [
cond_op_pred[x] for x in sorted(cond_op_pred.keys())
]
where_op_gt = [
cond_op_gt[x] for x in sorted(cond_op_gt.keys())
]
if where_op_pred != where_op_gt:
cond_op_err += 1
good = False
where_val_pred = [
cond_val_pred[x] for x in sorted(cond_val_pred.keys())
]
where_val_gt = [
cond_val_gt[x] for x in sorted(cond_val_gt.keys())
]
if where_val_pred != where_val_gt:
cond_val_err += 1
good = False
if not good:
tot_err += 1
return np.array(
(sel_num_err, sel_err, agg_err, cond_num_err, cond_col_err,
cond_op_err, cond_val_err, cond_rela_err)), tot_err
def gen_query(self, score, q, col, raw_q, reinforce=False, verbose=False):
def merge_tokens(tok_list, raw_tok_str):
"""
tok_list: list of string words in current cond
raw_tok_str: list of words in question
"""
tok_str = raw_tok_str.lower()
alphabet = 'abcdefghijklmnopqrstuvwxyz0123456789$('
special = {
'-LRB-': '(',
'-RRB-': ')',
'-LSB-': '[',
'-RSB-': ']',
'``': '"',
'\'\'': '"',
'--': u'\u2013'
}
ret = ''
double_quote_appear = 0
tok_list = [x for gx in tok_list for x in gx]
for raw_tok in tok_list:
if not raw_tok:
continue
tok = special.get(raw_tok, raw_tok)
if tok == '"':
double_quote_appear = 1 - double_quote_appear
if len(ret) == 0:
pass
elif len(ret) > 0 and ret + ' ' + tok in tok_str:
ret = ret + ' '
elif len(ret) > 0 and ret + tok in tok_str:
pass
elif tok == '"':
if double_quote_appear:
ret = ret + ' '
elif tok[0] not in alphabet:
pass
elif (ret[-1] not in ['(', '/', u'\u2013', '#', '$', '&']) \
and (ret[-1] != '"' or not double_quote_appear):
ret = ret + ' '
ret = ret + tok
return ret.strip()
sel_cond_score, agg_score, cond_op_str_score, where_rela_score = score
sel_num_score, cond_num_score, sel_score, cond_col_score \
= [x.data.cpu().numpy() for x in sel_cond_score]
cond_op_score, cond_str_score = [
x.data.cpu().numpy() for x in cond_op_str_score
]
# [64,4,6], [64,14], ..., [64,4]
#sel_num_score = sel_num_score.data.cpu().numpy()
#sel_score = sel_score.data.cpu().numpy()
agg_score = agg_score.data.cpu().numpy()
where_rela_score = where_rela_score.data.cpu().numpy()
ret_queries = []
B = len(agg_score)
for b in range(B):
cur_query = {}
cur_query['sel'] = []
cur_query['agg'] = []
sel_num = np.argmax(sel_num_score[b])
max_col_idxes = np.argsort(-sel_score[b])[:sel_num]
# find the most-probable columns' indexes
max_agg_idxes = np.argsort(-agg_score[b])[:sel_num]
cur_query['sel'].extend([int(i) for i in max_col_idxes])
cur_query['agg'].extend([i[0] for i in max_agg_idxes])
cur_query['cond_conn_op'] = np.argmax(where_rela_score[b])
cur_query['conds'] = []
cond_num = np.argmax(cond_num_score[b])
all_toks = ['<BEG>'] + q[b] + ['<END>']
max_idxes = np.argsort(-cond_col_score[b])[:cond_num]
for idx in range(cond_num):
cur_cond = []
cur_cond.append(max_idxes[idx]) # where-col
cur_cond.append(np.argmax(cond_op_score[b][idx])) # where-op
cur_cond_str_toks = []
for str_score in cond_str_score[b][idx]:
str_tok = np.argmax(str_score[:len(all_toks)])
str_val = all_toks[str_tok]
if str_val == '<END>':
break
cur_cond_str_toks.append(str_val)
cur_cond.append(merge_tokens(cur_cond_str_toks, raw_q[b]))
cur_query['conds'].append(cur_cond)
ret_queries.append(cur_query)
return ret_queries