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