nextrec 0.1.3__py3-none-any.whl → 0.1.7__py3-none-any.whl

nextrec/basic/layers.py

@@ -50,14 +50,7 @@ __all__ = [
 
 
 class PredictionLayer(nn.Module):
-    _CLASSIFICATION_TASKS = {
-        "classification",
-        "binary",
-        "ctr",
-        "ranking",
-        "match",
-        "matching",
-    }
+    _CLASSIFICATION_TASKS = {"classification", "binary", "ctr", "ranking", "match", "matching"}
     _REGRESSION_TASKS = {"regression", "continuous"}
     _MULTICLASS_TASKS = {"multiclass", "softmax"}
 
@@ -220,9 +213,7 @@ class EmbeddingLayer(nn.Module):
                 elif feature.combiner == "concat":
                     pooling_layer = ConcatPooling()
                 else:
-                    raise ValueError(
-                        f"Unknown combiner for {feature.name}: {feature.combiner}"
-                    )
+                    raise ValueError(f"Unknown combiner for {feature.name}: {feature.combiner}")
 
                 feature_mask = InputMask()(x, feature, seq_input)
                 sparse_embeds.append(pooling_layer(seq_emb, feature_mask).unsqueeze(1))
@@ -254,9 +245,7 @@ class EmbeddingLayer(nn.Module):
 
             if target_dim is not None:
                 aligned_dense = [
-                    emb.unsqueeze(1)
-                    for emb in dense_embeds
-                    if emb.shape[-1] == target_dim
+                    emb.unsqueeze(1) for emb in dense_embeds if emb.shape[-1] == target_dim
                 ]
                 output_embeddings.extend(aligned_dense)
 
@@ -268,9 +257,7 @@ class EmbeddingLayer(nn.Module):
 
         return torch.cat(output_embeddings, dim=1)
 
-    def _project_dense(
-        self, feature: DenseFeature, x: dict[str, torch.Tensor]
-    ) -> torch.Tensor:
+    def _project_dense(self, feature: DenseFeature, x: dict[str, torch.Tensor]) -> torch.Tensor:
         if feature.name not in x:
             raise KeyError(f"Dense feature '{feature.name}' is missing from input.")
 
@@ -293,7 +280,6 @@ class EmbeddingLayer(nn.Module):
     def _compute_output_dim(self):
         return
 
-
 class InputMask(nn.Module):
     """Utility module to build sequence masks for pooling layers."""
 
@@ -303,9 +289,9 @@ class InputMask(nn.Module):
     def forward(self, x, fea, seq_tensor=None):
         values = seq_tensor if seq_tensor is not None else x[fea.name]
         if fea.padding_idx is not None:
-            mask = values.long() != fea.padding_idx
+            mask = (values.long() != fea.padding_idx)
         else:
-            mask = values.long() != 0
+            mask = (values.long() != 0)
         if mask.dim() == 1:
             mask = mask.unsqueeze(-1)
         return mask.unsqueeze(1).float()
@@ -333,7 +319,7 @@ class ConcatPooling(nn.Module):
         super().__init__()
 
     def forward(self, x, mask=None):
-        return x.flatten(start_dim=1, end_dim=2)
+        return x.flatten(start_dim=1, end_dim=2) 
 
 
 class AveragePooling(nn.Module):
@@ -367,9 +353,7 @@ class SumPooling(nn.Module):
 class MLP(nn.Module):
     """Stacked fully connected layers used in the deep component."""
 
-    def __init__(
-        self, input_dim, output_layer=True, dims=None, dropout=0, activation="relu"
-    ):
+    def __init__(self, input_dim, output_layer=True, dims=None, dropout=0, activation="relu"):
         super().__init__()
         if dims is None:
             dims = []
@@ -396,7 +380,7 @@ class FM(nn.Module):
         self.reduce_sum = reduce_sum
 
     def forward(self, x):
-        square_of_sum = torch.sum(x, dim=1) ** 2
+        square_of_sum = torch.sum(x, dim=1)**2
         sum_of_square = torch.sum(x**2, dim=1)
         ix = square_of_sum - sum_of_square
         if self.reduce_sum:
@@ -415,16 +399,7 @@ class CIN(nn.Module):
         prev_dim, fc_input_dim = input_dim, 0
         for i in range(self.num_layers):
             cross_layer_size = cin_size[i]
-            self.conv_layers.append(
-                torch.nn.Conv1d(
-                    input_dim * prev_dim,
-                    cross_layer_size,
-                    1,
-                    stride=1,
-                    dilation=1,
-                    bias=True,
-                )
-            )
+            self.conv_layers.append(torch.nn.Conv1d(input_dim * prev_dim, cross_layer_size, 1, stride=1, dilation=1, bias=True))
             if self.split_half and i != self.num_layers - 1:
                 cross_layer_size //= 2
             prev_dim = cross_layer_size
@@ -446,7 +421,6 @@ class CIN(nn.Module):
             xs.append(x)
         return self.fc(torch.sum(torch.cat(xs, dim=1), 2))
 
-
 class CrossLayer(nn.Module):
     """Single cross layer used in DCN (Wang et al., 2017)."""
 
@@ -466,12 +440,8 @@ class CrossNetwork(nn.Module):
     def __init__(self, input_dim, num_layers):
         super().__init__()
         self.num_layers = num_layers
-        self.w = torch.nn.ModuleList(
-            [torch.nn.Linear(input_dim, 1, bias=False) for _ in range(num_layers)]
-        )
-        self.b = torch.nn.ParameterList(
-            [torch.nn.Parameter(torch.zeros((input_dim,))) for _ in range(num_layers)]
-        )
+        self.w = torch.nn.ModuleList([torch.nn.Linear(input_dim, 1, bias=False) for _ in range(num_layers)])
+        self.b = torch.nn.ParameterList([torch.nn.Parameter(torch.zeros((input_dim,))) for _ in range(num_layers)])
 
     def forward(self, x):
         """
@@ -483,30 +453,21 @@ class CrossNetwork(nn.Module):
             x = x0 * xw + self.b[i] + x
         return x
 
-
 class CrossNetV2(nn.Module):
     """Vector-wise cross network proposed in DCN V2 (Wang et al., 2021)."""
-
     def __init__(self, input_dim, num_layers):
         super().__init__()
         self.num_layers = num_layers
-        self.w = torch.nn.ModuleList(
-            [
-                torch.nn.Linear(input_dim, input_dim, bias=False)
-                for _ in range(num_layers)
-            ]
-        )
-        self.b = torch.nn.ParameterList(
-            [torch.nn.Parameter(torch.zeros((input_dim,))) for _ in range(num_layers)]
-        )
+        self.w = torch.nn.ModuleList([torch.nn.Linear(input_dim, input_dim, bias=False) for _ in range(num_layers)])
+        self.b = torch.nn.ParameterList([torch.nn.Parameter(torch.zeros((input_dim,))) for _ in range(num_layers)])
+
 
     def forward(self, x):
         x0 = x
         for i in range(self.num_layers):
-            x = x0 * self.w[i](x) + self.b[i] + x
+            x =x0*self.w[i](x) + self.b[i] + x
         return x
 
-
 class CrossNetMix(nn.Module):
     """Mixture of low-rank cross experts from DCN V2 (Wang et al., 2021)."""
 
@@ -516,46 +477,18 @@ class CrossNetMix(nn.Module):
         self.num_experts = num_experts
 
         # U: (input_dim, low_rank)
-        self.u_list = torch.nn.ParameterList(
-            [
-                nn.Parameter(
-                    nn.init.xavier_normal_(
-                        torch.empty(num_experts, input_dim, low_rank)
-                    )
-                )
-                for i in range(self.num_layers)
-            ]
-        )
+        self.u_list = torch.nn.ParameterList([nn.Parameter(nn.init.xavier_normal_(
+            torch.empty(num_experts, input_dim, low_rank))) for i in range(self.num_layers)])
         # V: (input_dim, low_rank)
-        self.v_list = torch.nn.ParameterList(
-            [
-                nn.Parameter(
-                    nn.init.xavier_normal_(
-                        torch.empty(num_experts, input_dim, low_rank)
-                    )
-                )
-                for i in range(self.num_layers)
-            ]
-        )
+        self.v_list = torch.nn.ParameterList([nn.Parameter(nn.init.xavier_normal_(
+            torch.empty(num_experts, input_dim, low_rank))) for i in range(self.num_layers)])
         # C: (low_rank, low_rank)
-        self.c_list = torch.nn.ParameterList(
-            [
-                nn.Parameter(
-                    nn.init.xavier_normal_(torch.empty(num_experts, low_rank, low_rank))
-                )
-                for i in range(self.num_layers)
-            ]
-        )
-        self.gating = nn.ModuleList(
-            [nn.Linear(input_dim, 1, bias=False) for i in range(self.num_experts)]
-        )
-
-        self.bias = torch.nn.ParameterList(
-            [
-                nn.Parameter(nn.init.zeros_(torch.empty(input_dim, 1)))
-                for i in range(self.num_layers)
-            ]
-        )
+        self.c_list = torch.nn.ParameterList([nn.Parameter(nn.init.xavier_normal_(
+            torch.empty(num_experts, low_rank, low_rank))) for i in range(self.num_layers)])
+        self.gating = nn.ModuleList([nn.Linear(input_dim, 1, bias=False) for i in range(self.num_experts)])
+
+        self.bias = torch.nn.ParameterList([nn.Parameter(nn.init.zeros_(
+            torch.empty(input_dim, 1))) for i in range(self.num_layers)])
 
     def forward(self, x):
         x_0 = x.unsqueeze(2)  # (bs, in_features, 1)
@@ -570,9 +503,7 @@ class CrossNetMix(nn.Module):
 
                 # (2) E(x_l)
                 # project the input x_l to $\mathbb{R}^{r}$
-                v_x = torch.matmul(
-                    self.v_list[i][expert_id].t(), x_l
-                )  # (bs, low_rank, 1)
+                v_x = torch.matmul(self.v_list[i][expert_id].t(), x_l)  # (bs, low_rank, 1)
 
                 # nonlinear activation in low rank space
                 v_x = torch.tanh(v_x)
@@ -580,9 +511,7 @@ class CrossNetMix(nn.Module):
                 v_x = torch.tanh(v_x)
 
                 # project back to $\mathbb{R}^{d}$
-                uv_x = torch.matmul(
-                    self.u_list[i][expert_id], v_x
-                )  # (bs, in_features, 1)
+                uv_x = torch.matmul(self.u_list[i][expert_id], v_x)  # (bs, in_features, 1)
 
                 dot_ = uv_x + self.bias[i]
                 dot_ = x_0 * dot_  # Hadamard-product
@@ -590,78 +519,53 @@ class CrossNetMix(nn.Module):
                 output_of_experts.append(dot_.squeeze(2))
 
             # (3) mixture of low-rank experts
-            output_of_experts = torch.stack(
-                output_of_experts, 2
-            )  # (bs, in_features, num_experts)
-            gating_score_experts = torch.stack(
-                gating_score_experts, 1
-            )  # (bs, num_experts, 1)
+            output_of_experts = torch.stack(output_of_experts, 2)  # (bs, in_features, num_experts)
+            gating_score_experts = torch.stack(gating_score_experts, 1)  # (bs, num_experts, 1)
             moe_out = torch.matmul(output_of_experts, gating_score_experts.softmax(1))
             x_l = moe_out + x_l  # (bs, in_features, 1)
 
         x_l = x_l.squeeze()  # (bs, in_features)
         return x_l
 
-
 class SENETLayer(nn.Module):
     """Squeeze-and-Excitation block adopted by FiBiNET (Huang et al., 2019)."""
 
     def __init__(self, num_fields, reduction_ratio=3):
         super(SENETLayer, self).__init__()
-        reduced_size = max(1, int(num_fields / reduction_ratio))
-        self.mlp = nn.Sequential(
-            nn.Linear(num_fields, reduced_size, bias=False),
-            nn.ReLU(),
-            nn.Linear(reduced_size, num_fields, bias=False),
-            nn.ReLU(),
-        )
-
+        reduced_size = max(1, int(num_fields/ reduction_ratio))
+        self.mlp = nn.Sequential(nn.Linear(num_fields, reduced_size, bias=False),
+                                 nn.ReLU(),
+                                 nn.Linear(reduced_size, num_fields, bias=False),
+                                 nn.ReLU())
     def forward(self, x):
         z = torch.mean(x, dim=-1, out=None)
         a = self.mlp(z)
-        v = x * a.unsqueeze(-1)
+        v = x*a.unsqueeze(-1)
         return v
 
-
 class BiLinearInteractionLayer(nn.Module):
     """Bilinear feature interaction from FiBiNET (Huang et al., 2019)."""
 
-    def __init__(self, input_dim, num_fields, bilinear_type="field_interaction"):
+    def __init__(self, input_dim, num_fields, bilinear_type = "field_interaction"):
         super(BiLinearInteractionLayer, self).__init__()
         self.bilinear_type = bilinear_type
         if self.bilinear_type == "field_all":
             self.bilinear_layer = nn.Linear(input_dim, input_dim, bias=False)
         elif self.bilinear_type == "field_each":
-            self.bilinear_layer = nn.ModuleList(
-                [nn.Linear(input_dim, input_dim, bias=False) for i in range(num_fields)]
-            )
+            self.bilinear_layer = nn.ModuleList([nn.Linear(input_dim, input_dim, bias=False) for i in range(num_fields)])
         elif self.bilinear_type == "field_interaction":
-            self.bilinear_layer = nn.ModuleList(
-                [
-                    nn.Linear(input_dim, input_dim, bias=False)
-                    for i, j in combinations(range(num_fields), 2)
-                ]
-            )
+            self.bilinear_layer = nn.ModuleList([nn.Linear(input_dim, input_dim, bias=False) for i,j in combinations(range(num_fields), 2)])
         else:
             raise NotImplementedError()
 
     def forward(self, x):
         feature_emb = torch.split(x, 1, dim=1)
         if self.bilinear_type == "field_all":
-            bilinear_list = [
-                self.bilinear_layer(v_i) * v_j
-                for v_i, v_j in combinations(feature_emb, 2)
-            ]
+            bilinear_list = [self.bilinear_layer(v_i)*v_j for v_i, v_j in combinations(feature_emb, 2)]
         elif self.bilinear_type == "field_each":
-            bilinear_list = [
-                self.bilinear_layer[i](feature_emb[i]) * feature_emb[j]
-                for i, j in combinations(range(len(feature_emb)), 2)
-            ]
+            bilinear_list = [self.bilinear_layer[i](feature_emb[i])*feature_emb[j] for i,j in combinations(range(len(feature_emb)), 2)]
         elif self.bilinear_type == "field_interaction":
-            bilinear_list = [
-                self.bilinear_layer[i](v[0]) * v[1]
-                for i, v in enumerate(combinations(feature_emb, 2))
-            ]
+            bilinear_list = [self.bilinear_layer[i](v[0])*v[1] for i,v in enumerate(combinations(feature_emb, 2))]
         return torch.cat(bilinear_list, dim=1)
 
 
@@ -674,23 +578,17 @@ class MultiInterestSA(nn.Module):
         self.interest_num = interest_num
         if hidden_dim == None:
             self.hidden_dim = self.embedding_dim * 4
-        self.W1 = torch.nn.Parameter(
-            torch.rand(self.embedding_dim, self.hidden_dim), requires_grad=True
-        )
-        self.W2 = torch.nn.Parameter(
-            torch.rand(self.hidden_dim, self.interest_num), requires_grad=True
-        )
-        self.W3 = torch.nn.Parameter(
-            torch.rand(self.embedding_dim, self.embedding_dim), requires_grad=True
-        )
+        self.W1 = torch.nn.Parameter(torch.rand(self.embedding_dim, self.hidden_dim), requires_grad=True)
+        self.W2 = torch.nn.Parameter(torch.rand(self.hidden_dim, self.interest_num), requires_grad=True)
+        self.W3 = torch.nn.Parameter(torch.rand(self.embedding_dim, self.embedding_dim), requires_grad=True)
 
     def forward(self, seq_emb, mask=None):
-        H = torch.einsum("bse, ed -> bsd", seq_emb, self.W1).tanh()
+        H = torch.einsum('bse, ed -> bsd', seq_emb, self.W1).tanh()
         if mask != None:
-            A = torch.einsum("bsd, dk -> bsk", H, self.W2) + -1.0e9 * (1 - mask.float())
+            A = torch.einsum('bsd, dk -> bsk', H, self.W2) + -1.e9 * (1 - mask.float())
             A = F.softmax(A, dim=1)
         else:
-            A = F.softmax(torch.einsum("bsd, dk -> bsk", H, self.W2), dim=1)
+            A = F.softmax(torch.einsum('bsd, dk -> bsk', H, self.W2), dim=1)
         A = A.permute(0, 2, 1)
         multi_interest_emb = torch.matmul(A, seq_emb)
         return multi_interest_emb
@@ -699,15 +597,7 @@ class MultiInterestSA(nn.Module):
 class CapsuleNetwork(nn.Module):
     """Dynamic routing capsule network used in MIND (Li et al., 2019)."""
 
-    def __init__(
-        self,
-        embedding_dim,
-        seq_len,
-        bilinear_type=2,
-        interest_num=4,
-        routing_times=3,
-        relu_layer=False,
-    ):
+    def __init__(self, embedding_dim, seq_len, bilinear_type=2, interest_num=4, routing_times=3, relu_layer=False):
         super(CapsuleNetwork, self).__init__()
         self.embedding_dim = embedding_dim  # h
         self.seq_len = seq_len  # s
@@ -717,24 +607,13 @@ class CapsuleNetwork(nn.Module):
 
         self.relu_layer = relu_layer
         self.stop_grad = True
-        self.relu = nn.Sequential(
-            nn.Linear(self.embedding_dim, self.embedding_dim, bias=False), nn.ReLU()
-        )
+        self.relu = nn.Sequential(nn.Linear(self.embedding_dim, self.embedding_dim, bias=False), nn.ReLU())
         if self.bilinear_type == 0:  # MIND
             self.linear = nn.Linear(self.embedding_dim, self.embedding_dim, bias=False)
         elif self.bilinear_type == 1:
-            self.linear = nn.Linear(
-                self.embedding_dim, self.embedding_dim * self.interest_num, bias=False
-            )
+            self.linear = nn.Linear(self.embedding_dim, self.embedding_dim * self.interest_num, bias=False)
         else:
-            self.w = nn.Parameter(
-                torch.Tensor(
-                    1,
-                    self.seq_len,
-                    self.interest_num * self.embedding_dim,
-                    self.embedding_dim,
-                )
-            )
+            self.w = nn.Parameter(torch.Tensor(1, self.seq_len, self.interest_num * self.embedding_dim, self.embedding_dim))
             nn.init.xavier_uniform_(self.w)
 
     def forward(self, item_eb, mask):
@@ -745,15 +624,11 @@ class CapsuleNetwork(nn.Module):
             item_eb_hat = self.linear(item_eb)
         else:
             u = torch.unsqueeze(item_eb, dim=2)
-            item_eb_hat = torch.sum(self.w[:, : self.seq_len, :, :] * u, dim=3)
+            item_eb_hat = torch.sum(self.w[:, :self.seq_len, :, :] * u, dim=3)
 
-        item_eb_hat = torch.reshape(
-            item_eb_hat, (-1, self.seq_len, self.interest_num, self.embedding_dim)
-        )
+        item_eb_hat = torch.reshape(item_eb_hat, (-1, self.seq_len, self.interest_num, self.embedding_dim))
         item_eb_hat = torch.transpose(item_eb_hat, 1, 2).contiguous()
-        item_eb_hat = torch.reshape(
-            item_eb_hat, (-1, self.interest_num, self.seq_len, self.embedding_dim)
-        )
+        item_eb_hat = torch.reshape(item_eb_hat, (-1, self.interest_num, self.seq_len, self.embedding_dim))
 
         if self.stop_grad:
             item_eb_hat_iter = item_eb_hat.detach()
@@ -761,47 +636,34 @@ class CapsuleNetwork(nn.Module):
             item_eb_hat_iter = item_eb_hat
 
         if self.bilinear_type > 0:
-            capsule_weight = torch.zeros(
-                item_eb_hat.shape[0],
-                self.interest_num,
-                self.seq_len,
-                device=item_eb.device,
-                requires_grad=False,
-            )
+            capsule_weight = torch.zeros(item_eb_hat.shape[0],
+                                         self.interest_num,
+                                         self.seq_len,
+                                         device=item_eb.device,
+                                         requires_grad=False)
         else:
-            capsule_weight = torch.randn(
-                item_eb_hat.shape[0],
-                self.interest_num,
-                self.seq_len,
-                device=item_eb.device,
-                requires_grad=False,
-            )
+            capsule_weight = torch.randn(item_eb_hat.shape[0],
+                                         self.interest_num,
+                                         self.seq_len,
+                                         device=item_eb.device,
+                                         requires_grad=False)
 
         for i in range(self.routing_times):  # 动态路由传播3次
             atten_mask = torch.unsqueeze(mask, 1).repeat(1, self.interest_num, 1)
             paddings = torch.zeros_like(atten_mask, dtype=torch.float)
 
             capsule_softmax_weight = F.softmax(capsule_weight, dim=-1)
-            capsule_softmax_weight = torch.where(
-                torch.eq(atten_mask, 0), paddings, capsule_softmax_weight
-            )
+            capsule_softmax_weight = torch.where(torch.eq(atten_mask, 0), paddings, capsule_softmax_weight)
             capsule_softmax_weight = torch.unsqueeze(capsule_softmax_weight, 2)
 
             if i < 2:
-                interest_capsule = torch.matmul(
-                    capsule_softmax_weight, item_eb_hat_iter
-                )
+                interest_capsule = torch.matmul(capsule_softmax_weight, item_eb_hat_iter)
                 cap_norm = torch.sum(torch.square(interest_capsule), -1, True)
                 scalar_factor = cap_norm / (1 + cap_norm) / torch.sqrt(cap_norm + 1e-9)
                 interest_capsule = scalar_factor * interest_capsule
 
-                delta_weight = torch.matmul(
-                    item_eb_hat_iter,
-                    torch.transpose(interest_capsule, 2, 3).contiguous(),
-                )
-                delta_weight = torch.reshape(
-                    delta_weight, (-1, self.interest_num, self.seq_len)
-                )
+                delta_weight = torch.matmul(item_eb_hat_iter, torch.transpose(interest_capsule, 2, 3).contiguous())
+                delta_weight = torch.reshape(delta_weight, (-1, self.interest_num, self.seq_len))
                 capsule_weight = capsule_weight + delta_weight
             else:
                 interest_capsule = torch.matmul(capsule_softmax_weight, item_eb_hat)
@@ -809,9 +671,7 @@ class CapsuleNetwork(nn.Module):
                 scalar_factor = cap_norm / (1 + cap_norm) / torch.sqrt(cap_norm + 1e-9)
                 interest_capsule = scalar_factor * interest_capsule
 
-        interest_capsule = torch.reshape(
-            interest_capsule, (-1, self.interest_num, self.embedding_dim)
-        )
+        interest_capsule = torch.reshape(interest_capsule, (-1, self.interest_num, self.embedding_dim))
 
         if self.relu_layer:
             interest_capsule = self.relu(interest_capsule)
@@ -823,18 +683,18 @@ class FFM(nn.Module):
     """Field-aware Factorization Machine (Juan et al., 2016)."""
 
     def __init__(self, num_fields, reduce_sum=True):
-        super().__init__()
+        super().__init__()        
         self.num_fields = num_fields
         self.reduce_sum = reduce_sum
 
     def forward(self, x):
         # compute (non-redundant) second order field-aware feature crossings
         crossed_embeddings = []
-        for i in range(self.num_fields - 1):
-            for j in range(i + 1, self.num_fields):
-                crossed_embeddings.append(x[:, i, j, :] * x[:, j, i, :])
+        for i in range(self.num_fields-1):
+            for j in range(i+1, self.num_fields):
+                crossed_embeddings.append(x[:, i, j, :] *  x[:, j, i, :])        
         crossed_embeddings = torch.stack(crossed_embeddings, dim=1)
-
+        
         # if reduce_sum is true, the crossing operation is effectively inner product, other wise Hadamard-product
         if self.reduce_sum:
             crossed_embeddings = torch.sum(crossed_embeddings, dim=-1, keepdim=True)
@@ -845,57 +705,49 @@ class CEN(nn.Module):
     """Field-attentive interaction network from FAT-DeepFFM (Wang et al., 2020)."""
 
     def __init__(self, embed_dim, num_field_crosses, reduction_ratio):
-        super().__init__()
-
+        super().__init__()        
+        
         # convolution weight (Eq.7 FAT-DeepFFM)
-        self.u = torch.nn.Parameter(
-            torch.rand(num_field_crosses, embed_dim), requires_grad=True
-        )
+        self.u = torch.nn.Parameter(torch.rand(num_field_crosses, embed_dim), requires_grad=True)
 
         # two FC layers that computes the field attention
-        self.mlp_att = MLP(
-            num_field_crosses,
-            dims=[num_field_crosses // reduction_ratio, num_field_crosses],
-            output_layer=False,
-            activation="relu",
-        )
-
-    def forward(self, em):
+        self.mlp_att = MLP(num_field_crosses, dims=[num_field_crosses//reduction_ratio, num_field_crosses], output_layer=False, activation="relu")
+        
+
+    def forward(self, em):  
         # compute descriptor vector (Eq.7 FAT-DeepFFM), output shape [batch_size, num_field_crosses]
         d = F.relu((self.u.squeeze(0) * em).sum(-1))
-
-        # compute field attention (Eq.9), output shape [batch_size, num_field_crosses]
-        s = self.mlp_att(d)
+        
+        # compute field attention (Eq.9), output shape [batch_size, num_field_crosses]    
+        s = self.mlp_att(d)                             
 
         # rescale original embedding with field attention (Eq.10), output shape [batch_size, num_field_crosses, embed_dim]
-        aem = s.unsqueeze(-1) * em
+        aem = s.unsqueeze(-1) * em                 
         return aem.flatten(start_dim=1)
 
 
 class MultiHeadSelfAttention(nn.Module):
     """Multi-head self-attention layer from AutoInt (Song et al., 2019)."""
-
+    
     def __init__(self, embedding_dim, num_heads=2, dropout=0.0, use_residual=True):
         super().__init__()
         if embedding_dim % num_heads != 0:
-            raise ValueError(
-                f"embedding_dim ({embedding_dim}) must be divisible by num_heads ({num_heads})"
-            )
-
+            raise ValueError(f"embedding_dim ({embedding_dim}) must be divisible by num_heads ({num_heads})")
+        
         self.embedding_dim = embedding_dim
         self.num_heads = num_heads
         self.head_dim = embedding_dim // num_heads
         self.use_residual = use_residual
-
+        
         self.W_Q = nn.Linear(embedding_dim, embedding_dim, bias=False)
         self.W_K = nn.Linear(embedding_dim, embedding_dim, bias=False)
         self.W_V = nn.Linear(embedding_dim, embedding_dim, bias=False)
-
+        
         if self.use_residual:
             self.W_Res = nn.Linear(embedding_dim, embedding_dim, bias=False)
-
+        
         self.dropout = nn.Dropout(dropout)
-
+        
     def forward(self, x):
         """
         Args:
@@ -904,47 +756,37 @@ class MultiHeadSelfAttention(nn.Module):
             output: [batch_size, num_fields, embedding_dim]
         """
         batch_size, num_fields, _ = x.shape
-
+        
         # Linear projections
         Q = self.W_Q(x)  # [batch_size, num_fields, embedding_dim]
         K = self.W_K(x)
         V = self.W_V(x)
-
+        
         # Split into multiple heads: [batch_size, num_heads, num_fields, head_dim]
-        Q = Q.view(batch_size, num_fields, self.num_heads, self.head_dim).transpose(
-            1, 2
-        )
-        K = K.view(batch_size, num_fields, self.num_heads, self.head_dim).transpose(
-            1, 2
-        )
-        V = V.view(batch_size, num_fields, self.num_heads, self.head_dim).transpose(
-            1, 2
-        )
-
+        Q = Q.view(batch_size, num_fields, self.num_heads, self.head_dim).transpose(1, 2)
+        K = K.view(batch_size, num_fields, self.num_heads, self.head_dim).transpose(1, 2)
+        V = V.view(batch_size, num_fields, self.num_heads, self.head_dim).transpose(1, 2)
+        
         # Attention scores
-        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.head_dim**0.5)
+        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.head_dim ** 0.5)
         attention_weights = F.softmax(scores, dim=-1)
         attention_weights = self.dropout(attention_weights)
-
+        
         # Apply attention to values
-        attention_output = torch.matmul(
-            attention_weights, V
-        )  # [batch_size, num_heads, num_fields, head_dim]
-
+        attention_output = torch.matmul(attention_weights, V)  # [batch_size, num_heads, num_fields, head_dim]
+        
         # Concatenate heads
         attention_output = attention_output.transpose(1, 2).contiguous()
-        attention_output = attention_output.view(
-            batch_size, num_fields, self.embedding_dim
-        )
-
+        attention_output = attention_output.view(batch_size, num_fields, self.embedding_dim)
+        
         # Residual connection
         if self.use_residual:
             output = attention_output + self.W_Res(x)
         else:
             output = attention_output
-
+            
         output = F.relu(output)
-
+        
         return output
 
 
@@ -953,31 +795,25 @@ class AttentionPoolingLayer(nn.Module):
     Attention pooling layer for DIN/DIEN
     Computes attention weights between query (candidate item) and keys (user behavior sequence)
     """
-
-    def __init__(
-        self,
-        embedding_dim,
-        hidden_units=[80, 40],
-        activation="sigmoid",
-        use_softmax=True,
-    ):
+    
+    def __init__(self, embedding_dim, hidden_units=[80, 40], activation='sigmoid', use_softmax=True):
         super().__init__()
         self.embedding_dim = embedding_dim
         self.use_softmax = use_softmax
-
+        
         # Build attention network
         # Input: [query, key, query-key, query*key] -> 4 * embedding_dim
         input_dim = 4 * embedding_dim
         layers = []
-
+        
         for hidden_unit in hidden_units:
             layers.append(nn.Linear(input_dim, hidden_unit))
             layers.append(activation_layer(activation))
             input_dim = hidden_unit
-
+        
         layers.append(nn.Linear(input_dim, 1))
         self.attention_net = nn.Sequential(*layers)
-
+    
     def forward(self, query, keys, keys_length=None, mask=None):
         """
         Args:
@@ -989,52 +825,48 @@ class AttentionPoolingLayer(nn.Module):
             output: [batch_size, embedding_dim] - attention pooled representation
         """
         batch_size, seq_len, emb_dim = keys.shape
-
+        
         # Expand query to match sequence length: [batch_size, seq_len, embedding_dim]
         query_expanded = query.unsqueeze(1).expand(-1, seq_len, -1)
-
+        
         # Compute attention features: [query, key, query-key, query*key]
-        attention_input = torch.cat(
-            [query_expanded, keys, query_expanded - keys, query_expanded * keys], dim=-1
-        )  # [batch_size, seq_len, 4*embedding_dim]
-
+        attention_input = torch.cat([
+            query_expanded,
+            keys,
+            query_expanded - keys,
+            query_expanded * keys
+        ], dim=-1)  # [batch_size, seq_len, 4*embedding_dim]
+        
         # Compute attention scores
-        attention_scores = self.attention_net(
-            attention_input
-        )  # [batch_size, seq_len, 1]
-
+        attention_scores = self.attention_net(attention_input)  # [batch_size, seq_len, 1]
+        
         # Apply mask if provided
         if mask is not None:
             attention_scores = attention_scores.masked_fill(mask == 0, -1e9)
-
+        
         # Apply softmax to get attention weights
         if self.use_softmax:
-            attention_weights = F.softmax(
-                attention_scores, dim=1
-            )  # [batch_size, seq_len, 1]
+            attention_weights = F.softmax(attention_scores, dim=1)  # [batch_size, seq_len, 1]
         else:
             attention_weights = attention_scores
-
+        
         # Weighted sum of keys
-        output = torch.sum(
-            attention_weights * keys, dim=1
-        )  # [batch_size, embedding_dim]
-
+        output = torch.sum(attention_weights * keys, dim=1)  # [batch_size, embedding_dim]
+        
         return output
 
 
 class DynamicGRU(nn.Module):
     """Dynamic GRU unit with auxiliary loss path from DIEN (Zhou et al., 2019)."""
-
     """
     GRU with dynamic routing for DIEN
     """
-
+    
     def __init__(self, input_size, hidden_size, bias=True):
         super().__init__()
         self.input_size = input_size
         self.hidden_size = hidden_size
-
+        
         # GRU parameters
         self.weight_ih = nn.Parameter(torch.randn(3 * hidden_size, input_size))
         self.weight_hh = nn.Parameter(torch.randn(3 * hidden_size, hidden_size))
@@ -1042,16 +874,16 @@ class DynamicGRU(nn.Module):
             self.bias_ih = nn.Parameter(torch.randn(3 * hidden_size))
             self.bias_hh = nn.Parameter(torch.randn(3 * hidden_size))
         else:
-            self.register_parameter("bias_ih", None)
-            self.register_parameter("bias_hh", None)
-
+            self.register_parameter('bias_ih', None)
+            self.register_parameter('bias_hh', None)
+        
         self.reset_parameters()
-
+    
     def reset_parameters(self):
         std = 1.0 / (self.hidden_size) ** 0.5
         for weight in self.parameters():
             weight.data.uniform_(-std, std)
-
+    
     def forward(self, x, att_scores=None):
         """
         Args:
@@ -1062,61 +894,60 @@ class DynamicGRU(nn.Module):
             hidden: [batch_size, hidden_size] - final hidden state
         """
         batch_size, seq_len, _ = x.shape
-
+        
         # Initialize hidden state
         h = torch.zeros(batch_size, self.hidden_size, device=x.device)
-
+        
         outputs = []
         for t in range(seq_len):
             x_t = x[:, t, :]  # [batch_size, input_size]
-
+            
             # GRU computation
             gi = F.linear(x_t, self.weight_ih, self.bias_ih)
             gh = F.linear(h, self.weight_hh, self.bias_hh)
             i_r, i_i, i_n = gi.chunk(3, 1)
             h_r, h_i, h_n = gh.chunk(3, 1)
-
+            
             resetgate = torch.sigmoid(i_r + h_r)
             inputgate = torch.sigmoid(i_i + h_i)
             newgate = torch.tanh(i_n + resetgate * h_n)
             h = newgate + inputgate * (h - newgate)
-
+            
             outputs.append(h.unsqueeze(1))
-
+        
         output = torch.cat(outputs, dim=1)  # [batch_size, seq_len, hidden_size]
-
+        
         return output, h
 
 
 class AUGRU(nn.Module):
     """Attention-aware GRU update gate used in DIEN (Zhou et al., 2019)."""
-
     """
     Attention-based GRU for DIEN
     Uses attention scores to weight the update of hidden states
     """
-
+    
     def __init__(self, input_size, hidden_size, bias=True):
         super().__init__()
         self.input_size = input_size
         self.hidden_size = hidden_size
-
+        
         self.weight_ih = nn.Parameter(torch.randn(3 * hidden_size, input_size))
         self.weight_hh = nn.Parameter(torch.randn(3 * hidden_size, hidden_size))
         if bias:
             self.bias_ih = nn.Parameter(torch.randn(3 * hidden_size))
             self.bias_hh = nn.Parameter(torch.randn(3 * hidden_size))
         else:
-            self.register_parameter("bias_ih", None)
-            self.register_parameter("bias_hh", None)
-
+            self.register_parameter('bias_ih', None)
+            self.register_parameter('bias_hh', None)
+        
         self.reset_parameters()
-
+    
     def reset_parameters(self):
         std = 1.0 / (self.hidden_size) ** 0.5
         for weight in self.parameters():
             weight.data.uniform_(-std, std)
-
+    
     def forward(self, x, att_scores):
         """
         Args:
@@ -1127,28 +958,28 @@ class AUGRU(nn.Module):
             hidden: [batch_size, hidden_size] - final hidden state
         """
         batch_size, seq_len, _ = x.shape
-
+        
         h = torch.zeros(batch_size, self.hidden_size, device=x.device)
-
+        
         outputs = []
         for t in range(seq_len):
             x_t = x[:, t, :]  # [batch_size, input_size]
             att_t = att_scores[:, t, :]  # [batch_size, 1]
-
+            
             gi = F.linear(x_t, self.weight_ih, self.bias_ih)
             gh = F.linear(h, self.weight_hh, self.bias_hh)
             i_r, i_i, i_n = gi.chunk(3, 1)
             h_r, h_i, h_n = gh.chunk(3, 1)
-
+            
             resetgate = torch.sigmoid(i_r + h_r)
             inputgate = torch.sigmoid(i_i + h_i)
             newgate = torch.tanh(i_n + resetgate * h_n)
-
+            
             # Use attention score to control update
             h = (1 - att_t) * h + att_t * newgate
-
+            
             outputs.append(h.unsqueeze(1))
-
+        
         output = torch.cat(outputs, dim=1)
-
-        return output, h
+        
+        return output, h