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

torch_rechub/models/matching/youtube_sbc.py

@@ -1,98 +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
+"""
+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 numpy as np
+import torch
+import torch.nn.functional as F
+
+from ...basic.layers import MLP, EmbeddingLayer
+
+
+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)
+        # Sampling Bias Corrected, using broadcast. (batch_size, batch_size)
+        scores = pred - torch.log(sample_weight)
+
+        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:
+            # (batch_size, 1 + self.n_neg)
+            scores = scores[self.index0, self.index1]
+
+        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
+        # [batch_size, num_features*deep_dims]
+        input_user = self.embedding(x, self.user_features, squeeze_dim=True)
+        # [batch_size, user_params["dims"][-1]]
+        user_embedding = self.user_mlp(input_user)
+        return user_embedding
+
+    def item_tower(self, x):
+        if self.mode == "user":
+            return None
+        # [batch_size, num_features*embed_dim]
+        input_item = self.embedding(x, self.item_features, squeeze_dim=True)
+        # [batch_size, item_params["dims"][-1]]
+        item_embedding = self.item_mlp(input_item)
+        return item_embedding

torch_rechub/models/multi_task/__init__.py

@@ -1,5 +1,7 @@
-from .shared_bottom import SharedBottom
-from .esmm import ESMM
-from .mmoe import MMOE
-from .ple import PLE
-from .aitm import AITM
+__all__ = ['SharedBottom', 'ESMM', 'MMOE', 'PLE', 'AITM']
+
+from .aitm import AITM
+from .esmm import ESMM
+from .mmoe import MMOE
+from .ple import PLE
+from .shared_bottom import SharedBottom

torch_rechub/models/multi_task/aitm.py

@@ -1,84 +1,83 @@
-"""
-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
+"""
+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
+            # [batch_size,1,bottom_dims[-1]]
+            info = self.info_gates[i - 1](input_towers[i - 1]).unsqueeze(1)
+            # [batch_size, 2, bottom_dims[-1]]
+            ait_input = torch.cat([input_towers[i].unsqueeze(1), info], dim=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,55 +1,56 @@
-"""
-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)
+"""
+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)

torch_rechub/models/multi_task/mmoe.py

@@ -1,58 +1,58 @@
-"""
-Date: create on 04/05/2022
-References: 
-    paper: (KDD'2018) Modeling Task Relationships in Multi-task Learning with Multi-gate Mixture-of-Experts
-    url: https://dl.acm.org/doi/pdf/10.1145/3219819.3220007
-Authors: Mincai Lai, laimincai@shanghaitech.edu.cn
-"""
-
-import torch
-import torch.nn as nn
-
-from ...basic.layers import MLP, EmbeddingLayer, PredictionLayer
-
-
-class MMOE(nn.Module):
-    """Multi-gate Mixture-of-Experts model.
-
-    Args:
-        features (list): the list of `Feature Class`, training by the expert and tower module.
-        task_types (list): types of tasks, only support `["classfication", "regression"]`.
-        n_expert (int): the number of expert net.
-        expert_params (dict): the params of all the 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, task_types, n_expert, expert_params, tower_params_list):
-        super().__init__()
-        self.features = features
-        self.task_types = task_types
-        self.n_task = len(task_types)
-        self.n_expert = n_expert
-        self.embedding = EmbeddingLayer(features)
-        self.input_dims = sum([fea.embed_dim for fea in features])
-        self.experts = nn.ModuleList(
-            MLP(self.input_dims, output_layer=False, **expert_params) for i in range(self.n_expert))
-        self.gates = nn.ModuleList(
-            MLP(self.input_dims, output_layer=False, **{
-                "dims": [self.n_expert],
-                "activation": "softmax"
-            }) for i in range(self.n_task))  #n_gate = n_task
-        self.towers = nn.ModuleList(MLP(expert_params["dims"][-1], **tower_params_list[i]) for i in range(self.n_task))
-        self.predict_layers = nn.ModuleList(PredictionLayer(task_type) for task_type in task_types)
-
-    def forward(self, x):
-        embed_x = self.embedding(x, self.features, squeeze_dim=True)  #[batch_size, input_dims]
-        expert_outs = [expert(embed_x).unsqueeze(1) for expert in self.experts
-                      ]  #expert_out[i]: [batch_size, 1, expert_dims[-1]]
-        expert_outs = torch.cat(expert_outs, dim=1)  #[batch_size, n_expert, expert_dims[-1]]
-        gate_outs = [gate(embed_x).unsqueeze(-1) for gate in self.gates]  #gate_out[i]: [batch_size, n_expert, 1]
-
-        ys = []
-        for gate_out, tower, predict_layer in zip(gate_outs, self.towers, self.predict_layers):
-            expert_weight = torch.mul(gate_out, expert_outs)  #[batch_size, n_expert, expert_dims[-1]]
-            expert_pooling = torch.sum(expert_weight, dim=1)  #[batch_size, expert_dims[-1]]
-            tower_out = tower(expert_pooling)  #[batch_size, 1]
-            y = predict_layer(tower_out)  #logit -> proba
-            ys.append(y)
-        return torch.cat(ys, dim=1)
+"""
+Date: create on 04/05/2022
+References:
+    paper: (KDD'2018) Modeling Task Relationships in Multi-task Learning with Multi-gate Mixture-of-Experts
+    url: https://dl.acm.org/doi/pdf/10.1145/3219819.3220007
+Authors: Mincai Lai, laimincai@shanghaitech.edu.cn
+"""
+
+import torch
+import torch.nn as nn
+
+from ...basic.layers import MLP, EmbeddingLayer, PredictionLayer
+
+
+class MMOE(nn.Module):
+    """Multi-gate Mixture-of-Experts model.
+
+    Args:
+        features (list): the list of `Feature Class`, training by the expert and tower module.
+        task_types (list): types of tasks, only support `["classfication", "regression"]`.
+        n_expert (int): the number of expert net.
+        expert_params (dict): the params of all the 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, task_types, n_expert, expert_params, tower_params_list):
+        super().__init__()
+        self.features = features
+        self.task_types = task_types
+        self.n_task = len(task_types)
+        self.n_expert = n_expert
+        self.embedding = EmbeddingLayer(features)
+        self.input_dims = sum([fea.embed_dim for fea in features])
+        self.experts = nn.ModuleList(MLP(self.input_dims, output_layer=False, **expert_params) for i in range(self.n_expert))
+        self.gates = nn.ModuleList(MLP(self.input_dims, output_layer=False, **{"dims": [self.n_expert], "activation": "softmax"}) for i in range(self.n_task))  # n_gate = n_task
+        self.towers = nn.ModuleList(MLP(expert_params["dims"][-1], **tower_params_list[i]) for i in range(self.n_task))
+        self.predict_layers = nn.ModuleList(PredictionLayer(task_type) for task_type in task_types)
+
+    def forward(self, x):
+        # [batch_size, input_dims]
+        embed_x = self.embedding(x, self.features, squeeze_dim=True)
+        # expert_out[i]: [batch_size, 1, expert_dims[-1]]
+        expert_outs = [expert(embed_x).unsqueeze(1) for expert in self.experts]
+        # [batch_size, n_expert, expert_dims[-1]]
+        expert_outs = torch.cat(expert_outs, dim=1)
+        # gate_out[i]: [batch_size, n_expert, 1]
+        gate_outs = [gate(embed_x).unsqueeze(-1) for gate in self.gates]
+
+        ys = []
+        for gate_out, tower, predict_layer in zip(gate_outs, self.towers, self.predict_layers):
+            # [batch_size, n_expert, expert_dims[-1]]
+            expert_weight = torch.mul(gate_out, expert_outs)
+            # [batch_size, expert_dims[-1]]
+            expert_pooling = torch.sum(expert_weight, dim=1)
+            tower_out = tower(expert_pooling)  # [batch_size, 1]
+            y = predict_layer(tower_out)  # logit -> proba
+            ys.append(y)
+        return torch.cat(ys, dim=1)