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

torch_rechub/models/matching/sine.py

@@ -0,0 +1,151 @@
+"""
+Date: created on 03/07/2022
+References:
+    paper: Sparse-Interest Network for Sequential Recommendation
+    url: https://arxiv.org/abs/2102.09267
+    code: https://github.com/Qiaoyut/SINE/blob/master/model.py
+Authors: Bo Kang, klinux@live.com
+"""
+
+import torch
+import torch.nn.functional as F
+from torch import einsum
+
+class SINE(torch.nn.Module):
+    """The match model was proposed in `Sparse-Interest Network for Sequential Recommendation` paper.
+
+    Args:        
+        history_features (list[str]): training history feature names, this is for indexing the historical sequences from input dictionary
+        item_features (list[str]): item feature names, this is for indexing the items from input dictionary
+        neg_item_features (list[str]): neg item feature names, this for indexing negative items from input dictionary
+        num_items (int): number of items in the data
+        embedding_dim (int): dimensionality of the embeddings        
+        hidden_dim (int): dimensionality of the hidden layer in self attention modules
+        num_concept (int): number of concept, also called conceptual prototypes 
+        num_intention (int): number of (user) specific intentions out of the concepts
+        seq_max_len (int): max sequence length of input item sequence
+        num_heads （int): number of attention heads in self attention modules, default to 1
+        temperature (float): temperature factor in the similarity measure, default to 1.0        
+    """
+    def __init__(self, history_features, item_features, neg_item_features, num_items, embedding_dim, hidden_dim, num_concept, num_intention, seq_max_len, num_heads=1, temperature=1.0):
+        super().__init__()        
+        self.item_features = item_features
+        self.history_features = history_features
+        self.neg_item_features = neg_item_features
+        self.temperature = temperature
+        self.num_concept = num_concept        
+        self.num_intention = num_intention
+        self.seq_max_len = seq_max_len
+
+        std = 1e-4
+        self.item_embedding = torch.nn.Embedding(num_items, embedding_dim)
+        torch.nn.init.normal_(self.item_embedding.weight, 0, std)
+        self.concept_embedding = torch.nn.Embedding(num_concept, embedding_dim)
+        torch.nn.init.normal_(self.concept_embedding.weight, 0, std)
+        self.position_embedding = torch.nn.Embedding(seq_max_len, embedding_dim)
+        torch.nn.init.normal_(self.position_embedding.weight, 0, std)
+    
+
+        self.w_1 = torch.nn.Parameter(torch.rand(embedding_dim, hidden_dim), requires_grad=True)        
+        self.w_2 = torch.nn.Parameter(torch.rand(hidden_dim, num_heads), requires_grad=True)
+
+        self.w_3 = torch.nn.Parameter(torch.rand(embedding_dim, embedding_dim), requires_grad=True)
+
+        self.w_k1 = torch.nn.Parameter(torch.rand(embedding_dim, hidden_dim), requires_grad=True)
+        self.w_k2 = torch.nn.Parameter(torch.rand(hidden_dim, num_intention), requires_grad=True)
+
+        self.w_4 = torch.nn.Parameter(torch.rand(embedding_dim, hidden_dim), requires_grad=True)
+        self.w_5 = torch.nn.Parameter(torch.rand(hidden_dim, num_heads), 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
+        
+        y = torch.mul(user_embedding, item_embedding).sum(dim=-1)
+
+        # # compute covariance regularizer
+        # M = torch.cov(self.concept_embedding.weight, correction=0)
+        # l_c = (torch.norm(M, p='fro')**2 - torch.norm(torch.diag(M), p='fro')**2)/2
+
+        return y
+
+    def user_tower(self, x):
+        if self.mode == "item":
+            return None
+        
+        # sparse interests extraction
+        ## user specific historical item embedding X^u
+        hist_item = x[self.history_features[0]]
+        x_u = self.item_embedding(hist_item) + self.position_embedding.weight.unsqueeze(0)
+        x_u_mask = (x[self.history_features[0]] > 0).long()
+
+        ## user specific conceptual prototypes C^u
+        ### attention a
+        h_1 = einsum('bse, ed -> bsd', x_u, self.w_1).tanh()
+        a_hist = F.softmax(einsum('bsd, dh -> bsh', h_1, self.w_2) + -1.e9 * (1 - x_u_mask.unsqueeze(-1).float()), dim=1) 
+
+        ### virtual concept vector z_u
+        z_u = einsum("bse, bsh -> be", x_u, a_hist)
+
+        ### similarity between user's concept vector and entire conceptual prototypes s^u
+        s_u = einsum("be, te -> bt", z_u, self.concept_embedding.weight)
+        s_u_top_k = torch.topk(s_u, self.num_intention)
+
+        ### final C^u
+        c_u =  einsum("bk, bke -> bke", torch.sigmoid(s_u_top_k.values), self.concept_embedding(s_u_top_k.indices))
+
+        ## user intention assignment P_{k|t}        
+        p_u = F.softmax(einsum("bse, bke -> bks", F.normalize(x_u @ self.w_3, dim=-1),  F.normalize(c_u, p=2, dim=-1)), dim=1)
+
+        ## attention weighing P_{t|k}
+        h_2 = einsum('bse, ed -> bsd', x_u, self.w_k1).tanh()
+        a_concept_k = F.softmax(einsum('bsd, dk -> bsk', h_2, self.w_k2) + -1.e9 * (1 - x_u_mask.unsqueeze(-1).float()), dim=1) 
+
+        ## multiple interests encoding \phi_\theta^k(x^u)
+        phi_u = einsum("bks, bse -> bke", p_u * a_concept_k.permute(0, 2, 1), x_u)
+
+
+        # adaptive interest aggregation
+        ## intention aware input behavior \hat{X^u}
+        x_u_hat = einsum('bks, bke -> bse', p_u, c_u)
+
+        ## user's next intention C^u_{apt}
+        h_3 = einsum('bse, ed -> bsd', x_u_hat, self.w_4).tanh()
+        c_u_apt = F.normalize(
+                    einsum(
+                        "bs, bse -> be",
+                        F.softmax(einsum('bsd, dh -> bsh', h_3, self.w_5).reshape(-1, self.seq_max_len) + -1.e9 * (1 - x_u_mask.float()), dim=1),
+                        x_u_hat
+                    ), -1)
+
+        ## aggregation weights e_k^u
+        e_u = F.softmax(einsum('be, bke -> bk', c_u_apt, phi_u) / self.temperature, dim=1)         
+
+
+        # final user representation v^u
+        v_u = einsum('bk, bke -> be', e_u, phi_u)
+
+        if self.mode == "user":
+            return v_u
+        return v_u.unsqueeze(1)
+
+    def item_tower(self, x):
+        if self.mode == "user":
+            return None
+        pos_embedding = self.item_embedding(x[self.item_features[0]]).unsqueeze(1)
+        if self.mode == "item":  #inference embedding mode
+            return pos_embedding.squeeze(1)  #[batch_size, embed_dim]
+        neg_embeddings = self.item_embedding(x[self.neg_item_features[0]]).squeeze(1)  #[batch_size, n_neg_items, embed_dim]        
+
+        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/stamp.py

@@ -0,0 +1,83 @@
+"""
+Date: created on 17/09/2022
+References:
+    paper: STAMP: Short-Term Attention/Memory Priority Model for Session-based Recommendation
+    url: https://dl.acm.org/doi/10.1145/3219819.3219950
+    official Tensorflow implementation: https://github.com/uestcnlp/STAMP
+Authors: Bo Kang, klinux@live.com
+"""
+
+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+
+
+class STAMP(nn.Module):
+    def __init__(self, item_history_feature, weight_std, emb_std):
+        super(STAMP, self).__init__()        
+        
+        # item embedding layer
+        self.item_history_feature = item_history_feature
+        n_items, item_emb_dim, = item_history_feature.vocab_size, item_history_feature.embed_dim
+        self.item_emb = nn.Embedding(n_items, item_emb_dim, padding_idx=0)
+        
+        # weights and biases for attention computation
+        self.w_0 = nn.Parameter(torch.zeros(item_emb_dim, 1))        
+        self.w_1_t = nn.Parameter(torch.zeros(item_emb_dim, item_emb_dim))
+        self.w_2_t = nn.Parameter(torch.zeros(item_emb_dim, item_emb_dim))        
+        self.w_3_t = nn.Parameter(torch.zeros(item_emb_dim, item_emb_dim))        
+        self.b_a = nn.Parameter(torch.zeros(item_emb_dim))
+        self._init_parameter_weights(weight_std)
+
+        # mlp layers
+        self.f_s = nn.Sequential(nn.Tanh(), nn.Linear(item_emb_dim, item_emb_dim))
+        self.f_t = nn.Sequential(nn.Tanh(), nn.Linear(item_emb_dim, item_emb_dim))        
+        self.emb_std = emb_std
+        self.apply(self._init_module_weights)
+
+    
+    def _init_parameter_weights(self, weight_std):
+        nn.init.normal_(self.w_0, std=weight_std)
+        nn.init.normal_(self.w_1_t, std=weight_std)
+        nn.init.normal_(self.w_2_t, std=weight_std)
+        nn.init.normal_(self.w_3_t, std=weight_std)
+
+
+    def _init_module_weights(self, module):        
+        if isinstance(module, nn.Linear):            
+            module.weight.data.normal_(std=self.emb_std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):            
+            module.weight.data.normal_(std=self.emb_std)
+        
+
+    def forward(self, input_dict):        
+        # Index the embeddings for the items in the session        
+        input = input_dict[self.item_history_feature.name]
+        value_mask = (input != 0).unsqueeze(-1)
+        value_counts = value_mask.sum(dim=1, keepdim=True).squeeze(-1)                
+        item_emb_batch = self.item_emb(input) * value_mask
+
+        # Index the embeddings of the latest clicked items
+        x_t = self.item_emb(torch.gather(input, 1, value_counts-1))        
+        
+        # Eq. 2, user's general interest in the current session
+        m_s = ((item_emb_batch).sum(1) / value_counts).unsqueeze(1)
+        
+        # Eq. 7, compute attention coefficient
+        a = F.normalize(torch.exp(torch.sigmoid(item_emb_batch@self.w_1_t + x_t@self.w_2_t + m_s@self.w_3_t + self.b_a) @ self.w_0) * value_mask, p=1, dim=1)
+
+        # Eq. 8, compute user's attention-based interests
+        m_a = (a * item_emb_batch).sum(1) + m_s.squeeze(1)
+
+        # Eq. 3, compute the output state of the general interest
+        h_s = self.f_s(m_a)
+
+        # Eq. 9, compute the output state of the short-term interest
+        h_t = self.f_t(x_t).squeeze(1)
+
+        # Eq. 4, compute candidate scores
+        z =  h_s * h_t @ self.item_emb.weight.T        
+        
+        return z

torch_rechub/models/matching/youtube_dnn.py

@@ -0,0 +1,71 @@
+"""
+Date: create on 23/05/2022
+References: 
+    paper: (RecSys'2016) Deep Neural Networks for YouTube Recommendations
+    url: https://dl.acm.org/doi/10.1145/2959100.2959190
+Authors: Mincai Lai, laimincai@shanghaitech.edu.cn
+"""
+
+import torch
+import torch.nn.functional as F
+from ...basic.layers import MLP, EmbeddingLayer
+
+
+class YoutubeDNN(torch.nn.Module):
+    """The match model mentioned in `Deep Neural Networks for YouTube Recommendations` paper.
+    It's a DSSM 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.
+        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.
+        user_params (dict): the params of the User Tower module, keys include:`{"dims":list, "activation":str, "dropout":float, "output_layer":bool`}.
+        temperature (float): temperature factor for similarity score, default to 1.0.
+    """
+
+    def __init__(self, user_features, item_features, neg_item_feature, user_params, temperature=1.0):
+        super().__init__()
+        self.user_features = user_features
+        self.item_features = item_features
+        self.neg_item_feature = neg_item_feature
+        self.temperature = temperature
+        self.user_dims = sum([fea.embed_dim for fea in user_features])
+        self.embedding = EmbeddingLayer(user_features + item_features)
+        self.user_mlp = MLP(self.user_dims, output_layer=False, **user_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=2)
+        y = y / self.temperature
+        return 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).unsqueeze(1)  #[batch_size, 1, embed_dim]
+        user_embedding = F.normalize(user_embedding, p=2, dim=2)
+        if self.mode == "user":
+            return user_embedding.squeeze(1)  #inference embedding mode -> [batch_size, 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=2)
+        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=2)
+        return torch.cat((pos_embedding, neg_embeddings), dim=1)  #[batch_size, 1+n_neg_items, embed_dim]

torch_rechub/models/matching/youtube_sbc.py

@@ -0,0 +1,98 @@
+"""
+Date: create on 23/05/2022
+References: 
+    paper: (RecSys'2019) Sampling-Bias-Corrected Neural Modeling for Large Corpus Item Recommendations
+    url: https://dl.acm.org/doi/10.1145/3298689.3346996
+Authors: Mincai Lai, laimincai@shanghaitech.edu.cn
+"""
+
+import torch
+import torch.nn.functional as F
+from ...basic.layers import MLP, EmbeddingLayer
+import numpy as np
+
+
+class YoutubeSBC(torch.nn.Module):
+    """Sampling-Bias-Corrected Neural Modeling for Matching by Youtube. 
+    It's a DSSM match model trained by In-batch softmax loss on list-wise samples, and add sample debias module.
+
+    Args:
+        user_features (list[Feature Class]): training by the user tower module.
+        item_features (list[Feature Class]): training by the item tower module.
+        sample_weight_feature (list[Feature Class]): used for sampling bias corrected in training.
+        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`}.
+        batch_size (int): same as batch size of DataLoader, used in in-batch sampling
+        n_neg (int): the number of negative sample for every positive sample, default to 3. Note it's must smaller than batch_size.
+        temperature (float): temperature factor for similarity score, default to 1.0.
+    """
+
+    def __init__(self,
+                 user_features,
+                 item_features,
+                 sample_weight_feature,
+                 user_params,
+                 item_params,
+                 batch_size,
+                 n_neg=3,
+                 temperature=1.0):
+        super().__init__()
+        self.user_features = user_features
+        self.item_features = item_features
+        self.sample_weight_feature = sample_weight_feature
+        self.n_neg = n_neg
+        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.batch_size = batch_size
+        self.embedding = EmbeddingLayer(user_features + item_features + sample_weight_feature)
+        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
+
+        # in-batch sampling index
+        self.index0 = np.repeat(np.arange(batch_size), n_neg + 1)
+        self.index1 = np.concatenate([np.arange(i, i + n_neg + 1) for i in range(batch_size)])
+        self.index1[np.where(self.index1 >= batch_size)] -= batch_size
+
+    def forward(self, x):
+        user_embedding = self.user_tower(x)  # (batch_size, embedding_dim)
+        item_embedding = self.item_tower(x)  # (batch_size, embedding_dim)
+        if self.mode == "user":
+            return user_embedding
+        if self.mode == "item":
+            return item_embedding
+
+        # pred[i, j] means predicted score that user_i give to item_j
+        pred = torch.cosine_similarity(user_embedding.unsqueeze(1), item_embedding, dim=2)  # (batch_size, batch_size)
+
+        # get sample weight of items in this batch
+        sample_weight = self.embedding(x, self.sample_weight_feature, squeeze_dim=True).squeeze(1)  # (batch_size)
+        scores = pred - torch.log(sample_weight)  #Sampling Bias Corrected, using broadcast. (batch_size, batch_size)
+
+        if user_embedding.shape[0] * (self.n_neg + 1) != self.index0.shape[0]:  # last batch
+            batch_size = user_embedding.shape[0]
+            index0 = self.index0[:batch_size * (self.n_neg + 1)]
+            index1 = self.index1[:batch_size * (self.n_neg + 1)]
+            index0[np.where(index0 >= batch_size)] -= batch_size
+            index1[np.where(index1 >= batch_size)] -= batch_size
+            scores = scores[index0, index1]  # (batch_size, 1 + self.n_neg)
+        else:
+            scores = scores[self.index0, self.index1]  # (batch_size, 1 + self.n_neg)
+
+        scores = scores / self.temperature
+        return scores.view(-1, self.n_neg + 1)  #(batch_size, 1 + self.n_neg)
+
+    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]]
+        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]]
+        return item_embedding

torch_rechub/models/multi_task/__init__.py

@@ -1,4 +1,5 @@
-from .shared_bottom import SharedBottom
-from .esmm import ESMM
-from .mmoe import MMOE
-from .ple import PLE
+from .shared_bottom import SharedBottom
+from .esmm import ESMM
+from .mmoe import MMOE
+from .ple import PLE
+from .aitm import AITM

torch_rechub/models/multi_task/aitm.py

@@ -0,0 +1,84 @@
+"""
+Date: create on 14/05/2022
+References: 
+    paper: (KDD'2021) Modeling the Sequential Dependence among Audience Multi-step Conversions with Multi-task Learning in Targeted Display Advertising
+    url: https://arxiv.org/abs/2105.08489
+    code: https://github.com/adtalos/AITM-torch
+Authors: Mincai Lai, laimincai@shanghaitech.edu.cn
+"""
+
+import torch
+import torch.nn as nn
+
+from ...basic.layers import MLP, EmbeddingLayer
+
+
+class AITM(nn.Module):
+    """ Adaptive Information Transfer Multi-task (AITM) framework. 
+        all the task type must be binary classificatioon.
+
+    Args:
+        features (list[Feature Class]): training by the whole module.
+        n_task (int): the number of binary classificatioon task.
+        bottom_params (dict): the params of all the botwer expert module, keys include:`{"dims":list, "activation":str, "dropout":float}`.
+        tower_params_list (list): the list of tower params dict, the keys same as expert_params.
+    """
+
+    def __init__(self, features, n_task, bottom_params, tower_params_list):
+        super().__init__()
+        self.features = features
+        self.n_task = n_task
+        self.input_dims = sum([fea.embed_dim for fea in features])
+        self.embedding = EmbeddingLayer(features)
+
+        self.bottoms = nn.ModuleList(
+            MLP(self.input_dims, output_layer=False, **bottom_params) for i in range(self.n_task))
+        self.towers = nn.ModuleList(MLP(bottom_params["dims"][-1], **tower_params_list[i]) for i in range(self.n_task))
+
+        self.info_gates = nn.ModuleList(
+            MLP(bottom_params["dims"][-1], output_layer=False, dims=[bottom_params["dims"][-1]])
+            for i in range(self.n_task - 1))
+        self.aits = nn.ModuleList(AttentionLayer(bottom_params["dims"][-1]) for _ in range(self.n_task - 1))
+
+    def forward(self, x):
+        embed_x = self.embedding(x, self.features, squeeze_dim=True)  #[batch_size, *]
+        input_towers = [self.bottoms[i](embed_x) for i in range(self.n_task)]  #[i]:[batch_size, bottom_dims[-1]]
+        for i in range(1, self.n_task):  #for task 1:n-1
+            info = self.info_gates[i - 1](input_towers[i - 1]).unsqueeze(1)  #[batch_size,1,bottom_dims[-1]]
+            ait_input = torch.cat([input_towers[i].unsqueeze(1), info], dim=1)  #[batch_size, 2, bottom_dims[-1]]
+            input_towers[i] = self.aits[i - 1](ait_input)
+
+        ys = []
+        for input_tower, tower in zip(input_towers, self.towers):
+            y = tower(input_tower)
+            ys.append(torch.sigmoid(y))
+        return torch.cat(ys, dim=1)
+
+
+class AttentionLayer(nn.Module):
+    """attention for info tranfer
+
+    Args:
+        dim (int): attention dim
+    
+    Shape:
+        Input: (batch_size, 2, dim)
+        Output: (batch_size, dim)
+    """
+
+    def __init__(self, dim=32):
+        super().__init__()
+        self.dim = dim
+        self.q_layer = nn.Linear(dim, dim, bias=False)
+        self.k_layer = nn.Linear(dim, dim, bias=False)
+        self.v_layer = nn.Linear(dim, dim, bias=False)
+        self.softmax = nn.Softmax(dim=1)
+
+    def forward(self, x):
+        Q = self.q_layer(x)
+        K = self.k_layer(x)
+        V = self.v_layer(x)
+        a = torch.sum(torch.mul(Q, K), -1) / torch.sqrt(torch.tensor(self.dim))
+        a = self.softmax(a)
+        outputs = torch.sum(torch.mul(torch.unsqueeze(a, -1), V), dim=1)
+        return outputs

torch_rechub/models/multi_task/esmm.py

@@ -1,45 +1,55 @@
-"""
-Date: create on 04/05/2022
-References: 
-    paper: (SIGIR'2018) Entire Space Multi-Task Model: An Effective Approach for Estimating Post-Click Conversion Rate 
-    url: https://arxiv.org/abs/1804.07931
-Authors: Mincai Lai, laimincai@shanghaitech.edu.cn
-"""
-
-import torch
-import torch.nn as nn
-
-from ...basic.layers import MLP, EmbeddingLayer
-
-
-class ESMM(nn.Module):
-    """Entire Space Multi-Task Model
-
-    Args:
-        user_features (list): the list of `Feature Class`, training by shared bottom and tower module. It means the user features.
-        item_features (list): the list of `Feature Class`, training by shared bottom and tower module. It means the item features.
-        cvr_params (dict): the params of the CVR Tower module, keys include:`{"dims":list, "activation":str, "dropout":float, "output_layer":bool`}
-        ctr_params (dict): the params of the CTR Tower module, keys include:`{"dims":list, "activation":str, "dropout":float, "output_layer":bool`}
-    """
-
-    def __init__(self, user_features, item_features, cvr_params, ctr_params):
-        super().__init__()
-        self.user_features = user_features
-        self.item_features = item_features
-        self.embedding = EmbeddingLayer(user_features + item_features)
-        self.tower_dims = user_features[0].embed_dim + item_features[0].embed_dim  #the dims after Field-wise Pooling Layer
-        self.tower_cvr = MLP(self.tower_dims, **cvr_params)
-        self.tower_ctr = MLP(self.tower_dims, **ctr_params)
-
-    def forward(self, x):
-        embed_user_features = self.embedding(x, self.user_features, squeeze_dim=False).sum(dim=1)  #[batch_size, embed_dim]
-        embed_item_features = self.embedding(x, self.item_features, squeeze_dim=False).sum(dim=1)  #[batch_size, embed_dim]
-        input_tower = torch.cat((embed_user_features, embed_item_features), dim=1)
-        cvr_logit = self.tower_cvr(input_tower)
-        ctr_logit = self.tower_ctr(input_tower)
-        cvr_pred = torch.sigmoid(cvr_logit)
-        ctr_pred = torch.sigmoid(ctr_logit)
-        ctcvr_pred = torch.mul(cvr_pred, cvr_pred)
-
-        ys = [cvr_pred, ctr_pred, ctcvr_pred]
-        return torch.cat(ys, dim=1)
+"""
+Date: create on 04/05/2022
+References: 
+    paper: (SIGIR'2018) Entire Space Multi-Task Model: An Effective Approach for Estimating Post-Click Conversion Rate 
+    url: https://arxiv.org/abs/1804.07931
+Authors: Mincai Lai, laimincai@shanghaitech.edu.cn
+"""
+
+import torch
+import torch.nn as nn
+
+from ...basic.layers import MLP, EmbeddingLayer
+
+
+class ESMM(nn.Module):
+    """Entire Space Multi-Task Model
+
+    Args:
+        user_features (list): the list of `Feature Class`, training by shared bottom and tower module. It means the user features.
+        item_features (list): the list of `Feature Class`, training by shared bottom and tower module. It means the item features.
+        cvr_params (dict): the params of the CVR Tower module, keys include:`{"dims":list, "activation":str, "dropout":float`}
+        ctr_params (dict): the params of the CTR Tower module, keys include:`{"dims":list, "activation":str, "dropout":float`}
+    """
+
+    def __init__(self, user_features, item_features, cvr_params, ctr_params):
+        super().__init__()
+        self.user_features = user_features
+        self.item_features = item_features
+        self.embedding = EmbeddingLayer(user_features + item_features)
+        self.tower_dims = len(user_features) * user_features[0].embed_dim + len(item_features) * item_features[0].embed_dim
+        self.tower_cvr = MLP(self.tower_dims, **cvr_params)
+        self.tower_ctr = MLP(self.tower_dims, **ctr_params)
+
+    def forward(self, x):
+        # # Field-wise Pooling Layer for user and item
+        # embed_user_features = self.embedding(x, self.user_features, squeeze_dim=False).sum(dim=1)  #[batch_size, embed_dim]
+        # embed_item_features = self.embedding(x, self.item_features, squeeze_dim=False).sum(dim=1)  #[batch_size, embed_dim]
+
+        # Here we concat all the features instead of field-wise pooling them
+        # [batch_size, num_features, embed_dim] --> [batch_size, num_features * embed_dim]
+        _batch_size = self.embedding(x, self.user_features, squeeze_dim=False).shape[0]
+        embed_user_features = self.embedding(x, self.user_features, squeeze_dim=False).reshape(_batch_size, -1)
+        embed_item_features = self.embedding(x, self.item_features, squeeze_dim=False).reshape(_batch_size, -1)
+
+        # print('embed_user_features', embed_user_features.shape)
+
+        input_tower = torch.cat((embed_user_features, embed_item_features), dim=1)
+        cvr_logit = self.tower_cvr(input_tower)
+        ctr_logit = self.tower_ctr(input_tower)
+        cvr_pred = torch.sigmoid(cvr_logit)
+        ctr_pred = torch.sigmoid(ctr_logit)
+        ctcvr_pred = torch.mul(ctr_pred, cvr_pred)
+
+        ys = [cvr_pred, ctr_pred, ctcvr_pred]
+        return torch.cat(ys, dim=1)