@@ -0,0 +1,2257 @@
+import os%0Aimport pickle%0Aimport json%0Aimport numpy as np%0Aimport sys%0Afrom ogb.lsc import WikiKG90MDataset, WikiKG90MEvaluator%0Aimport pdb%0Afrom collections import defaultdict%0Aimport torch.nn.functional as F%0Aimport torch%0A%0A# python dump_test_by_best_valid.py $CKPT_PATH $NUM_PROC%0Aif __name__ == '__main__':%0A    path = sys.argv%5B1%5D%0A    num_proc = int(sys.argv%5B2%5D)%0A%0A    all_file_names = os.listdir(path)%0A    test_file_names = %5Bname for name in all_file_names if '.pkl' in name and 'test' in name%5D%0A    valid_file_names = %5Bname for name in all_file_names if '.pkl' in name and 'valid' in name%5D%0A    steps = %5Bint(name.split('.')%5B0%5D.split('_')%5B-1%5D) for name in valid_file_names if 'valid_0' in name%5D%0A    steps.sort()%0A    evaluator = WikiKG90MEvaluator()%0A    device = torch.device('cpu')%0A%0A    all_test_dicts = %5B%5D%0A    best_valid_mrr = -1%0A    best_valid_idx = -1%0A%0A    for i, step in enumerate(steps):%0A        valid_result_dict = defaultdict(lambda: defaultdict(list))%0A        test_result_dict = defaultdict(lambda: defaultdict(list))%0A        for proc in range(num_proc):%0A            valid_result_dict_proc = torch.load(os.path.join(path, %22valid_%7B%7D_%7B%7D.pkl%22.format(proc, step)), map_location=device)%0A            test_result_dict_proc = torch.load(os.path.join(path, %22test_%7B%7D_%7B%7D.pkl%22.format(proc, step)), map_location=device)%0A            for result_dict_proc, result_dict in zip(%5Bvalid_result_dict_proc, test_result_dict_proc%5D, %5Bvalid_result_dict, test_result_dict%5D):%0A                for key in result_dict_proc%5B'h,r-%3Et'%5D:%0A                    result_dict%5B'h,r-%3Et'%5D%5Bkey%5D.append(result_dict_proc%5B'h,r-%3Et'%5D%5Bkey%5D.numpy())%0A        for result_dict in %5Bvalid_result_dict, test_result_dict%5D:%0A            for key in result_dict%5B'h,r-%3Et'%5D:%0A                result_dict%5B'h,r-%3Et'%5D%5Bkey%5D = np.concatenate(result_dict%5B'h,r-%3Et'%5D%5Bkey%5D, 0)%0A                    %0A        all_test_dicts.append(test_result_dict)%0A        metrics = evaluator.eval(valid_result_dict)%0A        metric = 'mrr'%0A        print(%22valid-%7B%7D at step %7B%7D: %7B%7D%22.format(metric, step, metrics%5Bmetric%5D))%0A        if metrics%5Bmetric%5D %3E best_valid_mrr:%0A            best_valid_mrr = metrics%5Bmetric%5D%0A            best_valid_idx = i%0A%0A    best_test_dict = all_test_dicts%5Bbest_valid_idx%5D%0A    evaluator.save_test_submission(best_test_dict, path)%0A%0A%0A%0A

update dump test wikikg