load_used=args.train_emb, use_small=USE_SMALL)
if args.baseline:
model = Seq2SQL(word_emb,
N_word=N_word,
gpu=GPU,
trainable_emb=args.train_emb)
assert not args.train_emb, "Seq2SQL can\'t train embedding."
else:
model = SQLNet(word_emb,
N_word=N_word,
use_ca=args.ca,
gpu=GPU,
trainable_emb=args.train_emb)
assert not args.rl, "SQLNet can\'t do reinforcement learning."
optimizer = torch.optim.Adam(model.parameters(),
lr=learning_rate,
weight_decay=0)
if args.train_emb:
agg_m, sel_m, cond_m, agg_e, sel_e, cond_e = best_model_name(args)
else:
agg_m, sel_m, cond_m = best_model_name(args)
if args.rl or args.train_emb: # Load pretrained model.
agg_lm, sel_lm, cond_lm = best_model_name(args, for_load=True)
print("Loading from %s" % agg_lm)
model.agg_pred.load_state_dict(torch.load(agg_lm))
print("Loading from %s" % sel_lm)
model.sel_pred.load_state_dict(torch.load(sel_lm))
print("Loading from %s" % cond_lm)
print '==> loading word embedding'
word_emb = load_word_emb('glove/glove.%dB.%dd.txt'%(B_word,N_word), load_used=args.train_emb, use_small=USE_SMALL)
# import pickle
# with open('glove/word_emb42B.pkl', 'rb') as f:
# pickle.load(f)
# word_emb = pickle.load(f)
print 'time to load word emb: ' + str(time() - tic)
# build sqlnet model
if not args.baseline:
tic = time()
print '==> loading sqlnet constructor'
model = SQLNet(word_emb, N_word=N_word, use_ca=args.ca, gpu=GPU, trainable_emb = args.train_emb)
print 'time to load sqlnet constructor: ' + str(time() - tic)
assert not args.rl, "SQLNet can\'t do reinforcement learning."
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay = 0)
if args.train_emb:
agg_m, sel_m, cond_m, agg_e, sel_e, cond_e = best_model_name(args)
# load model
agg_m, sel_m, cond_m = best_model_name(args, savedstr='_pretrain_wikisql')
print('==> best model names:', agg_m, sel_m, cond_m)
print "Loading from %s"%agg_m
model.agg_pred.load_state_dict(torch.load(agg_m))
print "Loading from %s"%sel_m
model.sel_pred.load_state_dict(torch.load(sel_m))
print "Loading from %s"%cond_m
model.cond_pred.load_state_dict(torch.load(cond_m))
if args.rl or args.train_emb:
print('train_emb is on, so loading best_model')