def __init__(self, vocab_src, vocab_tgt, config):
super(CoAEVNMT, self).__init__()
self.vocab_src = vocab_src
self.vocab_tgt = vocab_tgt
if config["share_embeddings"]:
self.emb_src = nn.Embedding(vocab_src.size(),
config["emb_size"],
padding_idx=vocab_src[PAD_TOKEN])
self.emb_tgt = nn.Embedding(vocab_tgt.size(),
config["emb_size"],
padding_idx=vocab_tgt[PAD_TOKEN])
else:
self.emb_src = None
self.emb_tgt = None
self.model_xy = aevnmt_utils.create_model(self.vocab_src,
self.vocab_tgt,
config,
emb_src=self.emb_src,
emb_tgt=self.emb_tgt)
self.model_yx = aevnmt_utils.create_model(self.vocab_tgt,
self.vocab_src,
config,
emb_src=self.emb_tgt,
emb_tgt=self.emb_src)
self.config = config
def create_models(vocab_src, vocab_tgt, config):
if config["model_type"] == "conmt":
model_xy = cond_nmt_utils.create_model(vocab_src, vocab_tgt, config)
model_yx = cond_nmt_utils.create_model(vocab_tgt, vocab_src, config)
bi_train_fn = conmt_utils.bi_train_fn
mono_train_fn = conmt_utils.mono_train_fn
validate_fn = cond_nmt_utils.validate
elif config["model_type"] == "coaevnmt":
model_xy = aevnmt_utils.create_model(vocab_src, vocab_tgt, config)
model_yx = aevnmt_utils.create_model(vocab_tgt, vocab_src, config)
bi_train_fn = coaevnmt_utils.bi_train_fn
mono_train_fn = coaevnmt_utils.mono_train_fn
validate_fn = aevnmt_utils.validate
else:
raise ValueError("Unknown model type: {}".format(config["model_type"]))
print("Model xy: ", model_xy)
print("Model yx: ", model_yx)
return model_xy, model_yx, bi_train_fn, mono_train_fn, validate_fn
def create_model(vocab_src, vocab_tgt, config):
if config["model_type"] == "cond_nmt":
model = cond_nmt_utils.create_model(vocab_src, vocab_tgt, config)
train_fn = cond_nmt_utils.train_step
validate_fn = cond_nmt_utils.validate
elif config["model_type"] == "aevnmt":
model = aevnmt_utils.create_model(vocab_src, vocab_tgt, config)
train_fn = aevnmt_utils.train_step
validate_fn = aevnmt_utils.validate
else:
raise ValueError("Unknown model type: {}".format(config["model_type"]))
print("Model: ", model)
return model, train_fn, validate_fn