print('NDCG@5 - ', ndcg_5)
print('NDCG@10 - ', ndcg_10)
if __name__ == "__main__":
# Load the saved pre-trained model
dictionary = helper.load_object(args.save_path + 'dictionary.p')
embeddings_index = helper.load_word_embeddings(args.word_vectors_directory, args.word_vectors_file,
dictionary.word2idx)
model = DRMM(dictionary, embeddings_index, args)
if 'CUDA_VISIBLE_DEVICES' in os.environ:
cuda_visible_devices = [int(x) for x in os.environ['CUDA_VISIBLE_DEVICES'].split(',')]
if len(cuda_visible_devices) > 1:
model = torch.nn.DataParallel(model, device_ids=cuda_visible_devices)
if args.cuda:
model = model.cuda()
checkpoint = helper.load_from_checkpoint(os.path.join(args.save_path, 'model_best.pth.tar'), args.cuda)
model.load_state_dict(checkpoint['state_dict'])
model.eval()
test_corpus = data.Corpus(args.tokenize, args.max_query_length, args.max_doc_length)
test_corpus.parse(args.data + 'test.txt', args.max_example)
print('test set size = ', len(test_corpus.data))
test_batches = helper.batchify(test_corpus.data, args.batch_size)
print('number of test batches = ', len(test_batches))
test_ranking(model, test_batches)
model = CDSSM(dictionary, args)
optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()),
args.lr)
best_loss = sys.maxsize
param_dict = helper.count_parameters(model)
print('#parameters = ', numpy.sum(list(param_dict.values())))
if args.cuda:
model = model.cuda()
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = helper.load_from_checkpoint(args.resume, args.cuda)
args.start_epoch = checkpoint['epoch']
best_loss = checkpoint['best_loss']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
# ###############################################################################
# # Train the model
# ###############################################################################
train = train.Train(model, optimizer, dictionary, args, best_loss)
train.train_epochs(train_corpus, dev_corpus, args.start_epoch, args.epochs)