def make_batch(self, seq, gpu):
tokens = [self.vocab.start_token] + seq + [self.vocab.stop_token]
idxs = torch.LongTensor([[self.vocab.index(t)] for t in tokens])
lengths = torch.LongTensor([len(tokens)])
encoder_inputs = Variable(idxs,
lengths,
batch_dim=1,
length_dim=0,
pad_value=0)
return {'encoder_input': encoder_inputs.cuda(gpu)}
def init_state(self, batch_size, encoder_state):
output = Variable(
torch.LongTensor([self.vocab.start_index] * batch_size)\
.view(1, -1),
lengths=torch.LongTensor([1] * batch_size),
length_dim=0, batch_dim=1, pad_value=self.vocab.pad_index)
if str(encoder_state.device) != "cpu":
output = output.cuda(encoder_state.device)
return {"output": output, "decoder_state": encoder_state}
def make_generator_inputs(self, data):
tokens = [self.source_vocab.start_token] + data["source"] \
+ [self.source_vocab.stop_token]
inputs = Variable(torch.LongTensor(
[[self.source_vocab[t] for t in tokens]]).t(),
lengths=torch.LongTensor([len(tokens)]),
length_dim=0,
batch_dim=1,
pad_value=self.source_vocab.pad_index)
if self._gpu > -1:
inputs = inputs.cuda(self._gpu)
return {"source_inputs": inputs}
def _batch_labels(self, labels_batch):
input_tokens = torch.LongTensor(
[[self.source_vocab[tok] for tok in self._labels2input(labels)]
for labels in labels_batch])
length = input_tokens.size(1)
inputs = Variable(input_tokens.t(),
lengths=torch.LongTensor([length] *
len(labels_batch)),
length_dim=0,
batch_dim=1,
pad_value=-1)
if self._gpu > -1:
inputs = inputs.cuda(self._gpu)
return {"source_inputs": inputs}
def init_state(self, batch_size, device, source_indices,
encoder_state=None):
output = Variable(
torch.LongTensor([self.vocab.start_index] * batch_size)\
.view(1, -1),
lengths=torch.LongTensor([1] * batch_size),
length_dim=0, batch_dim=1, pad_value=self.vocab.pad_index)
if str(device) != "cpu":
output = output.cuda(device)
mask = source_indices.mask.t()
return {"target": output, "decoder_state": encoder_state,
"inputs": output, 'pointer_mask': mask}
def search2inputs(search_outputs, vocab, gpu):
input_lens = [len(x) + 1 for x in search_outputs]
max_len = max(input_lens)
inputs = []
for out in search_outputs:
out = [vocab.start_token] + out
if len(out) < max_len:
out = out + [vocab.pad_token] * (max_len - len(out))
inputs.append([vocab[t] for t in out])
inputs = torch.LongTensor(inputs)
input_lens = torch.LongTensor(input_lens)
inputs = Variable(
inputs.t(),
lengths=input_lens,
batch_dim=1, length_dim=0, pad_value=vocab.pad_index)
if gpu > -1:
inputs = inputs.cuda(gpu)
return {"inputs": inputs}
def init_state_context(self, encoder_state):
batch_size = encoder_state["state"].size(1) * self.samples
output = Variable(torch.LongTensor([self.vocab.start_index] *
batch_size).view(1, -1),
lengths=torch.LongTensor([1] * batch_size),
length_dim=0,
batch_dim=1,
pad_value=self.vocab.pad_index)
if str(encoder_state["state"].device) != "cpu":
output = output.cuda(encoder_state["state"].device)
layers = encoder_state["state"].size(0)
decoder_state = encoder_state["state"].unsqueeze(2)\
.repeat(1, 1, self.samples, 1).view(layers, batch_size, -1)
search_state = {"output": output, "decoder_state": decoder_state}
encoder_output = encoder_state["output"].repeat_batch_dim(self.samples)
context = {"encoder_output": encoder_output}
return search_state, context
def _make_classifier_inputs(self, batch):
lens = []
clf_inputs = []
for out in batch:
clf_inputs.append([self.input_vocab.start_index] +
[self.input_vocab[t] for t in out[:-1]] +
[self.input_vocab.stop_index])
lens.append(len(out) + 1)
lens = torch.LongTensor(lens)
max_len = lens.max().item()
clf_inputs = torch.LongTensor([
inp + [self.input_vocab.pad_index] * (max_len - len(inp))
for inp in clf_inputs
]).t()
clf_inputs = Variable(clf_inputs,
lengths=lens,
batch_dim=1,
length_dim=0,
pad_value=self.input_vocab.pad_index)
if self.gpu > -1:
return clf_inputs.cuda(self.gpu)
else:
return clf_inputs