-
Notifications
You must be signed in to change notification settings - Fork 3
/
run_per_relation_evaluation.py
79 lines (59 loc) · 2.61 KB
/
run_per_relation_evaluation.py
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
import os
import sys
from Analysis_Utils import get_evaluator_from_config, parse_directory
from Loading_Utils import load_checkpoint
import json
def evaluation_per_relation(model_dir, config_path, sub_dir=None):
checkpoint_path, analysis_path = parse_directory(model_dir, sub_dir)
config, dataset, trained_model = load_checkpoint(config_path, checkpoint_path)
evaluator = get_evaluator_from_config(config)
result_per_relation = {}
# Perform an evaluation for each relation:
for relation_name in dataset.relation_to_id.keys():
# Filter training, testing and validation set by the name of the relation
training_for_curr_rel = dataset.training.new_with_restriction(
relations={relation_name},
)
validation_for_curr_rel = dataset.validation.new_with_restriction(
relations={relation_name},
)
testing_for_curr_rel = dataset.testing.new_with_restriction(
relations={relation_name},
)
# Perform the evaluation per relation
test_result_of_curr_relation = evaluator.evaluate(
model=trained_model,
mapped_triples=testing_for_curr_rel.mapped_triples,
additional_filter_triples=[
training_for_curr_rel.mapped_triples,
validation_for_curr_rel.mapped_triples,
],
batch_size=config['evaluator_kwargs']['batch_size'],
)
result_per_relation[relation_name] = test_result_of_curr_relation.to_flat_dict()
return result_per_relation
def main(**kwargs):
experiment_dir = './Benchmarking'
if 'model_dir' in kwargs.keys():
model_dir = kwargs['model_dir']
else:
raise Exception('The <model_dir> parameter needs to be specified.')
if 'sub_dir' in kwargs.keys():
sub_dir = kwargs['sub_dir']
else:
raise Exception('The <sub_dir> parameter needs to be specified.')
if 'config_path' in kwargs.keys():
config_path = kwargs['config_path']
else:
raise Exception('The <config_path> parameter needs to be specified.')
if 'gpu' in kwargs.keys():
gpu = kwargs['gpu']
os.environ['CUDA_VISIBLE_DEVICES'] = gpu
res = evaluation_per_relation(model_dir, config_path, sub_dir)
saving_directory = os.path.join(experiment_dir, 'per_relation_result', sub_dir)
if not (os.path.exists(saving_directory)):
os.makedirs(saving_directory)
with open(os.path.join(saving_directory, 'result.json'), 'w') as outfile:
json.dump(res, outfile)
if __name__ == '__main__':
main(**dict(arg.split('=') for arg in sys.argv[1:])) # kwargs