#!/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 + ['</s>'] * (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()