def make_loss_compute(vocab_size):
weight = torch.ones(vocab_size)
weight[data.NULL_ID] = 0
criterion = torch.nn.NLLLoss(weight, size_average=False)
if nu.gpu_available():
criterion = criterion.cuda()
return criterion
def build_model(opt, vocab_size):
embeddings = nn.Embedding(vocab_size,
opt.word_vec_size,
padding_idx=data.NULL_ID)
if opt.position_encoding:
pe = modules.PositionalEncoding(opt.dropout, opt.word_vec_size)
embeddings = nn.Sequential(embeddings, pe)
else:
embeddings = nn.Sequential(embeddings)
if opt.model_type == 'transformer':
encoder = TransformerEncoder(embeddings, opt.transformer_enc_layers,
opt.head_count,
opt.transformer_hidden_size, opt.dropout)
decoder = TransformerDecoder(embeddings, opt.transformer_dec_layers,
opt.head_count,
opt.transformer_hidden_size, opt.dropout)
elif opt.model_type == 'rnn':
encoder = RNNEncoder(embeddings, opt.num_layers, opt.rnn_size,
opt.bidirectional_encoder, opt.dropout)
decoder = InputFeedRNNDecoder(embeddings, opt.num_layers,
opt.bidirectional_encoder, opt.rnn_size,
opt.attn_type, opt.dropout)
model = NMTModel(encoder, decoder)
if nu.gpu_available():
model = model.cuda()
return model
def build_discriminator(opt):
model = Discriminator(opt.rnn_size)
if nu.gpu_available():
model = model.cuda()
return model