Coverage for hopwise/model/knowledge_graph_embedding_recommender/tucker.py: 83%

1# @Time : 2024/11/21 

2# @Author : Alessandro Soccol 

3# @Email : alessandro.soccol@unica.it 

4 

5"""TuckER 

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

7Reference: 

8 Balažević et al. "TuckER: Tensor Factorization for Knowledge Graph Completion." in EMNLP/IJCNLP 2019. 

9 

10Reference code: 

11 https://github.com/ibalazevic/TuckER 

12""" 

13 

14import torch 

15from torch import nn 

16 

17from hopwise.model.abstract_recommender import KnowledgeRecommender 

18from hopwise.model.init import xavier_normal_initialization 

19from hopwise.utils import InputType 

20 

21 

22class TuckER(KnowledgeRecommender): 

23 r"""TuckER relies on Tucker Decomposition. It handles entity and relation embeddings of independent dimension 

24 and jointly learns a share core W. 

25 

26 Note: 

27 In this version, we sample recommender data and knowledge data separately, and put them together for training. 

28 """ 

29 

30 input_type = InputType.PAIRWISE 

31 

32 def __init__(self, config, dataset): 

33 super().__init__(config, dataset) 

34 

35 # load parameters info 

36 self.embedding_size = config["embedding_size"] 

37 self.device = config["device"] 

38 self.ui_relation = dataset.field2token_id["relation_id"][dataset.ui_relation] 

39 self.label_smoothing = config["label_smoothing"] 

40 self.input_dropout = config["input_dropout"] 

41 self.input_dropout1 = config["input_dropout1"] 

42 self.input_dropout2 = config["input_dropout2"] 

43 

44 # define layers and loss 

45 self.user_embedding = nn.Embedding(self.n_users + self.n_items, self.embedding_size) 

46 self.entity_embedding = nn.Embedding(self.n_entities, self.embedding_size) 

47 

48 self.relation_embedding = nn.Embedding(self.n_relations, self.embedding_size) 

49 

50 self.weights = torch.nn.Parameter( 

51 torch.tensor( 

52 torch.empty(self.embedding_size, self.embedding_size, self.embedding_size).uniform_(-1, 1), 

53 requires_grad=True, 

54 ) 

55 ) 

56 

57 self.input_dropout = torch.nn.Dropout(self.input_dropout) 

58 self.hidden_dropout1 = torch.nn.Dropout(self.input_dropout1) 

59 self.hidden_dropout2 = torch.nn.Dropout(self.input_dropout2) 

60 self.bn0 = torch.nn.BatchNorm1d(self.embedding_size) 

61 self.bn1 = torch.nn.BatchNorm1d(self.embedding_size) 

62 

63 # Loss 

64 self.loss = nn.BCELoss() 

65 

66 # Parameters initialization 

67 self.apply(xavier_normal_initialization) 

68 

69 def forward(self, h, r, embeddings): 

70 x = self.bn0(h) 

71 x = self.input_dropout(x) 

72 x = x.view(-1, 1, h.size(1)) 

73 

74 w_mat = torch.mm(r, self.weights.view(r.size(1), -1)) 

75 w_mat = w_mat.view(-1, h.size(1), h.size(1)) 

76 w_mat = self.hidden_dropout1(w_mat) 

77 

78 x = torch.bmm(x, w_mat) 

79 x = x.view(-1, h.size(1)) 

80 x = self.bn1(x) 

81 x = self.hidden_dropout2(x) 

82 x = torch.mm(x, embeddings.weight.transpose(1, 0)) 

83 pred = torch.sigmoid(x) 

84 return pred 

85 

86 def _get_rec_embeddings(self, user): 

87 relation_users = torch.tensor([self.ui_relation] * user.shape[0], device=self.device) 

88 user_e = self.user_embedding(user) 

89 r_e = self.relation_embedding(relation_users) 

90 return user_e, r_e 

91 

92 def _get_kg_embeddings(self, h, r): 

93 h = self.entity_embedding(h) 

94 r = self.relation_embedding(r) 

95 return h, r 

96 

97 def calculate_loss(self, interaction): 

98 user = interaction[self.USER_ID] 

99 

100 item = interaction[self.ITEM_ID] 

101 

102 head = interaction[self.HEAD_ENTITY_ID] 

103 

104 relation = interaction[self.RELATION_ID] 

105 

106 tail = interaction[self.TAIL_ENTITY_ID] 

107 

108 head_e, relation_e = self._get_kg_embeddings(head, relation) 

109 user_e, rec_r_e = self._get_rec_embeddings(user) 

110 

111 item_new = torch.zeros((item.size(0), self.n_users + self.n_items), device=self.device) 

112 item_new[:, item + self.n_users] = 1.0 

113 

114 tail_new = torch.zeros((tail.size(0), self.n_entities), device=self.device) 

115 tail_new[:, tail] = 1.0 

116 

117 if self.label_smoothing: 

118 item_new = ((1.0 - self.label_smoothing) * item_new) + (1.0 / self.n_items) 

119 tail_new = ((1.0 - self.label_smoothing) * tail_new) + (1.0 / self.n_entities) 

120 

121 score_users = self.forward(user_e, rec_r_e, self.user_embedding) 

122 score_kg = self.forward(head_e, relation_e, self.entity_embedding) 

123 

124 loss_rec = self.loss(score_users, item_new) 

125 loss_kg = self.loss(score_kg, tail_new) 

126 

127 return loss_rec + loss_kg 

128 

129 def predict(self, interaction): 

130 user = interaction[self.USER_ID] 

131 item = interaction[self.ITEM_ID] 

132 

133 relation_users = torch.tensor([self.ui_relation] * user.shape[0], device=self.device) 

134 user_e = self.user_embedding(user) 

135 r_e = self.relation_embedding(relation_users) 

136 

137 score = self.forward(user_e, r_e, self.user_embedding) 

138 

139 score = score[torch.arange(user.size(0)), item] 

140 return score 

141 

142 def predict_kg(self, interaction): 

143 head = interaction[self.HEAD_ENTITY_ID] 

144 relation = interaction[self.RELATION_ID] 

145 tail = interaction[self.TAIL_ENTITY_ID] 

146 

147 head_e = self.entity_embedding(head) 

148 r_e = self.relation_embedding(relation) 

149 

150 score = self.forward(head_e, r_e, self.entity_embedding) 

151 score = score[torch.arange(head.size(0)), tail] 

152 return score 

153 

154 def full_sort_predict(self, interaction): 

155 user = interaction[self.USER_ID] 

156 relation_users = torch.tensor([self.ui_relation] * user.shape[0], device=self.device) 

157 user_e = self.user_embedding(user) 

158 r_e = self.relation_embedding(relation_users) 

159 

160 score = self.forward(user_e, r_e, self.user_embedding) 

161 score = score[:, self.n_users :] 

162 return score 

163 

164 # def full_sort_predict_kg(self, interaction): 

165 # head = interaction[self.HEAD_ENTITY_ID] 

166 # relation = interaction[self.RELATION_ID] 

167 

168 # head_e = self.entity_embedding(head) 

169 # r_e = self.relation_embedding(relation) 

170 

171 # score = self.forward(head_e, r_e, self.entity_embedding) 

172 # return score