I’m trying to deploy this model for generating SQL from text and establish an interactive interface where users can input natural language queries and receive SQL outputs. You will find below main.py and model.py. this model is taken from TREQS repository
this is main.py
import argparse
import torch
from model import modelABS
from LeafNATS.utils.utils import str2bool
from end2end_large import End2EndBase
parser = argparse.ArgumentParser()
'''
Use in the framework and cannot remove.
'''
parser.add_argument('--task', default='train', help='train | validate | test | evaluate')
parser.add_argument('--data_dir', default='mimicsql_data/mimicsql_natural', help='directory that store the data.')
parser.add_argument('--file_train', default='train.json', help='Training')
parser.add_argument('--file_val', default='dev.json', help='validation')
parser.add_argument('--file_test', default='test.json', help='test data')
parser.add_argument('--n_epoch', type=int, default=20, help='number of epochs.')
parser.add_argument('--batch_size', type=int, default=16, help='batch size.')
parser.add_argument('--checkpoint', type=int, default=100, help='How often you want to save model?')
parser.add_argument('--val_num_batch', type=int, default=100, help='how many batches')
parser.add_argument('--nbestmodel', type=int, default=5, help='How many models you want to keep?')
parser.add_argument('--continue_training', type=str2bool, default=True, help='Do you want to continue?')
parser.add_argument('--train_base_model', type=str2bool, default=False, help='True: Use Pretrained Param | False: Transfer Learning')
parser.add_argument('--use_move_avg', type=str2bool, default=False, help='move average')
parser.add_argument('--use_optimal_model', type=str2bool, default=True, help='Do you want to use the best model?')
parser.add_argument('--model_optimal_key', default='0,0', help='epoch,batch')
parser.add_argument('--is_lower', type=str2bool, default=True, help='convert all tokens to lower case?')
'''
User specified parameters.
'''
parser.add_argument('--device', default=torch.device("cuda:0"), help='device')
parser.add_argument('--file_vocab', default='vocab', help='file store training vocabulary.')
parser.add_argument('--max_vocab_size', type=int, default=50000, help='max number of words in the vocabulary.')
parser.add_argument('--word_minfreq', type=int, default=5, help='min word frequency')
parser.add_argument('--emb_dim', type=int, default=128, help='source embedding dimension')
parser.add_argument('--src_hidden_dim', type=int, default=256, help='encoder hidden dimension')
parser.add_argument('--trg_hidden_dim', type=int, default=256, help='decoder hidden dimension')
parser.add_argument('--src_seq_len', type=int, default=400, help='length of source documents.')
parser.add_argument('--trg_seq_len', type=int, default=100, help='length of target documents.')
parser.add_argument('--nLayers', type=int, default=1, help='Encoder RNN layers')
parser.add_argument('--learning_rate', type=float, default=0.0005, help='learning rate.')
parser.add_argument('--grad_clip', type=float, default=2.0, help='clip the gradient norm.')
parser.add_argument('--file_output', default='output.json', help='test output file')
parser.add_argument('--beam_size', type=int, default=5, help='beam size.')
parser.add_argument('--copy_words', type=str2bool, default=True, help='Do you want to copy words?')
# scheduler
parser.add_argument('--step_size', type=int, default=2, help='---')
parser.add_argument('--step_decay', type=float, default=0.8, help='---')
args = parser.parse_args()
if args.task == 'train' or args.task == 'validate' or args.task == 'test':
from model import modelABS
model = modelABS(args)
if args.task == "train":
model.train()
if args.task == "validate":
model.validate()
if args.task == "test":
model.test()
if args.task == "evaluate":
from LeafNATS.eval_scripts.eval_pyrouge_v2 import run_pyrouge
run_pyrouge(args)
and this is model.py
import os
import time
import torch
from torch.autograd import Variable
from seq2sql.model_seq2seq_base import modelSeq2SeqBase
from LeafNATS.data.seq2sql.process_batch_cqa_v1 import process_batch
from LeafNATS.modules.embedding.nats_embedding import natsEmbedding
from LeafNATS.modules.encoder.encoder_rnn import EncoderRNN
from LeafNATS.modules.encoder2decoder.nats_encoder2decoder import natsEncoder2Decoder
from LeafNATS.modules.attention.nats_attention_encoder import AttentionEncoder
from LeafNATS.modules.attention.nats_attention_decoder import AttentionDecoder
# from LeafNATS.utils.utils import *
class modelABS(modelSeq2SeqBase):
def __init__(self, args):
super().__init__(args=args)
def build_scheduler(self, optimizer):
'''
Schedule Learning Rate
'''
scheduler = torch.optim.lr_scheduler.StepLR(
optimizer=optimizer, step_size=self.args.step_size,
gamma=self.args.step_decay)
return scheduler
def build_batch(self, batch_id):
'''
get batch data
'''
output = process_batch(
batch_id=batch_id,
path_=os.path.join('..', 'nats_results'),
fkey_=self.args.task,
batch_size=self.args.batch_size,
vocab2id=self.batch_data['vocab2id'],
max_lens=[self.args.src_seq_len, self.args.trg_seq_len])
self.batch_data['ext_id2oov'] = output['ext_id2oov']
self.batch_data['src_var'] = output['src_var'].to(self.args.device)
self.batch_data['batch_size'] = self.batch_data['src_var'].size(0)
self.batch_data['src_seq_len'] = self.batch_data['src_var'].size(1)
self.batch_data['src_mask_pad'] = output['src_mask_pad'].to(self.args.device)
if self.args.task == 'train' or self.args.task == 'validate':
self.batch_data['trg_input'] = output['trg_input_var'].to(self.args.device)
# different from seq2seq models.
self.batch_data['trg_output'] = output['trg_output_var'].to(self.args.device)
self.batch_data['trg_seq_len'] = self.batch_data['trg_input'].size(1)
else:
self.batch_data['src_mask_unk'] = output['src_mask_unk'].to(self.args.device)
self.batch_data['src_txt'] = output['src_txt']
self.batch_data['trg_txt'] = output['trg_txt']
self.batch_data['trg_seq_len'] = 1
def build_models(self):
'''
build all models.
in this model source and target share embeddings
'''
self.train_models['embedding'] = natsEmbedding(
vocab_size = self.batch_data['vocab_size'],
emb_dim = self.args.emb_dim,
share_emb_weight = True
).to(self.args.device)
self.train_models['encoder'] = EncoderRNN(
self.args.emb_dim, self.args.src_hidden_dim,
self.args.nLayers, 'lstm',
device = self.args.device
).to(self.args.device)
self.train_models['encoder2decoder'] = natsEncoder2Decoder(
src_hidden_size = self.args.src_hidden_dim,
trg_hidden_size = self.args.trg_hidden_dim,
rnn_network = 'lstm',
device = self.args.device
).to(self.args.device)
self.train_models['decoderRNN'] = torch.nn.LSTMCell(
self.args.emb_dim+self.args.trg_hidden_dim,
self.args.trg_hidden_dim
).to(self.args.device)
self.train_models['attnEncoder'] = AttentionEncoder(
self.args.src_hidden_dim,
self.args.trg_hidden_dim,
attn_method='luong_general',
repetition='temporal'
).to(self.args.device)
self.train_models['attnDecoder'] = AttentionDecoder(
self.args.trg_hidden_dim,
attn_method='luong_general'
).to(self.args.device)
self.train_models['wrapDecoder'] = torch.nn.Linear(
self.args.src_hidden_dim*2+self.args.trg_hidden_dim*2,
self.args.trg_hidden_dim, bias=True
).to(self.args.device)
self.train_models['genPrb'] = torch.nn.Linear(
self.args.emb_dim+self.args.src_hidden_dim*2+self.args.trg_hidden_dim, 1
).to(self.args.device)
# decoder to vocab
self.train_models['decoder2proj'] = torch.nn.Linear(
self.args.trg_hidden_dim, self.args.emb_dim, bias=False
).to(self.args.device)
def build_encoder(self):
'''
Encoder Pipeline
self.pipe_data = {
'encoder': {},
'decoderA': {}}
'decoderB': {'accu_attn': [], 'last_word': word}}
'''
src_emb = self.train_models['embedding'].get_embedding(
self.batch_data['src_var'])
src_enc, hidden_encoder = self.train_models['encoder'](src_emb)
trg_hidden0 = self.train_models['encoder2decoder'](hidden_encoder)
# set up pipe_data pass to decoder
self.pipe_data['encoder'] = {}
self.pipe_data['encoder']['src_emb'] = src_emb
self.pipe_data['encoder']['src_enc'] = src_enc
self.pipe_data['decoderB'] = {}
self.pipe_data['decoderB']['hidden'] = trg_hidden0
self.pipe_data['decoderB']['h_attn'] = Variable(torch.zeros(
self.batch_data['batch_size'], self.args.trg_hidden_dim
)).to(self.args.device)
self.pipe_data['decoderB']['past_attn'] = Variable(torch.ones(
self.batch_data['batch_size'], self.batch_data['src_seq_len']
)/float(self.batch_data['src_seq_len'])).to(self.args.device)
self.pipe_data['decoderB']['past_dech'] = Variable(torch.zeros(
1, 1)).to(self.args.device)
self.pipe_data['decoderB']['accu_attn'] = []
self.pipe_data['decoderFF'] = {}
self.pipe_data['decoderFF']['h_attn'] = []
self.pipe_data['decoderFF']['attn'] = []
self.pipe_data['decoderFF']['genPrb'] = []
# when training get target embedding at the same time.
if self.args.task == 'train' or self.args.task == 'validate':
trg_emb = self.train_models['embedding'].get_embedding(
self.batch_data['trg_input'])
self.pipe_data['decoderFF']['trg_seq_emb'] = trg_emb
def build_decoder_one_step(self, k=0):
'''
Decoder one-step
'''
# embedding at current decoding step
if self.args.task == 'train' or self.args.task == 'validate':
self.pipe_data['decoderA'] = self.pipe_data['decoderB']
word_emb = self.pipe_data['decoderFF']['trg_seq_emb'][:, k]
else:
word_emb = self.train_models['embedding'].get_embedding(
self.pipe_data['decoderA']['last_word'])
h_attn = self.pipe_data['decoderA']['h_attn']
dec_input = torch.cat((word_emb, h_attn), 1)
hidden = self.pipe_data['decoderA']['hidden']
past_attn = self.pipe_data['decoderA']['past_attn']
accu_attn = self.pipe_data['decoderA']['accu_attn']
past_dech = self.pipe_data['decoderA']['past_dech']
hidden = self.train_models['decoderRNN'](dec_input, hidden)
ctx_enc, attn, attn_ee = self.train_models['attnEncoder'](
hidden[0], self.pipe_data['encoder']['src_enc'],
past_attn, self.batch_data['src_mask_pad'])
# temporal attention
past_attn = past_attn + attn_ee
# decoder attention
if k == 0:
ctx_dec = Variable(torch.zeros(
self.batch_data['batch_size'], self.args.trg_hidden_dim
)).to(self.args.device)
else:
ctx_dec, _ = self.train_models['attnDecoder'](
hidden[0], past_dech)
past_dech = past_dech.transpose(0, 1) # seqL*batch*hidden
dec_idx = past_dech.size(0)
if k == 0:
past_dech = hidden[0].unsqueeze(0) # seqL*batch*hidden
past_dech = past_dech.transpose(0, 1) # batch*seqL*hidden
else:
past_dech = past_dech.contiguous().view(
-1, self.args.trg_hidden_dim) # seqL*batch**hidden
past_dech = torch.cat((past_dech, hidden[0]), 0) # (seqL+1)*batch**hidden
past_dech = past_dech.view(
dec_idx+1, self.batch_data['batch_size'], self.args.trg_hidden_dim
) # (seqL+1)*batch*hidden
past_dech = past_dech.transpose(0, 1) # batch*(seqL+1)*hidden
# wrap up.
h_attn = self.train_models['wrapDecoder'](torch.cat((ctx_enc, ctx_dec, hidden[0]), 1))
# pointer generator
pt_input = torch.cat((word_emb, hidden[0], ctx_enc), 1)
genPrb = torch.sigmoid(self.train_models['genPrb'](pt_input))
# setup piped_data
self.pipe_data['decoderB'] = {}
self.pipe_data['decoderB']['h_attn'] = h_attn
self.pipe_data['decoderB']['past_attn'] = past_attn
self.pipe_data['decoderB']['hidden'] = hidden
self.pipe_data['decoderB']['past_dech'] = past_dech
self.pipe_data['decoderB']['accu_attn'] = [a for a in accu_attn]
self.pipe_data['decoderB']['accu_attn'].append(attn)
if self.args.task == 'train' or self.args.task == 'validate':
self.pipe_data['decoderFF']['h_attn'].append(h_attn)
self.pipe_data['decoderFF']['attn'].append(attn)
self.pipe_data['decoderFF']['genPrb'].append(genPrb)
if k == self.batch_data['trg_seq_len']-1:
self.pipe_data['decoderFF']['h_attn'] = \
torch.cat(self.pipe_data['decoderFF']['h_attn'], 0).view(
self.batch_data['trg_seq_len'],
self.batch_data['batch_size'],
self.args.trg_hidden_dim).transpose(0,1)
self.pipe_data['decoderFF']['attn'] = \
torch.cat(self.pipe_data['decoderFF']['attn'], 0).view(
self.batch_data['trg_seq_len'],
self.batch_data['batch_size'],
self.args.src_seq_len).transpose(0,1)
self.pipe_data['decoderFF']['genPrb'] = \
torch.cat(self.pipe_data['decoderFF']['genPrb'], 0).view(
self.batch_data['trg_seq_len'],
self.batch_data['batch_size']).transpose(0,1)
else:
self.pipe_data['decoderFF']['h_attn'] = h_attn
self.pipe_data['decoderFF']['attn'] = attn.unsqueeze(0)
self.pipe_data['decoderFF']['genPrb'] = genPrb
def build_vocab_distribution(self):
'''
Data flow from input to output.
'''
trg_out = self.pipe_data['decoderFF']['h_attn']
trg_out = self.train_models['decoder2proj'](trg_out)
trg_out = self.train_models['embedding'].get_decode2vocab(trg_out)
trg_out = trg_out.view(
self.batch_data['batch_size'], self.batch_data['trg_seq_len'], -1)
prb = torch.softmax(trg_out, dim=2)
vocab_size = self.batch_data['vocab_size']
batch_size = self.batch_data['batch_size']
# trg_seq_len = self.batch_data['trg_seq_len']
src_seq_len = self.batch_data['src_seq_len']
# pointer-generator calculate index matrix
pt_idx = Variable(torch.FloatTensor(torch.zeros(1, 1, 1))).to(self.args.device)
pt_idx = pt_idx.repeat(batch_size, src_seq_len, vocab_size)
pt_idx.scatter_(2, self.batch_data['src_var'].unsqueeze(2), 1.0)
p_gen = self.pipe_data['decoderFF']['genPrb']
attn_ = self.pipe_data['decoderFF']['attn']
prb_output = p_gen.unsqueeze(2)*prb + \
(1.0-p_gen.unsqueeze(2))*torch.bmm(attn_, pt_idx)
return prb_output + 1e-20
def build_pipelines(self):
'''
Build pipeline from input to output.
Output is loss.
Input is word one-hot encoding.
'''
self.build_encoder()
for k in range(self.args.trg_seq_len):
self.build_decoder_one_step(k)
prb = self.build_vocab_distribution()
pad_mask = torch.ones(self.batch_data['vocab_size']).to(self.args.device)
pad_mask[self.batch_data['vocab2id']['<pad>']] = 0
self.loss_criterion = torch.nn.NLLLoss(pad_mask).to(self.args.device)
prb = torch.log(prb)
loss = self.loss_criterion(
prb.reshape(-1, self.batch_data['vocab_size']),
self.batch_data['trg_output'].view(-1))
return loss