torch-rechub 0.0.1__py3-none-any.whl → 0.0.3__py3-none-any.whl

torch_rechub/basic/loss_func.py

@@ -0,0 +1,34 @@
+import torch
+import torch.functional as F
+
+
+class HingeLoss(torch.nn.Module):
+    """Hinge Loss for pairwise learning.
+    reference: https://github.com/ustcml/RecStudio/blob/main/recstudio/model/loss_func.py
+
+    """
+
+    def __init__(self, margin=2, num_items=None):
+        super().__init__()
+        self.margin = margin
+        self.n_items = num_items
+
+    def forward(self, pos_score, neg_score):
+        loss = torch.maximum(torch.max(neg_score, dim=-1).values - pos_score + self.margin, torch.tensor([0]).type_as(pos_score))
+        if self.n_items is not None:
+            impostors = neg_score - pos_score.view(-1, 1) + self.margin > 0
+            rank = torch.mean(impostors, -1) * self.n_items
+            return torch.mean(loss * torch.log(rank + 1))
+        else:
+            return torch.mean(loss)
+
+
+class BPRLoss(torch.nn.Module):
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, pos_score, neg_score):
+        loss = torch.mean(-(pos_score - neg_score).sigmoid().log(), dim=-1)
+        return loss
+        #loss = -torch.mean(F.logsigmoid(pos_score - torch.max(neg_score, dim=-1))) need v1.10

torch_rechub/basic/metaoptimizer.py

@@ -0,0 +1,72 @@
+"""The metaoptimizer module, it provides a class MetaBalance 
+MetaBalance is used to scale the gradient and balance the gradient of each task
+Authors: Qida Dong, dongjidan@126.com
+"""
+import torch
+from torch.optim.optimizer import Optimizer
+
+
+class MetaBalance(Optimizer):
+    """MetaBalance Optimizer
+       This method is used to scale the gradient and balance the gradient of each task
+
+    Args:
+        parameters (list): the parameters of model
+        relax_factor (float, optional): the relax factor of gradient scaling (default: 0.7)
+        beta (float, optional): the coefficient of moving average (default: 0.9)
+		"""
+
+    def __init__(self, parameters, relax_factor=0.7, beta=0.9):
+
+        if relax_factor < 0. or relax_factor >= 1.:
+            raise ValueError(f'Invalid relax_factor: {relax_factor}, it should be 0. <= relax_factor < 1.')
+        if beta < 0. or beta >= 1.:
+            raise ValueError(f'Invalid beta: {beta}, it should be 0. <= beta < 1.')
+        rel_beta_dict = {'relax_factor': relax_factor, 'beta': beta}
+        super(MetaBalance, self).__init__(parameters, rel_beta_dict)
+
+    @torch.no_grad()
+    def step(self, losses):
+        """_summary_
+        Args:
+            losses (_type_): _description_
+
+        Raises:
+            RuntimeError: _description_
+        """
+
+        for idx, loss in enumerate(losses):
+            loss.backward(retain_graph=True)
+            for group in self.param_groups:
+                for gp in group['params']:
+                    if gp.grad is None:
+                        # print('breaking')
+                        break
+                    if gp.grad.is_sparse:
+                        raise RuntimeError('MetaBalance does not support sparse gradients')
+                    # store the result of moving average
+                    state = self.state[gp]
+                    if len(state) == 0:
+                        for i in range(len(losses)):
+                            if i == 0:
+                                gp.norms = [0]
+                            else:
+                                gp.norms.append(0)
+                    # calculate the moving average
+                    beta = group['beta']
+                    gp.norms[idx] = gp.norms[idx] * beta + (1 - beta) * torch.norm(gp.grad)
+                    # scale the auxiliary gradient
+                    relax_factor = group['relax_factor']
+                    gp.grad = gp.grad * gp.norms[0] / (gp.norms[idx] + 1e-5) * relax_factor + gp.grad * (1. - relax_factor)
+                    # store the gradient of each auxiliary task in state
+                    if idx == 0:
+                        state['sum_gradient'] = torch.zeros_like(gp.data)
+                        state['sum_gradient'] += gp.grad
+                    else:
+                        state['sum_gradient'] += gp.grad
+
+                    if gp.grad is not None:
+                        gp.grad.detach_()
+                        gp.grad.zero_()
+                    if idx == len(losses) - 1:
+                        gp.grad = state['sum_gradient']

torch_rechub/basic/metric.py

@@ -0,0 +1,250 @@
+"""The metric module, it is used to provide some metrics for recommenders.
+Available function:
+- auc_score: compute AUC
+- gauc_score: compute GAUC
+- log_loss: compute LogLoss
+- topk_metrics: compute topk metrics contains 'ndcg', 'mrr', 'recall', 'hit'
+Authors: Qida Dong, dongjidan@126.com
+"""
+from sklearn.metrics import roc_auc_score
+import numpy as np
+from collections import defaultdict
+
+
+def auc_score(y_true, y_pred):
+
+	return roc_auc_score(y_true, y_pred)
+
+
+def get_user_pred(y_true, y_pred, users):
+	"""divide the result into different group by user id
+
+	Args:
+		y_true (array): all true labels of the data
+		y_pred (array): the predicted score
+		users (array): user id 
+
+	Return:
+		user_pred (dict): {userid: values}, key is user id and value is the labels and scores of each user
+	"""
+	user_pred = {}
+	for i, u in enumerate(users):
+		if u not in user_pred:
+			user_pred[u] = {'y_true': [y_true[i]], 'y_pred': [y_pred[i]]}
+		else:
+			user_pred[u]['y_true'].append(y_true[i])
+			user_pred[u]['y_pred'].append(y_pred[i])
+
+	return user_pred
+
+
+def gauc_score(y_true, y_pred, users, weights=None):
+	"""compute GAUC
+
+	Args: 
+		y_true (array): dim(N, ), all true labels of the data
+		y_pred (array): dim(N, ), the predicted score
+		users (array): dim(N, ), user id 
+		weight (dict): {userid: weight_value}, it contains weights for each group. 
+				if it is None, the weight is equal to the number
+				of times the user is recommended
+	Return:
+		score: float, GAUC
+	"""
+	assert len(y_true) == len(y_pred) and len(y_true) == len(users)
+
+	user_pred = get_user_pred(y_true, y_pred, users)
+	score = 0
+	num = 0
+	for u in user_pred.keys():
+		auc = auc_score(user_pred[u]['y_true'], user_pred[u]['y_pred'])
+		if weights is None:
+			user_weight = len(user_pred[u]['y_true'])
+		else:
+			user_weight = weights[u]
+		auc *= user_weight
+		num += user_weight
+		score += auc
+	return score / num
+
+
+
+def ndcg_score(y_true, y_pred, topKs=None):
+	if topKs is None:
+		topKs = [5]
+	result = topk_metrics(y_true, y_pred, topKs)
+	return result['NDCG']
+	
+
+
+def hit_score(y_true, y_pred, topKs=None):
+	if topKs is None:
+		topKs = [5]
+	result = topk_metrics(y_true, y_pred, topKs)
+	return result['Hit']
+
+
+def mrr_score(y_true, y_pred, topKs=None):
+	if topKs is None:
+		topKs = [5]
+	result = topk_metrics(y_true, y_pred, topKs)
+	return result['MRR']
+
+
+def recall_score(y_true, y_pred, topKs=None):
+	if topKs is None:
+		topKs = [5]
+	result = topk_metrics(y_true, y_pred, topKs)
+	return result['Recall']
+
+
+def precision_score(y_true, y_pred, topKs=None):
+	if topKs is None:
+		topKs = [5]
+	result = topk_metrics(y_true, y_pred, topKs)
+	return result['Precision']
+
+
+def topk_metrics(y_true, y_pred, topKs=None):
+	"""choice topk metrics and compute it
+	the metrics contains 'ndcg', 'mrr', 'recall', 'precision' and 'hit'
+
+	Args:
+		y_true (dict): {userid, item_ids}, the key is user id and the value is the list that contains the items the user interacted
+		y_pred (dict): {userid, item_ids}, the key is user id and the value is the list that contains the items recommended  
+		topKs (list or tuple): if you want to get top5 and top10, topKs=(5, 10)
+
+	Return:
+		results (dict): {metric_name: metric_values}, it contains five metrics, 'ndcg', 'recall', 'mrr', 'hit', 'precision'
+
+	"""
+	if topKs is None:
+		topKs = [5]
+	assert len(y_true) == len(y_pred)
+
+	if not isinstance(topKs, (tuple, list)):
+		raise ValueError('topKs wrong, it should be tuple or list')
+	
+	pred_array = []
+	true_array = []
+	for u in y_true.keys():
+		pred_array.append(y_pred[u])
+		true_array.append(y_true[u])
+	ndcg_result = []
+	mrr_result = []
+	hit_result = []
+	precision_result = []
+	recall_result = []
+	for idx in range(len(topKs)):
+		ndcgs = 0
+		mrrs = 0
+		hits = 0
+		precisions = 0
+		recalls = 0
+		gts = 0
+		for i in range(len(true_array)):
+			if len(true_array[i]) != 0:
+				mrr_tmp = 0
+				mrr_flag = True
+				hit_tmp = 0
+				dcg_tmp = 0
+				idcg_tmp = 0
+				for j in range(topKs[idx]):
+					if pred_array[i][j] in true_array[i]:
+						hit_tmp += 1.
+						if mrr_flag:
+							mrr_flag = False
+							mrr_tmp = 1. / (1 + j)
+						dcg_tmp += 1. / (np.log2(j + 2))
+					if j < len(true_array[i]):
+						idcg_tmp += 1. / (np.log2(j + 2))
+				gts += len(true_array[i])
+				hits += hit_tmp
+				mrrs += mrr_tmp
+				recalls += hit_tmp / len(true_array[i])
+				precisions += hit_tmp / topKs[idx]
+				if idcg_tmp != 0:
+					ndcgs += dcg_tmp / idcg_tmp
+		hit_result.append(round(hits / gts, 4))
+		mrr_result.append(round(mrrs / len(pred_array), 4))
+		recall_result.append(round(recalls / len(pred_array), 4))
+		precision_result.append(round(precisions / len(pred_array), 4))
+		ndcg_result.append(round(ndcgs / len(pred_array), 4))
+
+	results = defaultdict(list)
+	for idx in range(len(topKs)):
+
+		output = f'NDCG@{topKs[idx]}: {ndcg_result[idx]}'
+		results['NDCG'].append(output)
+
+		output = f'MRR@{topKs[idx]}: {mrr_result[idx]}'
+		results['MRR'].append(output)
+
+		output = f'Recall@{topKs[idx]}: {recall_result[idx]}'
+		results['Recall'].append(output)
+
+		output = f'Hit@{topKs[idx]}: {hit_result[idx]}'
+		results['Hit'].append(output)
+
+		output = f'Precision@{topKs[idx]}: {precision_result[idx]}'
+		results['Precision'].append(output)
+	return results
+
+def log_loss(y_true, y_pred):
+	score = y_true * np.log(y_pred) + (1 - y_true) * np.log(1 - y_pred)
+	return -score.sum() / len(y_true)
+
+def Coverage(y_pred, all_items, topKs=None):
+	"""compute the coverage
+	This method measures the diversity of the recommended items 
+	and the ability to explore the long-tailed items
+	Arg:
+		y_pred (dict): {userid, item_ids}, the key is user id and the value is the list that contains the items recommended  
+		all_items (set): all unique items
+	Return:
+		result (list[float]): the list of coverage scores
+	"""
+	if topKs is None:
+		topKs = [5]
+	result = []
+	for k in topKs:
+		rec_items = set([])
+		for u in y_pred.keys():
+			tmp_items = set(y_pred[u][:k])
+			rec_items = rec_items | tmp_items
+		score = len(rec_items) * 1. / len(all_items)
+		score = round(score, 4)
+		result.append(f'Coverage@{k}: {score}')
+	return result
+
+
+# print(Coverage({'0':[0,1,2],'1':[1,3,4]}, {0,1,2,3,4,5}, [2,3]))
+
+# pred = np.array([  0.3, 0.2, 0.5, 0.9, 0.7, 0.31, 0.8, 0.1, 0.4, 0.6])
+# label = np.array([   1,   0,   0,   1,   0,   0,    1,   0,   0,   1])
+# users_id = np.array([ 2,   1,   0,   2,   1,   0,    0,   2,   1,   1])
+
+# print('auc: ', auc_score(label, pred))
+# print('gauc: ', gauc_score(label, pred, users_id))
+# print('log_loss: ', log_loss(label, pred))
+
+# for mt in ['ndcg', 'mrr', 'recall', 'hit','s']:
+# 	tm = topk_metrics(y_true, y_pred, users_id, 3, metric_type=mt)
+# 	print(f'{mt}: {tm}')
+# y_pred = {'0': [0, 1], '1': [0, 1], '2': [2, 3]}
+# y_true = {'0': [1, 2], '1': [0, 1, 2], '2': [2, 3]}
+# out = topk_metrics(y_true, y_pred, topKs=(1,2))
+# ndcgs = ndcg_score(y_true,y_pred, topKs=(1,2))
+# print(out)
+# print(ndcgs)
+
+# ground_truth, match_res = np.load("C:\\Users\\dongj\\Desktop/res.npy", allow_pickle=True)
+# print(len(ground_truth),len(match_res))
+# out = topk_metrics(y_true=ground_truth, y_pred=match_res, topKs=[50])
+# print(out)
+
+if __name__ == "__main__":
+	y_pred = {'0': [0, 1], '1': [0, 1], '2': [2, 3]}
+	y_true = {'0': [1, 2], '1': [0, 1, 2], '2': [2, 3]}
+	out = topk_metrics(y_true, y_pred, topKs=(1,2))
+	print(out)

torch_rechub/models/matching/__init__.py

@@ -0,0 +1,11 @@
+from .dssm import DSSM
+from .youtube_dnn import YoutubeDNN
+from .youtube_sbc import YoutubeSBC
+from .dssm_facebook import FaceBookDSSM
+from .gru4rec import GRU4Rec
+from .comirec import ComirecSA, ComirecDR
+from .mind import MIND
+from .narm import NARM
+from .stamp import STAMP
+from .sasrec import SASRec
+from .sine import SINE

torch_rechub/models/matching/comirec.py

@@ -0,0 +1,188 @@
+"""
+Date: create on 07/06/2022
+References: 
+    paper: Controllable Multi-Interest Framework for Recommendation
+    url: https://arxiv.org/pdf/2005.09347.pdf
+    code: https://github.com/ShiningCosmos/pytorch_ComiRec/blob/main/ComiRec.py
+Authors: Kai Wang, 306178200@qq.com
+"""
+
+import torch
+
+from ...basic.layers import MLP, EmbeddingLayer, MultiInterestSA, CapsuleNetwork
+from torch import nn
+import torch.nn.functional as F
+
+
+class ComirecSA(torch.nn.Module):
+    """The match model mentioned in `Controllable Multi-Interest Framework for Recommendation` paper.
+    It's a ComirecSA match model trained by global softmax loss on list-wise samples.
+    Note in origin paper, it's without item dnn tower and train item embedding directly.
+
+    Args:
+        user_features (list[Feature Class]): training by the user tower module.
+        history_features (list[Feature Class]): training history
+        item_features (list[Feature Class]): training by the embedding table, it's the item id feature.
+        neg_item_feature (list[Feature Class]): training by the embedding table, it's the negative items id feature.
+        temperature (float): temperature factor for similarity score, default to 1.0.
+        interest_num （int): interest num
+    """
+
+    def __init__(self, user_features, history_features, item_features, neg_item_feature, temperature=1.0, interest_num=4):
+        super().__init__()
+        self.user_features = user_features
+        self.item_features = item_features
+        self.history_features = history_features
+        self.neg_item_feature = neg_item_feature
+        self.temperature = temperature
+        self.interest_num = interest_num
+        self.user_dims = sum([fea.embed_dim for fea in user_features+history_features])
+
+        self.embedding = EmbeddingLayer(user_features + item_features + history_features)
+        self.multi_interest_sa = MultiInterestSA(embedding_dim=self.history_features[0].embed_dim, interest_num=self.interest_num)
+        self.convert_user_weight = nn.Parameter(torch.rand(self.user_dims, self.history_features[0].embed_dim), requires_grad=True)
+        self.mode = None
+
+    def forward(self, x):
+        user_embedding = self.user_tower(x)
+        item_embedding = self.item_tower(x)
+        if self.mode == "user":
+            return user_embedding
+        if self.mode == "item":
+            return item_embedding
+
+        pos_item_embedding = item_embedding[:,0,:]
+        dot_res = torch.bmm(user_embedding, pos_item_embedding.squeeze(1).unsqueeze(-1))
+        k_index = torch.argmax(dot_res, dim=1)
+        best_interest_emb = torch.rand(user_embedding.shape[0], user_embedding.shape[2]).to(user_embedding.device)
+        for k in range(user_embedding.shape[0]):
+            best_interest_emb[k, :] = user_embedding[k, k_index[k], :]
+        best_interest_emb = best_interest_emb.unsqueeze(1)
+
+        y = torch.mul(best_interest_emb, item_embedding).sum(dim=1)
+
+        return y
+
+    def user_tower(self, x):
+        if self.mode == "item":
+            return None
+        input_user = self.embedding(x, self.user_features, squeeze_dim=True).unsqueeze(1)  #[batch_size, num_features*deep_dims]
+        input_user = input_user.expand([input_user.shape[0], self.interest_num, input_user.shape[-1]])
+
+        history_emb = self.embedding(x, self.history_features).squeeze(1)
+        mask = self.gen_mask(x)
+        mask = mask.unsqueeze(-1).float()
+        multi_interest_emb = self.multi_interest_sa(history_emb,mask)
+
+        input_user = torch.cat([input_user,multi_interest_emb],dim=-1)
+
+        # user_embedding = self.user_mlp(input_user).unsqueeze(1)  #[batch_size, interest_num, embed_dim]
+        user_embedding = torch.matmul(input_user,self.convert_user_weight)
+        user_embedding = F.normalize(user_embedding, p=2, dim=-1)  # L2 normalize
+        if self.mode == "user":
+            return user_embedding  #inference embedding mode -> [batch_size, interest_num, embed_dim]
+        return user_embedding
+
+    def item_tower(self, x):
+        if self.mode == "user":
+            return None
+        pos_embedding = self.embedding(x, self.item_features, squeeze_dim=False)  #[batch_size, 1, embed_dim]
+        pos_embedding = F.normalize(pos_embedding, p=2, dim=-1)  # L2 normalize
+        if self.mode == "item":  #inference embedding mode
+            return pos_embedding.squeeze(1)  #[batch_size, embed_dim]
+        neg_embeddings = self.embedding(x, self.neg_item_feature,
+                                        squeeze_dim=False).squeeze(1)  #[batch_size, n_neg_items, embed_dim]
+        neg_embeddings = F.normalize(neg_embeddings, p=2, dim=-1)  # L2 normalize
+        return torch.cat((pos_embedding, neg_embeddings), dim=1)  #[batch_size, 1+n_neg_items, embed_dim]
+
+    def gen_mask(self, x):
+        his_list = x[self.history_features[0].name]
+        mask = (his_list > 0).long()
+        return mask
+
+class ComirecDR(torch.nn.Module):
+    """The match model mentioned in `Controllable Multi-Interest Framework for Recommendation` paper.
+    It's a ComirecDR match model trained by global softmax loss on list-wise samples.
+    Note in origin paper, it's without item dnn tower and train item embedding directly.
+
+    Args:
+        user_features (list[Feature Class]): training by the user tower module.
+        history_features (list[Feature Class]): training history
+        item_features (list[Feature Class]): training by the embedding table, it's the item id feature.
+        neg_item_feature (list[Feature Class]): training by the embedding table, it's the negative items id feature.
+        max_length (int): max sequence length of input item sequence
+        temperature (float): temperature factor for similarity score, default to 1.0.
+        interest_num （int): interest num
+    """
+
+    def __init__(self, user_features, history_features, item_features, neg_item_feature, max_length, temperature=1.0, interest_num=4):
+        super().__init__()
+        self.user_features = user_features
+        self.item_features = item_features
+        self.history_features = history_features
+        self.neg_item_feature = neg_item_feature
+        self.temperature = temperature
+        self.interest_num = interest_num
+        self.max_length = max_length
+        self.user_dims = sum([fea.embed_dim for fea in user_features+history_features])
+
+        self.embedding = EmbeddingLayer(user_features + item_features + history_features)
+        self.capsule = CapsuleNetwork(self.history_features[0].embed_dim,self.max_length,bilinear_type=2,interest_num=self.interest_num)
+        self.convert_user_weight = nn.Parameter(torch.rand(self.user_dims, self.history_features[0].embed_dim), requires_grad=True)
+        self.mode = None
+
+    def forward(self, x):
+        user_embedding = self.user_tower(x)
+        item_embedding = self.item_tower(x)
+        if self.mode == "user":
+            return user_embedding
+        if self.mode == "item":
+            return item_embedding
+
+        pos_item_embedding = item_embedding[:,0,:]
+        dot_res = torch.bmm(user_embedding, pos_item_embedding.squeeze(1).unsqueeze(-1))
+        k_index = torch.argmax(dot_res, dim=1)
+        best_interest_emb = torch.rand(user_embedding.shape[0], user_embedding.shape[2]).to(user_embedding.device)
+        for k in range(user_embedding.shape[0]):
+            best_interest_emb[k, :] = user_embedding[k, k_index[k], :]
+        best_interest_emb = best_interest_emb.unsqueeze(1)
+
+        y = torch.mul(best_interest_emb, item_embedding).sum(dim=1)
+
+        return y
+
+    def user_tower(self, x):
+        if self.mode == "item":
+            return None
+        input_user = self.embedding(x, self.user_features, squeeze_dim=True).unsqueeze(1)  #[batch_size, num_features*deep_dims]
+        input_user = input_user.expand([input_user.shape[0], self.interest_num, input_user.shape[-1]])
+
+        history_emb = self.embedding(x, self.history_features).squeeze(1)
+        mask = self.gen_mask(x)
+        multi_interest_emb = self.capsule(history_emb,mask)
+
+        input_user = torch.cat([input_user,multi_interest_emb],dim=-1)
+
+        # user_embedding = self.user_mlp(input_user).unsqueeze(1)  #[batch_size, interest_num, embed_dim]
+        user_embedding = torch.matmul(input_user,self.convert_user_weight)
+        user_embedding = F.normalize(user_embedding, p=2, dim=-1)  # L2 normalize
+        if self.mode == "user":
+            return user_embedding  #inference embedding mode -> [batch_size, interest_num, embed_dim]
+        return user_embedding
+
+    def item_tower(self, x):
+        if self.mode == "user":
+            return None
+        pos_embedding = self.embedding(x, self.item_features, squeeze_dim=False)  #[batch_size, 1, embed_dim]
+        pos_embedding = F.normalize(pos_embedding, p=2, dim=-1)  # L2 normalize
+        if self.mode == "item":  #inference embedding mode
+            return pos_embedding.squeeze(1)  #[batch_size, embed_dim]
+        neg_embeddings = self.embedding(x, self.neg_item_feature,
+                                        squeeze_dim=False).squeeze(1)  #[batch_size, n_neg_items, embed_dim]
+        neg_embeddings = F.normalize(neg_embeddings, p=2, dim=-1)  # L2 normalize
+        return torch.cat((pos_embedding, neg_embeddings), dim=1)  #[batch_size, 1+n_neg_items, embed_dim]
+
+    def gen_mask(self, x):
+        his_list = x[self.history_features[0].name]
+        mask = (his_list > 0).long()
+        return mask

torch_rechub/models/matching/dssm.py

@@ -0,0 +1,66 @@
+"""
+Date: create on 12/05/2022, update on 20/05/2022
+References: 
+    paper: (CIKM'2013) Learning Deep Structured Semantic Models for Web Search using Clickthrough Data
+    url: https://posenhuang.github.io/papers/cikm2013_DSSM_fullversion.pdf
+    code: https://github.com/bbruceyuan/DeepMatch-Torch/blob/main/deepmatch_torch/models/dssm.py
+Authors: Mincai Lai, laimincai@shanghaitech.edu.cn
+"""
+
+import torch
+import torch.nn.functional as F
+from ...basic.layers import MLP, EmbeddingLayer
+
+
+class DSSM(torch.nn.Module):
+    """Deep Structured Semantic Model
+
+    Args:
+        user_features (list[Feature Class]): training by the user tower module.
+        item_features (list[Feature Class]): training by the item tower module.
+        temperature (float): temperature factor for similarity score, default to 1.0.
+        user_params (dict): the params of the User Tower module, keys include:`{"dims":list, "activation":str, "dropout":float, "output_layer":bool`}.
+        item_params (dict): the params of the Item Tower module, keys include:`{"dims":list, "activation":str, "dropout":float, "output_layer":bool`}.
+    """
+
+    def __init__(self, user_features, item_features, user_params, item_params, temperature=1.0):
+        super().__init__()
+        self.user_features = user_features
+        self.item_features = item_features
+        self.temperature = temperature
+        self.user_dims = sum([fea.embed_dim for fea in user_features])
+        self.item_dims = sum([fea.embed_dim for fea in item_features])
+
+        self.embedding = EmbeddingLayer(user_features + item_features)
+        self.user_mlp = MLP(self.user_dims, output_layer=False, **user_params)
+        self.item_mlp = MLP(self.item_dims, output_layer=False, **item_params)
+        self.mode = None
+
+    def forward(self, x):
+        user_embedding = self.user_tower(x)
+        item_embedding = self.item_tower(x)
+        if self.mode == "user":
+            return user_embedding
+        if self.mode == "item":
+            return item_embedding
+
+        # calculate cosine score
+        y = torch.mul(user_embedding, item_embedding).sum(dim=1)
+        # y = y / self.temperature
+        return torch.sigmoid(y)
+
+    def user_tower(self, x):
+        if self.mode == "item":
+            return None
+        input_user = self.embedding(x, self.user_features, squeeze_dim=True)  #[batch_size, num_features*deep_dims]
+        user_embedding = self.user_mlp(input_user)  #[batch_size, user_params["dims"][-1]]
+        user_embedding = F.normalize(user_embedding, p=2, dim=1)  # L2 normalize
+        return user_embedding
+
+    def item_tower(self, x):
+        if self.mode == "user":
+            return None
+        input_item = self.embedding(x, self.item_features, squeeze_dim=True)  #[batch_size, num_features*embed_dim]
+        item_embedding = self.item_mlp(input_item)  #[batch_size, item_params["dims"][-1]]
+        item_embedding = F.normalize(item_embedding, p=2, dim=1)
+        return item_embedding