Merge pull request #21 from RUCAIBox/master

get latest code

data/utils.py

@@ -1,5 +1,6 @@
 from .dataloader import *
 from config import EvalSetting
+from utils import ModelType
 
 def data_preparation(config, model, dataset):
     es = EvalSetting(config)
@@ -35,7 +36,7 @@ def data_preparation(config, model, dataset):
     return train_data, test_data, valid_data
 
 def dataloader_construct(name, config, eval_setting, dataset,
-                         dl_type='general', dl_format='pointwise',
+                         dl_type=ModelType.GENERAL, dl_format='pointwise',
                          batch_size=1, shuffle=False):
     if not isinstance(dataset, list):
         dataset = [dataset]
@@ -50,7 +51,7 @@ def dataloader_construct(name, config, eval_setting, dataset,
     print(eval_setting)
     print('batch_size = {}, shuffle = {}\n'.format(batch_size, shuffle))
 
-    if dl_type == 'general':
+    if dl_type == ModelType.GENERAL:
         DataLoader = GeneralDataLoader
     else:
         raise NotImplementedError('dl_type [{}] has not been implemented'.format(dl_type))

evaluator/evaluator.py

@@ -11,8 +11,8 @@
 metric_name = {metric.lower() : metric for metric in ['Hit', 'Recall', 'MRR', 'AUC', 'Precision', 'NDCG']}
 
 # These metrics are typical in topk recommendations
-topk_metric = {'hit', 'recall', 'precision', 'ndcg'}
-other_metric = {'auc', 'mrr'}
+topk_metric = {'hit', 'recall', 'precision', 'ndcg', 'mrr'}
+other_metric = {'auc'}
 
 class AbstractEvaluator(metaclass=abc.ABCMeta):
     """The abstract class of the evaluation module, its subclasses must implement their functions

model/general_recommender/bprmf.py

@@ -15,14 +15,15 @@
 
 from model.abstract_recommender import AbstractRecommender
 from model.loss import BPRLoss
+from utils import ModelType
 
 
 class BPRMF(AbstractRecommender):
 
     def __init__(self, config, dataset):
         super(BPRMF, self).__init__()
 
-        self.type = 'general'
+        self.type = ModelType.GENERAL
 
         self.USER_ID = config['USER_ID_FIELD']
         self.ITEM_ID = config['ITEM_ID_FIELD']

utils/__init__.py

@@ -1,3 +1,4 @@
 from .metrics import *
 from .logger import Logger
 from .utils import get_local_time, ensure_dir
+from .enum_type import *

utils/enum_type.py

@@ -0,0 +1,8 @@
+from enum import Enum
+
+class ModelType(Enum):
+    GENERAL = 1
+    SEQUENTIAL = 2
+    CONTEXT = 3
+    KNOWLEDGE = 4
+    SOCIAL = 5

utils/metrics.py

@@ -12,16 +12,17 @@ def hit(rank, label, k):
     """
     return int(any(rank[label] <= k))
 
-def mrr(rank, label, k=None):
+def mrr(rank, label, k):
     """The MRR (also known as mean reciprocal rank) is a statistic measure for evaluating any process that produces a list
     of possible responses to a sample of queries, ordered by probability of correctness. 
 
     url:https://en.wikipedia.org/wiki/Mean_reciprocal_rank
 
     """
     ground_truth_ranks = rank[label]
-    if ground_truth_ranks.all():
-        return (1 / ground_truth_ranks.min())
+    ground_truth_at_k = ground_truth_ranks[ground_truth_ranks <= k]
+    if ground_truth_at_k.shape[0] > 0:
+        return (1 / ground_truth_at_k.min())
     return 0
 
 def recall(rank, label, k):