def initialize_models(voc):
print('Building encoder and decoder ...')
# initialize word embeddings
embedding = nn.Embedding(voc.num_words, c.HIDDEN_SIZE)
if c.LOAD_FILENAME:
embedding.load_state_dict(embedding_sd)
# initialize encoder & decoder models
encoder = EncoderRNN(c.HIDDEN_SIZE, embedding, c.ENCODER_N_LAYERS,
c.DROPOUT)
decoder = LuongAttnDecoderRNN(c.ATTN_MODEL, embedding, c.HIDDEN_SIZE,
voc.num_words, c.DECODER_N_LAYERS, c.DROPOUT)
if c.LOAD_FILENAME:
encoder.load_state_dict(encoder_sd)
decoder.load_state_dict(decoder_sd)
# Use appropriate device
encoder = encoder.to(c.DEVICE)
decoder = decoder.to(c.DEVICE)
print('Models built and ready to go!')
return embedding, encoder, decoder
encoder_sd = checkpoint['en']
decoder_sd = checkpoint['de']
encoder_optimizer_sd = checkpoint['en_opt']
decoder_optimizer_sd = checkpoint['de_opt']
embedding_sd = checkpoint['embedding']
voc.__dict__ = checkpoint['voc_dict']
# Initialize word embeddings
embedding = nn.Embedding(voc.num_words, hidden_size)
embedding.load_state_dict(embedding_sd)
# Initialize encoder & decoder models
encoder = EncoderRNN(hidden_size, embedding, encoder_n_layers, dropout)
decoder = LuongAttnDecoderRNN(attn_model, embedding, hidden_size,
voc.num_words, decoder_n_layers, dropout)
encoder.load_state_dict(encoder_sd)
decoder.load_state_dict(decoder_sd)
encoder = encoder.to(device)
decoder = decoder.to(device)
# Set dropout layers to eval mode
encoder.eval()
decoder.eval()
# Initialize search module
searcher = GreedySearchDecoder(encoder, decoder)
# Begin chatting (uncomment and run the following line to begin)
evaluateInput(encoder, decoder, searcher, voc)
target_embedding = Embedding(filename=args.glove_filename, embedding_size=embedding_size).load_embedding_plus_keyword(
train_dataset.tgt_vocab)
else:
train_embedding = Embedding(filename=args.glove_filename, embedding_size=embedding_size).load_embedding(train_dataset.src_vocab)
target_embedding = Embedding(filename=args.glove_filename, embedding_size=embedding_size).load_embedding(train_dataset.tgt_vocab)
encoder.embedding.weight.data.copy_(train_embedding)
decoder.embedding.weight.data.copy_(target_embedding)
if opts.fixed_embeddings:
encoder.embedding.weight.requires_grad = False
decoder.embedding.weight.requires_grad = False
else:
decoder.embedding.weight.requires_grad = True
print("emb end")
if LOAD_CHECKPOINT:
encoder.load_state_dict(checkpoint['encoder_state_dict'])
decoder.load_state_dict(checkpoint['decoder_state_dict'])
if USE_CUDA:
encoder.cuda()
decoder.cuda()
FINE_TUNE = True
if FINE_TUNE:
encoder.embedding.weight.requires_grad = True
print('='*100)
print('Model log:\n')
print(encoder)
print(decoder)
print('- Encoder input embedding requires_grad={}'.format(encoder.embedding.weight.requires_grad))
print('- Decoder input embedding requires_grad={}'.format(decoder.embedding.weight.requires_grad))
print('- Decoder output embedding requires_grad={}'.format(decoder.W_s.weight.requires_grad))
print('='*100 + '\n')