waynesun 4 years ago
parent
commit
710220adbe

+ 5 - 9
sqlnet/model/modules/aggregator_predict.py

@@ -4,7 +4,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 from torch.autograd import Variable
 import numpy as np
-from net_utils import run_lstm, col_name_encode
+from sqlnet.model.modules.net_utils import run_lstm, col_name_encode
 
 
 
@@ -13,17 +13,13 @@ class AggPredictor(nn.Module):
         super(AggPredictor, self).__init__()
         self.use_ca = use_ca
 
-        self.agg_lstm = nn.LSTM(input_size=N_word, hidden_size=N_h/2,
-                num_layers=N_depth, batch_first=True,
-                dropout=0.3, bidirectional=True)
+        self.agg_lstm = nn.LSTM(input_size=N_word, hidden_size=int(N_h/2), num_layers=N_depth, batch_first=True, dropout=0.3, bidirectional=True)
         if use_ca:
-            print "Using column attention on aggregator predicting"
-            self.agg_col_name_enc = nn.LSTM(input_size=N_word,
-                    hidden_size=N_h/2, num_layers=N_depth,
-                    batch_first=True, dropout=0.3, bidirectional=True)
+            print ("Using column attention on aggregator predicting")
+            self.agg_col_name_enc = nn.LSTM(input_size=N_word, hidden_size=int(N_h/2), num_layers=N_depth, batch_first=True, dropout=0.3, bidirectional=True)
             self.agg_att = nn.Linear(N_h, N_h)
         else:
-            print "Not using column attention on aggregator predicting"
+            print ("Not using column attention on aggregator predicting")
             self.agg_att = nn.Linear(N_h, 1)
         self.agg_out = nn.Sequential(nn.Linear(N_h, N_h), nn.Tanh(), nn.Linear(N_h, 6))
         self.softmax = nn.Softmax(dim=-1)

+ 10 - 13
sqlnet/model/modules/select_number.py

@@ -4,7 +4,8 @@ import torch.nn as nn
 import torch.nn.functional as F
 from torch.autograd import Variable
 import numpy as np
-from net_utils import run_lstm, col_name_encode
+
+from sqlnet.model.modules.net_utils import run_lstm, col_name_encode
 
 class SelNumPredictor(nn.Module):
     def __init__(self, N_word, N_h, N_depth, use_ca):
@@ -12,20 +13,18 @@ class SelNumPredictor(nn.Module):
         self.N_h = N_h
         self.use_ca = use_ca
 
-        self.sel_num_lstm = nn.LSTM(input_size=N_word, hidden_size=N_h/2,
-                                    num_layers=N_depth, batch_first=True,
-                                    dropout=0.3, bidirectional=True)
+        self.sel_num_lstm = nn.LSTM(input_size=N_word, hidden_size=int(N_h/2), num_layers=N_depth, batch_first=True, dropout=0.3, bidirectional=True)
+
         self.sel_num_att = nn.Linear(N_h, 1)
         self.sel_num_col_att = nn.Linear(N_h, 1)
-        self.sel_num_out = nn.Sequential(nn.Linear(N_h, N_h),
-                                         nn.Tanh(), nn.Linear(N_h,4))
+        self.sel_num_out = nn.Sequential(nn.Linear(N_h, N_h), nn.Tanh(), nn.Linear(N_h,4))
         self.softmax = nn.Softmax(dim=-1)
         self.sel_num_col2hid1 = nn.Linear(N_h, 2 * N_h)
         self.sel_num_col2hid2 = nn.Linear(N_h, 2 * N_h)
 
 
         if self.use_ca:
-            print "Using column attention on select number predicting"
+            print ("Using column attention on select number predicting")
 
     def forward(self, x_emb_var, x_len, col_inp_var, col_name_len, col_len, col_num):
         B = len(x_len)
@@ -42,11 +41,10 @@ class SelNumPredictor(nn.Module):
                 num_col_att_val[idx, num:] = -1000000
         num_col_att = self.softmax(num_col_att_val)
         K_num_col = (e_num_col * num_col_att.unsqueeze(2)).sum(1)
-        sel_num_h1 = self.sel_num_col2hid1(K_num_col).view(B, 4, self.N_h/2).transpose(0,1).contiguous()
-        sel_num_h2 = self.sel_num_col2hid2(K_num_col).view(B, 4, self.N_h/2).transpose(0,1).contiguous()
+        sel_num_h1 = self.sel_num_col2hid1(K_num_col).view((B, 4, self.N_h//2)).transpose(0,1).contiguous()
+        sel_num_h2 = self.sel_num_col2hid2(K_num_col).view((B, 4, self.N_h//2)).transpose(0,1).contiguous()
 
-        h_num_enc, _ = run_lstm(self.sel_num_lstm, x_emb_var, x_len,
-                                hidden=(sel_num_h1, sel_num_h2))
+        h_num_enc, _ = run_lstm(self.sel_num_lstm, x_emb_var, x_len,hidden=(sel_num_h1, sel_num_h2))
 
         num_att_val = self.sel_num_att(h_num_enc).squeeze()
         for idx, num in enumerate(x_len):
@@ -54,8 +52,7 @@ class SelNumPredictor(nn.Module):
                 num_att_val[idx, num:] = -1000000
         num_att = self.softmax(num_att_val)
 
-        K_sel_num = (h_num_enc * num_att.unsqueeze(2).expand_as(
-            h_num_enc)).sum(1)
+        K_sel_num = (h_num_enc * num_att.unsqueeze(2).expand_as(h_num_enc)).sum(1)
         sel_num_score = self.sel_num_out(K_sel_num)
         return sel_num_score
 

+ 5 - 9
sqlnet/model/modules/selection_predict.py

@@ -4,25 +4,21 @@ import torch.nn as nn
 import torch.nn.functional as F
 from torch.autograd import Variable
 import numpy as np
-from net_utils import run_lstm, col_name_encode
+from sqlnet.model.modules.net_utils import run_lstm, col_name_encode
 
 class SelPredictor(nn.Module):
     def __init__(self, N_word, N_h, N_depth, max_tok_num, use_ca):
         super(SelPredictor, self).__init__()
         self.use_ca = use_ca
         self.max_tok_num = max_tok_num
-        self.sel_lstm = nn.LSTM(input_size=N_word, hidden_size=N_h/2,
-                num_layers=N_depth, batch_first=True,
-                dropout=0.3, bidirectional=True)
+        self.sel_lstm = nn.LSTM(input_size=N_word, hidden_size=int(N_h/2), num_layers=N_depth, batch_first=True, dropout=0.3, bidirectional=True)
         if use_ca:
-            print "Using column attention on selection predicting"
+            print ("Using column attention on selection predicting")
             self.sel_att = nn.Linear(N_h, N_h)
         else:
-            print "Not using column attention on selection predicting"
+            print ("Not using column attention on selection predicting")
             self.sel_att = nn.Linear(N_h, 1)
-        self.sel_col_name_enc = nn.LSTM(input_size=N_word, hidden_size=N_h/2,
-                num_layers=N_depth, batch_first=True,
-                dropout=0.3, bidirectional=True)
+        self.sel_col_name_enc = nn.LSTM(input_size=N_word, hidden_size=int(N_h/2), num_layers=N_depth, batch_first=True, dropout=0.3, bidirectional=True)
         self.sel_out_K = nn.Linear(N_h, N_h)
         self.sel_out_col = nn.Linear(N_h, N_h)
         self.sel_out = nn.Sequential(nn.Tanh(), nn.Linear(N_h, 1))

+ 16 - 36
sqlnet/model/modules/sqlnet_condition_predict.py

@@ -4,7 +4,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 from torch.autograd import Variable
 import numpy as np
-from net_utils import run_lstm, col_name_encode
+from sqlnet.model.modules.net_utils import run_lstm, col_name_encode
 
 class SQLNetCondPredictor(nn.Module):
     def __init__(self, N_word, N_h, N_depth, max_col_num, max_tok_num, use_ca, gpu):
@@ -15,59 +15,41 @@ class SQLNetCondPredictor(nn.Module):
         self.gpu = gpu
         self.use_ca = use_ca
 
-        self.cond_num_lstm = nn.LSTM(input_size=N_word, hidden_size=N_h/2,
-                num_layers=N_depth, batch_first=True,
-                dropout=0.3, bidirectional=True)
+        self.cond_num_lstm = nn.LSTM(input_size=N_word, hidden_size=int(N_h/2), num_layers=N_depth, batch_first=True, dropout=0.3, bidirectional=True)
         self.cond_num_att = nn.Linear(N_h, 1)
         self.cond_num_out = nn.Sequential(nn.Linear(N_h, N_h),
                 nn.Tanh(), nn.Linear(N_h, 5))
-        self.cond_num_name_enc = nn.LSTM(input_size=N_word, hidden_size=N_h/2,
-                num_layers=N_depth, batch_first=True,
-                dropout=0.3, bidirectional=True)
+        self.cond_num_name_enc = nn.LSTM(input_size=N_word, hidden_size=int(N_h/2), num_layers=N_depth, batch_first=True, dropout=0.3, bidirectional=True)
         self.cond_num_col_att = nn.Linear(N_h, 1)
         self.cond_num_col2hid1 = nn.Linear(N_h, 2*N_h)
         self.cond_num_col2hid2 = nn.Linear(N_h, 2*N_h)
 
-        self.cond_col_lstm = nn.LSTM(input_size=N_word, hidden_size=N_h/2,
-                num_layers=N_depth, batch_first=True,
-                dropout=0.3, bidirectional=True)
+        self.cond_col_lstm = nn.LSTM(input_size=N_word, hidden_size=int(N_h/2), num_layers=N_depth, batch_first=True, dropout=0.3, bidirectional=True)
         if use_ca:
-            print "Using column attention on where predicting"
+            print ("Using column attention on where predicting")
             self.cond_col_att = nn.Linear(N_h, N_h)
         else:
-            print "Not using column attention on where predicting"
+            print ("Not using column attention on where predicting")
             self.cond_col_att = nn.Linear(N_h, 1)
-        self.cond_col_name_enc = nn.LSTM(input_size=N_word, hidden_size=N_h/2,
-                num_layers=N_depth, batch_first=True,
-                dropout=0.3, bidirectional=True)
+        self.cond_col_name_enc = nn.LSTM(input_size=N_word, hidden_size=int(N_h/2), num_layers=N_depth, batch_first=True, dropout=0.3, bidirectional=True)
         self.cond_col_out_K = nn.Linear(N_h, N_h)
         self.cond_col_out_col = nn.Linear(N_h, N_h)
         self.cond_col_out = nn.Sequential(nn.ReLU(), nn.Linear(N_h, 1))
 
-        self.cond_op_lstm = nn.LSTM(input_size=N_word, hidden_size=N_h/2,
-                num_layers=N_depth, batch_first=True,
-                dropout=0.3, bidirectional=True)
+        self.cond_op_lstm = nn.LSTM(input_size=N_word, hidden_size=int(N_h/2), num_layers=N_depth, batch_first=True, dropout=0.3, bidirectional=True)
         if use_ca:
             self.cond_op_att = nn.Linear(N_h, N_h)
         else:
             self.cond_op_att = nn.Linear(N_h, 1)
         self.cond_op_out_K = nn.Linear(N_h, N_h)
-        self.cond_op_name_enc = nn.LSTM(input_size=N_word, hidden_size=N_h/2,
-                num_layers=N_depth, batch_first=True,
-                dropout=0.3, bidirectional=True)
+        self.cond_op_name_enc = nn.LSTM(input_size=N_word, hidden_size=int(N_h/2), num_layers=N_depth, batch_first=True, dropout=0.3, bidirectional=True)
         self.cond_op_out_col = nn.Linear(N_h, N_h)
         self.cond_op_out = nn.Sequential(nn.Linear(N_h, N_h), nn.Tanh(),
                 nn.Linear(N_h, 4))
 
-        self.cond_str_lstm = nn.LSTM(input_size=N_word, hidden_size=N_h/2,
-                num_layers=N_depth, batch_first=True,
-                dropout=0.3, bidirectional=True)
-        self.cond_str_decoder = nn.LSTM(input_size=self.max_tok_num,
-                hidden_size=N_h, num_layers=N_depth,
-                batch_first=True, dropout=0.3)
-        self.cond_str_name_enc = nn.LSTM(input_size=N_word, hidden_size=N_h/2,
-                num_layers=N_depth, batch_first=True,
-                dropout=0.3, bidirectional=True)
+        self.cond_str_lstm = nn.LSTM(input_size=N_word, hidden_size=int(N_h/2), num_layers=N_depth, batch_first=True, dropout=0.3, bidirectional=True)
+        self.cond_str_decoder = nn.LSTM(input_size=self.max_tok_num, hidden_size=N_h, num_layers=N_depth, batch_first=True, dropout=0.3)
+        self.cond_str_name_enc = nn.LSTM(input_size=N_word, hidden_size=int(N_h/2), num_layers=N_depth, batch_first=True, dropout=0.3, bidirectional=True)
         self.cond_str_out_g = nn.Linear(N_h, N_h)
         self.cond_str_out_h = nn.Linear(N_h, N_h)
         self.cond_str_out_col = nn.Linear(N_h, N_h)
@@ -78,7 +60,7 @@ class SQLNetCondPredictor(nn.Module):
 
     def gen_gt_batch(self, split_tok_seq):
         B = len(split_tok_seq)
-        max_len = max([max([len(tok) for tok in tok_seq]+[0]) for 
+        max_len = max([max([len(tok) for tok in tok_seq]+[0]) for
             tok_seq in split_tok_seq]) - 1 # The max seq len in the batch.
         if max_len < 1:
             max_len = 1
@@ -121,10 +103,8 @@ class SQLNetCondPredictor(nn.Module):
                 num_col_att_val[idx, num:] = -100
         num_col_att = self.softmax(num_col_att_val)
         K_num_col = (e_num_col * num_col_att.unsqueeze(2)).sum(1)
-        cond_num_h1 = self.cond_num_col2hid1(K_num_col).view(
-                B, 4, self.N_h/2).transpose(0, 1).contiguous()
-        cond_num_h2 = self.cond_num_col2hid2(K_num_col).view(
-                B, 4, self.N_h/2).transpose(0, 1).contiguous()
+        cond_num_h1 = self.cond_num_col2hid1(K_num_col).view(B, 4, self.N_h//2).transpose(0, 1).contiguous()
+        cond_num_h2 = self.cond_num_col2hid2(K_num_col).view(B, 4, self.N_h//2).transpose(0, 1).contiguous()
 
         h_num_enc, _ = run_lstm(self.cond_num_lstm, x_emb_var, x_len,
                 hidden=(cond_num_h1, cond_num_h2))
@@ -185,7 +165,7 @@ class SQLNetCondPredictor(nn.Module):
         h_op_enc, _ = run_lstm(self.cond_op_lstm, x_emb_var, x_len)
         col_emb = []
         for b in range(B):
-            cur_col_emb = torch.stack([e_cond_col[b, x] 
+            cur_col_emb = torch.stack([e_cond_col[b, x]
                 for x in chosen_col_gt[b]] + [e_cond_col[b, 0]] *
                 (4 - len(chosen_col_gt[b])))  # Pad the columns to maximum (4)
             col_emb.append(cur_col_emb)

+ 5 - 7
sqlnet/model/modules/where_relation.py

@@ -4,7 +4,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 from torch.autograd import Variable
 import numpy as np
-from net_utils import run_lstm, col_name_encode
+from sqlnet.model.modules.net_utils import run_lstm, col_name_encode
 
 class WhereRelationPredictor(nn.Module):
     def __init__(self, N_word, N_h, N_depth, use_ca):
@@ -12,9 +12,7 @@ class WhereRelationPredictor(nn.Module):
         self.N_h = N_h
         self.use_ca = use_ca
 
-        self.where_rela_lstm = nn.LSTM(input_size=N_word, hidden_size=N_h/2,
-                                    num_layers=N_depth, batch_first=True,
-                                    dropout=0.3, bidirectional=True)
+        self.where_rela_lstm = nn.LSTM(input_size=N_word, hidden_size=int(N_h/2), num_layers=N_depth, batch_first=True, dropout=0.3, bidirectional=True)
         self.where_rela_att = nn.Linear(N_h, 1)
         self.where_rela_col_att = nn.Linear(N_h, 1)
         self.where_rela_out = nn.Sequential(nn.Linear(N_h, N_h), nn.Tanh(), nn.Linear(N_h,3))
@@ -23,7 +21,7 @@ class WhereRelationPredictor(nn.Module):
         self.col2hid2 = nn.Linear(N_h, 2 * N_h)
 
         if self.use_ca:
-            print "Using column attention on where relation predicting"
+            print ("Using column attention on where relation predicting")
 
     def forward(self, x_emb_var, x_len, col_inp_var, col_name_len, col_len, col_num):
         B = len(x_len)
@@ -40,8 +38,8 @@ class WhereRelationPredictor(nn.Module):
                 col_att_val[idx, num:] = -1000000
         num_col_att = self.softmax(col_att_val)
         K_num_col = (e_num_col * num_col_att.unsqueeze(2)).sum(1)
-        h1 = self.col2hid1(K_num_col).view(B, 4, self.N_h/2).transpose(0,1).contiguous()
-        h2 = self.col2hid2(K_num_col).view(B, 4, self.N_h/2).transpose(0,1).contiguous()
+        h1 = self.col2hid1(K_num_col).view(B, 4, self.N_h//2).transpose(0,1).contiguous()
+        h2 = self.col2hid2(K_num_col).view(B, 4, self.N_h//2).transpose(0,1).contiguous()
 
         h_enc, _ = run_lstm(self.where_rela_lstm, x_emb_var, x_len, hidden=(h1, h2))
 

+ 13 - 22
sqlnet/model/modules/word_embedding.py

@@ -15,15 +15,14 @@ class WordEmbedding(nn.Module):
         self.SQL_TOK = SQL_TOK
 
         if trainable:
-            print "Using trainable embedding"
+            print ("Using trainable embedding")
             self.w2i, word_emb_val = word_emb
             self.embedding = nn.Embedding(len(self.w2i), N_word)
             self.embedding.weight = nn.Parameter(
                     torch.from_numpy(word_emb_val.astype(np.float32)))
         else:
             self.word_emb = word_emb
-            print "Using fixed embedding"
-
+            print ("Using fixed embedding")
 
     def gen_x_batch(self, q, col):
         B = len(q)
@@ -31,24 +30,17 @@ class WordEmbedding(nn.Module):
         val_len = np.zeros(B, dtype=np.int64)
         for i, (one_q, one_col) in enumerate(zip(q, col)):
             if self.trainable:
-                q_val = map(lambda x:self.w2i.get(x, 0), one_q)
-            else:
-                q_val = map(lambda x:self.word_emb.get(x, np.zeros(self.N_word, dtype=np.float32)), one_q)
-            if self.our_model:
-                if self.trainable:
-                    val_embs.append([1] + q_val + [2])  #<BEG> and <END>
-                else:
-                    val_embs.append([np.zeros(self.N_word, dtype=np.float32)] + q_val + [np.zeros(self.N_word, dtype=np.float32)])  #<BEG> and <END>
-                val_len[i] = 1 + len(q_val) + 1
+                q_val = [self.w2i.get(x,0) for x in one_q]
+                val_embs.append([1] + q_val + [2])  #<BEG> and <END>
             else:
-                one_col_all = [x for toks in one_col for x in toks+[',']]
-                if self.trainable:
-                    col_val = map(lambda x:self.w2i.get(x, 0), one_col_all)
-                    val_embs.append( [0 for _ in self.SQL_TOK] + col_val + [0] + q_val+ [0])
-                else:
-                    col_val = map(lambda x:self.word_emb.get(x, np.zeros(self.N_word, dtype=np.float32)), one_col_all)
-                    val_embs.append( [np.zeros(self.N_word, dtype=np.float32) for _ in self.SQL_TOK] + col_val + [np.zeros(self.N_word, dtype=np.float32)] + q_val+ [np.zeros(self.N_word, dtype=np.float32)])
-                val_len[i] = len(self.SQL_TOK) + len(col_val) + 1 + len(q_val) + 1
+                # print (i)
+                # print ([x.encode('utf-8') for x in one_q])
+                q_val = [self.word_emb.get(x, np.zeros(self.N_word, dtype=np.float32)) for x in one_q]
+                # print (q_val)
+                # print ("#"*60)
+                val_embs.append([np.zeros(self.N_word, dtype=np.float32)] + q_val + [np.zeros(self.N_word, dtype=np.float32)])  #<BEG> and <END>
+            # exit(0)
+            val_len[i] = len(q_val) + 2
         max_len = max(val_len)
 
         if self.trainable:
@@ -93,8 +85,7 @@ class WordEmbedding(nn.Module):
             if self.trainable:
                 val = [self.w2i.get(x, 0) for x in one_str]
             else:
-                val = [self.word_emb.get(x, np.zeros(
-                    self.N_word, dtype=np.float32)) for x in one_str]
+                val = [self.word_emb.get(x, np.zeros(self.N_word, dtype=np.float32)) for x in one_str]
             val_embs.append(val)
             val_len[i] = len(val)
         max_len = max(val_len)

+ 20 - 17
sqlnet/model/sqlnet.py

@@ -3,12 +3,12 @@ import torch.nn as nn
 import torch.nn.functional as F
 from torch.autograd import Variable
 import numpy as np
-from modules.word_embedding import WordEmbedding
-from modules.aggregator_predict import AggPredictor
-from modules.selection_predict import SelPredictor
-from modules.sqlnet_condition_predict import SQLNetCondPredictor
-from modules.select_number import SelNumPredictor
-from modules.where_relation import WhereRelationPredictor
+from sqlnet.model.modules.word_embedding import WordEmbedding
+from sqlnet.model.modules.aggregator_predict import AggPredictor
+from sqlnet.model.modules.selection_predict import SelPredictor
+from sqlnet.model.modules.sqlnet_condition_predict import SQLNetCondPredictor
+from sqlnet.model.modules.select_number import SelNumPredictor
+from sqlnet.model.modules.where_relation import WhereRelationPredictor
 
 
 class SQLNet(nn.Module):
@@ -140,7 +140,8 @@ class SQLNet(nn.Module):
         loss = 0
 
         # Evaluate select number
-        sel_num_truth = map(lambda x:x[0], truth_num)
+        # sel_num_truth = map(lambda x:x[0], truth_num)
+        sel_num_truth = [x[0] for x in truth_num]
         sel_num_truth = torch.from_numpy(np.array(sel_num_truth))
         if self.gpu:
             sel_num_truth = Variable(sel_num_truth.cuda())
@@ -173,20 +174,21 @@ class SQLNet(nn.Module):
                 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)
+        sel_agg_pred = agg_score[b, :len(truth_num[b][1])]
+        loss += (self.CE(sel_agg_pred, sel_agg_truth_var)) / len(truth_num)
 
         cond_num_score, cond_col_score, cond_op_score, cond_str_score = cond_score
 
         # Evaluate the number of conditions
-        cond_num_truth = map(lambda x:x[3], truth_num)
+        # cond_num_truth = map(lambda x:x[3], truth_num)
+        cond_num_truth = [x[3] for x in 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
+                print ("cond_num_truth_var error")
+                print (data)
                 exit(0)
         else:
             cond_num_truth_var = Variable(data)
@@ -224,8 +226,8 @@ class SQLNet(nn.Module):
             try:
                 loss += (self.CE(cond_op_pred, cond_op_truth_var) / len(truth_num))
             except:
-                print cond_op_pred
-                print cond_op_truth_var
+                print (cond_op_pred)
+                print (cond_op_truth_var)
                 exit(0)
 
         #Evaluate the strings of conditions
@@ -245,14 +247,15 @@ class SQLNet(nn.Module):
                         / (len(gt_where) * len(gt_where[b])))
 
         # Evaluate condition relationship, and / or
-        where_rela_truth = map(lambda x:x[6], truth_num)
+        # where_rela_truth = map(lambda x:x[6], truth_num)
+        where_rela_truth = [x[6] for x in 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
+                print ("where_rela_truth error")
+                print (data)
                 exit(0)
         else:
             where_rela_truth = Variable(data)

+ 18 - 13
sqlnet/utils.py

@@ -1,5 +1,5 @@
 import json
-from lib.dbengine import DBEngine
+from sqlnet.lib.dbengine import DBEngine
 import numpy as np
 from tqdm import tqdm
 
@@ -12,20 +12,20 @@ def load_data(sql_paths, table_paths, use_small=False):
     table_data = {}
 
     for SQL_PATH in sql_paths:
-        with open(SQL_PATH) as inf:
+        with open(SQL_PATH, encoding='utf-8') as inf:
             for idx, line in enumerate(inf):
                 sql = json.loads(line.strip())
                 if use_small and idx >= 1000:
                     break
                 sql_data.append(sql)
-        print "Loaded %d data from %s" % (len(sql_data), SQL_PATH)
+        print ("Loaded %d data from %s" % (len(sql_data), SQL_PATH))
 
     for TABLE_PATH in table_paths:
-        with open(TABLE_PATH) as inf:
+        with open(TABLE_PATH, encoding='utf-8') as inf:
             for line in inf:
                 tab = json.loads(line.strip())
                 table_data[tab[u'id']] = tab
-        print "Loaded %d data from %s" % (len(table_data), TABLE_PATH)
+        print ("Loaded %d data from %s" % (len(table_data), TABLE_PATH))
 
     ret_sql_data = []
     for sql in sql_data:
@@ -35,7 +35,7 @@ def load_data(sql_paths, table_paths, use_small=False):
     return ret_sql_data, table_data
 
 def load_dataset(toy=False, use_small=False, mode='train'):
-    print "Loading dataset"
+    print ("Loading dataset")
     dev_sql, dev_table = load_data('data/val/val.json', 'data/val/val.tables.json', use_small=use_small)
     dev_db = 'data/val/val.db'
     if mode == 'train':
@@ -107,6 +107,7 @@ def to_batch_query(sql_data, idxes, st, ed):
 def epoch_train(model, optimizer, batch_size, sql_data, table_data):
     model.train()
     perm=np.random.permutation(len(sql_data))
+    perm = list(range(len(sql_data)))
     cum_loss = 0.0
     for st in tqdm(range(len(sql_data)//batch_size+1)):
         ed = (st+1)*batch_size if (st+1)*batch_size < len(perm) else len(perm)
@@ -173,7 +174,7 @@ def epoch_acc(model, batch_size, sql_data, table_data, db_path):
             one_err, tot_err = model.check_acc(raw_data, pred_queries, query_gt)
         except:
             badcase += 1
-            print 'badcase', badcase
+            print ('badcase', badcase)
             continue
         one_acc_num += (ed-st-one_err)
         tot_acc_num += (ed-st-tot_err)
@@ -191,10 +192,14 @@ def epoch_acc(model, batch_size, sql_data, table_data, db_path):
 
 def load_word_emb(file_name):
     print ('Loading word embedding from %s'%file_name)
-    ret = {}
-    with open(file_name) as inf:
-        for idx, line in enumerate(inf):
-            info = line.strip().split(' ')
-            if info[0].lower() not in ret:
-                ret[info[0].decode('utf-8')] = np.array(map(lambda x:float(x), info[1:]))
+    f = open(file_name)
+    ret = json.load(f)
+    f.close()
+    # ret = {}
+    # with open(file_name, encoding='latin') as inf:
+    #     ret = json.load(inf)
+    #     for idx, line in enumerate(inf):
+    #         info = line.strip().split(' ')
+    #         if info[0].lower() not in ret:
+    #             ret[info[0]] = np.array([float(x) for x in info[1:]])
     return ret