default=1.0)
parser.add_argument(
'--negative_sampling',
dest='negative_sampling',
type=str,
help='choose unit or bern to generate negative examples',
default='bern')
parser.add_argument(
'--score_func',
dest='score_func',
type=str,
default='l1',
help='choose l1 or l2 to calculate distance of vectors')
args = parser.parse_args()
print(args)
KG = KnowledgeGraph(data_dir=args.data_dir,
negative_sampling=args.negative_sampling)
model = TransE(num_entity=KG.num_entity,
num_relation=KG.num_relation,
learning_rate=args.learning_rate,
batch_size=args.batch_size,
num_epochs=args.num_epochs,
margin=args.margin,
dimension=args.dimension,
score_func=args.score_func)
model.compile()
tp, tn = KG.get_training_data()
train_model(model, tp, tn)
model.save_embeddings()
test_model(model, KG.get_test_data())
