# This code builds on the AWD-LSTM codebase

# (https://github.com/salesforce/awd-lstm-lm).

#

# groc is free software: you can redistribute it and/or modify

# it under the terms of the GNU General Public License version 3 as

# published by the Free Software Foundation.

#

# groc is distributed in the hope that it will be useful,

# but WITHOUT ANY WARRANTY; without even the implied warranty of

# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the

# GNU General Public License for more details.

#

# You should have received a copy of the GNU General Public License

# along with groc. If not, see http://www.gnu.org/licenses/

import argparse

import os, shutil

import sys

import hashlib

import time

import math

import numpy as np

import torch

import torch.nn as nn

import data

import model

from util import *

import IPython as ipy

parser = argparse.ArgumentParser(description='PyTorch PennTreeBank RNN/LSTM Language Model')

parser.add_argument('--data', type=str, default='data/penn/',

help='location of the data corpus')

parser.add_argument('--rnn_type', type=str, default='LSTM',

help='type of recurrent net (LSTM)')

parser.add_argument('--emsize', type=int, default=400,

help='size of word embeddings')

parser.add_argument('--nhid', type=int, default=1150,

help='number of hidden units per layer')

parser.add_argument('--nlayers', type=int, default=3,

help='number of layers')

parser.add_argument('--lr', type=float, default=1e-03,

help='initial learning rate')

parser.add_argument('--beta0', type=float, default=0.9,

help='initial learning rate')

parser.add_argument('--clip', type=float, default=0.25,

help='gradient clipping')

parser.add_argument('--epochs', type=int, default=8000,

help='upper epoch limit')

parser.add_argument('--batch_size', type=int, default=80, metavar='N',

help='batch size')

parser.add_argument('--bptt', type=int, default=70,

help='sequence length')

parser.add_argument('--dropout', type=float, default=0.4,

help='dropout applied to layers (0 = no dropout)')

parser.add_argument('--dropouth', type=float, default=0.3,

help='dropout for rnn layers (0 = no dropout)')

parser.add_argument('--dropouti', type=float, default=0.65,

help='dropout for input embedding layers (0 = no dropout)')

parser.add_argument('--dropoute', type=float, default=0.1,

help='dropout to remove words from embedding layer (0 = no dropout)')

parser.add_argument('--joint_emb', type=int, default=None,

help='whether to use joint embedding or not')

parser.add_argument('--fullsoftmax', action='store_true',

help='whether to use full softmax or not')

parser.add_argument('--joint_emb_depth', type=int, default=1,

help='depth of the joint embedding')

parser.add_argument('--joint_dropout', type=float, default=0.2,

help='dropout for joint embedding layers')

parser.add_argument('--joint_emb_dense', action='store_true', default=False,

help='add residuals to all previous joint embedding projections')

parser.add_argument('--joint_emb_activation', type=str, default='Sigmoid',

help='activation function for the joint embedding layers')

parser.add_argument('--joint_emb_dual', action='store_true', default=False,

help='whether to use projection on both input and output side or not')

parser.add_argument('--joint_locked_dropout', action='store_true', default=False,

help='whether to use locked dropout or not for the joint space')

parser.add_argument('--joint_residual_prev', action='store_true', default=False,

help='whether to use residual connection to previous layer')

parser.add_argument('--joint_noresid', action='store_true', default=False,

help='disable residual connections')

parser.add_argument('--seed', type=int, default=1111,

help='random seed')

parser.add_argument('--nonmono', type=int, default=5,

help='random seed')

parser.add_argument('--cuda', action="store_true", default=False,

help='use CUDA device')

parser.add_argument('--cuda_device', type=int, default=-1,

help='set CUDA device')

parser.add_argument('--log-interval', type=int, default=200, metavar='N',

help='report interval')

randomhash = ''.join(str(time.time()).split('.'))

parser.add_argument('--save', type=str, default=randomhash+'.pt',

help='path to save the final model')

parser.add_argument('--alpha', type=float, default=2,

help='alpha L2 regularization on RNN activation (alpha = 0 means no regularization)')

parser.add_argument('--beta', type=float, default=1,

help='beta slowness regularization applied on RNN activiation (beta = 0 means no regularization)')

parser.add_argument('--wdecay', type=float, default=1.2e-6,

help='weight decay applied to all weights')

parser.add_argument('--init', type=float, default=0.1,

help='range for initialization')

parser.add_argument('--resume', action="store_true", default=False,

help='resume or not')

parser.add_argument('--char_emb', action="store_true", default=False,

help='use character embedding or not')

parser.add_argument('--rel_emb', action="store_true", default=False,

help='use relation embedding or not')

parser.add_argument('--def_emb', action="store_true", default=False,

help='use definition embedding or not')

parser.add_argument('--combine', type=str, default="add",

help='how to combine different forms (add | concat | multiply)')

parser.add_argument('--predict_bias', action="store_true", default=False,

help='use the bias estimator or not')

parser.add_argument('--bias_drop', type=float, default=0.2,

help='dropout for bias estimator (0 = no dropout)')

parser.add_argument('--bias_out', type=int, default=1,

help='output dim for bias estimator')

parser.add_argument('--bias_activation', type=str, default="tanh",

help='activation for bias estimator')

parser.add_argument('--div', type=int, default=500,

help='number to divide the vocab size for batch computation')

parser.add_argument('--offset', type=int, default=None,

help='offset for reducing compute for the char emb')

parser.add_argument('--hdepth', type=int, default=1,

help='character emb highway depth')

parser.add_argument('--max_deflen', type=int, default=10,

help='maximum definition length')

parser.add_argument('--max_rellen', type=int, default=3,

help='maximum relation length')

parser.add_argument('--max_charlen', type=int, default=20,

help='maximum char length')

parser.add_argument('--defenc', type=str, default=None,

help='encoder for description and relations (lstm | highway)')

parser.add_argument('--char_emsize', type=int, default=16,

help='size of character embeddings')

parser.add_argument('--char_activation', type=str, default="selu",

help='character embedding activation (applies to convolutions and highway network)')

parser.add_argument('--optimizer', type=str, default='adam',

help='optimizer to use (sgd, adam)')

parser.add_argument('--when', nargs="+", type=int, default=[-1],

help='When (which epochs) to divide the learning rate by 10 - accepts multiple')

parser.add_argument('--cnnsoftmax', action="store_true",

default=False, help='whether to use cnn softmax or not')

parser.add_argument('--cnncorr', type=int, default=None,

help='whether to use cnn correction or not and its dimensionality')

parser.add_argument('--bilinear', action="store_true",

default=False, help='whether to use bilinear output embedding not')

parser.add_argument('--adaptiveoutputs', action="store_true",

default=False, help='whether to use adaptive output embedding or not')

parser.add_argument('--adaptiveoutputs_tied', action="store_true",

default=False, help='whether to use tied adaptive output embedding or not')

parser.add_argument('--adaptivecutoffs', type=str,

default="[5000]", help='cutoffs for adaptive outputs')

parser.add_argument('--output_dropout', type=float,

default=0.0, help='locked dropout for outputs')

parser.add_argument('--char_update_ratio', type=float,

default=1.0, help='portion of regular updates versus detached ones for the grounded embedding. ')

parser.add_argument('--char_nohighways', action="store_true",

help='do not use highways after the convolutional network. ')

parser.add_argument('--coverage', type=float, default=1.0,

help='control vocabulary coverage from the external knowledge base. ')

parser.add_argument('--finetune', action="store_true",

help="fine-tune an existing model on new data (with a new vocab)")

parser.add_argument('--finetune_save', default=None,

help="save fine-tuned model and logs to this directory (not the original")

args = parser.parse_args()

args.tied = True

if args.joint_emb is not None:

args.tied = False

if args.fullsoftmax:

args.tied = False

if args.adaptiveoutputs:

args.tied = False

if args.finetune:

args.resume = True

save_dir = args.finetune_save

load_dir = args.save

else:

save_dir = args.save

load_dir = args.save

if args.finetune or not args.resume:

try:

create_exp_dir(save_dir, scripts_to_save=['main.py', 'model.py'])

except:

res = input("Directory exists! Try again with --resume. ")

exit(0)

with open(os.path.join(save_dir, 'replicate.sh'), 'w') as f:

f.write('python ' + ' '.join(sys.argv))

def logging(s, print_=True, log_=True):

f_log.write(str(s) + '\n')

# Set the random seed manually for reproducibility.

np.random.seed(args.seed)

torch.manual_seed(args.seed)

#torch.backends.cudnn.deterministic = True

#torch.backends.cudnn.benchmark = False

if torch.cuda.is_available():

if not args.cuda:

logging("WARNING: You have a CUDA device, so you should probably run with --cuda and --cuda_device [device_id]")

else:

torch.cuda.set_device(int(args.cuda_device))

torch.cuda.manual_seed(args.seed)

###############################################################################

# Load data

###############################################################################

def model_save(model, criterion, optimizer, save):

save_checkpoint(model, criterion, optimizer, save)

def model_load(fn):

return model, criterion, optimizer

def corpus_load(corpus_path, use_unk=False):

return corpus

corpus = data.Corpus(args.data)

eval_batch_size = 30

test_batch_size = 2

def own_batchify(data):

return data.reshape((-1, 800))

def get_batch_padded(source, batch_size):

yield source[b*batch_size:(b+1)*batch_size, :-1].permute(1,0).cuda(), source[b*batch_size:(b+1)*batch_size, 1:].permute(1,0).cuda()

train_data = own_batchify(corpus.train)

val_data = own_batchify(corpus.valid)

test_data = own_batchify(corpus.test)

ntokens = len(corpus.dictionary)

args.ntoken = ntokens

if not args.finetune:

char_arr, rel_arr, def_arr = get_external_knowledge(args, corpus)

###############################################################################

# Build the model

###############################################################################

criterion = None

if args.resume:

logging('Resuming model ...')

model, criterion, optimizer = model_load(load_dir)

optimizer.param_groups[0]['lr'] = args.lr

model.dropouti, model.dropouth, model.dropout, args.dropoute = args.dropouti, args.dropouth, args.dropout, args.dropoute

else:

model = model.RNNModel( args, char_arr=char_arr, rel_arr=rel_arr, def_arr=def_arr, dict=corpus.dictionary)

criterion = load_criterion(args, ntokens, logging)

if args.cuda:

model = model.cuda()

criterion = criterion.cuda()

params = list(model.parameters()) + list(criterion.parameters())

total_params = sum(x.size()[0] * x.size()[1] if len(x.size()) > 1 else x.size()[0] for x in params if x.size() and x.requires_grad)

logging('Args: %s' % args)

logging('Model total parameters: %d ' % total_params)

logging('Model vocab size: %d ' % len(model.dict.word2idx))

logging(model)

if args.finetune:

char_arr, rel_arr, def_arr = get_external_knowledge(model.H, corpus)

model.change_embedding_vocab(char_arr, rel_arr, def_arr, corpus.dictionary)

params = list(model.parameters()) + list(criterion.parameters())

total_params = sum(x.size()[0] * x.size()[1] if len(x.size()) > 1 else x.size()[0] for x in params if x.size() and x.requires_grad)

logging('Model total parameters post-vocab-change: %d ' % total_params)

logging('Model vocab size post-vocab-change: %d ' % len(model.dict.word2idx))

logging(model)

###############################################################################

# Training code

###############################################################################

def evaluate(data_source, batch_size=10):

return total_loss / (data_source[1:, :] != corpus.dictionary.word2idx[corpus.padding]).sum().item()

def train():

# i += seq_len

# Loop over epochs.

lr = args.lr

best_val_loss = []

stored_loss = 100000000

steps_wo_inc = 0

total_steps_wo_inc = 0

# At any point you can hit Ctrl + C to break out of training early.

try:

optimizer = None

# Ensure the optimizer is optimizing params, which includes both the model's weights as well as the criterion's weight (i.e. Adaptive Softmax)

if args.optimizer == 'sgd':

optimizer = torch.optim.SGD(params, lr=args.lr, weight_decay=args.wdecay)

if args.optimizer == 'adam':

optimizer = torch.optim.Adam(params, lr=args.lr, betas=(args.beta0, 0.999), eps=1e-09, weight_decay=args.wdecay)

if args.optimizer == 'asgd':

optimizer = torch.optim.ASGD(model.parameters(), lr=args.lr, t0=0, lambd=0., weight_decay=args.wdecay)

for epoch in range(1, args.epochs+1):

epoch_start_time = time.time()

train()

if 't0' in optimizer.param_groups[0]:

tmp = {}

for j, prm in enumerate(model.parameters()):

tmp[prm] = prm.data.clone()

curstate = optimizer.state[prm]

prm.data = curstate['ax'].clone()

val_loss2 = evaluate(val_data, eval_batch_size)

logging('-' * 89)

logging('| end of epoch {:3d} | time: {:5.2f}s | valid loss {:5.2f} | '

'valid ppl {:8.2f} | valid bpc {:8.3f}'.format(

epoch, (time.time() - epoch_start_time), val_loss2, math.exp(val_loss2), val_loss2 / math.log(2)))

logging('-' * 89)

if val_loss2 < stored_loss:

model_save(model, criterion, optimizer, save_dir)

logging('Saving Averaged!')

stored_loss = val_loss2

for prm in model.parameters():

prm.data = tmp[prm].clone()

else:

val_loss = evaluate(val_data, eval_batch_size)

logging('-' * 89)

logging('| end of epoch {:3d} | time: {:5.2f}s | valid loss {:5.2f} | '

'valid ppl {:8.2f} | valid bpc {:8.3f}'.format(

epoch, (time.time() - epoch_start_time), val_loss, math.exp(val_loss), val_loss / math.log(2)))

logging('-' * 89)

if np.isnan(val_loss):

logging('Exiting...(loss has nan value)')

exit(-1)

if val_loss < stored_loss:

model_save(model, criterion, optimizer, save_dir)

logging('Saving model (new best validation)')

stored_loss = val_loss

steps_wo_inc = 0

total_steps_wo_inc = 0

if args.optimizer == 'sgd' and 't0' not in optimizer.param_groups[0] and (len(best_val_loss)>args.nonmono and val_loss > min(best_val_loss[:-args.nonmono])):

logging('Switching to ASGD')

optimizer = torch.optim.ASGD(model.parameters(), lr=args.lr, t0=0, lambd=0., weight_decay=args.wdecay)

if val_loss > stored_loss:

steps_wo_inc += 1

if steps_wo_inc == 4:

logging('Dividing learning rate by 10')

optimizer.param_groups[0]['lr'] /= 10.

steps_wo_inc = 0

total_steps_wo_inc += 1

if total_steps_wo_inc == 8:

logging('Exiting...')

break

if epoch in args.when:

logging('Saving model before learning rate decreased')

model_save(model, criterion, optimizer, '{}.e{}'.format(save_dir, epoch))

logging('Dividing learning rate by 10')

optimizer.param_groups[0]['lr'] /= 10.

best_val_loss.append(val_loss)

except KeyboardInterrupt:

logging('-' * 89)

logging('Exiting from training early')

# Load the best saved model.

# even if we resumed, we want to load the model saved on _this_ run, so use save_dir

model, criterion, optimizer = model_load(save_dir)

# Run on test data.

test_loss = evaluate(test_data, test_batch_size)

logging('=' * 89)

logging('| End of training | test loss {:5.2f} | test ppl {:8.2f} | test bpc {:8.3f}'.format(

test_loss, math.exp(test_loss), test_loss / math.log(2)))

logging('=' * 89)