class LanguageModelAgent(Agent):
""" Agent which trains an RNN on a language modeling task.
It is adapted from the language model featured in Pytorch's examples repo
here: <https://github.com/pytorch/examples/tree/master/word_language_model>.
"""
@staticmethod
def dictionary_class():
return DictionaryAgent
@staticmethod
def add_cmdline_args(argparser):
"""Add command-line arguments specifically for this agent."""
argparser.set_defaults(batch_sort=False)
agent = argparser.add_argument_group('Language Model Arguments')
agent.add_argument(
'--init-model',
type=str,
default=None,
help='load dict/features/weights/opts from this file')
agent.add_argument('-hs',
'--hiddensize',
type=int,
default=200,
help='size of the hidden layers')
agent.add_argument('-esz',
'--embeddingsize',
type=int,
default=200,
help='size of the token embeddings')
agent.add_argument('-nl',
'--numlayers',
type=int,
default=2,
help='number of hidden layers')
agent.add_argument('-dr',
'--dropout',
type=float,
default=0.2,
help='dropout rate')
agent.add_argument('-clip',
'--gradient-clip',
type=float,
default=0.25,
help='gradient clipping')
agent.add_argument('--no-cuda',
action='store_true',
default=False,
help='disable GPUs even if available')
agent.add_argument(
'-rnn',
'--rnn-class',
default='LSTM',
help='type of recurrent net (RNN_TANH, RNN_RELU, LSTM, GRU)')
agent.add_argument('-sl',
'--seq-len',
type=int,
default=35,
help='sequence length')
agent.add_argument('-tied',
'--emb-tied',
action='store_true',
help='tie the word embedding and softmax weights')
agent.add_argument('-seed',
'--random-seed',
type=int,
default=1111,
help='random seed')
agent.add_argument('--gpu',
type=int,
default=-1,
help='which GPU device to use')
agent.add_argument('-tr',
'--truncate-pred',
type=int,
default=50,
help='truncate predictions')
agent.add_argument('-rf',
'--report-freq',
type=float,
default=0.1,
help='report frequency of prediction during eval')
agent.add_argument('-pt',
'--person-tokens',
type='bool',
default=True,
help='append person1 and person2 tokens to text')
# learning rate parameters
agent.add_argument('-lr',
'--learningrate',
type=float,
default=20,
help='initial learning rate')
agent.add_argument(
'-lrf',
'--lr-factor',
type=float,
default=1.0,
help='mutliply learning rate by this factor when the \
validation loss does not decrease')
agent.add_argument('-lrp',
'--lr-patience',
type=int,
default=10,
help='wait before decreasing learning rate')
agent.add_argument('-lrm',
'--lr-minimum',
type=float,
default=0.1,
help='minimum learning rate')
agent.add_argument(
'-sm',
'--sampling-mode',
type='bool',
default=False,
help='sample when generating tokens instead of taking \
the max and do not produce UNK token (when bs=1)')
LanguageModelAgent.dictionary_class().add_cmdline_args(argparser)
return agent
def __init__(self, opt, shared=None):
"""Set up model if shared params not set, otherwise no work to do."""
super().__init__(opt, shared)
opt = self.opt # there is a deepcopy in the init
self.metrics = {
'loss': 0,
'num_tokens': 0,
'lmloss': 0,
'lm_num_tokens': 0
}
self.states = {}
# check for cuda
self.use_cuda = not opt.get('no_cuda') and torch.cuda.is_available()
self.batchsize = opt.get('batchsize', 1)
self.use_person_tokens = opt.get('person_tokens', True)
self.sampling_mode = opt.get('sampling_mode', False)
if shared:
# set up shared properties
self.opt = shared['opt']
opt = self.opt
self.dict = shared['dict']
self.model = shared['model']
self.metrics = shared['metrics']
# get NULL token and END token
self.NULL_IDX = self.dict[self.dict.null_token]
self.END_IDX = self.dict[self.dict.end_token]
if 'states' in shared:
self.states = shared['states']
if self.use_person_tokens:
# add person1 and person2 tokens
self.dict.add_to_dict(self.dict.tokenize("PERSON1"))
self.dict.add_to_dict(self.dict.tokenize("PERSON2"))
else:
# this is not a shared instance of this class, so do full init
if self.use_cuda:
print('[ Using CUDA ]')
torch.cuda.set_device(opt['gpu'])
init_model = None
# check first for 'init_model' for loading model from file
if opt.get('init_model') and os.path.isfile(opt['init_model']):
init_model = opt['init_model']
# next check for 'model_file', this would override init_model
if opt.get('model_file') and os.path.isfile(opt['model_file']):
init_model = opt['model_file']
# for backwards compatibility: will only be called for older models
# for which .opt file does not exist
if (init_model is not None
and not os.path.isfile(init_model + '.opt')):
new_opt = self.load_opt(init_model)
# load model parameters if available
print('[ Setting opt from {} ]'.format(init_model))
# since .opt file does not exist, save one for future use
print("Saving opt file at:", init_model + ".opt")
with open(init_model + ".opt", 'wb') as handle:
pickle.dump(new_opt,
handle,
protocol=pickle.HIGHEST_PROTOCOL)
opt = self.override_opt(new_opt)
if ((init_model is not None
and os.path.isfile(init_model + '.dict'))
or opt['dict_file'] is None):
opt['dict_file'] = init_model + '.dict'
# load dictionary and basic tokens & vectors
self.dict = DictionaryAgent(opt)
self.id = 'LanguageModel'
# get NULL token and END token
self.NULL_IDX = self.dict[self.dict.null_token]
self.END_IDX = self.dict[self.dict.end_token]
if self.use_person_tokens:
# add person1 and person2 tokens
self.dict.add_to_dict(self.dict.tokenize("PERSON1"))
self.dict.add_to_dict(self.dict.tokenize("PERSON2"))
# set model
self.model = RNNModel(opt, len(self.dict))
if init_model is not None:
self.load(init_model)
if self.use_cuda:
self.model.cuda()
self.next_observe = []
self.next_batch = []
self.is_training = True
self.clip = opt.get('gradient_clip', 0.25)
# set up criteria
self.criterion = nn.CrossEntropyLoss(ignore_index=self.NULL_IDX,
size_average=False)
if self.use_cuda:
# push to cuda
self.criterion.cuda()
# init hidden state
self.hidden = self.model.init_hidden(self.batchsize)
# init tensor of end tokens
self.ends = torch.LongTensor(
[self.END_IDX for _ in range(self.batchsize)])
if self.use_cuda:
self.ends = self.ends.cuda()
# set up model and learning rate scheduler parameters
self.lr = opt['learningrate']
self.optimizer = torch.optim.SGD(self.model.parameters(), lr=self.lr)
self.best_val_loss = self.states.get('best_val_loss', None)
self.lr_factor = opt['lr_factor']
if self.lr_factor < 1.0:
self.lr_patience = opt['lr_patience']
self.lr_min = opt['lr_minimum']
self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
self.optimizer,
factor=self.lr_factor,
verbose=True,
patience=self.lr_patience,
min_lr=self.lr_min)
# initial step for scheduler if self.best_val_loss is initialized
if self.best_val_loss is not None:
self.scheduler.step(self.best_val_loss)
else:
self.scheduler = None
self.reset()
def override_opt(self, new_opt):
"""Set overridable opts from loaded opt file.
Print out each added key and each overriden key.
Only override args specific to the model.
"""
model_args = {
'hiddensize', 'embeddingsize', 'numlayers', 'dropout', 'seq_len',
'emb_tied', 'truncate_pred', 'report_freq', 'person_tokens',
'learningrate'
}
for k, v in new_opt.items():
if k not in model_args:
# skip non-model args
continue
if k not in self.opt:
print('Adding new option [ {k}: {v} ]'.format(k=k, v=v))
elif self.opt[k] != v:
print('Overriding option [ {k}: {old} => {v}]'.format(
k=k, old=self.opt[k], v=v))
self.opt[k] = v
return self.opt
def parse(self, text):
"""Convert string to token indices."""
return self.dict.txt2vec(text)
def zero_grad(self):
"""Zero out optimizer."""
self.optimizer.zero_grad()
def update_params(self):
"""Do one optimization step."""
torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.clip)
self.optimizer.step()
def reset(self):
"""Reset observation and episode_done."""
self.observation = None
self.reset_metrics()
def reset_metrics(self):
self.metrics.clear()
self.metrics['loss'] = 0
self.metrics['lmloss'] = 0
self.metrics['num_tokens'] = 0
self.metrics['lm_num_tokens'] = 0
def report(self):
m = {}
if self.metrics['num_tokens'] > 0:
m['loss'] = self.metrics['loss'] / self.metrics['num_tokens']
m['ppl'] = math.exp(m['loss'])
if self.metrics['lm_num_tokens'] > 0:
m['lmloss'] = self.metrics['lmloss'] / self.metrics['lm_num_tokens']
m['lmppl'] = math.exp(m['lmloss'])
for k, v in m.items():
# clean up: rounds to sigfigs and converts tensors to floats
m[k] = round_sigfigs(v, 4)
return m
def share(self):
"""Share internal states between parent and child instances."""
shared = super().share()
shared['opt'] = self.opt
shared['dict'] = self.dict
shared['NULL_IDX'] = self.NULL_IDX
shared['END_IDX'] = self.END_IDX
shared['model'] = self.model
if self.opt.get('numthreads', 1) > 1:
if type(self.metrics) == dict:
# move metrics and model to shared memory
self.metrics = SharedTable(self.metrics)
self.model.share_memory()
shared['states'] = { # only need to pass optimizer states
'optimizer': self.optimizer.state_dict(),
}
shared['metrics'] = self.metrics
return shared
def observe(self, observation):
"""Save observation for act.
If multiple observations are from the same episode, concatenate them.
"""
# shallow copy observation (deep copy can be expensive)
obs = observation.copy()
seq_len = self.opt['seq_len']
is_training = True
if 'labels' not in obs:
is_training = False
if is_training:
if 'text' in obs:
if self.use_person_tokens:
obs['text'] = 'PERSON1 ' + obs['text']
vec = self.parse(obs['text'])
vec.append(self.END_IDX)
self.next_observe += vec
if 'labels' in obs:
if self.use_person_tokens:
labels = [
'PERSON2 ' + label for label in obs['labels']
if label != ''
]
obs['labels'] = tuple(labels)
vec = self.parse(obs['labels'][0])
vec.append(self.END_IDX)
self.next_observe += vec
if len(self.next_observe) < (seq_len + 1):
# not enough to return to make a batch
# we handle this case in vectorize
# labels indicates that we are training
self.observation = {'labels': ''}
return self.observation
else:
vecs_to_return = []
total = len(self.next_observe) // (seq_len + 1)
for _ in range(total):
observe = self.next_observe[:(seq_len + 1)]
self.next_observe = self.next_observe[(seq_len + 1):]
vecs_to_return.append(observe)
dict_to_return = {
'text': '',
'labels': '',
'text2vec': vecs_to_return
}
self.observation = dict_to_return
return dict_to_return
else:
if 'text' in obs:
if self.use_person_tokens:
obs['text'] = 'PERSON1 ' + obs['text']
if 'eval_labels' in obs:
if self.use_person_tokens:
eval_labels = [
'PERSON2 ' + label for label in obs['eval_labels']
if label != ''
]
obs['eval_labels'] = tuple(eval_labels)
self.observation = obs
return obs
def repackage_hidden(self, h):
"""Wraps hidden states in new Variables, to detach them from their history."""
if isinstance(h, Variable):
return Variable(h.data)
else:
return tuple(self.repackage_hidden(v) for v in h)
def get_target_loss(self, data, hidden, targets):
"""Calculates the loss with respect to the targets, token by token,
where each output token is conditioned on either the input or the
previous target token.
"""
loss = 0.0
bsz = data.size(0)
# during interactive mode, when no targets exist, we return 0
if targets is None:
return loss
# feed in inputs without end token
output, hidden = self.model(data.transpose(0, 1), hidden)
self.hidden = self.repackage_hidden(hidden)
# feed in end tokens
output, hidden = self.model(Variable(self.ends[:bsz].view(1, bsz)),
self.hidden)
self.hidden = self.repackage_hidden(hidden)
output_flat = output.view(-1, len(self.dict))
loss += self.criterion(output_flat, targets.select(1, 0).view(-1)).data
for i in range(1, targets.size(1)):
output, hidden = self.model(targets.select(1, i - 1).view(1, bsz),
self.hidden,
no_pack=True)
self.hidden = self.repackage_hidden(hidden)
output_flat = output.view(-1, len(self.dict))
loss += self.criterion(output_flat,
targets.select(1, i).view(-1)).data
return loss
def get_predictions(self, data):
"""Generates predictions word by word until we either reach the end token
or some max length (opt['truncate_pred']).
"""
token_list = []
bsz = data.size(0)
done = [False for _ in range(bsz)]
total_done = 0
hidden = self.model.init_hidden(bsz)
i = 0
word_idx = None
while total_done < bsz and i <= self.opt['truncate_pred']:
if i == 0:
# feed in input without end tokens
output, hidden = self.model(data.transpose(0, 1), hidden)
hidden = self.repackage_hidden(hidden)
# feed in end tokens
output, hidden = self.model(
Variable(self.ends[:bsz].view(1, bsz)), hidden)
else:
output, hidden = self.model(Variable(word_idx.view(1, bsz)),
hidden,
no_pack=True)
hidden = self.repackage_hidden(hidden)
word_weights = output.squeeze().data.exp()
if bsz > 1:
_, word_idx = torch.max(word_weights, 1)
else:
if self.sampling_mode:
unk_idx = self.dict[self.dict.unk_token]
# make word_weights have smaller norm so that calculated
# norm does not blow up
word_weights = word_weights.div(1e10)
# make word_weights have L2 norm 1
ww_norm = torch.norm(word_weights, p=2)
word_weights = word_weights.div(ww_norm)
# square distribution
word_weights = torch.mul(word_weights, word_weights)
# sample distribution
word_idx = torch.multinomial(word_weights, 1)
# do not produce UNK token
while word_idx == unk_idx:
word_idx = torch.multinomial(word_weights, 1)
else:
_, word_idx = torch.max(word_weights, 0)
# mark end indices for items in batch
word_idx = word_idx.view(-1)
for k in range(word_idx.size(0)):
if not done[k]:
if int(word_idx[k]) == self.END_IDX:
done[k] = True
total_done += 1
token_list.append(word_idx.view(bsz, 1))
i += 1
return torch.cat(token_list, 1)
def predict(self,
data,
hidden,
targets=None,
is_training=True,
y_lens=None):
"""Produce a prediction from our model."""
output = None
predictions = None
if is_training:
self.model.train()
self.zero_grad()
output, hidden = self.model(data, hidden)
loss = self.criterion(output.view(-1, len(self.dict)),
targets.view(-1))
# save loss to metrics
target_tokens = targets.ne(self.NULL_IDX).float().sum().item()
self.metrics['lmloss'] += loss.double().item()
self.metrics['lm_num_tokens'] += target_tokens
# average loss per token
loss /= target_tokens
loss.backward(retain_graph=True)
self.update_params()
else:
self.model.eval()
predictions = self.get_predictions(data)
bsz = data.size(0)
if bsz != self.batchsize:
self.hidden = self.model.init_hidden(bsz)
if targets is not None:
loss = self.get_target_loss(data, self.hidden, targets)
self.metrics['loss'] += loss
self.metrics['num_tokens'] += sum(y_lens)
return output, hidden, predictions
def vectorize(self, observations, seq_len, is_training):
"""Convert a list of observations into input & target tensors."""
labels = None
valid_inds = None
y_lens = None
if is_training:
for obs in observations:
if obs:
if 'text2vec' in obs:
self.next_batch += obs['text2vec']
if len(self.next_batch) <= self.batchsize:
return None, None, None, None, None
else:
data_list = []
targets_list = []
# total is the number of batches
total = len(self.next_batch) // self.batchsize
for _ in range(total):
batch = self.next_batch[:self.batchsize]
self.next_batch = self.next_batch[self.batchsize:]
source = torch.LongTensor(batch).t().contiguous()
data = Variable(source[:seq_len])
targets = Variable(source[1:])
if self.use_cuda:
data = data.cuda()
targets = targets.cuda()
data_list.append(data)
targets_list.append(targets)
else:
# here we get valid examples and pad them with zeros
xs, ys, labels, valid_inds, _, y_lens = PaddingUtils.pad_text(
observations,
self.dict,
end_idx=self.END_IDX,
null_idx=self.NULL_IDX)
if self.use_cuda:
if xs is not None:
xs = Variable(torch.LongTensor(xs)).cuda()
if ys is not None:
ys = Variable(torch.LongTensor(ys)).cuda()
else:
if xs is not None:
xs = Variable(torch.LongTensor(xs))
if ys is not None:
ys = Variable(torch.LongTensor(ys))
data_list = [xs]
targets_list = [ys]
return data_list, targets_list, labels, valid_inds, y_lens
def batch_act(self, observations):
batch_reply = [{'id': self.getID()} for _ in range(len(observations))]
if any(['labels' in obs for obs in observations]):
# if we are starting a new training epoch, reinitialize hidden
if not self.is_training:
self.hidden = self.model.init_hidden(self.batchsize)
self.is_training = True
data_list, targets_list, _c, _v, y_lens = self.vectorize(
observations, self.opt['seq_len'], self.is_training)
else:
# if we just finished training, reinitialize hidden
if self.is_training:
self.hidden = self.model.init_hidden(self.batchsize)
self.is_training = False
data_list, targets_list, labels, valid_inds, y_lens = self.vectorize(
observations, self.opt['seq_len'], self.is_training)
if data_list is None:
# not enough data to batch act yet, return empty responses
return batch_reply
batch_reply = []
# during evaluation, len(data_list) is always 1
# during training, len(dat_list) >= 0: vectorize returns a list
# containing all batches available at the time it is called
for i in range(len(data_list)):
temp_dicts = [{
'id': self.getID()
} for _ in range(len(observations))]
# ignore case when we do not return any valid indices
if data_list[i] is not None:
output, hidden, predictions = self.predict(
data_list[i], self.hidden, targets_list[i],
self.is_training, y_lens)
self.hidden = self.repackage_hidden(hidden)
if predictions is not None:
# map predictions back to the right order
PaddingUtils.map_predictions(
predictions.cpu(),
valid_inds,
temp_dicts,
observations,
self.dict,
self.END_IDX,
report_freq=self.opt['report_freq'])
batch_reply += temp_dicts
# for prediction metrics computations, we get rid of PERSON1 and PERSON2 tokens
if not self.is_training:
for reply in batch_reply:
if 'text' in reply:
reply['text'] = reply['text'].replace('PERSON1 ', '')
reply['text'] = reply['text'].replace('PERSON2 ', '')
return batch_reply
def act(self):
# call batch_act with this batch of one
return self.batch_act([self.observation])[0]
def save(self, path=None):
"""Save model parameters if model_file is set."""
path = self.opt.get('model_file', None) if path is None else path
if path and hasattr(self, 'model'):
model = {}
model['model'] = self.model.state_dict()
model['opt'] = self.opt
model['best_val_loss'] = self.best_val_loss
with open(path, 'wb') as write:
torch.save(model, write)
# save opt file
with open(path + ".opt", 'wb') as handle:
pickle.dump(self.opt, handle, protocol=pickle.HIGHEST_PROTOCOL)
def shutdown(self):
"""Save the state of the model when shutdown."""
path = self.opt.get('model_file', None)
if path is not None:
self.save(path + '.shutdown_state')
super().shutdown()
def receive_metrics(self, metrics_dict):
if 'loss' in metrics_dict and self.scheduler is not None:
self.scheduler.step(metrics_dict['loss'])
def load_opt(self, path):
"""Return opt, states."""
states = torch.load(path, map_location=lambda cpu, _: cpu)
return states['opt']
def load(self, path):
"""Load model states."""
if os.path.isfile(path):
# load model parameters if available
print('[ Loading existing model params from {} ]'.format(path))
self.states = torch.load(path, map_location=lambda cpu, _: cpu)
self.model.load_state_dict(self.states['model'])
class LanguageModelAgent(Agent):
""" Agent which trains an RNN on a language modeling task.
It is adapted from the language model featured in Pytorch's examples repo
here: <https://github.com/pytorch/examples/tree/master/word_language_model>.
"""
@staticmethod
def dictionary_class():
return DictionaryAgent
@staticmethod
def add_cmdline_args(argparser):
"""Add command-line arguments specifically for this agent."""
argparser.set_defaults(batch_sort=False)
LanguageModelAgent.dictionary_class().add_cmdline_args(argparser)
agent = argparser.add_argument_group('Language Model Arguments')
agent.add_argument('-hs',
'--hiddensize',
type=int,
default=200,
help='size of the hidden layers')
agent.add_argument('-esz',
'--embeddingsize',
type=int,
default=200,
help='size of the token embeddings')
agent.add_argument('-nl',
'--numlayers',
type=int,
default=2,
help='number of hidden layers')
agent.add_argument('-lr',
'--learningrate',
type=float,
default=20,
help='initial learning rate')
agent.add_argument('-dr',
'--dropout',
type=float,
default=0.2,
help='dropout rate')
agent.add_argument('-clip',
'--gradient-clip',
type=float,
default=0.25,
help='gradient clipping')
agent.add_argument('--no-cuda',
action='store_true',
default=False,
help='disable GPUs even if available')
agent.add_argument(
'-rnn',
'--rnn-class',
default='LSTM',
help='type of recurrent net (RNN_TANH, RNN_RELU, LSTM, GRU)')
agent.add_argument('-sl',
'--seq-len',
type=int,
default=35,
help='sequence length')
agent.add_argument('-tied',
'--emb-tied',
action='store_true',
help='tie the word embedding and softmax weights')
agent.add_argument('-seed',
'--random-seed',
type=int,
default=1111,
help='random seed')
agent.add_argument('--gpu',
type=int,
default=-1,
help='which GPU device to use')
agent.add_argument('-tr',
'--truncate-pred',
type=int,
default=50,
help='truncate predictions')
agent.add_argument('-rf',
'--report-freq',
type=float,
default=0.1,
help='report frequency of prediction during eval')
agent.add_argument('-pt',
'--person-tokens',
type=bool,
default=True,
help='append person1 and person2 tokens to text')
agent.add_argument(
'-lrf',
'--lr-factor',
type=float,
default=0.5,
help='mutliply learning rate by this factor when the \
validation loss does not decrease')
def __init__(self, opt, shared=None):
"""Set up model if shared params not set, otherwise no work to do."""
super().__init__(opt, shared)
opt = self.opt # there is a deepcopy in the init
self.states = {}
# check for cuda
self.use_cuda = not opt.get('no_cuda') and torch.cuda.is_available()
self.batchsize = opt.get('batchsize', 1)
self.use_person_tokens = opt.get('person_tokens', True)
if shared:
# set up shared properties
self.dict = shared['dict']
if 'model' in shared:
# model is shared during hogwild
self.model = shared['model']
self.states = shared['states']
# get NULL token and END token
self.NULL_IDX = self.dict[self.dict.null_token]
self.END_IDX = self.dict[self.dict.end_token]
if self.use_person_tokens:
# add person1 and person2 tokens
self.dict.add_to_dict(self.dict.tokenize("PERSON1"))
self.dict.add_to_dict(self.dict.tokenize("PERSON2"))
else:
# this is not a shared instance of this class, so do full init
if self.use_cuda:
print('[ Using CUDA ]')
torch.cuda.set_device(opt['gpu'])
if opt.get('model_file') and os.path.isfile(opt['model_file']):
# load model parameters if available
print('Loading existing model params from ' +
opt['model_file'])
new_opt, self.states = self.load(opt['model_file'])
# override model-specific options with stored ones
opt = self.override_opt(new_opt)
if opt['dict_file'] is None and opt.get('model_file'):
# set default dict-file if not set
opt['dict_file'] = opt['model_file'] + '.dict'
# load dictionary and basic tokens & vectors
self.dict = DictionaryAgent(opt)
self.id = 'LanguageModel'
# get NULL token and END token
self.NULL_IDX = self.dict[self.dict.null_token]
self.END_IDX = self.dict[self.dict.end_token]
if self.use_person_tokens:
# add person1 and person2 tokens
self.dict.add_to_dict(self.dict.tokenize("PERSON1"))
self.dict.add_to_dict(self.dict.tokenize("PERSON2"))
# set model
self.model = RNNModel(opt, len(self.dict))
if self.states:
# set loaded states if applicable
self.model.load_state_dict(self.states['model'])
if self.use_cuda:
self.model.cuda()
self.next_observe = []
self.next_batch = []
self.is_training = True
if hasattr(self, 'model'):
# if model was built, do more setup
self.clip = opt.get('gradient_clip', 0.25)
# set up criteria
self.criterion = nn.CrossEntropyLoss(ignore_index=self.NULL_IDX)
if self.use_cuda:
# push to cuda
self.criterion.cuda()
# set up criterion for eval: we do not want to average over size
self.eval_criterion = nn.CrossEntropyLoss(
ignore_index=self.NULL_IDX, size_average=False)
if self.use_cuda:
# push to cuda
self.eval_criterion.cuda()
# init hidden state
self.hidden = self.model.init_hidden(self.batchsize)
# init tensor of end tokens
self.ends = torch.LongTensor(
[self.END_IDX for _ in range(self.batchsize)])
if self.use_cuda:
self.ends = self.ends.cuda()
# set up optimizer
self.lr = opt['learningrate']
self.lr_factor = opt['lr_factor']
self.best_val_loss = None
self.reset()
def override_opt(self, new_opt):
"""Set overridable opts from loaded opt file.
Print out each added key and each overriden key.
Only override args specific to the model.
"""
model_args = {
'hiddensize', 'embeddingsize', 'numlayers', 'dropout', 'seq_len',
'emb_tied'
}
for k, v in new_opt.items():
if k not in model_args:
# skip non-model args
continue
if k not in self.opt:
print('Adding new option [ {k}: {v} ]'.format(k=k, v=v))
elif self.opt[k] != v:
print('Overriding option [ {k}: {old} => {v}]'.format(
k=k, old=self.opt[k], v=v))
self.opt[k] = v
return self.opt
def parse(self, text):
"""Convert string to token indices."""
return self.dict.txt2vec(text)
def zero_grad(self):
"""Zero out optimizer."""
self.model.zero_grad()
def update_params(self):
"""Do one optimization step."""
torch.nn.utils.clip_grad_norm(self.model.parameters(), self.clip)
for p in self.model.parameters():
p.data.add_(-self.lr, p.grad.data)
def reset(self):
"""Reset observation and episode_done."""
self.observation = None
def share(self):
"""Share internal states between parent and child instances."""
shared = super().share()
shared['dict'] = self.dict
shared['NULL_IDX'] = self.NULL_IDX
shared['END_IDX'] = self.END_IDX
if self.opt.get('numthreads', 1) > 1:
shared['model'] = self.model
self.model.share_memory()
shared['states'] = self.states
return shared
def observe(self, observation):
"""Save observation for act.
If multiple observations are from the same episode, concatenate them.
"""
#shallow copy observation (deep copy can be expensive)
obs = observation.copy()
seq_len = self.opt['seq_len']
is_training = True
if 'labels' not in obs:
is_training = False
if is_training:
if 'text' in obs:
if self.use_person_tokens:
obs['text'] = 'PERSON1 ' + obs['text']
vec = self.parse(obs['text'])
vec.append(self.END_IDX)
self.next_observe += vec
if 'labels' in obs:
if self.use_person_tokens:
labels = [
'PERSON2 ' + label for label in obs['labels']
if label != ''
]
obs['labels'] = tuple(labels)
vec = self.parse(obs['labels'][0])
vec.append(self.END_IDX)
self.next_observe += vec
if len(self.next_observe) < (seq_len + 1):
# not enough to return to make a batch
# we handle this case in vectorize
# labels indicates that we are training
self.observation = {'labels': ''}
return self.observation
else:
vecs_to_return = []
total = len(self.next_observe) // (seq_len + 1)
for _ in range(total):
observe = self.next_observe[:(seq_len + 1)]
self.next_observe = self.next_observe[(seq_len + 1):]
vecs_to_return.append(observe)
dict_to_return = {
'text': '',
'labels': '',
'text2vec': vecs_to_return
}
self.observation = dict_to_return
return dict_to_return
else:
if 'text' in obs:
if self.use_person_tokens:
obs['text'] = 'PERSON1 ' + obs['text']
if 'eval_labels' in obs:
if self.use_person_tokens:
eval_labels = [
'PERSON2 ' + label for label in obs['eval_labels']
if label != ''
]
obs['eval_labels'] = tuple(eval_labels)
self.observation = obs
return obs
def repackage_hidden(self, h):
"""Wraps hidden states in new Variables, to detach them from their history."""
if type(h) == Variable:
return Variable(h.data)
else:
return tuple(self.repackage_hidden(v) for v in h)
def get_target_loss(self, data, hidden, targets, y_lens):
"""Calculates the loss with respect to the targets, token by token,
where each output token is conditioned on either the input or the
previous target token.
"""
loss = 0.0
bsz = data.size(0)
# during interactive mode, when no targets exist, we return 0
if targets is None:
return loss
# feed in inputs without end token
output, hidden = self.model(data.transpose(0, 1), hidden)
self.hidden = self.repackage_hidden(hidden)
# feed in end tokens
output, hidden = self.model(Variable(self.ends[:bsz].view(1, bsz)),
self.hidden)
self.hidden = self.repackage_hidden(hidden)
output_flat = output.view(-1, len(self.dict))
loss += self.eval_criterion(output_flat,
targets.select(1, 0).view(-1)).data
for i in range(1, targets.size(1)):
output, hidden = self.model(targets.select(1, i - 1).view(1, bsz),
self.hidden,
no_pack=True)
self.hidden = self.repackage_hidden(hidden)
output_flat = output.view(-1, len(self.dict))
loss += self.eval_criterion(output_flat,
targets.select(1, i).view(-1)).data
return loss / float(sum(y_lens))
def get_predictions(self, data):
"""Generates predictions word by word until we either reach the end token
or some max length (opt['truncate_pred']).
"""
token_list = []
bsz = data.size(0)
done = [False for _ in range(bsz)]
total_done = 0
hidden = self.model.init_hidden(bsz)
i = 0
while total_done < bsz and i <= self.opt['truncate_pred']:
if i == 0:
# feed in input without end tokens
output, hidden = self.model(data.transpose(0, 1), hidden)
hidden = self.repackage_hidden(hidden)
# feed in end tokens
output, hidden = self.model(
Variable(self.ends[:bsz].view(1, bsz)), hidden)
else:
output, hidden = self.model(Variable(word_idx.view(1, bsz)),
hidden,
no_pack=True)
hidden = self.repackage_hidden(hidden)
word_weights = output.squeeze().data.exp()
if bsz > 1:
value, word_idx = torch.max(word_weights, 1)
else:
value, word_idx = torch.max(word_weights, 0)
# mark end indices for items in batch
for k in range(word_idx.size(0)):
if not done[k]:
if int(word_idx[k]) == self.END_IDX:
done[k] = True
total_done += 1
token_list.append(word_idx.view(bsz, 1))
i += 1
if token_list:
return torch.cat(token_list, 1)
else:
return None
def predict(self,
data,
hidden,
targets=None,
is_training=True,
y_lens=None):
"""Produce a prediction from our model.
"""
loss_dict = None
output = None
predictions = None
if is_training:
self.model.train()
self.zero_grad()
output, hidden = self.model(data, hidden)
loss = self.criterion(output.view(-1, len(self.dict)),
targets.view(-1))
loss.backward(retain_graph=True)
self.update_params()
loss_dict = {'lmloss': loss.data}
loss_dict['lmppl'] = math.exp(loss.data)
else:
self.model.eval()
predictions = self.get_predictions(data)
loss_dict = {}
bsz = data.size(0)
if bsz != self.batchsize:
self.hidden = self.model.init_hidden(bsz)
loss = self.get_target_loss(data, self.hidden, targets, y_lens)
loss_dict['loss'] = loss
loss_dict['ppl'] = math.exp(loss)
return output, hidden, loss_dict, predictions
def vectorize(self, observations, seq_len, is_training):
"""Convert a list of observations into input & target tensors."""
labels = None
valid_inds = None
y_lens = None
if is_training:
for obs in observations:
if obs:
if 'text2vec' in obs:
self.next_batch += obs['text2vec']
if len(self.next_batch) <= self.batchsize:
return None, None, None, None, None
else:
data_list = []
targets_list = []
# total is the number of batches
total = len(self.next_batch) // self.batchsize
for i in range(total):
batch = self.next_batch[:self.batchsize]
self.next_batch = self.next_batch[self.batchsize:]
source = torch.LongTensor(batch).t().contiguous()
data = Variable(source[:seq_len])
targets = Variable(source[1:])
if self.use_cuda:
data = data.cuda()
targets = targets.cuda()
data_list.append(data)
targets_list.append(targets)
else:
# here we get valid examples and pad them with zeros
xs, ys, labels, valid_inds, _, y_lens = PaddingUtils.pad_text(
observations, self.dict, self.END_IDX, self.NULL_IDX)
if self.use_cuda:
xs = Variable(xs).cuda()
if ys is not None:
ys = Variable(ys).cuda()
else:
xs = Variable(xs)
if ys is not None:
ys = Variable(ys)
data_list = [xs]
targets_list = [ys]
return data_list, targets_list, labels, valid_inds, y_lens
def batch_act(self, observations):
batch_reply = [{'id': self.getID()} for _ in range(len(observations))]
if any(['labels' in obs for obs in observations]):
# if we are starting a new training epoch, reinitialize hidden
if self.is_training == False:
self.hidden = self.model.init_hidden(self.batchsize)
self.is_training = True
data_list, targets_list, _, _, y_lens = self.vectorize(
observations, self.opt['seq_len'], self.is_training)
else:
# if we just finished training, reinitialize hidden
if self.is_training == True:
self.hidden = self.model.init_hidden(self.batchsize)
self.is_training = False
data_list, targets_list, labels, valid_inds, y_lens = self.vectorize(
observations, self.opt['seq_len'], self.is_training)
if data_list is None:
# not enough data to batch act yet, return empty responses
return batch_reply
batch_reply = []
# during evaluation, len(data_list) is always 1
# during training, len(dat_list) >= 0: vectorize returns a list containing all batches available at the time it is called
for i in range(len(data_list)):
temp_dicts = [{
'id': self.getID()
} for _ in range(len(observations))]
output, hidden, loss_dict, predictions = self.predict(
data_list[i], self.hidden, targets_list[i], self.is_training,
y_lens)
self.hidden = self.repackage_hidden(hidden)
if predictions is not None:
# map predictions back to the right order
PaddingUtils.map_predictions(
predictions,
valid_inds,
temp_dicts,
observations,
self.dict,
self.END_IDX,
report_freq=self.opt['report_freq'])
if loss_dict is not None:
if 'metrics' in temp_dicts[0]:
for k, v in loss_dict.items():
temp_dicts[0]['metrics'][k] = v
else:
temp_dicts[0]['metrics'] = loss_dict
batch_reply += temp_dicts
return batch_reply
def act(self):
# call batch_act with this batch of one
return self.batch_act([self.observation])[0]
def save(self, path=None):
"""Save model parameters if model_file is set."""
path = self.opt.get('model_file', None) if path is None else path
if path and hasattr(self, 'model'):
model = {}
model['model'] = self.model.state_dict()
model['opt'] = self.opt
with open(path, 'wb') as write:
torch.save(model, write)
def shutdown(self):
"""Save the state of the model when shutdown."""
path = self.opt.get('model_file', None)
if path is not None:
self.save(path + '.shutdown_state')
super().shutdown()
def receive_metrics(self, metrics_dict):
if 'loss' in metrics_dict:
if self.best_val_loss is None:
self.best_val_loss = metrics_dict['loss']
else:
if metrics_dict['loss'] > self.best_val_loss:
self.lr *= self.lr_factor
print("Updating learning rate: lr =", self.lr)
def load(self, path):
"""Return opt and model states."""
with open(path, 'rb') as read:
states = torch.load(read)
return states['opt'], states
class MemnnFeedbackAgent(Agent):
"""
Memory Network agent for question answering that supports reward-based learning
(RBI), forward prediction (FP), and imitation learning (IM).
For more details on settings see: https://arxiv.org/abs/1604.06045.
Models settings 'FP', 'RBI', 'RBI+FP', and 'IM_feedback' assume that
feedback and reward for the current example immediatly follow the query
(add ':feedback' argument when specifying task name).
python examples/train_model.py --setting 'FP'
-m "projects.memnn_feedback.agent.memnn_feedback:MemnnFeedbackAgent"
-t "projects.memnn_feedback.tasks.dbll_babi.agents:taskTeacher:3_p0.5:feedback"
"""
@staticmethod
def add_cmdline_args(argparser):
DictionaryAgent.add_cmdline_args(argparser)
arg_group = argparser.add_argument_group('MemNN Arguments')
arg_group.add_argument('-lr',
'--learning-rate',
type=float,
default=0.01,
help='learning rate')
arg_group.add_argument('--embedding-size',
type=int,
default=128,
help='size of token embeddings')
arg_group.add_argument('--hops',
type=int,
default=3,
help='number of memory hops')
arg_group.add_argument('--mem-size',
type=int,
default=100,
help='size of memory')
arg_group.add_argument(
'--time-features',
type='bool',
default=True,
help='use time features for memory embeddings',
)
arg_group.add_argument(
'--position-encoding',
type='bool',
default=False,
help='use position encoding instead of bag of words embedding',
)
arg_group.add_argument(
'-clip',
'--gradient-clip',
type=float,
default=0.2,
help='gradient clipping using l2 norm',
)
arg_group.add_argument('--output',
type=str,
default='rank',
help='type of output (rank|generate)')
arg_group.add_argument(
'--rnn-layers',
type=int,
default=2,
help='number of hidden layers in RNN decoder for generative output',
)
arg_group.add_argument(
'--dropout',
type=float,
default=0.1,
help='dropout probability for RNN decoder training',
)
arg_group.add_argument('--optimizer',
default='sgd',
help='optimizer type (sgd|adam)')
arg_group.add_argument(
'--no-cuda',
action='store_true',
default=False,
help='disable GPUs even if available',
)
arg_group.add_argument('--gpu',
type=int,
default=-1,
help='which GPU device to use')
arg_group.add_argument(
'--setting',
type=str,
default='IM',
help='choose among IM, IM_feedback, RBI, FP, RBI+FP',
)
arg_group.add_argument(
'--num-feedback-cands',
type=int,
default=6,
help='number of feedback candidates',
)
arg_group.add_argument(
'--single_embedder',
type='bool',
default=False,
help='number of embedding matrices in the model',
)
def __init__(self, opt, shared=None):
super().__init__(opt, shared)
opt['cuda'] = not opt['no_cuda'] and torch.cuda.is_available()
if opt['cuda']:
print('[ Using CUDA ]')
torch.cuda.device(opt['gpu'])
if not shared:
self.id = 'MemNN'
self.dict = DictionaryAgent(opt)
self.decoder = None
if opt['output'] == 'generate' or opt['output'] == 'g':
self.decoder = Decoder(
opt['embedding_size'],
opt['embedding_size'],
opt['rnn_layers'],
opt,
self.dict,
)
elif opt['output'] != 'rank' and opt['output'] != 'r':
raise NotImplementedError('Output type not supported.')
if 'FP' in opt['setting']:
# add extra beta-word to indicate learner's answer
self.beta_word = 'betaword'
self.dict.add_to_dict([self.beta_word])
self.model = MemNN(opt, self.dict)
optim_params = [
p for p in self.model.parameters() if p.requires_grad
]
lr = opt['learning_rate']
if opt['optimizer'] == 'sgd':
self.optimizers = {'memnn': optim.SGD(optim_params, lr=lr)}
if self.decoder is not None:
self.optimizers['decoder'] = optim.SGD(
self.decoder.parameters(), lr=lr)
elif opt['optimizer'] == 'adam':
self.optimizers = {'memnn': optim.Adam(optim_params, lr=lr)}
if self.decoder is not None:
self.optimizers['decoder'] = optim.Adam(
self.decoder.parameters(), lr=lr)
else:
raise NotImplementedError('Optimizer not supported.')
if opt['cuda']:
self.model.share_memory()
if self.decoder is not None:
self.decoder.cuda()
if opt.get('model_file') and os.path.isfile(opt['model_file']):
print('Loading existing model parameters from ' +
opt['model_file'])
self.load(opt['model_file'])
else:
if 'model' in shared:
# model is shared during hogwild
self.model = shared['model']
self.dict = shared['dict']
self.decoder = shared['decoder']
self.optimizers = shared['optimizer']
if 'FP' in opt['setting']:
self.beta_word = shared['betaword']
if hasattr(self, 'model'):
self.opt = opt
self.mem_size = opt['mem_size']
self.loss_fn = CrossEntropyLoss()
self.gradient_clip = opt.get('gradient_clip', 0.2)
self.model_setting = opt['setting']
if 'FP' in opt['setting']:
self.feedback_cands = set([])
self.num_feedback_cands = opt['num_feedback_cands']
self.longest_label = 1
self.END = self.dict.end_token
self.END_TENSOR = torch.LongTensor(self.dict.parse(self.END))
self.START = self.dict.start_token
self.START_TENSOR = torch.LongTensor(self.dict.parse(self.START))
self.reset()
self.last_cands, self.last_cands_list = None, None
def share(self):
# Share internal states between parent and child instances
shared = super().share()
if self.opt.get('numthreads', 1) > 1:
shared['model'] = self.model
self.model.share_memory()
shared['optimizer'] = self.optimizers
shared['dict'] = self.dict
shared['decoder'] = self.decoder
if 'FP' in self.model_setting:
shared['betaword'] = self.beta_word
return shared
def observe(self, observation):
observation = copy.copy(observation)
# extract feedback for forward prediction
# IM setting - no feedback provided in the dataset
if self.opt['setting'] != 'IM':
if 'text' in observation:
split = observation['text'].split('\n')
feedback = split[-1]
observation['feedback'] = feedback
observation['text'] = '\n'.join(split[:-1])
if not self.episode_done:
# if the last example wasn't the end of an episode, then we need to
# recall what was said in that example
prev_dialogue = (self.observation['text']
if self.observation is not None else '')
# append answer and feedback (if available) given in the previous example to the previous dialog
if 'eval_labels' in self.observation:
prev_dialogue += '\n' + random.choice(
self.observation['eval_labels'])
elif 'labels' in self.observation:
prev_dialogue += '\n' + random.choice(
self.observation['labels'])
if 'feedback' in self.observation:
prev_dialogue += '\n' + self.observation['feedback']
observation['text'] = prev_dialogue + '\n' + observation['text']
self.observation = observation
self.episode_done = observation['episode_done']
return observation
def reset(self):
# reset observation and episode_done
self.observation = None
self.episode_done = True
def backward(self, loss, retain_graph=False):
# zero out optimizer and take one optimization step
for o in self.optimizers.values():
o.zero_grad()
loss.backward(retain_graph=retain_graph)
torch.nn.utils.clip_grad_norm(self.model.parameters(),
self.gradient_clip)
for o in self.optimizers.values():
o.step()
def parse_cands(self, cand_answers):
"""Returns:
cand_answers = tensor (vector) of token indices for answer candidates
cand_answers_lengths = tensor (vector) with lengths of each answer candidate
"""
parsed_cands = [to_tensors(c, self.dict) for c in cand_answers]
cand_answers_tensor = torch.cat([x[1] for x in parsed_cands])
max_cands_len = max([len(c) for c in cand_answers])
cand_answers_lengths = torch.LongTensor(len(cand_answers),
max_cands_len).zero_()
for i in range(len(cand_answers)):
if len(parsed_cands[i][0]) > 0:
cand_answers_lengths[
i, -len(parsed_cands[i][0]):] = parsed_cands[i][0]
cand_answers_tensor = Variable(cand_answers_tensor)
cand_answers_lengths = Variable(cand_answers_lengths)
return cand_answers_tensor, cand_answers_lengths
def get_cand_embeddings_with_added_beta(self, cands, selected_answer_inds):
# add beta_word to the candidate selected by the learner to indicate learner's answer
cand_answers_with_beta = copy.deepcopy(cands)
for i in range(len(cand_answers_with_beta)):
cand_answers_with_beta[i][
selected_answer_inds[i]] += ' ' + self.beta_word
# get candidate embeddings after adding beta_word to the selected candidate
(
cand_answers_tensor_with_beta,
cand_answers_lengths_with_beta,
) = self.parse_cands(cand_answers_with_beta)
cands_embeddings_with_beta = self.model.answer_embedder(
cand_answers_lengths_with_beta, cand_answers_tensor_with_beta)
if self.opt['cuda']:
cands_embeddings_with_beta = cands_embeddings_with_beta.cuda()
return cands_embeddings_with_beta
def predict(self, xs, answer_cands, ys=None, feedback_cands=None):
is_training = ys is not None
if is_training and 'FP' not in self.model_setting:
# Subsample to reduce training time
answer_cands = [
list(set(random.sample(c, min(32, len(c))) + self.labels))
for c in answer_cands
]
else:
# rank all cands to increase accuracy
answer_cands = [list(set(c)) for c in answer_cands]
self.model.train(mode=is_training)
# Organize inputs for network (see contents of xs and ys in batchify method)
inputs = [Variable(x, volatile=is_training) for x in xs]
if self.decoder:
output_embeddings = self.model(*inputs)
self.decoder.train(mode=is_training)
output_lines, loss = self.decode(output_embeddings, ys)
predictions = self.generated_predictions(output_lines)
self.backward(loss)
return predictions
scores = None
if is_training:
label_inds = [
cand_list.index(self.labels[i])
for i, cand_list in enumerate(answer_cands)
]
if 'FP' in self.model_setting:
if len(feedback_cands) == 0:
print(
'FP is not training... waiting for negative feedback examples'
)
else:
cand_answers_embs_with_beta = self.get_cand_embeddings_with_added_beta(
answer_cands, label_inds)
scores, forward_prediction_output = self.model(
*inputs, answer_cands, cand_answers_embs_with_beta)
fp_scores = self.model.get_score(feedback_cands,
forward_prediction_output,
forward_predict=True)
feedback_label_inds = [
cand_list.index(self.feedback_labels[i])
for i, cand_list in enumerate(feedback_cands)
]
if self.opt['cuda']:
feedback_label_inds = Variable(
torch.cuda.LongTensor(feedback_label_inds))
else:
feedback_label_inds = Variable(
torch.LongTensor(feedback_label_inds))
loss_fp = self.loss_fn(fp_scores, feedback_label_inds)
if loss_fp.data[0] > 100000:
raise Exception("Loss might be diverging. Loss:",
loss_fp.data[0])
self.backward(loss_fp, retain_graph=True)
if self.opt['cuda']:
label_inds = Variable(torch.cuda.LongTensor(label_inds))
else:
label_inds = Variable(torch.LongTensor(label_inds))
if scores is None:
output_embeddings = self.model(*inputs)
scores = self.model.get_score(answer_cands, output_embeddings)
predictions = self.ranked_predictions(answer_cands, scores)
if is_training:
update_params = True
# don't perform regular training if in FP mode
if self.model_setting == 'FP':
update_params = False
elif 'RBI' in self.model_setting:
if len(self.rewarded_examples_inds) == 0:
update_params = False
else:
self.rewarded_examples_inds = torch.LongTensor(
self.rewarded_examples_inds)
if self.opt['cuda']:
self.rewarded_examples_inds = self.rewarded_examples_inds.cuda(
)
# use only rewarded examples for training
loss = self.loss_fn(
scores[self.rewarded_examples_inds, :],
label_inds[self.rewarded_examples_inds],
)
else:
# regular IM training
loss = self.loss_fn(scores, label_inds)
if update_params:
self.backward(loss)
return predictions
def ranked_predictions(self, cands, scores):
_, inds = scores.data.sort(descending=True, dim=1)
return [[cands[i][j] for j in r if j < len(cands[i])]
for i, r in enumerate(inds)]
def decode(self, output_embeddings, ys=None):
batchsize = output_embeddings.size(0)
hn = output_embeddings.unsqueeze(0).expand(self.opt['rnn_layers'],
batchsize,
output_embeddings.size(1))
x = self.model.answer_embedder(Variable(torch.LongTensor([1])),
Variable(self.START_TENSOR))
xes = x.unsqueeze(1).expand(x.size(0), batchsize, x.size(1))
loss = 0
output_lines = [[] for _ in range(batchsize)]
done = [False for _ in range(batchsize)]
total_done = 0
idx = 0
while (total_done < batchsize) and idx < self.longest_label:
# keep producing tokens until we hit END or max length for each ex
if self.opt['cuda']:
xes = xes.cuda()
hn = hn.contiguous()
preds, scores = self.decoder(xes, hn)
if ys is not None:
y = Variable(ys[0][:, idx])
temp_y = y.cuda() if self.opt['cuda'] else y
loss += self.loss_fn(scores, temp_y)
else:
y = preds
# use the true token as the next input for better training
xes = self.model.answer_embedder(
Variable(torch.LongTensor(preds.numel()).fill_(1)),
y).unsqueeze(0)
for b in range(batchsize):
if not done[b]:
token = self.dict.vec2txt(preds.data[b])
if token == self.END:
done[b] = True
total_done += 1
else:
output_lines[b].append(token)
idx += 1
return output_lines, loss
def generated_predictions(self, output_lines):
return [[
' '.join(c for c in o
if c != self.END and c != self.dict.null_token)
] for o in output_lines]
def parse(self, text):
"""Returns:
query = tensor (vector) of token indices for query
query_length = length of query
memory = tensor (matrix) where each row contains token indices for a memory
memory_lengths = tensor (vector) with lengths of each memory
"""
sp = text.split('\n')
query_sentence = sp[-1]
query = self.dict.txt2vec(query_sentence)
query = torch.LongTensor(query)
query_length = torch.LongTensor([len(query)])
sp = sp[:-1]
sentences = []
for s in sp:
sentences.extend(s.split('\t'))
if len(sentences) == 0:
sentences.append(self.dict.null_token)
num_mems = min(self.mem_size, len(sentences))
memory_sentences = sentences[-num_mems:]
memory = [self.dict.txt2vec(s) for s in memory_sentences]
memory = [torch.LongTensor(m) for m in memory]
memory_lengths = torch.LongTensor([len(m) for m in memory])
memory = torch.cat(memory)
return (query, memory, query_length, memory_lengths)
def batchify(self, obs):
"""Returns:
xs = [memories, queries, memory_lengths, query_lengths]
ys = [labels, label_lengths] (if available, else None)
cands = list of candidates for each example in batch
valid_inds = list of indices for examples with valid observations
"""
exs = [ex for ex in obs if 'text' in ex]
valid_inds = [i for i, ex in enumerate(obs) if 'text' in ex]
if not exs:
return [None] * 5
if 'RBI' in self.model_setting:
self.rewarded_examples_inds = [
i for i, ex in enumerate(obs)
if 'text' in ex and ex.get('reward', 0) > 0
]
parsed = [self.parse(ex['text']) for ex in exs]
queries = torch.cat([x[0] for x in parsed])
memories = torch.cat([x[1] for x in parsed])
query_lengths = torch.cat([x[2] for x in parsed])
memory_lengths = torch.LongTensor(len(exs), self.mem_size).zero_()
for i in range(len(exs)):
if len(parsed[i][3]) > 0:
memory_lengths[i, -len(parsed[i][3]):] = parsed[i][3]
xs = [memories, queries, memory_lengths, query_lengths]
ys = None
self.labels = [
random.choice(ex['labels']) for ex in exs if 'labels' in ex
]
if len(self.labels) == len(exs):
parsed = [self.dict.txt2vec(l) for l in self.labels]
parsed = [torch.LongTensor(p) for p in parsed]
label_lengths = torch.LongTensor([len(p)
for p in parsed]).unsqueeze(1)
self.longest_label = max(self.longest_label, label_lengths.max())
padded = [
torch.cat((
p,
torch.LongTensor(self.longest_label - len(p)).fill_(
self.END_TENSOR[0]),
)) for p in parsed
]
labels = torch.stack(padded)
ys = [labels, label_lengths]
feedback_cands = []
if 'FP' in self.model_setting:
self.feedback_labels = [
ex['feedback'] for ex in exs
if 'feedback' in ex and ex['feedback'] is not None
]
self.feedback_cands = self.feedback_cands | set(
self.feedback_labels)
if (len(self.feedback_labels) == len(exs)
and len(self.feedback_cands) > self.num_feedback_cands):
feedback_cands = [
list(
set(
random.sample(self.feedback_cands,
self.num_feedback_cands) +
[feedback])) for feedback in self.feedback_labels
]
cands = [
ex['label_candidates'] for ex in exs if 'label_candidates' in ex
]
# Use words in dict as candidates if no candidates are provided
if len(cands) < len(exs):
cands = build_cands(exs, self.dict)
# Avoid rebuilding candidate list every batch if its the same
if self.last_cands != cands:
self.last_cands = cands
self.last_cands_list = [list(c) for c in cands]
cands = self.last_cands_list
return xs, ys, cands, valid_inds, feedback_cands
def batch_act(self, observations):
batchsize = len(observations)
batch_reply = [{'id': self.getID()} for _ in range(batchsize)]
xs, ys, cands, valid_inds, feedback_cands = self.batchify(observations)
if xs is None or len(xs[1]) == 0:
return batch_reply
# Either train or predict
predictions = self.predict(xs, cands, ys, feedback_cands)
for i in range(len(valid_inds)):
batch_reply[valid_inds[i]]['text'] = predictions[i][0]
batch_reply[valid_inds[i]]['text_candidates'] = predictions[i]
return batch_reply
def act(self):
return self.batch_act([self.observation])[0]
def save(self, path=None):
path = self.opt.get('model_file', None) if path is None else path
if path:
checkpoint = {}
checkpoint['memnn'] = self.model.state_dict()
checkpoint['memnn_optim'] = self.optimizers['memnn'].state_dict()
if self.decoder is not None:
checkpoint['decoder'] = self.decoder.state_dict()
checkpoint['decoder_optim'] = self.optimizers[
'decoder'].state_dict()
checkpoint['longest_label'] = self.longest_label
with open(path, 'wb') as write:
torch.save(checkpoint, write)
def load(self, path):
with open(path, 'rb') as read:
checkpoint = torch.load(read)
self.model.load_state_dict(checkpoint['memnn'])
self.optimizers['memnn'].load_state_dict(checkpoint['memnn_optim'])
if self.decoder is not None:
self.decoder.load_state_dict(checkpoint['decoder'])
self.optimizers['decoder'].load_state_dict(
checkpoint['decoder_optim'])
self.longest_label = checkpoint['longest_label']
import pickle
import os
from parlai.core.dict import DictionaryAgent
path = 'dat/MovieTriples_Dataset.tar'
with open(os.path.join(path, 'Training.dict.pkl'), 'rb') as data_file:
dictionary = pickle.load(data_file)
parlai_dict = DictionaryAgent({'vocab_size': 10004})
dictionary = sorted(dictionary, key=lambda x: x[1])
print(dictionary[:10])
for word in dictionary:
# print(word[0])
parlai_dict.add_to_dict([word[0]])
parlai_dict.freq[word[0]] = word[2]
# print(word)
# print(parlai_dict)
# parlai_dict.add_to_dict(['hello'])
parlai_dict.save('test_hred.dict', sort=True)
