Coverage for hopwise/evaluator/evaluator.py: 78%

1# @Time : 2021/6/25 

2# @Author : Zhichao Feng 

3# @email : fzcbupt@gmail.com 

4 

5"""hopwise.evaluator.evaluator 

6##################################### 

7""" 

8 

9from collections import OrderedDict 

10 

11from hopwise.evaluator.collector import DataStruct 

12from hopwise.evaluator.register import metrics_dict 

13 

14 

15class Evaluator: 

16 """Evaluator is used to check parameter correctness, and summarize the results of all metrics.""" 

17 

18 def __init__(self, config): 

19 self.config = config 

20 self.metrics = [metric.lower() for metric in self.config["metrics"]] 

21 self.metric_class = {} 

22 

23 for metric in self.metrics: 

24 self.metric_class[metric] = metrics_dict[metric](self.config) 

25 

26 def evaluate(self, dataobject: DataStruct): 

27 """Calculate all the metrics. It is called at the end of each epoch 

28 

29 Args: 

30 dataobject (DataStruct): It contains all the information needed for metrics. 

31 

32 Returns: 

33 collections.OrderedDict: such as ``{'hit@20': 0.3824, 'recall@20': 0.0527, 'hit@10': 0.3153, 'recall@10': 0.0329, 'gauc': 0.9236}`` 

34 

35 """ # noqa: E501 

36 

37 result_dict = OrderedDict() 

38 for metric in self.metrics: 

39 metric_val = self.metric_class[metric].calculate_metric(dataobject) 

40 result_dict.update(metric_val) 

41 return result_dict 

42 

43 

44class Evaluator_KG(Evaluator): 

45 """Evaluator KG extends the Evaluator class for link prediction tasks.""" 

46 

47 def __init__(self, config): 

48 super().__init__(config) 

49 self.metrics = [metric.lower() for metric in self.config["metrics_lp"]] 

50 self.metric_class = {} 

51 

52 for metric in self.metrics: 

53 self.metric_class[metric] = metrics_dict[metric](self.config)