Training and generation processes for neural encoder-decoder machine translation.
September 1, 2015 ยท View on GitHub
#!/usr/bin/python3
import datetime import sys import math import numpy as np from argparse import ArgumentParser from collections import defaultdict
from chainer import FunctionSet, Variable, functions, optimizers from chainer import cuda
def trace(*args): print(datetime.datetime.now(), '...', *args, file=sys.stderr) sys.stderr.flush()
class Vocabulary: def init(self): pass
def __len__(self):
return self.__size
def stoi(self, s):
return self.__stoi[s]
def itos(self, i):
return self.__itos[i]
@staticmethod
def new(line_gen, size):
self = Vocabulary()
self.__size = size
word_freq = defaultdict(lambda: 0)
for line in line_gen:
words = line.split()
for word in words:
word_freq[word] += 1
self.__stoi = defaultdict(lambda: 0)
self.__stoi['<unk>'] = 0
self.__stoi['<s>'] = 1
self.__stoi['</s>'] = 2
self.__itos = [''] * self.__size
self.__itos[0] = '<unk>'
self.__itos[1] = '<s>'
self.__itos[2] = '</s>'
for i, (k, v) in zip(range(self.__size - 3), sorted(word_freq.items(), key=lambda x: -x[1])):
self.__stoi[k] = i + 3
self.__itos[i + 3] = k
return self
def save(self, fp):
print(self.__size, file=fp)
for i in range(self.__size):
print(self.__itos[i], file=fp)
@staticmethod
def load(line_gen):
self = Vocabulary()
self.__size = int(next(line_gen))
self.__stoi = defaultdict(lambda: 0)
self.__itos = [''] * self.__size
for i in range(self.__size):
s = next(line_gen).strip()
if s:
self.__stoi[s] = i
self.__itos[i] = s
return self
class wrapper: @staticmethod def make_var(array, dtype=np.float32): #return Variable(np.array(array, dtype=dtype)) return Variable(cuda.to_gpu(np.array(array, dtype=dtype)))
@staticmethod
def get_data(variable):
#return variable.data
return cuda.to_cpu(variable.data)
@staticmethod
def zeros(shape, dtype=np.float32):
#return Variable(np.zeros(shape, dtype=dtype))
return Variable(cuda.zeros(shape, dtype=dtype))
@staticmethod
def make_model(**kwargs):
#return FunctionSet(**kwargs)
return FunctionSet(**kwargs).to_gpu()
class EncoderDecoderModel: def init(self): pass
def __new_opt(self):
#return optimizers.SGD(lr=0.01)
return optimizers.AdaGrad(lr=0.01)
def __to_cpu(self):
self.__model.to_cpu()
self.__opt = self.__new_opt()
self.__opt.setup(self.__model)
def __to_gpu(self):
self.__model.to_gpu()
self.__opt = self.__new_opt()
self.__opt.setup(self.__model)
def __make_model(self):
self.__model = wrapper.make_model(
# encoder
w_xi = functions.EmbedID(len(self.__src_vocab), self.__n_src_embed),
w_ip = functions.Linear(self.__n_src_embed, 4 * self.__n_hidden),
w_pp = functions.Linear(self.__n_hidden, 4 * self.__n_hidden),
# decoder
w_pq = functions.Linear(self.__n_hidden, 4 * self.__n_hidden),
w_qj = functions.Linear(self.__n_hidden, self.__n_trg_embed),
w_jy = functions.Linear(self.__n_trg_embed, len(self.__trg_vocab)),
w_yq = functions.EmbedID(len(self.__trg_vocab), 4 * self.__n_hidden),
w_qq = functions.Linear(self.__n_hidden, 4 * self.__n_hidden),
)
self.__to_gpu()
@staticmethod
def new(src_vocab, trg_vocab, n_src_embed, n_trg_embed, n_hidden):
self = EncoderDecoderModel()
self.__src_vocab = src_vocab
self.__trg_vocab = trg_vocab
self.__n_src_embed = n_src_embed
self.__n_trg_embed = n_trg_embed
self.__n_hidden = n_hidden
self.__make_model()
return self
def save(self, fp):
vtos = lambda v: ' '.join('%.8e' % x for x in v)
fprint = lambda x: print(x, file=fp)
def print_embed(f):
for row in f.W: fprint(vtos(row))
def print_linear(f):
for row in f.W: fprint(vtos(row))
fprint(vtos(f.b))
self.__src_vocab.save(fp)
self.__trg_vocab.save(fp)
fprint(self.__n_src_embed)
fprint(self.__n_trg_embed)
fprint(self.__n_hidden)
self.__to_cpu()
print_embed(self.__model.w_xi)
print_linear(self.__model.w_ip)
print_linear(self.__model.w_pp)
print_linear(self.__model.w_pq)
print_linear(self.__model.w_qj)
print_linear(self.__model.w_jy)
print_embed(self.__model.w_yq)
print_linear(self.__model.w_qq)
self.__to_gpu()
@staticmethod
def load(line_gen):
loadv = lambda tp: [tp(x) for x in next(line_gen).split()]
def copyv(tp, row):
data = loadv(tp)
for i in range(len(data)):
row[i] = data[i]
def copy_embed(f):
for row in f.W: copyv(float, row)
def copy_linear(f):
for row in f.W: copyv(float, row)
copyv(float, f.b)
self = EncoderDecoderModel()
self.__src_vocab = Vocabulary.load(line_gen)
self.__trg_vocab = Vocabulary.load(line_gen)
self.__n_src_embed = loadv(int)[0]
self.__n_trg_embed = loadv(int)[0]
self.__n_hidden = loadv(int)[0]
self.__make_model()
self.__to_cpu()
copy_embed(self.__model.w_xi)
copy_linear(self.__model.w_ip)
copy_linear(self.__model.w_pp)
copy_linear(self.__model.w_pq)
copy_linear(self.__model.w_qj)
copy_linear(self.__model.w_jy)
copy_embed(self.__model.w_yq)
copy_linear(self.__model.w_qq)
self.__to_gpu()
return self
def __predict(self, is_training, src_batch, trg_batch = None, generation_limit = None):
m = self.__model
tanh = functions.tanh
lstm = functions.lstm
batch_size = len(src_batch)
src_len = len(src_batch[0])
src_stoi = self.__src_vocab.stoi
trg_stoi = self.__trg_vocab.stoi
trg_itos = self.__trg_vocab.itos
s_c = wrapper.zeros((batch_size, self.__n_hidden))
# encoding
s_x = wrapper.make_var([src_stoi('</s>') for _ in range(batch_size)], dtype=np.int32)
s_i = tanh(m.w_xi(s_x))
s_c, s_p = lstm(s_c, m.w_ip(s_i))
for l in reversed(range(src_len)):
s_x = wrapper.make_var([src_stoi(src_batch[k][l]) for k in range(batch_size)], dtype=np.int32)
s_i = tanh(m.w_xi(s_x))
s_c, s_p = lstm(s_c, m.w_ip(s_i) + m.w_pp(s_p))
s_c, s_q = lstm(s_c, m.w_pq(s_p))
hyp_batch = [[] for _ in range(batch_size)]
# decoding
if is_training:
accum_loss = wrapper.zeros(())
trg_len = len(trg_batch[0])
for l in range(trg_len):
s_j = tanh(m.w_qj(s_q))
r_y = m.w_jy(s_j)
s_t = wrapper.make_var([trg_stoi(trg_batch[k][l]) for k in range(batch_size)], dtype=np.int32)
accum_loss += functions.softmax_cross_entropy(r_y, s_t)
output = wrapper.get_data(r_y).argmax(1)
for k in range(batch_size):
hyp_batch[k].append(trg_itos(output[k]))
#s_y = wrapper.make_var(output, dtype=np.int32)
#s_c, s_q = lstm(s_c, m.w_yq(s_y) + m.w_qq(s_q))
s_c, s_q = lstm(s_c, m.w_yq(s_t) + m.w_qq(s_q))
return hyp_batch, accum_loss
else:
while len(hyp_batch[0]) < generation_limit:
s_j = tanh(m.w_qj(s_q))
r_y = m.w_jy(s_j)
output = wrapper.get_data(r_y).argmax(1)
for k in range(batch_size):
hyp_batch[k].append(trg_itos(output[k]))
if all(hyp_batch[k][-1] == '</s>' for k in range(batch_size)): break
s_y = wrapper.make_var(output, dtype=np.int32)
s_c, s_q = lstm(s_c, m.w_yq(s_y) + m.w_qq(s_q))
return hyp_batch
def train(self, src_batch, trg_batch):
self.__opt.zero_grads()
hyp_batch, accum_loss = self.__predict(True, src_batch, trg_batch=trg_batch)
accum_loss.backward()
self.__opt.clip_grads(10)
self.__opt.update()
return hyp_batch, wrapper.get_data(accum_loss).reshape(())
def predict(self, src_batch, generation_limit):
return self.__predict(False, src_batch, generation_limit=generation_limit)
def parse_args(): def_vocab = 32768 def_embed_in = 256 def_embed_out = 256 def_hidden = 512 def_epoch = 100 def_minibatch = 256 def_generation_limit = 256
p = ArgumentParser(description='Encoder-decoder trainslation model trainer')
p.add_argument('mode', help='\'train\' or \'test\'')
p.add_argument('source', help='[in] source corpus')
p.add_argument('target', help='[in/out] target corpus')
p.add_argument('model', help='[in/out] model file')
p.add_argument('--vocab', default=def_vocab, metavar='INT', type=int,
help='vocabulary size (default: %d)' % def_vocab)
p.add_argument('--embed-in', default=def_embed_in, metavar='INT', type=int,
help='input embedding layer size (default: %d)' % def_embed_in)
p.add_argument('--embed-out', default=def_embed_out, metavar='INT', type=int,
help='output embedding layer size (default: %d)' % def_embed_out)
p.add_argument('--hidden', default=def_hidden, metavar='INT', type=int,
help='hidden layer size (default: %d)' % def_hidden)
p.add_argument('--epoch', default=def_epoch, metavar='INT', type=int,
help='number of training epoch (default: %d)' % def_epoch)
p.add_argument('--minibatch', default=def_minibatch, metavar='INT', type=int,
help='minibatch size (default: %d)' % def_minibatch)
p.add_argument('--generation-limit', default=def_generation_limit, metavar='INT', type=int,
help='maximum number of words to be generated for test input')
args = p.parse_args()
# check args
try:
if args.mode not in ['train', 'test']: raise ValueError('you must set mode = \'train\' or \'test\'')
if args.vocab < 1: raise ValueError('you must set --vocab >= 1')
if args.embed_in < 1: raise ValueError('you must set --embed-in >= 1')
if args.embed_out < 1: raise ValueError('you must set --embed-out >= 1')
if args.hidden < 1: raise ValueError('you must set --hidden >= 1')
if args.epoch < 1: raise ValueError('you must set --epoch >= 1')
if args.minibatch < 1: raise ValueError('you must set --minibatch >= 1')
if args.generation_limit < 1: raise ValueError('you must set --generation-limit >= 1')
except Exception as ex:
p.print_usage(file=sys.stderr)
print(ex, file=sys.stderr)
sys.exit()
return args
def generate_parallel_sorted(src_filename, trg_filename, batch_size): with open(src_filename) as fsrc, open(trg_filename) as ftrg: batch = [] try: while True: for i in range(batch_size): batch.append((next(fsrc).split(), next(ftrg).split())) batch = sorted(batch, key=lambda x: len(x[1])) for src, trg in batch: yield src, trg batch = [] except StopIteration: if batch: batch = sorted(batch, key=lambda x: len(x[1])) for src, trg in batch: yield src, trg except: raise
def normalize_batch(batch): max_len = max(len(x) for x in batch) return [x + [''] * (max_len - len(x) + 1) for x in batch]
def generate_train_minibatch(src_filename, trg_filename, batch_size): gen = iter(generate_parallel_sorted(src_filename, trg_filename, batch_size * 100)) src_batch = [] trg_batch = []
try:
while True:
for i in range(batch_size):
src, trg = next(gen)
src_batch.append(src)
trg_batch.append(trg)
yield normalize_batch(src_batch), normalize_batch(trg_batch)
src_batch = []
trg_batch = []
except StopIteration:
if src_batch and len(src_batch) == len(trg_batch):
yield normalize_batch(src_batch), normalize_batch(trg_batch)
except:
raise
def generate_test_minibatch(src_filename, batch_size): with open(src_filename) as fsrc: batch = [] try: while True: for i in range(batch_size): batch.append(next(fsrc).split()) yield normalize_batch(batch) batch = [] except StopIteration: if batch: yield normalize_batch(batch) except: raise
def train_model(args): trace('making vocaburaries ...') with open(args.source) as fp: src_vocab = Vocabulary.new(fp, args.vocab) with open(args.target) as fp: trg_vocab = Vocabulary.new(fp, args.vocab)
trace('start training ...')
model = EncoderDecoderModel.new(src_vocab, trg_vocab, args.embed_in, args.embed_out, args.hidden)
for epoch in range(args.epoch):
trace('epoch %d/%d: ' % (epoch + 1, args.epoch))
log_ppl = 0.0
trained = 0
for src_batch, trg_batch in generate_train_minibatch(args.source, args.target, args.minibatch):
hyp_batch, loss = model.train(src_batch, trg_batch)
K = len(src_batch)
for k in range(K):
trace('epoch %3d/%3d, sample %8d' % (epoch + 1, args.epoch, trained + k + 1))
trace(' src = ' + ' '.join([x if x != '</s>' else '*' for x in src_batch[k]]))
trace(' trg = ' + ' '.join([x if x != '</s>' else '*' for x in trg_batch[k]]))
trace(' hyp = ' + ' '.join([x if x != '</s>' else '*' for x in hyp_batch[k]]))
trained += K
with open(args.model + '.%03d' % (epoch + 1), 'w') as fp: model.save(fp)
trace('finished.')
def test_model(args): trace('loading model ...') with open(args.model) as fp: model = EncoderDecoderModel.load(fp)
trace('generating translation ...')
generated = 0
with open(args.target, 'w') as fp:
for src_batch in generate_test_minibatch(args.source, args.minibatch):
K = len(src_batch)
trace('sample %8d - %8d ...' % (generated + 1, generated + K))
hyp_batch = model.predict(src_batch, args.generation_limit)
for hyp in hyp_batch:
hyp.append('</s>')
hyp = hyp[:hyp.index('</s>')]
print(' '.join(hyp), file=fp)
generated += K
trace('finished.')
def main(): args = parse_args()
trace('initializing CUDA ...')
cuda.init()
if args.mode == 'train': train_model(args)
elif args.mode == 'test': test_model(args)
if name == 'main': main()