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

torch_rechub/models/matching/dssm.py

@@ -0,0 +1,72 @@
+"""
+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
+        # [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)  # L2 normalize
+        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)
+        item_embedding = F.normalize(item_embedding, p=2, dim=1)
+        return item_embedding

torch_rechub/models/matching/dssm_facebook.py

@@ -0,0 +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
+        # [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

@@ -0,0 +1,87 @@
+"""
+Date: create on 12/19/2024
+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
+
+
+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.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
+
+    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
+        # [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
+        # [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

torch_rechub/models/matching/gru4rec.py

@@ -0,0 +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
+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

@@ -0,0 +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
+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

torch_rechub/models/matching/narm.py

@@ -0,0 +1,82 @@
+"""
+Date: created on 06/09/2022
+References:
+    paper: Neural Attentive Session-based Recommendation
+    url: http://arxiv.org/abs/1711.04725
+    official Theano implementation: https://github.com/lijingsdu/sessionRec_NARM
+    another Pytorch implementation: https://github.com/Wang-Shuo/Neural-Attentive-Session-Based-Recommendation-PyTorch
+Authors: Bo Kang, klinux@live.com
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.utils.rnn as rnn_utils
+from torch import sigmoid
+from torch.nn import GRU, Dropout, Embedding, Parameter
+
+
+class NARM(nn.Module):
+
+    def __init__(self, item_history_feature, hidden_dim, emb_dropout_p, session_rep_dropout_p):
+        super(NARM, self).__init__()
+
+        # item embedding layer
+        self.item_history_feature = item_history_feature
+        self.item_emb = Embedding(item_history_feature.vocab_size, item_history_feature.embed_dim, padding_idx=0)
+
+        # embedding dropout layer
+        self.emb_dropout = Dropout(emb_dropout_p)
+
+        # gru unit
+        self.gru = GRU(input_size=item_history_feature.embed_dim, hidden_size=hidden_dim)
+
+        # attention projection matrices
+        self.a_1, self.a_2 = Parameter(torch.randn(hidden_dim, hidden_dim)), Parameter(torch.randn(hidden_dim, hidden_dim))
+
+        # attention context vector
+        self.v = Parameter(torch.randn(hidden_dim, 1))
+
+        # session representation dropout layer
+        self.session_rep_dropout = Dropout(session_rep_dropout_p)
+
+        # bilinear projection matrix
+        self.b = Parameter(torch.randn(item_history_feature.embed_dim, hidden_dim * 2))
+
+    def forward(self, input_dict):
+        # Eq. 1-4, index item embeddings and pass through gru
+        # # Fetch the embeddings for items in the session
+        input = input_dict[self.item_history_feature.name]
+        value_mask = (input != 0)
+        value_counts = value_mask.sum(dim=1, keepdim=False).to("cpu").detach()
+        embs = rnn_utils.pack_padded_sequence(self.emb_dropout(self.item_emb(input)), value_counts, batch_first=True, enforce_sorted=False)
+
+        # # compute hidden states at each time step
+        h, h_t = self.gru(embs)
+        h_t = h_t.permute(1, 0, 2)
+        h, _ = rnn_utils.pad_packed_sequence(h, batch_first=True)
+
+        # Eq. 5, set last hidden state of gru as the output of the global
+        # encoder
+        c_g = h_t.squeeze(1)
+
+        # Eq. 8, compute similarity between final hidden state and previous
+        # hidden states
+        q = sigmoid(h_t @ self.a_1.T + h @ self.a_2.T) @ self.v
+
+        # Eq. 7, compute attention
+        alpha = torch.exp(q) * value_mask.unsqueeze(-1)
+        alpha /= alpha.sum(dim=1, keepdim=True)
+
+        # Eq. 6, compute the output of the local encoder
+        c_l = (alpha * h).sum(1)
+
+        # Eq. 9, compute session representation by concatenating user
+        # sequential behavior (global) and main purpose in the current session
+        # (local)
+        c = self.session_rep_dropout(torch.hstack((c_g, c_l)))
+
+        # Eq. 10, compute bilinear similarity between current session and each
+        # candidate items
+        s = c @ self.b.T @ self.item_emb.weight.T
+
+        return s

torch_rechub/models/matching/sasrec.py

@@ -0,0 +1,143 @@
+"""
+Date: create on 2022/5/8, update on 2022/5/8
+References:
+    paper: (ICDM'2018) Self-attentive sequential recommendation
+    url: https://arxiv.org/pdf/1808.09781.pdf
+    code: https://github.com/kang205/SASRec
+Authors: Yuchen Wang, 615922749@qq.com
+"""
+import numpy as np
+import torch
+import torch.nn as nn
+
+from torch_rechub.basic.features import DenseFeature, SequenceFeature, SparseFeature
+from torch_rechub.basic.layers import MLP, EmbeddingLayer
+
+
+class SASRec(torch.nn.Module):
+    """SASRec: Self-Attentive Sequential Recommendation
+    Args:
+        features (list): the list of `Feature Class`. In sasrec, the features list needs to have three elements in order: user historical behavior sequence features, positive sample sequence, and negative sample sequence.
+        max_len: The length of the sequence feature.
+        num_blocks: The number of stacks of attention modules.
+        num_heads: The number of heads in MultiheadAttention.
+
+    """
+
+    def __init__(
+        self,
+        features,
+        max_len=50,
+        dropout_rate=0.5,
+        num_blocks=2,
+        num_heads=1,
+    ):
+        super(SASRec, self).__init__()
+
+        self.features = features
+
+        self.item_num = self.features[0].vocab_size
+        self.embed_dim = self.features[0].embed_dim
+
+        self.item_emb = EmbeddingLayer(self.features)
+        self.position_emb = torch.nn.Embedding(max_len, self.embed_dim)
+        self.emb_dropout = torch.nn.Dropout(p=dropout_rate)
+
+        self.attention_layernorms = torch.nn.ModuleList()
+        self.attention_layers = torch.nn.ModuleList()
+        self.forward_layernorms = torch.nn.ModuleList()
+        self.forward_layers = torch.nn.ModuleList()
+        self.last_layernorm = torch.nn.LayerNorm(self.embed_dim, eps=1e-8)
+
+        for _ in range(num_blocks):
+            new_attn_layernorm = torch.nn.LayerNorm(self.embed_dim, eps=1e-8)
+            self.attention_layernorms.append(new_attn_layernorm)
+
+            new_attn_layer = torch.nn.MultiheadAttention(self.embed_dim, num_heads, dropout_rate)
+            self.attention_layers.append(new_attn_layer)
+
+            new_fwd_layernorm = torch.nn.LayerNorm(self.embed_dim, eps=1e-8)
+            self.forward_layernorms.append(new_fwd_layernorm)
+
+            new_fwd_layer = PointWiseFeedForward(self.embed_dim, dropout_rate)
+            self.forward_layers.append(new_fwd_layer)
+
+    def seq_forward(self, x, embed_x_feature):
+        x = x['seq']
+
+        embed_x_feature *= self.features[0].embed_dim**0.5
+        embed_x_feature = embed_x_feature.squeeze()  # (bacth_size, max_len, embed_dim)
+
+        positions = np.tile(np.array(range(x.shape[1])), [x.shape[0], 1])
+
+        embed_x_feature += self.position_emb(torch.LongTensor(positions))
+        embed_x_feature = self.emb_dropout(embed_x_feature)
+
+        timeline_mask = torch.BoolTensor(x == 0)
+        embed_x_feature *= ~timeline_mask.unsqueeze(-1)
+
+        attention_mask = ~torch.tril(torch.ones((embed_x_feature.shape[1], embed_x_feature.shape[1]), dtype=torch.bool))
+
+        for i in range(len(self.attention_layers)):
+            embed_x_feature = torch.transpose(embed_x_feature, 0, 1)
+            Q = self.attention_layernorms[i](embed_x_feature)
+            mha_outputs, _ = self.attention_layers[i](Q, embed_x_feature, embed_x_feature, attn_mask=attention_mask)
+
+            embed_x_feature = Q + mha_outputs
+            embed_x_feature = torch.transpose(embed_x_feature, 0, 1)
+
+            embed_x_feature = self.forward_layernorms[i](embed_x_feature)
+            embed_x_feature = self.forward_layers[i](embed_x_feature)
+            embed_x_feature *= ~timeline_mask.unsqueeze(-1)
+
+        seq_output = self.last_layernorm(embed_x_feature)
+
+        return seq_output
+
+    def forward(self, x):
+        # (batch_size, 3, max_len, embed_dim)
+        embedding = self.item_emb(x, self.features)
+        # (batch_size, max_len, embed_dim)
+        seq_embed, pos_embed, neg_embed = embedding[:, 0], embedding[:, 1], embedding[:, 2]
+
+        # (batch_size, max_len, embed_dim)
+        seq_output = self.seq_forward(x, seq_embed)
+
+        pos_logits = (seq_output * pos_embed).sum(dim=-1)
+        neg_logits = (seq_output * neg_embed).sum(dim=-1)  # (batch_size, max_len)
+
+        return pos_logits, neg_logits
+
+
+class PointWiseFeedForward(torch.nn.Module):
+
+    def __init__(self, hidden_units, dropout_rate):
+        super(PointWiseFeedForward, self).__init__()
+
+        self.conv1 = torch.nn.Conv1d(hidden_units, hidden_units, kernel_size=1)
+        self.dropout1 = torch.nn.Dropout(p=dropout_rate)
+        self.relu = torch.nn.ReLU()
+        self.conv2 = torch.nn.Conv1d(hidden_units, hidden_units, kernel_size=1)
+        self.dropout2 = torch.nn.Dropout(p=dropout_rate)
+
+    def forward(self, inputs):
+        outputs = self.dropout2(self.conv2(self.relu(self.dropout1(self.conv1(inputs.transpose(-1, -2))))))
+        outputs = outputs.transpose(-1, -2)
+        outputs += inputs
+        return outputs
+
+
+if __name__ == '__main__':
+    seq = SequenceFeature('seq', vocab_size=17, embed_dim=7, pooling='concat')
+    pos = SequenceFeature('pos', vocab_size=17, embed_dim=7, pooling='concat', shared_with='seq')
+    neg = SequenceFeature('neg', vocab_size=17, embed_dim=7, pooling='concat', shared_with='seq')
+
+    seq = [seq, pos, neg]
+
+    hist_seq = torch.tensor([[1, 2, 3, 4], [2, 3, 7, 8]])
+    pos_seq = hist_seq
+    neg_seq = hist_seq
+
+    x = {'seq': hist_seq, 'pos': pos_seq, 'neg': neg_seq}
+    model = SASRec(features=seq)
+    print('out', model(x))