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

torch_rechub/models/multi_task/ple.py

@@ -1,130 +1,116 @@
-"""
-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
+"""
+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):
+        # [batch_size, input_dims]
+        embed_x = self.embedding(x, self.features, squeeze_dim=True)
+        ple_inputs = [embed_x] * (self.n_task + 1)
+        ple_outs = []
+        for i in range(self.n_level):
+            # ple_outs[i]: [batch_size, expert_dims[-1]]
+            ple_outs = self.cgc_layers[i](ple_inputs)
+            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
+        # the first layer expert dim is the input data dim other expert dim
+        input_dims = input_dims if cur_level == 1 else expert_params["dims"][-1]
+        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]])
+        # x_list[-1]: the input for shared experts
+        expert_shared_outs = [expert(x_list[-1]).unsqueeze(1) for expert in self.experts_shared]
+        # gate_out[i]: [batch_size, n_expert_specific+n_expert_shared, 1]
+        gate_specific_outs = [gate(x_list[i]).unsqueeze(-1) for i, gate in enumerate(self.gates_specific)]
+        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
+            # [batch_size, n_expert_specific+n_expert_shared, expert_dims[-1]]
+            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, expert_dims[-1]]
+            expert_pooling = torch.sum(expert_weight, dim=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]]
+            # [batch_size, n_expert_all, expert_dims[-1]]
+            expert_weight = torch.mul(gate_shared_out, expert_concat)
+            # [batch_size, expert_dims[-1]]
+            expert_pooling = torch.sum(expert_weight, dim=1)
+            cgc_outs.append(expert_pooling)  # length: n_task+1
+
+        return cgc_outs

torch_rechub/models/multi_task/shared_bottom.py

@@ -1,45 +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)
+            # regression->keep, binary classification->sigmoid
+            y = predict_layer(tower_out)
+            ys.append(y)
+        return torch.cat(ys, dim=1)

torch_rechub/models/ranking/__init__.py

@@ -1,11 +1,14 @@
-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
+__all__ = ['WideDeep', 'DeepFM', 'DCN', 'DCNv2', 'EDCN', 'AFM', 'FiBiNet', 'DeepFFM', 'BST', 'DIN', 'DIEN', 'FatDeepFFM', 'AutoInt']
+
+from .afm import AFM
+from .autoint import AutoInt
+from .bst import BST
+from .dcn import DCN
+from .dcn_v2 import DCNv2
+from .deepffm import DeepFFM, FatDeepFFM
+from .deepfm import DeepFM
+from .dien import DIEN
+from .din import DIN
+from .edcn import EDCN
+from .fibinet import FiBiNet
+from .widedeep import WideDeep

torch_rechub/models/ranking/afm.py

@@ -1,63 +1,65 @@
-"""
-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))
+"""
+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 torch import nn
+from torch.nn import Parameter, init
+
+from ...basic.layers import FM, LR, MLP, EmbeddingLayer
+
+
+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):
+        # [batch_size, num_fields, embed_dim]
+        input_fm = self.embedding(x, self.fm_features, squeeze_dim=False)
+
+        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/autoint.py

@@ -0,0 +1,102 @@
+"""
+Date: create on 14/11/2025
+References:
+    paper: (CIKM'2019) AutoInt: Automatic Feature Interaction Learning via Self-Attentive Neural Networks
+    url: https://arxiv.org/abs/1810.11921
+Authors: Yang Zhou, zyaztec@gmail.com
+"""
+
+import torch
+import torch.nn as nn
+
+from ...basic.layers import LR, MLP, EmbeddingLayer, InteractingLayer
+
+
+class AutoInt(torch.nn.Module):
+    """AutoInt Model
+
+    Args:
+        sparse_features (list): the list of `SparseFeature` Class
+        dense_features (list): the list of `DenseFeature` Class
+        num_layers (int): number of interacting layers
+        num_heads (int): number of attention heads
+        dropout (float): dropout rate for attention
+        mlp_params (dict): parameters for MLP, keys: {"dims":list, "activation":str,
+                                             "dropout":float, "output_layer":bool"}
+    """
+
+    def __init__(self, sparse_features, dense_features, num_layers=3, num_heads=2, dropout=0.0, mlp_params=None):
+        super(AutoInt, self).__init__()
+        self.sparse_features = sparse_features
+
+        self.dense_features = dense_features if dense_features is not None else []
+        embed_dims = [fea.embed_dim for fea in self.sparse_features]
+        self.embed_dim = embed_dims[0]
+        if len(self.sparse_features) == 0:
+            raise ValueError("AutoInt requires at least one sparse feature to determine embed_dim.")
+
+        # field nums = sparse + dense
+        self.num_sparse = len(self.sparse_features)
+        self.num_dense = len(self.dense_features)
+        self.num_fields = self.num_sparse + self.num_dense
+
+        # total dims = num_fields * embed_dim
+        self.dims = self.num_fields * self.embed_dim
+        self.num_layers = num_layers
+
+        self.sparse_embedding = EmbeddingLayer(self.sparse_features)
+
+        # dense feature embedding
+        self.dense_embeddings = nn.ModuleDict()
+        for fea in self.dense_features:
+            self.dense_embeddings[fea.name] = nn.Linear(1, self.embed_dim, bias=False)
+
+        self.interacting_layers = torch.nn.ModuleList([InteractingLayer(self.embed_dim, num_heads=num_heads, dropout=dropout, residual=True) for _ in range(num_layers)])
+
+        self.linear = LR(self.dims)
+
+        self.attn_linear = nn.Linear(self.dims, 1)
+
+        if mlp_params is not None:
+            self.use_mlp = True
+            self.mlp = MLP(self.dims, **mlp_params)
+        else:
+            self.use_mlp = False
+
+    def forward(self, x):
+        # sparse feature embedding: [B, num_sparse, embed_dim]
+        sparse_emb = self.sparse_embedding(x, self.sparse_features, squeeze_dim=False)
+
+        dense_emb_list = []
+        for fea in self.dense_features:
+            v = x[fea.name].float().view(-1, 1, 1)
+            dense_emb = self.dense_embeddings[fea.name](v)  # [B, 1, embed_dim]
+            dense_emb_list.append(dense_emb)
+
+        if len(dense_emb_list) > 0:
+            dense_emb = torch.cat(dense_emb_list, dim=1)  # [B, num_dense, d]
+            embed_x = torch.cat([sparse_emb, dense_emb], dim=1)  # [B, num_fields, d]
+        else:
+            embed_x = sparse_emb  # [B, num_sparse, d]
+
+        embed_x_flatten = embed_x.flatten(start_dim=1)  # [B, num_fields * embed_dim]
+
+        # Multi-head self-attention layers
+        attn_out = embed_x
+        for layer in self.interacting_layers:
+            attn_out = layer(attn_out)  # [B, num_fields, embed_dim]
+
+        # Attention linear
+        attn_out_flatten = attn_out.flatten(start_dim=1)  # [B, num_fields * embed_dim]
+        y_attn = self.attn_linear(attn_out_flatten)  # [B, 1]
+
+        # Linear part
+        y_linear = self.linear(embed_x_flatten)  # [B, 1]
+
+        # Deep MLP
+        y = y_attn + y_linear
+        if self.use_mlp:
+            y_deep = self.mlp(embed_x_flatten)  # [B, 1]
+            y = y + y_deep
+
+        return torch.sigmoid(y.squeeze(1))