def __init__(self, ernie, config):
""" Model which Based on the PaddleNLP PretrainedModel
Note:
1. the ernie must be the first argument.
2. must set self.XX = ernie to load weights.
3. the self.config keyword is taken by PretrainedModel class.
Args:
ernie (nn.Layer): the submodule layer of ernie model.
config (Dict): the config file
"""
super(ErnieSageForLinkPrediction, self).__init__()
self.config_file = config
self.ernie = ernie
self.encoder = Encoder.factory(self.config_file, self.ernie)
self.loss_func = LossFactory(self.config_file)
def forward(self, features):
num_nodes, num_edges, edges, node_feat_index, node_feat_term_ids, user_index, \
pos_item_index, neg_item_index, user_real_index, pos_item_real_index = features
node_feat = {"index": node_feat_index, "term_ids": node_feat_term_ids}
graph_wrapper = BatchGraphWrapper(num_nodes, num_edges, edges,
node_feat)
encoder = Encoder.factory(self.hparam)
outputs = encoder([graph_wrapper],
[user_index, pos_item_index, neg_item_index])
user_feat, pos_item_feat, neg_item_feat = outputs
# loss
if self.hparam.neg_type == "batch_neg":
neg_item_feat = pos_item_feat
if self.mode is propeller.RunMode.TRAIN:
return user_feat, pos_item_feat, neg_item_feat
elif self.mode is propeller.RunMode.PREDICT:
return user_feat, user_real_index
elif self.mode is propeller.RunMode.EVAL:
return user_feat, pos_item_feat, neg_item_feat
