| import argparse |
| import os |
|
|
| from openea.modules.load.kg import KG |
| from openea.modules.load.read import read_relation_triples, read_links |
|
|
| from mtranse import MTransEV2 |
| from utils import read_items, MyKGs |
|
|
| parser = argparse.ArgumentParser(description='NullEA') |
| parser.add_argument('--training_data', type=str, default='./DBP2.0/zh_en/') |
| parser.add_argument('--output', type=str, default='../../output/results/') |
| parser.add_argument('--dataset_division', type=str, default='splits') |
|
|
| parser.add_argument('--align_direction', type=str, default='left') |
|
|
| parser.add_argument('--init', type=str, default='xavier') |
| parser.add_argument('--alignment_module', type=str, default='mapping') |
| parser.add_argument('--search_module', type=str, default='greedy') |
| parser.add_argument('--neg_sampling', type=str, default='truncated', choices=['uniform', 'truncated']) |
|
|
| parser.add_argument('--dim', type=int, default=128) |
| parser.add_argument('--loss_norm', type=str, default='L2') |
| parser.add_argument('--ent_l2_norm', type=bool, default=True) |
| parser.add_argument('--rel_l2_norm', type=bool, default=True) |
| parser.add_argument('--batch_size', type=int, default=20480) |
|
|
| parser.add_argument('--embed_margin', type=float, default=1.5) |
| parser.add_argument('--mapping_margin', type=float, default=1.0) |
| parser.add_argument('--mapping_neg_num', type=int, default=1) |
| parser.add_argument('--mapping_neg_weight', type=int, default=0.1) |
|
|
| parser.add_argument('--neg_triple_num', type=int, default=1) |
| parser.add_argument('--truncated_epsilon', type=float, default=0.9) |
| parser.add_argument('--truncated_freq', type=int, default=10) |
|
|
| parser.add_argument('--learning_rate', type=float, default=0.001) |
| parser.add_argument('--optimizer', type=str, default='Adam') |
| parser.add_argument('--batch_threads_num', type=int, default=4) |
| parser.add_argument('--test_threads_num', type=int, default=1) |
| parser.add_argument('--max_epoch', type=int, default=200) |
| parser.add_argument('--eval_freq', type=int, default=10) |
|
|
| parser.add_argument('--ordered', type=bool, default=True) |
| parser.add_argument('--top_k', type=list, default=[1, 5, 10]) |
|
|
| parser.add_argument('--eval_norm', type=bool, default=True) |
| parser.add_argument('--start_valid', type=int, default=0) |
| parser.add_argument('--stop_metric', type=str, default='mrr') |
| parser.add_argument('--eval_metric', type=str, default='inner') |
|
|
| args = parser.parse_args() |
| print(args) |
|
|
|
|
| def read_kgs_from_folder(training_data_folder, division, mode, ordered, direction): |
| kg1_relation_triples, _, _ = read_relation_triples(training_data_folder + 'rel_triples_1') |
| kg2_relation_triples, _, _ = read_relation_triples(training_data_folder + 'rel_triples_2') |
|
|
| train_links = read_links(os.path.join(training_data_folder, division, 'train_links')) |
| valid_links = read_links(os.path.join(training_data_folder, division, 'valid_links')) |
| test_links = read_links(os.path.join(training_data_folder, division, 'test_links')) |
|
|
| train_unlinked_ent1 = read_items(os.path.join(training_data_folder, division, 'train_unlinked_ent1')) |
| valid_unlinked_ent1 = read_items(os.path.join(training_data_folder, division, 'valid_unlinked_ent1')) |
| test_unlinked_ent1 = read_items(os.path.join(training_data_folder, division, 'test_unlinked_ent1')) |
|
|
| train_unlinked_ent2 = read_items(os.path.join(training_data_folder, division, 'train_unlinked_ent2')) |
| valid_unlinked_ent2 = read_items(os.path.join(training_data_folder, division, 'valid_unlinked_ent2')) |
| test_unlinked_ent2 = read_items(os.path.join(training_data_folder, division, 'test_unlinked_ent2')) |
|
|
| kg1 = KG(kg1_relation_triples, set()) |
| kg2 = KG(kg2_relation_triples, set()) |
|
|
| if direction == "left": |
| two_kgs = MyKGs(kg1, kg2, train_links, test_links, |
| train_unlinked_ent1, valid_unlinked_ent1, test_unlinked_ent1, |
| train_unlinked_ent2, valid_unlinked_ent2, test_unlinked_ent2, |
| valid_links=valid_links, mode=mode, ordered=ordered) |
| else: |
| assert direction == "right" |
| train_links_rev = [(e2, e1) for e1, e2 in train_links] |
| test_links_rev = [(e2, e1) for e1, e2 in test_links] |
| valid_links_rev = [(e2, e1) for e1, e2 in valid_links] |
| two_kgs = MyKGs(kg2, kg1, train_links_rev, test_links_rev, |
| train_unlinked_ent2, valid_unlinked_ent2, test_unlinked_ent2, |
| train_unlinked_ent1, valid_unlinked_ent1, test_unlinked_ent1, |
| valid_links=valid_links_rev, mode=mode, ordered=ordered) |
| return two_kgs |
|
|
|
|
| if __name__ == '__main__': |
| kgs = read_kgs_from_folder(args.training_data, args.dataset_division, args.alignment_module, |
| args.ordered, args.align_direction) |
| model = MTransEV2() |
| model.set_args(args) |
| model.set_kgs(kgs) |
| model.init() |
| model.run() |
| model.test() |
|
|
