File size: 4,859 Bytes
4523bda | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | 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()
|