model = HEER(type_offset, config)
elif args.model == 'DistMult':
print('Loading DistMult model')
model = DistMult('embeddings/distmult/distmult.pt')
elif args.model == 'ComplEx':
print('Loading ComplEx model')
model = ComplEx('embeddings/complex/fb15k237_ComplEx.pt')
elif args.model == 'Node2Vec':
print('Loading Node2Vec model')
model = Node2Vec()
elif args.model == 'DeepWalk':
print('Loading DeepWalk model')
model = DeepWalk()
elif args.model == 'TransE':
print('Loading TransE model')
model = TransE('embeddings/transe/transe.pt')
else:
model = HEER(type_offset, config)
suffix = '_fb15k237_ko_0.07_eval.txt'
in_mapping = cPickle.load(
open('./datasets/FB15K237/fb15k237_ko_0.07_in_mapping.p'))
train_edge_dict = cPickle.load(
open('./datasets/FB15K237/fb15k237_train_edge_dict.p'))
node_type_dict = cPickle.load(
open('./datasets/FB15K237/fb15k237_node_type_dict.p'))
ko_dic = cPickle.load(open('./datasets/FB15K237/fb15k237_ko_dict.p'))
metrics = {
'predict_relation': {
'rr_sum': 0.0,
def evaluate_on_test(model_type, parameter_list, test_tg, saved_model_path, test_size, reverse_relation_dictionary):
tf.reset_default_graph()
with tf.Session() as session:
# Need to instantiate model again
print(parameter_list)
if model_type == TranslationModels.Trans_E:
model = TransE(*parameter_list)
elif model_type == TranslationModels.Trans_H:
model = TransH(*parameter_list)
elif model_type == TranslationModels.RESCAL:
model = RESCAL(*parameter_list)
elif model_type == TranslationModels.TEKE:
model = TEKE(*parameter_list)
model.create_graph()
saver = tf.train.Saver(model.variables())
saver.restore(session, saved_model_path)
test_batch_pos, _ = test_tg.next(test_size)
filter_triples = test_tg.all_triples
test_inpl = test_batch_pos[1, :]
test_inpr = test_batch_pos[0, :]
test_inpo = test_batch_pos[2, :]
scores_l, scores_r = model.scores(session, test_inpl, test_inpr, test_inpo)
errl, errr = ranking_error_triples(filter_triples, scores_l, scores_r, test_inpl, test_inpo, test_inpr)
# END OF SESSION
results = []
err_arr = np.asarray(errl + errr)
hits_10 = np.mean(err_arr <= 10) * 100
hits_3 = np.mean(err_arr <= 3) * 100
hits_1 = np.mean(err_arr <= 1) * 100
mean_rank = np.mean(err_arr)
mrr = np.mean(1.0 / err_arr)
print("Test Hits10: ", hits_10)
print("Test Hits3: ", hits_3)
print("Test Hits1: ", hits_1)
print("Test MRR: ", mrr)
print("Test MeanRank: ", mean_rank)
results.append(mean_rank)
results.append(mrr)
results.append(hits_10)
results.append(hits_3)
results.append(hits_1)
relation_results = dict()
for i in np.unique(test_inpo):
indices = np.argwhere(np.array(test_inpo) == i)
err_arr = np.concatenate([np.array(errl)[indices], np.array(errr)[indices]])
hits_10 = np.mean(err_arr <= 10) * 100
hits_3 = np.mean(err_arr <= 3) * 100
hits_1 = np.mean(err_arr <= 1) * 100
mean_rank = np.mean(err_arr)
mrr = np.mean(1.0 / err_arr)
relation_results[reverse_relation_dictionary[i]] = {'MeanRank': mean_rank, 'MRR' : mrr, 'Hits@10' : hits_10,
'Hits@3': hits_3, 'Hits@1': hits_1}
for k, v in relation_results.iteritems():
print(k, v)
return results, relation_results
shared=shared,
alpha=params.alpha)
else:
batch_size_sg = 0
num_sampled = 0
event_model = None
pre_train = False
# Model Selection
if model_type == TranslationModels.Trans_E:
param_list = [
num_entities, num_relations, params.embedding_size,
params.batch_size, batch_size_sg, num_sampled,
vocab_size, fnsim, params.learning_rate, event_model
]
model = TransE(*param_list)
elif model_type == TranslationModels.Trans_H:
param_list = [
num_entities, num_relations, params.embedding_size,
params.batch_size, batch_size_sg, num_sampled,
vocab_size, params.learning_rate, event_model,
params.lambd
]
model = TransH(*param_list)
elif model_type == TranslationModels.RESCAL:
param_list = [
num_entities, num_relations, params.embedding_size,
params.batch_size, batch_size_sg, num_sampled,
vocab_size, params.learning_rate, event_model,
params.lambd
]