nextrec 0.3.6__py3-none-any.whl → 0.4.2__py3-none-any.whl

nextrec/models/ranking/dcn_v2.py

@@ -5,25 +5,30 @@ Date: create on 09/11/2025
 import torch
 import torch.nn as nn
 
-from nextrec.basic.model import BaseModel
-from nextrec.basic.layers import EmbeddingLayer, MLP, PredictionLayer
-from nextrec.basic.features import DenseFeature, SparseFeature, SequenceFeature
 
 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)."""
 
@@ -33,18 +38,46 @@ 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)
@@ -59,7 +92,9 @@ 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)
@@ -67,7 +102,9 @@ 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
@@ -75,10 +112,14 @@ 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
+        return x_l

nextrec/models/ranking/deepfm.py

@@ -43,57 +43,60 @@ embedding，无需手工构造交叉特征即可端到端训练，常用于 CTR/
 - CTR/CVR 任务的常用强基线
 """
 
-import torch
 import torch.nn as nn
 
 from nextrec.basic.model import BaseModel
 from nextrec.basic.layers import FM, LR, EmbeddingLayer, MLP, PredictionLayer
 from nextrec.basic.features import DenseFeature, SparseFeature, SequenceFeature
 
+
 class DeepFM(BaseModel):
     @property
     def model_name(self):
         return "DeepFM"
 
     @property
-    def task_type(self):
+    def default_task(self):
         return "binary"
-    
-    def __init__(self,
-                 dense_features: list[DenseFeature]|list = [],
-                 sparse_features: list[SparseFeature]|list = [],
-                 sequence_features: list[SequenceFeature]|list = [],
-                 mlp_params: dict = {},
-                 target: list[str]|str = [],
-                 optimizer: str = "adam",
-                 optimizer_params: dict = {},
-                 loss: str | nn.Module | None = "bce",
-                 loss_params: dict | list[dict] | None = None,
-                 device: str = 'cpu',
-                 embedding_l1_reg=1e-6,
-                 dense_l1_reg=1e-5,
-                 embedding_l2_reg=1e-5,
-                 dense_l2_reg=1e-4, **kwargs):
-        
+
+    def __init__(
+        self,
+        dense_features: list[DenseFeature] | list = [],
+        sparse_features: list[SparseFeature] | list = [],
+        sequence_features: list[SequenceFeature] | list = [],
+        mlp_params: dict = {},
+        target: list[str] | str = [],
+        task: str | list[str] | None = None,
+        optimizer: str = "adam",
+        optimizer_params: dict = {},
+        loss: str | nn.Module | None = "bce",
+        loss_params: dict | list[dict] | None = None,
+        device: str = "cpu",
+        embedding_l1_reg=1e-6,
+        dense_l1_reg=1e-5,
+        embedding_l2_reg=1e-5,
+        dense_l2_reg=1e-4,
+        **kwargs,
+    ):
+
         super(DeepFM, self).__init__(
             dense_features=dense_features,
             sparse_features=sparse_features,
             sequence_features=sequence_features,
             target=target,
-            task=self.task_type,
+            task=task or self.default_task,
             device=device,
             embedding_l1_reg=embedding_l1_reg,
             dense_l1_reg=dense_l1_reg,
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
-            early_stop_patience=20,
-            **kwargs
+            **kwargs,
         )
 
         self.loss = loss
         if self.loss is None:
             self.loss = "bce"
-            
+
         self.fm_features = sparse_features + sequence_features
         self.deep_features = dense_features + sparse_features + sequence_features
         self.embedding = EmbeddingLayer(features=self.deep_features)
@@ -104,11 +107,18 @@ class DeepFM(BaseModel):
         self.linear = LR(fm_emb_dim_total)
         self.fm = FM(reduce_sum=True)
         self.mlp = MLP(input_dim=mlp_input_dim, **mlp_params)
-        self.prediction_layer = PredictionLayer(task_type=self.task_type)
+        self.prediction_layer = PredictionLayer(task_type=self.default_task)
 
         # Register regularization weights
-        self.register_regularization_weights(embedding_attr='embedding', include_modules=['linear', 'mlp'])
-        self.compile(optimizer=optimizer, optimizer_params=optimizer_params, loss=loss, loss_params=loss_params)
+        self.register_regularization_weights(
+            embedding_attr="embedding", include_modules=["linear", "mlp"]
+        )
+        self.compile(
+            optimizer=optimizer,
+            optimizer_params=optimizer_params,
+            loss=loss,
+            loss_params=loss_params,
+        )
 
     def forward(self, x):
         input_deep = self.embedding(x=x, features=self.deep_features, squeeze_dim=True)