def train_val_model(pipeline_cfg, model_cfg, train_cfg):
data_pipeline = DataPipeline(**pipeline_cfg)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if model_cfg['cxt_emb_pretrained'] is not None:
model_cfg['cxt_emb_pretrained'] = torch.load(
model_cfg['cxt_emb_pretrained'])
bidaf = BiDAF(word_emb=data_pipeline.word_type.vocab.vectors, **model_cfg)
ema = EMA(train_cfg['exp_decay_rate'])
for name, param in bidaf.named_parameters():
if param.requires_grad:
ema.register(name, param.data)
parameters = filter(lambda p: p.requires_grad, bidaf.parameters())
optimizer = optim.Adadelta(parameters, lr=train_cfg['lr'])
criterion = nn.CrossEntropyLoss()
result = {'best_f1': 0.0, 'best_model': None}
num_epochs = train_cfg['num_epochs']
for epoch in range(1, num_epochs + 1):
print('Epoch {}/{}'.format(epoch, num_epochs))
print('-' * 10)
for phase in ['train', 'val']:
val_answers = dict()
val_f1 = 0
val_em = 0
val_cnt = 0
val_r = 0
if phase == 'train':
bidaf.train()
else:
bidaf.eval()
backup_params = EMA(0)
for name, param in bidaf.named_parameters():
if param.requires_grad:
backup_params.register(name, param.data)
param.data.copy_(ema.get(name))
with torch.set_grad_enabled(phase == 'train'):
for batch_num, batch in enumerate(
data_pipeline.data_iterators[phase]):
optimizer.zero_grad()
p1, p2 = bidaf(batch)
loss = criterion(p1, batch.s_idx) + criterion(
p2, batch.e_idx)
if phase == 'train':
loss.backward()
optimizer.step()
for name, param in bidaf.named_parameters():
if param.requires_grad:
ema.update(name, param.data)
if batch_num % train_cfg['batch_per_disp'] == 0:
batch_loss = loss.item()
print('batch %d: loss %.3f' %
(batch_num, batch_loss))
if phase == 'val':
batch_size, c_len = p1.size()
val_cnt += batch_size
ls = nn.LogSoftmax(dim=1)
mask = (torch.ones(c_len, c_len) * float('-inf')).to(device).tril(-1). \
unsqueeze(0).expand(batch_size, -1, -1)
score = (ls(p1).unsqueeze(2) +
ls(p2).unsqueeze(1)) + mask
score, s_idx = score.max(dim=1)
score, e_idx = score.max(dim=1)
s_idx = torch.gather(s_idx, 1,
e_idx.view(-1, 1)).squeeze()
for i in range(batch_size):
answer = (s_idx[i], e_idx[i])
gt = (batch.s_idx[i], batch.e_idx[i])
val_f1 += f1_score(answer, gt)
val_em += exact_match_score(answer, gt)
val_r += r_score(answer, gt)
if phase == 'val':
val_f1 = val_f1 * 100 / val_cnt
val_em = val_em * 100 / val_cnt
val_r = val_r * 100 / val_cnt
print('Epoch %d: %s f1 %.3f | %s em %.3f | %s rouge %.3f' %
(epoch, phase, val_f1, phase, val_em, phase, val_r))
if val_f1 > result['best_f1']:
result['best_f1'] = val_f1
result['best_em'] = val_em
result['best_model'] = copy.deepcopy(bidaf.state_dict())
torch.save(result, train_cfg['ckpoint_file'])
# with open(train_cfg['val_answers'], 'w', encoding='utf-8') as f:
# print(json.dumps(val_answers), file=f)
for name, param in bidaf.named_parameters():
if param.requires_grad:
param.data.copy_(backup_params.get(name))
class SOLVER():
def __init__(self, args):
self.args = args
self.device = torch.device("cuda:{}".format(self.args.GPU) if torch.
cuda.is_available() else "cpu")
self.data = READ(self.args)
glove = self.data.WORD.vocab.vectors
char_size = len(self.data.CHAR.vocab)
self.model = BiDAF(self.args, char_size, glove).to(self.device)
self.optimizer = optim.Adadelta(self.model.parameters(),
lr=self.args.Learning_Rate)
self.ema = EMA(self.args.Exp_Decay_Rate)
if APEX_AVAILABLE: # Mixed Precision
self.model, self.optimizer = amp.initialize(self.model,
self.optimizer,
opt_level='O2')
for name, param in self.model.named_parameters():
if param.requires_grad:
self.ema.register(name, param.data)
self.parameters = filter(lambda p: p.requires_grad,
self.model.parameters())
def train(self):
criterion = nn.NLLLoss()
criterion = criterion.to(self.device)
self.model.train()
max_dev_em, max_dev_f1 = -1, -1
num_batches = len(self.data.train_iter)
logging.info("Begin Training")
self.model.zero_grad()
loss = 0.0
for epoch in range(self.args.Epoch):
self.model.train()
for i, batch in enumerate(self.data.train_iter):
i += 1
p1, p2 = self.model(batch)
batch_loss = criterion(
p1, batch.start_idx.to(self.device)) + criterion(
p2, batch.end_idx.to(self.device))
if APEX_AVAILABLE:
with amp.scale_loss(batch_loss,
self.optimizer) as scaled_loss:
scaled_loss.backward()
else:
batch_loss.backward()
loss = batch_loss.item()
self.optimizer.step()
del p1, p2, batch_loss
for name, param in self.model.named_parameters():
if param.requires_grad:
self.ema.update(name, param.data)
self.model.zero_grad()
logging.info("Epoch [{}/{}] Step [{}/{}] Train Loss {}".format(epoch+1, self.args.Epoch, \
i, int(num_batches) +1, round(loss,3)))
if epoch > 7:
if i % 100 == 0:
dev_em, dev_f1 = self.evaluate()
logging.info("Epoch [{}/{}] Dev EM {} Dev F1 {}".format(epoch + 1, self.args.Epoch, \
round(dev_em,3), round(dev_f1,3)))
self.model.train()
if dev_f1 > max_dev_f1:
max_dev_f1 = dev_f1
max_dev_em = dev_em
dev_em, dev_f1 = self.evaluate()
logging.info("Epoch [{}/{}] Dev EM {} Dev F1 {}".format(epoch + 1, self.args.Epoch, \
round(dev_em,3), round(dev_f1,3)))
self.model.train()
if dev_f1 > max_dev_f1:
max_dev_f1 = dev_f1
max_dev_em = dev_em
logging.info('Max Dev EM: {} Max Dev F1: {}'.format(
round(max_dev_em, 3), round(max_dev_f1, 3)))
def evaluate(self):
logging.info("Evaluating on Dev Dataset")
answers = dict()
self.model.eval()
temp_ema = EMA(0)
for name, param in self.model.named_parameters():
if param.requires_grad:
temp_ema.register(name, param.data)
param.data.copy_(self.ema.get(name))
with torch.no_grad():
for _, batch in enumerate(self.data.dev_iter):
p1, p2 = self.model(batch)
batch_size, _ = p1.size()
_, s_idx = p1.max(dim=1)
_, e_idx = p2.max(dim=1)
for i in range(batch_size):
qid = batch.qid[i]
answer = batch.c_word[0][i][s_idx[i]:e_idx[i] + 1]
answer = ' '.join(
[self.data.WORD.vocab.itos[idx] for idx in answer])
answers[qid] = answer
for name, param in self.model.named_parameters():
if param.requires_grad:
param.data.copy_(temp_ema.get(name))
results = evaluate(self.args, answers)
return results['exact_match'], results['f1']
def eval(context, question):
with open(os.path.join(config.data_dir, "train", "word2idx.pkl"), "rb") as wi, \
open(os.path.join(config.data_dir, "train", "char2idx.pkl"), "rb") as ci, \
open(os.path.join(config.data_dir, "train", "word_embeddings.pkl"), "rb") as wb, \
open(os.path.join(config.data_dir, "train", "char_embeddings.pkl"), "rb") as cb:
word2idx = pickle.load(wi)
char2idx = pickle.load(ci)
word_embedding_matrix = pickle.load(wb)
char_embedding_matrix = pickle.load(cb)
# transform them into Tensors
word_embedding_matrix = torch.from_numpy(
np.array(word_embedding_matrix)).type(torch.float32)
char_embedding_matrix = torch.from_numpy(
np.array(char_embedding_matrix)).type(torch.float32)
idx2word = dict([(y, x) for x, y in word2idx.items()])
context = clean_text(context)
context = [w for w in word_tokenize(context) if w]
question = clean_text(question)
question = [w for w in word_tokenize(question) if w]
if len(context) > config.max_len_context:
print("The context is too long. Maximum accepted length is",
config.max_len_context, "words.")
if max([len(w) for w in context]) > config.max_len_word:
print("Some words in the context are longer than", config.max_len_word,
"characters.")
if len(question) > config.max_len_question:
print("The question is too long. Maximum accepted length is",
config.max_len_question, "words.")
if max([len(w) for w in question]) > config.max_len_word:
print("Some words in the question are longer than",
config.max_len_word, "characters.")
if len(question) < 3:
print(
"The question is too short. It needs to be at least a three words question."
)
context_idx = np.zeros([config.max_len_context], dtype=np.int32)
question_idx = np.zeros([config.max_len_question], dtype=np.int32)
context_char_idx = np.zeros([config.max_len_context, config.max_len_word],
dtype=np.int32)
question_char_idx = np.zeros(
[config.max_len_question, config.max_len_word], dtype=np.int32)
# replace 0 values with word and char IDs
for j, word in enumerate(context):
if word in word2idx:
context_idx[j] = word2idx[word]
else:
context_idx[j] = 1
for k, char in enumerate(word):
if char in char2idx:
context_char_idx[j, k] = char2idx[char]
else:
context_char_idx[j, k] = 1
for j, word in enumerate(question):
if word in word2idx:
question_idx[j] = word2idx[word]
else:
question_idx[j] = 1
for k, char in enumerate(word):
if char in char2idx:
question_char_idx[j, k] = char2idx[char]
else:
question_char_idx[j, k] = 1
model = BiDAF(word_vectors=word_embedding_matrix,
char_vectors=char_embedding_matrix,
hidden_size=config.hidden_size,
drop_prob=config.drop_prob)
try:
if config.cuda:
model.load_state_dict(
torch.load(os.path.join(config.squad_models,
"model_final.pkl"))["state_dict"])
else:
model.load_state_dict(
torch.load(
os.path.join(config.squad_models, "model_final.pkl"),
map_location=lambda storage, loc: storage)["state_dict"])
print("Model weights successfully loaded.")
except:
pass
print(
"Model weights not found, initialized model with random weights.")
model.to(device)
model.eval()
with torch.no_grad():
context_idx, context_char_idx, question_idx, question_char_idx = torch.tensor(context_idx, dtype=torch.int64).unsqueeze(0).to(device),\
torch.tensor(context_char_idx, dtype=torch.int64).unsqueeze(0).to(device),\
torch.tensor(question_idx, dtype=torch.int64).unsqueeze(0).to(device),\
torch.tensor(question_char_idx, dtype=torch.int64).unsqueeze(0).to(device)
pred1, pred2 = model(context_idx, context_char_idx, question_idx,
question_char_idx)
starts, ends = discretize(pred1.exp(), pred2.exp(), 15, False)
prediction = " ".join(context[starts.item():ends.item() + 1])
return prediction
batch[4][:, 0].long().to(device),\
batch[4][:, 1].long().to(device)
optimizer.zero_grad()
pred1, pred2 = model(w_context, c_context, w_question, c_question)
loss = criterion(pred1, label1) + criterion(pred2, label2)
train_losses += loss.item()
loss.backward()
optimizer.step()
writer.add_scalars("train", {"loss": np.round(train_losses / len(train_dataloader), 2),
"epoch": epoch + 1})
print("Train loss of the model at epoch {} is: {}".format(epoch + 1, np.round(train_losses /
len(train_dataloader), 2)))
model.eval()
valid_losses = 0
valid_em = 0
valid_f1 = 0
n_samples = 0
with torch.no_grad():
for i, batch in enumerate(valid_dataloader):
w_context, c_context, w_question, c_question, labels = batch[0].long().to(device), \
batch[1].long().to(device), \
batch[2].long().to(device), \
batch[3].long().to(device), \
batch[4]
first_labels = torch.tensor([[int(a) for a in l.split("|")[0].split(" ")]
for l in labels], dtype=torch.int64).to(device)
pred1, pred2 = model(w_context, c_context, w_question, c_question)
def main(NMT_config):
### Load RL (global) configurations ###
config = parse_args()
### Load trained QA model ###
QA_checkpoint = torch.load(config.data_dir + config.QA_best_model)
QA_config = QA_checkpoint['config']
QA_mod = BiDAF(QA_config)
if QA_config.use_gpu:
QA_mod.cuda()
QA_mod.load_state_dict(QA_checkpoint['state_dict'])
### Load SQuAD dataset ###
data_filter = get_squad_data_filter(QA_config)
train_data = read_data(QA_config,
'train',
QA_config.load,
data_filter=data_filter)
dev_data = read_data(QA_config, 'dev', True, data_filter=data_filter)
update_config(QA_config, [train_data, dev_data])
print("Total vocabulary for training is %s" % QA_config.word_vocab_size)
# from all
word2vec_dict = train_data.shared[
'lower_word2vec'] if QA_config.lower_word else train_data.shared[
'word2vec']
# from filter-out set
word2idx_dict = train_data.shared['word2idx']
# filter-out set idx-vector
idx2vec_dict = {
word2idx_dict[word]: vec
for word, vec in word2vec_dict.items() if word in word2idx_dict
}
print("{}/{} unique words have corresponding glove vectors.".format(
len(idx2vec_dict), len(word2idx_dict)))
# <null> and <unk> do not have corresponding vector so random.
emb_mat = np.array([
idx2vec_dict[idx]
if idx in idx2vec_dict else np.random.multivariate_normal(
np.zeros(QA_config.word_emb_size), np.eye(QA_config.word_emb_size))
for idx in range(QA_config.word_vocab_size)
])
config.emb_mat = emb_mat
config.new_emb_mat = train_data.shared['new_emb_mat']
num_steps = int(
math.ceil(train_data.num_examples /
(QA_config.batch_size *
QA_config.num_gpus))) * QA_config.num_epochs
# offset for question mark
NMT_config.max_length = QA_config.ques_size_th - 1
NMT_config.batch_size = QA_config.batch_size
### Construct translator ###
translator = make_translator(NMT_config, report_score=True)
### Construct optimizer ###
optimizer = optim.SGD(filter(lambda p: p.requires_grad,
translator.model.parameters()),
lr=config.lr)
### Start RL training ###
count = 0
QA_mod.eval()
F1_eval = F1Evaluator(QA_config, QA_mod)
#eval_model(QA_mod, train_data, dev_data, QA_config, NMT_config, config, translator)
for i in range(config.n_episodes):
for batches in tqdm(train_data.get_multi_batches(
QA_config.batch_size,
QA_config.num_gpus,
num_steps=num_steps,
shuffle=True,
cluster=QA_config.cluster),
total=num_steps):
#for n, p in translator.model.named_parameters():
# print(n)
# print(p)
#print(p.requires_grad)
start = datetime.now()
to_input(batches[0][1].data['q'], config.RL_path + config.RL_file)
# obtain rewrite and log_prob
q, scores, log_prob = translator.translate(NMT_config.src_dir,
NMT_config.src,
NMT_config.tgt,
NMT_config.batch_size,
NMT_config.attn_debug)
q, cq = ref_query(q)
batches[0][1].data['q'] = q
batches[0][1].data['cq'] = cq
log_prob = torch.stack(log_prob).squeeze(-1)
#print(log_prob)
translator.model.zero_grad()
QA_mod(batches)
e = F1_eval.get_evaluation(batches, False, NMT_config, config,
translator)
reward = Variable(torch.FloatTensor(e.f1s), requires_grad=False)
#print(reward)
## Initial loss
loss = create_loss(log_prob, reward)
loss.backward()
optimizer.step()