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

torch_rechub/models/matching/dssm_facebook.py

@@ -1,79 +1,77 @@
-"""
-Date: create on 24/05/2022
-References: 
-    paper: (KDD'2020) Embedding-based Retrieval in Facebook Search
-    url: https://arxiv.org/abs/2006.11632
-Authors: Mincai Lai, laimincai@shanghaitech.edu.cn
-"""
-
-import torch
-import torch.nn.functional as F
-from ...basic.layers import MLP, EmbeddingLayer
-
-
-class FaceBookDSSM(torch.nn.Module):
-    """Embedding-based Retrieval in Facebook Search
-    It's a DSSM match model trained by hinge loss on pair-wise samples.
-
-    Args:
-        user_features (list[Feature Class]): training by the user tower module.
-        pos_item_features (list[Feature Class]): negative sample features, training by the item tower module.
-        neg_item_features (list[Feature Class]): positive sample features, 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,
-                 pos_item_features,
-                 neg_item_features,
-                 user_params,
-                 item_params,
-                 temperature=1.0):
-        super().__init__()
-        self.user_features = user_features
-        self.pos_item_features = pos_item_features
-        self.neg_item_features = neg_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 pos_item_features])
-
-        self.embedding = EmbeddingLayer(user_features + pos_item_features + neg_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)
-        pos_item_embedding, neg_item_embedding = self.item_tower(x)
-        if self.mode == "user":
-            return user_embedding
-        if self.mode == "item":
-            return pos_item_embedding
-
-        # calculate cosine score
-        pos_score = torch.mul(user_embedding, pos_item_embedding).sum(dim=1)
-        neg_score = torch.mul(user_embedding, neg_item_embedding).sum(dim=1)
-
-        return pos_score, neg_score
-
-    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)
-        return user_embedding
-
-    def item_tower(self, x):
-        if self.mode == "user":
-            return None, None
-        input_item_pos = self.embedding(x, self.pos_item_features, squeeze_dim=True)
-        if self.mode == "item":  #inference embedding mode, the zeros is just for placefolder
-            return self.item_mlp(input_item_pos), None
-        input_item_neg = self.embedding(x, self.neg_item_features, squeeze_dim=True)
-        pos_embedding, neg_embedding = self.item_mlp(input_item_pos), self.item_mlp(input_item_neg)
-        pos_embedding = F.normalize(pos_embedding, p=2, dim=1)
-        neg_embedding = F.normalize(neg_embedding, p=2, dim=1)
-        return pos_embedding, neg_embedding
+"""
+Date: create on 24/05/2022
+References:
+    paper: (KDD'2020) Embedding-based Retrieval in Facebook Search
+    url: https://arxiv.org/abs/2006.11632
+Authors: Mincai Lai, laimincai@shanghaitech.edu.cn
+"""
+
+import torch
+import torch.nn.functional as F
+
+from ...basic.layers import MLP, EmbeddingLayer
+
+
+class FaceBookDSSM(torch.nn.Module):
+    """Embedding-based Retrieval in Facebook Search
+    It's a DSSM match model trained by hinge loss on pair-wise samples.
+
+    Args:
+        user_features (list[Feature Class]): training by the user tower module.
+        pos_item_features (list[Feature Class]): negative sample features, training by the item tower module.
+        neg_item_features (list[Feature Class]): positive sample features, 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, pos_item_features, neg_item_features, user_params, item_params, temperature=1.0):
+        super().__init__()
+        self.user_features = user_features
+        self.pos_item_features = pos_item_features
+        self.neg_item_features = neg_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 pos_item_features])
+
+        self.embedding = EmbeddingLayer(user_features + pos_item_features + neg_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)
+        pos_item_embedding, neg_item_embedding = self.item_tower(x)
+        if self.mode == "user":
+            return user_embedding
+        if self.mode == "item":
+            return pos_item_embedding
+
+
+# calculate cosine score
+        pos_score = torch.mul(user_embedding, pos_item_embedding).sum(dim=1)
+        neg_score = torch.mul(user_embedding, neg_item_embedding).sum(dim=1)
+
+        return pos_score, neg_score
+
+    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)
+        user_embedding = F.normalize(user_embedding, p=2, dim=1)
+        return user_embedding
+
+    def item_tower(self, x):
+        if self.mode == "user":
+            return None, None
+        input_item_pos = self.embedding(x, self.pos_item_features, squeeze_dim=True)
+        if self.mode == "item":  # inference embedding mode, the zeros is just for placefolder
+            return self.item_mlp(input_item_pos), None
+        input_item_neg = self.embedding(x, self.neg_item_features, squeeze_dim=True)
+        pos_embedding, neg_embedding = self.item_mlp(input_item_pos), self.item_mlp(input_item_neg)
+        pos_embedding = F.normalize(pos_embedding, p=2, dim=1)
+        neg_embedding = F.normalize(neg_embedding, p=2, dim=1)
+        return pos_embedding, neg_embedding

torch_rechub/models/matching/dssm_senet.py

@@ -1,14 +1,15 @@
 """
 Date: create on 12/19/2024
-References: 
+References:
     url: https://zhuanlan.zhihu.com/p/358779957
 Authors: @1985312383
 """
 
 import torch
 import torch.nn.functional as F
+
+from ...basic.features import SequenceFeature, SparseFeature
 from ...basic.layers import MLP, EmbeddingLayer, SENETLayer
-from ...basic.features import SparseFeature, SequenceFeature
 
 
 class DSSM(torch.nn.Module):
@@ -33,8 +34,8 @@ class DSSM(torch.nn.Module):
         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.user_num_features = len([fea.embed_dim for fea in self.user_features if isinstance(fea, SparseFeature) or isinstance(fea, SequenceFeature) and fea.shared_with == None])
-        self.item_num_features = len([fea.embed_dim for fea in self.item_features if isinstance(fea, SparseFeature) or isinstance(fea, SequenceFeature) and fea.shared_with == None])
+        self.user_num_features = len([fea.embed_dim for fea in self.user_features if isinstance(fea, SparseFeature) or isinstance(fea, SequenceFeature) and fea.shared_with is None])
+        self.item_num_features = len([fea.embed_dim for fea in self.item_features if isinstance(fea, SparseFeature) or isinstance(fea, SequenceFeature) and fea.shared_with is None])
         self.user_senet = SENETLayer(self.user_num_features)
         self.item_senet = SENETLayer(self.item_num_features)
         self.mode = None
@@ -47,7 +48,8 @@ class DSSM(torch.nn.Module):
         if self.mode == "item":
             return item_embedding
 
-        # calculate cosine score
+
+# calculate cosine score
         y = torch.mul(user_embedding, item_embedding).sum(dim=1)
         y = y / self.temperature
         return torch.sigmoid(y)
@@ -55,21 +57,31 @@ class DSSM(torch.nn.Module):
     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 * embed_dim]
-        input_user = input_user.view(input_user.size(0), self.user_num_features, -1)  #[batch_size, num_features, embed_dim]
-        input_user = self.user_senet(input_user)  #[batch_size, num_features, embed_dim]
-        input_user = input_user.view(input_user.size(0), -1)   #[batch_size, num_features * embed_dim]
-        user_embedding = self.user_mlp(input_user)  #[batch_size, user_params["dims"][-1]]
+        # [batch_size, num_features * embed_dim]
+        input_user = self.embedding(x, self.user_features, squeeze_dim=True)
+        # [batch_size, num_features, embed_dim]
+        input_user = input_user.view(input_user.size(0), self.user_num_features, -1)
+        # [batch_size, num_features, embed_dim]
+        input_user = self.user_senet(input_user)
+        # [batch_size, num_features * embed_dim]
+        input_user = input_user.view(input_user.size(0), -1)
+        # [batch_size, user_params["dims"][-1]]
+        user_embedding = self.user_mlp(input_user)
         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]
-        input_item = input_item.view(input_item.size(0), self.item_num_features, -1)  #[batch_size, num_features, embed_dim]
-        input_item = self.item_senet(input_item)  #[batch_size, num_features, embed_dim]
-        input_item = input_item.view(input_item.size(0), -1)   #[batch_size, num_features * embed_dim]
-        item_embedding = self.item_mlp(input_item)  #[batch_size, item_params["dims"][-1]]
+        # [batch_size, num_features * embed_dim]
+        input_item = self.embedding(x, self.item_features, squeeze_dim=True)
+        # [batch_size, num_features, embed_dim]
+        input_item = input_item.view(input_item.size(0), self.item_num_features, -1)
+        # [batch_size, num_features, embed_dim]
+        input_item = self.item_senet(input_item)
+        # [batch_size, num_features * embed_dim]
+        input_item = input_item.view(input_item.size(0), -1)
+        # [batch_size, item_params["dims"][-1]]
+        item_embedding = self.item_mlp(input_item)
         item_embedding = F.normalize(item_embedding, p=2, dim=1)
-        return item_embedding
+        return item_embedding

torch_rechub/models/matching/gru4rec.py

@@ -1,87 +1,85 @@
-"""
-Date: create on 03/06/2022
-References: 
-    paper: SESSION-BASED RECOMMENDATIONS WITH RECURRENT NEURAL NETWORKS
-    url: http://arxiv.org/abs/1511.06939
-Authors: Kai Wang, 306178200@qq.com
-"""
-
-import torch
-
-from ...basic.layers import MLP, EmbeddingLayer
-from torch import nn
-import torch.nn.functional as F
-
-
-class GRU4Rec(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.
-        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.
-        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, history_features, item_features, neg_item_feature, user_params, temperature=1.0):
-        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.user_dims = sum([fea.embed_dim for fea in user_features+history_features])
-
-        self.embedding = EmbeddingLayer(user_features + item_features + history_features)
-        self.gru = nn.GRU(input_size = history_features[0].embed_dim,
-                          hidden_size = history_features[0].embed_dim,
-                          num_layers = user_params.get('num_layers',2),
-                          batch_first = True,
-                          bias = False)
-        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
-
-        y = torch.mul(user_embedding, 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)  #[batch_size, num_features*deep_dims]
-
-        history_emb = self.embedding(x, self.history_features).squeeze(1)
-        _, history_emb = self.gru(history_emb)
-        history_emb = history_emb[-1]
-
-        input_user = torch.cat([input_user,history_emb],dim=-1)
-
-        user_embedding = self.user_mlp(input_user).unsqueeze(1)  #[batch_size, 1, embed_dim]
-        user_embedding = F.normalize(user_embedding, p=2, dim=-1)  # L2 normalize
-        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=-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]
+"""
+Date: create on 03/06/2022
+References:
+    paper: SESSION-BASED RECOMMENDATIONS WITH RECURRENT NEURAL NETWORKS
+    url: http://arxiv.org/abs/1511.06939
+Authors: Kai Wang, 306178200@qq.com
+"""
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from ...basic.layers import MLP, EmbeddingLayer
+
+
+class GRU4Rec(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.
+        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.
+        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, history_features, item_features, neg_item_feature, user_params, temperature=1.0):
+        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.user_dims = sum([fea.embed_dim for fea in user_features + history_features])
+
+        self.embedding = EmbeddingLayer(user_features + item_features + history_features)
+        self.gru = nn.GRU(input_size=history_features[0].embed_dim, hidden_size=history_features[0].embed_dim, num_layers=user_params.get('num_layers', 2), batch_first=True, bias=False)
+        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
+
+        y = torch.mul(user_embedding, item_embedding).sum(dim=1)
+
+        return y
+
+    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)
+
+        history_emb = self.embedding(x, self.history_features).squeeze(1)
+        _, history_emb = self.gru(history_emb)
+        history_emb = history_emb[-1]
+
+        input_user = torch.cat([input_user, history_emb], dim=-1)
+
+        user_embedding = self.user_mlp(input_user).unsqueeze(1)  # [batch_size, 1, embed_dim]
+        user_embedding = F.normalize(user_embedding, p=2, dim=-1)  # L2 normalize
+        if self.mode == "user":
+            # inference embedding mode -> [batch_size, embed_dim]
+            return user_embedding.squeeze(1)
+        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
+        # [batch_size, 1+n_neg_items, embed_dim]
+        return torch.cat((pos_embedding, neg_embeddings), dim=1)

torch_rechub/models/matching/mind.py

@@ -1,101 +1,103 @@
-"""
-Date: create on 08/06/2022
-References:
-    paper: Multi-Interest Network with Dynamic Routing
-    url: https://arxiv.org/pdf/1904.08030v1
-    code: https://github.com/ShiningCosmos/pytorch_ComiRec/blob/main/MIND.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 MIND(torch.nn.Module):
-    """The match model mentioned in `Multi-Interest Network with Dynamic Routing` 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=0,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
+"""
+Date: create on 08/06/2022
+References:
+    paper: Multi-Interest Network with Dynamic Routing
+    url: https://arxiv.org/pdf/1904.08030v1
+    code: https://github.com/ShiningCosmos/pytorch_ComiRec/blob/main/MIND.py
+Authors: Kai Wang, 306178200@qq.com
+"""
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from ...basic.layers import MLP, CapsuleNetwork, EmbeddingLayer, MultiInterestSA
+
+
+class MIND(torch.nn.Module):
+    """The match model mentioned in `Multi-Interest Network with Dynamic Routing` 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=0, 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":
+            # inference embedding mode -> [batch_size, interest_num, embed_dim]
+            return user_embedding
+        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
+        # [batch_size, 1+n_neg_items, embed_dim]
+        return torch.cat((pos_embedding, neg_embeddings), dim=1)
+
+    def gen_mask(self, x):
+        his_list = x[self.history_features[0].name]
+        mask = (his_list > 0).long()
+        return mask