save_results = params.get("save_topk_result")
new_data = nnquery.get_topk_predictions(
reranker,
test_dataloader,
candidate_pool,
candidate_encoding,
params["silent"],
logger,
params["top_k"],
params.get("zeshel", None),
save_results,
)
if save_results:
save_data_path = os.path.join(
params['output_path'],
'candidates_%s_top%d.t7' % (params['mode'], params['top_k']))
torch.save(new_data, save_data_path)
if __name__ == "__main__":
parser = BlinkParser(add_model_args=True)
parser.add_eval_args()
args = parser.parse_args()
print(args)
params = args.__dict__
main(params)
logger.info("The training took {} minutes\n".format(execution_time))
# save the best models
logger.info(
"Best ctxt performance in epoch: {}".format(ctxt_best_epoch_idx))
best_ctxt_model_path = os.path.join(model_output_path,
"epoch_{}".format(ctxt_best_epoch_idx),
"ctxt")
logger.info(
"Best cand performance in epoch: {}".format(cand_best_epoch_idx))
best_cand_model_path = os.path.join(model_output_path,
"epoch_{}".format(cand_best_epoch_idx),
"cand")
copy_directory(best_ctxt_model_path,
os.path.join(model_output_path, "best_epoch", "ctxt"))
copy_directory(best_cand_model_path,
os.path.join(model_output_path, "best_epoch", "cand"))
if __name__ == "__main__":
parser = BlinkParser(add_model_args=True)
parser.add_training_args()
parser.add_joint_train_args()
args = parser.parse_args()
print(args)
params = args.__dict__
main(params)