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

torch_rechub/models/multi_task/mmoe.py

@@ -1,52 +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}, keep `{"output_layer":False}`.
-        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, **{**expert_params, **{"output_layer": False}}) for i in range(self.n_expert))
-        self.gates = nn.ModuleList(MLP(self.input_dims, **{"dims": [self.n_expert], "activation": "softmax", "output_layer": False}) 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):
+        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)

torch_rechub/models/multi_task/ple.py

@@ -1,104 +1,130 @@
-"""
-Date: create on 05/05/2022
-References: 
-    paper: (RecSys'2020) Progressive Layered Extraction (PLE): A Novel Multi-Task Learning (MTL) Model for Personalized Recommendations
-    url: https://dl.acm.org/doi/abs/10.1145/3383313.3412236
-Authors: Mincai Lai, laimincai@shanghaitech.edu.cn
-"""
-
-import torch
-import torch.nn as nn
-
-from ...basic.layers import MLP, EmbeddingLayer, PredictionLayer
-
-
-class PLE(nn.Module):
-    """Progressive Layered Extraction 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_level (int): the  number of CGC layer.
-        n_expert_specific (int): the number of task-specific expert net.
-        n_expert_shared (int): the number of task-shared expert net.
-        expert_params (dict): the params of all the expert module, keys include:`{"dims":list, "activation":str, "dropout":float}, keep `{"output_layer":False}`.
-        tower_params_list (list): the list of tower params dict, the keys same as expert_params.
-    """
-
-    def __init__(self, features, task_types, n_level, n_expert_specific, n_expert_shared, expert_params={"dims": [32, 16], "output_layer": False}, tower_params_list=[{"dims": [32, 16]}, {"dims": [32, 16]}]):
-        super().__init__()
-        self.features = features
-        self.n_task = len(task_types)
-        self.task_types = task_types
-        self.n_level = n_level
-        self.input_dims = sum([fea.embed_dim for fea in features])
-        self.embedding = EmbeddingLayer(features)
-        self.cgc_layers = nn.ModuleList(CGC(i + 1, n_level, self.n_task, n_expert_specific, n_expert_shared, self.input_dims, expert_params) for i in range(n_level))
-        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]
-        ple_inputs = [embed_x] * (self.n_task + 1)
-        ple_outs = []
-        for i in range(self.n_level):
-            ple_outs = self.cgc_layers[i](ple_inputs)  #ple_outs[i]: [batch_size, expert_dims[-1]]
-            ple_inputs = ple_outs
-        #predict
-        ys = []
-        for ple_out, tower, predict_layer in zip(ple_outs, self.towers, self.predict_layers):
-            tower_out = tower(ple_out)  #[batch_size, 1]
-            y = predict_layer(tower_out)  #logit -> proba
-            ys.append(y)
-        return torch.cat(ys, dim=1)
-
-
-class CGC(nn.Module):
-    """Customized Gate Control (CGC) Model mentioned in PLE paper.
-
-    Args:
-        cur_level (int): the current level of CGC in PLE.
-        n_level (int): the  number of CGC layer.
-        n_task (int): the number of tasks.
-        n_expert_specific (int): the number of task-specific expert net.
-        n_expert_shared (int): the number of task-shared expert net.
-        input_dims (int): the input dims of the xpert module in current CGC layer.
-        expert_params (dict): the params of all the expert module, keys include:`{"dims":list, "activation":str, "dropout":float}, keep `{"output_layer":False}`.
-    """
-
-    def __init__(self, cur_level, n_level, n_task, n_expert_specific, n_expert_shared, input_dims, expert_params={"dims": [32, 16], "output_layer": False}):
-        super().__init__()
-        self.cur_level = cur_level  # the CGC level of PLE
-        self.n_level = n_level
-        self.n_task = n_task
-        self.n_expert_specific = n_expert_specific
-        self.n_expert_shared = n_expert_shared
-        self.n_expert_all = n_expert_specific * self.n_task + n_expert_shared
-        input_dims = input_dims if cur_level == 1 else expert_params["dims"][-1]  #the first layer expert dim is the input data dim other expert dim
-        self.experts_specific = nn.ModuleList(MLP(input_dims, **expert_params) for _ in range(self.n_task * self.n_expert_specific))
-        self.experts_shared = nn.ModuleList(MLP(input_dims, **expert_params) for _ in range(self.n_expert_shared))
-        self.gates_specific = nn.ModuleList(MLP(input_dims, **{"dims": [self.n_expert_specific + self.n_expert_shared], "activation": "softmax", "output_layer": False}) for _ in range(self.n_task))  #n_gate_specific = n_task
-        if cur_level < n_level:
-            self.gate_shared = MLP(input_dims, **{"dims": [self.n_expert_all], "activation": "softmax", "output_layer": False})  #n_gate_specific = n_task
-
-    def forward(self, x_list):
-        expert_specific_outs = []  #expert_out[i]: [batch_size, 1, expert_dims[-1]]
-        for i in range(self.n_task):
-            expert_specific_outs.extend([expert(x_list[i]).unsqueeze(1) for expert in self.experts_specific[i * self.n_expert_specific:(i + 1) * self.n_expert_specific]])
-        expert_shared_outs = [expert(x_list[-1]).unsqueeze(1) for expert in self.experts_shared]  #x_list[-1]: the input for shared experts
-        gate_specific_outs = [gate(x_list[i]).unsqueeze(-1) for i, gate in enumerate(self.gates_specific)]  #gate_out[i]: [batch_size, n_expert_specific+n_expert_shared, 1]
-        cgc_outs = []
-        for i, gate_out in enumerate(gate_specific_outs):
-            cur_expert_list = expert_specific_outs[i * self.n_expert_specific:(i + 1) * self.n_expert_specific] + expert_shared_outs
-            expert_concat = torch.cat(cur_expert_list, dim=1)  #[batch_size, n_expert_specific+n_expert_shared, expert_dims[-1]]
-            expert_weight = torch.mul(gate_out, expert_concat)  #[batch_size, n_expert_specific+n_expert_shared, expert_dims[-1]]
-            expert_pooling = torch.sum(expert_weight, dim=1)  #[batch_size, expert_dims[-1]]
-            cgc_outs.append(expert_pooling)  #length: n_task
-        if self.cur_level < self.n_level:  #not the last layer
-            gate_shared_out = self.gate_shared(x_list[-1]).unsqueeze(-1)  #[batch_size, n_expert_all, 1]
-            expert_concat = torch.cat(expert_specific_outs + expert_shared_outs, dim=1)  #[batch_size, n_expert_all, expert_dims[-1]]
-            expert_weight = torch.mul(gate_shared_out, expert_concat)  #[batch_size, n_expert_all, expert_dims[-1]]
-            expert_pooling = torch.sum(expert_weight, dim=1)  #[batch_size, expert_dims[-1]]
-            cgc_outs.append(expert_pooling)  #length: n_task+1
-
-        return cgc_outs
+"""
+Date: create on 05/05/2022
+References: 
+    paper: (RecSys'2020) Progressive Layered Extraction (PLE): A Novel Multi-Task Learning (MTL) Model for Personalized Recommendations
+    url: https://dl.acm.org/doi/abs/10.1145/3383313.3412236
+Authors: Mincai Lai, laimincai@shanghaitech.edu.cn
+"""
+
+import torch
+import torch.nn as nn
+
+from ...basic.layers import MLP, EmbeddingLayer, PredictionLayer
+
+
+class PLE(nn.Module):
+    """Progressive Layered Extraction 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_level (int): the  number of CGC layer.
+        n_expert_specific (int): the number of task-specific expert net.
+        n_expert_shared (int): the number of task-shared 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_level, n_expert_specific, n_expert_shared, expert_params,
+                 tower_params_list):
+        super().__init__()
+        self.features = features
+        self.n_task = len(task_types)
+        self.task_types = task_types
+        self.n_level = n_level
+        self.input_dims = sum([fea.embed_dim for fea in features])
+        self.embedding = EmbeddingLayer(features)
+        self.cgc_layers = nn.ModuleList(
+            CGC(i + 1, n_level, self.n_task, n_expert_specific, n_expert_shared, self.input_dims, expert_params)
+            for i in range(n_level))
+        self.towers = nn.ModuleList(
+            MLP(expert_params["dims"][-1], output_layer=False, **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]
+        ple_inputs = [embed_x] * (self.n_task + 1)
+        ple_outs = []
+        for i in range(self.n_level):
+            ple_outs = self.cgc_layers[i](ple_inputs)  #ple_outs[i]: [batch_size, expert_dims[-1]]
+            ple_inputs = ple_outs
+        #predict
+        ys = []
+        for ple_out, tower, predict_layer in zip(ple_outs, self.towers, self.predict_layers):
+            tower_out = tower(ple_out)  #[batch_size, 1]
+            y = predict_layer(tower_out)  #logit -> proba
+            ys.append(y)
+        return torch.cat(ys, dim=1)
+
+
+class CGC(nn.Module):
+    """Customized Gate Control (CGC) Model mentioned in PLE paper.
+
+    Args:
+        cur_level (int): the current level of CGC in PLE.
+        n_level (int): the  number of CGC layer.
+        n_task (int): the number of tasks.
+        n_expert_specific (int): the number of task-specific expert net.
+        n_expert_shared (int): the number of task-shared expert net.
+        input_dims (int): the input dims of the xpert module in current CGC layer.
+        expert_params (dict): the params of all the expert module, keys include:`{"dims":list, "activation":str, "dropout":float}.
+    """
+
+    def __init__(self, cur_level, n_level, n_task, n_expert_specific, n_expert_shared, input_dims, expert_params):
+        super().__init__()
+        self.cur_level = cur_level  # the CGC level of PLE
+        self.n_level = n_level
+        self.n_task = n_task
+        self.n_expert_specific = n_expert_specific
+        self.n_expert_shared = n_expert_shared
+        self.n_expert_all = n_expert_specific * self.n_task + n_expert_shared
+        input_dims = input_dims if cur_level == 1 else expert_params["dims"][
+            -1]  #the first layer expert dim is the input data dim other expert dim
+        self.experts_specific = nn.ModuleList(
+            MLP(input_dims, output_layer=False, **expert_params) for _ in range(self.n_task * self.n_expert_specific))
+        self.experts_shared = nn.ModuleList(
+            MLP(input_dims, output_layer=False, **expert_params) for _ in range(self.n_expert_shared))
+        self.gates_specific = nn.ModuleList(
+            MLP(
+                input_dims, **{
+                    "dims": [self.n_expert_specific + self.n_expert_shared],
+                    "activation": "softmax",
+                    "output_layer": False
+                }) for _ in range(self.n_task))  #n_gate_specific = n_task
+        if cur_level < n_level:
+            self.gate_shared = MLP(input_dims, **{
+                "dims": [self.n_expert_all],
+                "activation": "softmax",
+                "output_layer": False
+            })  #n_gate_specific = n_task
+
+    def forward(self, x_list):
+        expert_specific_outs = []  #expert_out[i]: [batch_size, 1, expert_dims[-1]]
+        for i in range(self.n_task):
+            expert_specific_outs.extend([
+                expert(x_list[i]).unsqueeze(1)
+                for expert in self.experts_specific[i * self.n_expert_specific:(i + 1) * self.n_expert_specific]
+            ])
+        expert_shared_outs = [expert(x_list[-1]).unsqueeze(1) for expert in self.experts_shared
+                             ]  #x_list[-1]: the input for shared experts
+        gate_specific_outs = [gate(x_list[i]).unsqueeze(-1) for i, gate in enumerate(self.gates_specific)
+                             ]  #gate_out[i]: [batch_size, n_expert_specific+n_expert_shared, 1]
+        cgc_outs = []
+        for i, gate_out in enumerate(gate_specific_outs):
+            cur_expert_list = expert_specific_outs[i * self.n_expert_specific:(i + 1) *
+                                                   self.n_expert_specific] + expert_shared_outs
+            expert_concat = torch.cat(cur_expert_list,
+                                      dim=1)  #[batch_size, n_expert_specific+n_expert_shared, expert_dims[-1]]
+            expert_weight = torch.mul(gate_out,
+                                      expert_concat)  #[batch_size, n_expert_specific+n_expert_shared, expert_dims[-1]]
+            expert_pooling = torch.sum(expert_weight, dim=1)  #[batch_size, expert_dims[-1]]
+            cgc_outs.append(expert_pooling)  #length: n_task
+        if self.cur_level < self.n_level:  #not the last layer
+            gate_shared_out = self.gate_shared(x_list[-1]).unsqueeze(-1)  #[batch_size, n_expert_all, 1]
+            expert_concat = torch.cat(expert_specific_outs + expert_shared_outs,
+                                      dim=1)  #[batch_size, n_expert_all, expert_dims[-1]]
+            expert_weight = torch.mul(gate_shared_out, expert_concat)  #[batch_size, n_expert_all, expert_dims[-1]]
+            expert_pooling = torch.sum(expert_weight, dim=1)  #[batch_size, expert_dims[-1]]
+            cgc_outs.append(expert_pooling)  #length: n_task+1
+
+        return cgc_outs

torch_rechub/models/multi_task/shared_bottom.py

@@ -1,44 +1,45 @@
-"""
-Date: create on 04/05/2022
-Reference: 
-    paper: Caruana, R. (1997). Multitask learning. Machine learning, 28(1), 41-75.
-Authors: Mincai Lai, laimincai@shanghaitech.edu.cn
-"""
-
-import torch
-import torch.nn as nn
-
-from ...basic.layers import MLP, EmbeddingLayer, PredictionLayer
-
-
-class SharedBottom(nn.Module):
-    """Shared Bottom multi task model.
-
-    Args:
-        features (list): the list of `Feature Class`, training by the bottom and tower module.
-        task_types (list): types of tasks, only support `["classfication", "regression"]`.
-        bottom_params (dict): the params of the last MLP module, keys include:`{"dims":list, "activation":str, "dropout":float}, keep `{"output_layer":False}`.
-        tower_params_list (list): the list of tower params dict, the keys same as bottom_params.
-    """
-
-    def __init__(self, features, task_types, bottom_params, tower_params_list):
-        super().__init__()
-        self.features = features
-        self.task_types = task_types
-        self.embedding = EmbeddingLayer(features)
-        self.bottom_dims = sum([fea.embed_dim for fea in features])
-
-        self.bottom_mlp = MLP(self.bottom_dims, **{**bottom_params, **{"output_layer": False}})
-        self.towers = nn.ModuleList(MLP(bottom_params["dims"][-1], **tower_params_list[i]) for i in range(len(task_types)))
-        self.predict_layers = nn.ModuleList(PredictionLayer(task_type) for task_type in task_types)
-
-    def forward(self, x):
-        input_bottom = self.embedding(x, self.features, squeeze_dim=True)
-        x = self.bottom_mlp(input_bottom)
-
-        ys = []
-        for tower, predict_layer in zip(self.towers, self.predict_layers):
-            tower_out = tower(x)
-            y = predict_layer(tower_out)  #regression->keep, binary classification->sigmoid
-            ys.append(y)
-        return torch.cat(ys, dim=1)
+"""
+Date: create on 04/05/2022
+Reference: 
+    paper: Caruana, R. (1997). Multitask learning. Machine learning, 28(1), 41-75.
+Authors: Mincai Lai, laimincai@shanghaitech.edu.cn
+"""
+
+import torch
+import torch.nn as nn
+
+from ...basic.layers import MLP, EmbeddingLayer, PredictionLayer
+
+
+class SharedBottom(nn.Module):
+    """Shared Bottom multi task model.
+
+    Args:
+        features (list): the list of `Feature Class`, training by the bottom and tower module.
+        task_types (list): types of tasks, only support `["classfication", "regression"]`.
+        bottom_params (dict): the params of the last MLP module, keys include:`{"dims":list, "activation":str, "dropout":float}, keep `{"output_layer":False}`.
+        tower_params_list (list): the list of tower params dict, the keys same as bottom_params.
+    """
+
+    def __init__(self, features, task_types, bottom_params, tower_params_list):
+        super().__init__()
+        self.features = features
+        self.task_types = task_types
+        self.embedding = EmbeddingLayer(features)
+        self.bottom_dims = sum([fea.embed_dim for fea in features])
+
+        self.bottom_mlp = MLP(self.bottom_dims, **{**bottom_params, **{"output_layer": False}})
+        self.towers = nn.ModuleList(
+            MLP(bottom_params["dims"][-1], **tower_params_list[i]) for i in range(len(task_types)))
+        self.predict_layers = nn.ModuleList(PredictionLayer(task_type) for task_type in task_types)
+
+    def forward(self, x):
+        input_bottom = self.embedding(x, self.features, squeeze_dim=True)
+        x = self.bottom_mlp(input_bottom)
+
+        ys = []
+        for tower, predict_layer in zip(self.towers, self.predict_layers):
+            tower_out = tower(x)
+            y = predict_layer(tower_out)  #regression->keep, binary classification->sigmoid
+            ys.append(y)
+        return torch.cat(ys, dim=1)

torch_rechub/models/ranking/__init__.py

@@ -1,3 +1,11 @@
-from .widedeep import WideDeep
-from .deepfm import DeepFM
-from .din import DIN
+from .widedeep import WideDeep
+from .deepfm import DeepFM
+from .din import DIN
+from .dcn import DCN
+from .dcn_v2 import DCNv2
+from .edcn import EDCN
+from .deepffm import DeepFFM, FatDeepFFM
+from .fibinet import FiBiNet
+from .bst import BST
+from .afm import AFM
+from .dien import DIEN

torch_rechub/models/ranking/afm.py

@@ -0,0 +1,63 @@
+"""
+Date: create on 23/04/2024
+References: 
+    paper: (IJCAI'2017) Attentional Factorization Machines：Learning the Weight of Feature Interactions via Attention Networks
+    url: https://arxiv.org/abs/1708.04617
+Authors: Tao Fan, thisisevy@foxmail.com
+"""
+
+import torch
+from ...basic.layers import FM, MLP, LR, EmbeddingLayer
+from torch import nn
+from torch.nn import Parameter, init
+
+
+class AFM(nn.Module):
+    """Attentional Factorization Machine Model
+
+    Args:
+        fm_features (list): the list of `Feature Class`, training by the fm part module.
+        embed_dim (int): the dimension of input embedding.
+        t (int): the size of the hidden layer in the attention network.
+    """
+
+    def __init__(self, fm_features, embed_dim, t=64):
+        super(AFM, self).__init__()
+        self.fm_features = fm_features
+        self.embed_dim = embed_dim
+        self.fm_dims = sum([fea.embed_dim for fea in fm_features])
+        self.linear = LR(self.fm_dims)  # 1-odrder interaction
+        self.fm = FM(reduce_sum=False)  # 2-odrder interaction
+        self.embedding = EmbeddingLayer(fm_features)
+
+        # 注意力计算中的线性层
+        self.attention_liner = nn.Linear(self.embed_dim, t)
+        # AFM公式中的h
+        self.h = init.xavier_uniform_(Parameter(torch.empty(t, 1)))
+        # AFM公式中的p
+        self.p = init.xavier_uniform_(Parameter(torch.empty(self.embed_dim, 1)))
+
+    def attention(self, y_fm):
+        # embs: [ batch_size, k ]
+        # [ batch_size, t ]
+        y_fm = self.attention_liner(y_fm)
+        # [ batch_size, t ]
+        y_fm = torch.relu(y_fm)
+        # [ batch_size, 1 ]
+        y_fm = torch.matmul(y_fm, self.h)
+        # [ batch_size, 1 ]
+        atts = torch.softmax(y_fm, dim=1)
+        return atts
+
+    def forward(self, x):
+        input_fm = self.embedding(x, self.fm_features, squeeze_dim=False)  # [batch_size, num_fields, embed_dim]
+
+        y_linear = self.linear(input_fm.flatten(start_dim=1))
+        y_fm = self.fm(input_fm)
+        # 得到注意力
+        atts = self.attention(y_fm)
+        # [ batch_size, 1 ]
+        outs = torch.matmul(atts * y_fm, self.p)
+        # print(y_linear.size(), outs.size())
+        y = y_linear + outs
+        return torch.sigmoid(y.squeeze(1))

torch_rechub/models/ranking/bst.py

@@ -0,0 +1,63 @@
+"""
+Date: create on 26/02/2024, update on 30/04/2022
+References: 
+    paper: Behavior Sequence Transformer for E-commerce Recommendation in Alibaba
+    url: https://arxiv.org/pdf/1905.06874
+    code: https://github.com/jiwidi/Behavior-Sequence-Transformer-Pytorch/blob/master/pytorch_bst.ipynb
+Authors: Tao Fan, thisisevy@foxmail.com
+"""
+
+import torch
+import torch.nn as nn
+
+from ...basic.layers import EmbeddingLayer, MLP
+
+
+class BST(nn.Module):
+    """Behavior Sequence Transformer
+    Args:
+        features (list): the list of `Feature Class`. training by MLP. It means the user profile features and context features in origin paper, exclude history and target features.
+        history_features (list): the list of `Feature Class`,training by ActivationUnit. It means the user behaviour sequence features, eg.item id sequence, shop id sequence.
+        target_features (list): the list of `Feature Class`, training by ActivationUnit. It means the target feature which will execute target-attention with history feature.
+        mlp_params (dict): the params of the last MLP module, keys include:`{"dims":list, "activation":str, "dropout":float, "output_layer":bool`}.
+        nhead (int): the number of heads in the multi-head-attention models.
+        dropout (float): the dropout value in the multi-head-attention models.
+        num_layers (Any): the number of sub-encoder-layers in the encoder.
+    """
+
+    def __init__(self, features, history_features, target_features, mlp_params, nhead=8, dropout=0.2, num_layers=1):
+        super().__init__()
+        self.features = features
+        self.history_features = history_features
+        self.target_features = target_features
+        self.num_history_features = len(history_features)
+        self.embed_dim = target_features[0].embed_dim
+        self.seq_len = 50
+        # TODO 在 'torch_rechub.basic.features.SequenceFeature' 中加入seq_len属性
+        self.all_dims = (len(features) + len(history_features) * (self.seq_len + len(target_features))) * self.embed_dim
+        self.embedding = EmbeddingLayer(features + history_features + target_features)
+        self.encoder_layer = nn.TransformerEncoderLayer(d_model=self.embed_dim, nhead=nhead, dropout=dropout)
+        self.transformer_layers = nn.TransformerEncoder(self.encoder_layer, num_layers=num_layers)
+        self.mlp = MLP(self.all_dims, activation="leakyrelu",
+                       **mlp_params)  # # 定义模型，模型的参数需要我们之前的feature类，用于构建模型的输入层，mlp指定模型后续DNN的结构
+
+    def forward(self, x):
+        embed_x_features = self.embedding(x, self.features)  # (batch_size, num_features, emb_dim)
+        embed_x_history = self.embedding(x,
+                                         self.history_features)  # (batch_size, num_history_features, seq_length, emb_dim)
+        embed_x_target = self.embedding(x, self.target_features)  # (batch_size, num_target_features, emb_dim)
+        attention_pooling = []
+        for i in range(self.num_history_features):
+            attention_seq = self.transformer_layers(
+                torch.cat([torch.squeeze(embed_x_history[:, i, :, :], 1), embed_x_target], dim=1))
+            attention_pooling.append(attention_seq)  # (batch_size, seq_length + num_target_features, emb_dim)
+        attention_pooling = torch.cat(attention_pooling,
+                                      dim=1)  # (batch_size, num_history_features * (seq_length + num_target_features), emb_dim)
+
+        mlp_in = torch.cat([
+            attention_pooling.flatten(start_dim=1),
+            embed_x_features.flatten(start_dim=1)
+        ],
+            dim=1)  # (batch_size, N)
+        y = self.mlp(mlp_in)
+        return torch.sigmoid(y.squeeze(1))

torch_rechub/models/ranking/dcn.py

@@ -0,0 +1,38 @@
+"""
+Date: create on 12/05/2022
+References: 
+    paper: (AKDD'2017) Deep & Cross Network for Ad Click Predictions
+    url: https://arxiv.org/abs/1708.05123
+Authors: Mincai Lai, laimincai@shanghaitech.edu.cn
+"""
+
+import torch
+
+from ...basic.layers import LR, MLP, CrossNetwork, EmbeddingLayer
+
+
+class DCN(torch.nn.Module):
+    """Deep & Cross Network
+
+    Args:
+        features (list[Feature Class]): training by the whole module.
+        mlp_params (dict): the params of the last MLP module, keys include:`{"dims":list, "activation":str, "dropout":float, "output_layer":bool`}
+    """
+
+    def __init__(self, features, n_cross_layers, mlp_params):
+        super().__init__()
+        self.features = features
+        self.dims = sum([fea.embed_dim for fea in features])
+
+        self.embedding = EmbeddingLayer(features)
+        self.cn = CrossNetwork(self.dims, n_cross_layers)
+        self.mlp = MLP(self.dims, output_layer=False, **mlp_params)
+        self.linear = LR(self.dims + mlp_params["dims"][-1])
+
+    def forward(self, x):
+        embed_x = self.embedding(x, self.features, squeeze_dim=True)
+        cn_out = self.cn(embed_x)
+        mlp_out = self.mlp(embed_x)
+        x_stack = torch.cat([cn_out, mlp_out], dim=1)
+        y = self.linear(x_stack)
+        return torch.sigmoid(y.squeeze(1))