def init_model(config, mode):
print('Initializing %s model...' % mode)
if mode == 'disc':
if config.model == 'transe':
em = EmbeddingModel.TransE(config)
elif config.model == 'transd':
# config.embedding_size = config.embedding_size / 2
em = EmbeddingModel.TransD(config)
# config.embedding_size = config.embedding_size * 2
else:
raise ValueError('Unrecognized model type: %s' % config.model)
model = Discriminator.BaseModel(config, em)
elif mode == 'gen':
em = EmbeddingModel.DistMult(config)
model = GanGenerator(config, em)
elif mode == 'sn_gen':
em = EmbeddingModel.DistMult(config)
model = GanGenerator(config, em)
else:
raise ValueError('Unrecognized mode: %s' % config.mode)
model.build()
return model
