nextrec 0.4.25__py3-none-any.whl → 0.4.28__py3-none-any.whl

nextrec/models/multi_task/pepnet.py

@@ -1,11 +1,11 @@
 """
-Date: create on 09/11/2025
-Checkpoint: edit on 30/12/2025
+Date: create on 01/01/2026
+Checkpoint: edit on 01/01/2026
 Author: Yang Zhou, zyaztec@gmail.com
 Reference:
-[1] Yang et al. "PEPNet: Parameter and Embedding Personalized Network for Multi-Task Learning", 2021.
-[2] MMLRec-A-Unified-Multi-Task-and-Multi-Scenario-Learning-Benchmark-for-Recommendation:
-https://github.com/alipay/MMLRec-A-Unified-Multi-Task-and-Multi-Scenario-Learning-Benchmark-for-Recommendation/blob/main/model/pepnet.py
+- [1] Chang J, Zhang C, Hui Y, Leng D, Niu Y, Song Y, Gai K. PEPNet: Parameter and Embedding Personalized Network for Infusing with Personalized Prior Information. In: Proceedings of the 29th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining (KDD ’23), 2023.
+URL: https://arxiv.org/abs/2302.01115
+- [2] MMLRec-A-Unified-Multi-Task-and-Multi-Scenario-Learning-Benchmark-for-Recommendation: https://github.com/alipay/MMLRec-A-Unified-Multi-Task-and-Multi-Scenario-Learning-Benchmark-for-Recommendation/
 
 PEPNet (Parameter and Embedding Personalized Network) is a multi-task learning
 model that personalizes both input features and layer transformations with
@@ -58,12 +58,12 @@ from nextrec.basic.layers import EmbeddingLayer, GateMLP
 from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
 from nextrec.utils.model import select_features
-from nextrec.utils.types import ActivationName, TaskTypeName
+from nextrec.utils.types import TaskTypeName
 
 
-class PPNetBlock(nn.Module):
+class PPNet(nn.Module):
     """
-    PEPNet block with per-layer gates conditioned on task context.
+    PPNet: per-task tower with layer-wise gates conditioned on task context.
     """
 
     def __init__(
@@ -71,34 +71,42 @@ class PPNetBlock(nn.Module):
         input_dim: int,
         output_dim: int,
         gate_input_dim: int,
-        gate_hidden_dim: int | None,
-        hidden_units: list[int] | None = None,
-        hidden_activations: ActivationName | list[ActivationName] = "relu",
-        dropout_rates: float | list[float] = 0.0,
-        batch_norm: bool = False,
+        mlp_params: dict | None = None,
+        gate_mlp_params: dict | None = None,
         use_bias: bool = True,
-        gate_activation: ActivationName = "relu",
-        gate_dropout: float = 0.0,
-        gate_use_bn: bool = False,
     ) -> None:
         super().__init__()
-        hidden_units = hidden_units or []
+        mlp_params = mlp_params or {}
+        gate_mlp_params = gate_mlp_params or {}
 
-        if isinstance(dropout_rates, list):
-            if len(dropout_rates) != len(hidden_units):
+        mlp_params.setdefault("hidden_dims", [])
+        mlp_params.setdefault("activation", "relu")
+        mlp_params.setdefault("dropout", 0.0)
+        mlp_params.setdefault("norm_type", "none")
+
+        gate_mlp_params.setdefault("hidden_dim", None)
+        gate_mlp_params.setdefault("activation", "relu")
+        gate_mlp_params.setdefault("dropout", 0.0)
+        gate_mlp_params.setdefault("use_bn", False)
+
+        hidden_units = mlp_params["hidden_dims"]
+        norm_type = mlp_params["norm_type"]
+
+        if isinstance(mlp_params["dropout"], list):
+            if len(mlp_params["dropout"]) != len(hidden_units):
                 raise ValueError("dropout_rates length must match hidden_units length.")
-            dropout_list = dropout_rates
+            dropout_list = mlp_params["dropout"]
         else:
-            dropout_list = [dropout_rates] * len(hidden_units)
+            dropout_list = [mlp_params["dropout"]] * len(hidden_units)
 
-        if isinstance(hidden_activations, list):
-            if len(hidden_activations) != len(hidden_units):
+        if isinstance(mlp_params["activation"], list):
+            if len(mlp_params["activation"]) != len(hidden_units):
                 raise ValueError(
                     "hidden_activations length must match hidden_units length."
                 )
-            activation_list = hidden_activations
+            activation_list = mlp_params["activation"]
         else:
-            activation_list = [hidden_activations] * len(hidden_units)
+            activation_list = [mlp_params["activation"]] * len(hidden_units)
 
         self.gate_layers = nn.ModuleList()
         self.mlp_layers = nn.ModuleList()
@@ -108,7 +116,7 @@ class PPNetBlock(nn.Module):
             dense_layers: list[nn.Module] = [
                 nn.Linear(layer_units[idx], layer_units[idx + 1], bias=use_bias)
             ]
-            if batch_norm:
+            if norm_type == "batch_norm":
                 dense_layers.append(nn.BatchNorm1d(layer_units[idx + 1]))
             dense_layers.append(activation_layer(activation_list[idx]))
             if dropout_list[idx] > 0:
@@ -117,11 +125,11 @@ class PPNetBlock(nn.Module):
             self.gate_layers.append(
                 GateMLP(
                     input_dim=gate_input_dim,
-                    hidden_dim=gate_hidden_dim,
+                    hidden_dim=gate_mlp_params["hidden_dim"],
                     output_dim=layer_units[idx],
-                    activation=gate_activation,
-                    dropout=gate_dropout,
-                    use_bn=gate_use_bn,
+                    activation=gate_mlp_params["activation"],
+                    dropout=gate_mlp_params["dropout"],
+                    use_bn=gate_mlp_params["use_bn"],
                     scale_factor=2.0,
                 )
             )
@@ -130,11 +138,11 @@ class PPNetBlock(nn.Module):
         self.gate_layers.append(
             GateMLP(
                 input_dim=gate_input_dim,
-                hidden_dim=gate_hidden_dim,
+                hidden_dim=gate_mlp_params["hidden_dim"],
                 output_dim=layer_units[-1],
-                activation=gate_activation,
-                dropout=gate_dropout,
-                use_bn=gate_use_bn,
+                activation=gate_mlp_params["activation"],
+                dropout=gate_mlp_params["dropout"],
+                use_bn=gate_mlp_params["use_bn"],
                 scale_factor=1.0,
             )
         )
@@ -177,15 +185,9 @@ class PEPNet(BaseModel):
         sequence_features: list[SequenceFeature] | None = None,
         target: list[str] | str | None = None,
         task: TaskTypeName | list[TaskTypeName] | None = None,
-        dnn_hidden_units: list[int] | None = None,
-        dnn_activation: ActivationName = "relu",
-        dnn_dropout: float | list[float] = 0.0,
-        dnn_use_bn: bool = False,
-        feature_gate_hidden_dim: int = 128,
-        gate_hidden_dim: int | None = None,
-        gate_activation: ActivationName = "relu",
-        gate_dropout: float = 0.0,
-        gate_use_bn: bool = False,
+        mlp_params: dict | None = None,
+        feature_gate_mlp_params: dict | None = None,
+        gate_mlp_params: dict | None = None,
         domain_features: list[str] | str | None = None,
         user_features: list[str] | str | None = None,
         item_features: list[str] | str | None = None,
@@ -195,7 +197,24 @@ class PEPNet(BaseModel):
         dense_features = dense_features or []
         sparse_features = sparse_features or []
         sequence_features = sequence_features or []
-        dnn_hidden_units = dnn_hidden_units or [256, 128]
+        mlp_params = mlp_params or {}
+        feature_gate_mlp_params = feature_gate_mlp_params or {}
+        gate_mlp_params = gate_mlp_params or {}
+
+        mlp_params.setdefault("hidden_dims", [256, 128])
+        mlp_params.setdefault("activation", "relu")
+        mlp_params.setdefault("dropout", 0.0)
+        mlp_params.setdefault("norm_type", "none")
+
+        feature_gate_mlp_params.setdefault("hidden_dim", 128)
+        feature_gate_mlp_params.setdefault("activation", "relu")
+        feature_gate_mlp_params.setdefault("dropout", 0.0)
+        feature_gate_mlp_params.setdefault("use_bn", False)
+
+        gate_mlp_params.setdefault("hidden_dim", None)
+        gate_mlp_params.setdefault("activation", "relu")
+        gate_mlp_params.setdefault("dropout", 0.0)
+        gate_mlp_params.setdefault("use_bn", False)
 
         if target is None:
             target = []
@@ -203,24 +222,13 @@ class PEPNet(BaseModel):
             target = [target]
 
         self.nums_task = len(target) if target else 1
-        resolved_task = task
-        if resolved_task is None:
-            resolved_task = self.default_task
-        elif isinstance(resolved_task, str):
-            resolved_task = [resolved_task] * self.nums_task
-        elif len(resolved_task) == 1 and self.nums_task > 1:
-            resolved_task = resolved_task * self.nums_task
-        elif len(resolved_task) != self.nums_task:
-            raise ValueError(
-                f"Length of task ({len(resolved_task)}) must match number of targets ({self.nums_task})."
-            )
 
         super().__init__(
             dense_features=dense_features,
             sparse_features=sparse_features,
             sequence_features=sequence_features,
             target=target,
-            task=resolved_task,
+            task=task,
             **kwargs,
         )
 
@@ -266,30 +274,27 @@ class PEPNet(BaseModel):
         )
         task_dim = domain_dim + user_dim + item_dim
 
-        self.feature_gate = GateMLP(
+        # EPNet: shared feature-level gate (paper's EPNet).
+        self.epnet = GateMLP(
             input_dim=input_dim + domain_dim,
-            hidden_dim=feature_gate_hidden_dim,
+            hidden_dim=feature_gate_mlp_params["hidden_dim"],
             output_dim=input_dim,
-            activation=gate_activation,
-            dropout=gate_dropout,
-            use_bn=gate_use_bn,
+            activation=feature_gate_mlp_params["activation"],
+            dropout=feature_gate_mlp_params["dropout"],
+            use_bn=feature_gate_mlp_params["use_bn"],
+            scale_factor=2.0,
         )
 
-        self.ppn_blocks = nn.ModuleList(
+        # PPNet: per-task gated towers (paper's PPNet).
+        self.ppnet_blocks = nn.ModuleList(
             [
-                PPNetBlock(
+                PPNet(
                     input_dim=input_dim,
                     output_dim=1,
                     gate_input_dim=input_dim + task_dim,
-                    gate_hidden_dim=gate_hidden_dim,
-                    hidden_units=dnn_hidden_units,
-                    hidden_activations=dnn_activation,
-                    dropout_rates=dnn_dropout,
-                    batch_norm=dnn_use_bn,
+                    mlp_params=mlp_params,
+                    gate_mlp_params=gate_mlp_params,
                     use_bias=use_bias,
-                    gate_activation=gate_activation,
-                    gate_dropout=gate_dropout,
-                    gate_use_bn=gate_use_bn,
                 )
                 for _ in range(self.nums_task)
             ]
@@ -298,9 +303,9 @@ class PEPNet(BaseModel):
         self.prediction_layer = TaskHead(
             task_type=self.task, task_dims=[1] * self.nums_task
         )
-        self.grad_norm_shared_modules = ["embedding", "feature_gate"]
+        self.grad_norm_shared_modules = ["embedding", "epnet"]
         self.register_regularization_weights(
-            embedding_attr="embedding", include_modules=["feature_gate", "ppn_blocks"]
+            embedding_attr="embedding", include_modules=["epnet", "ppnet_blocks"]
         )
 
     def forward(self, x: dict[str, torch.Tensor]) -> torch.Tensor:
@@ -325,11 +330,11 @@ class PEPNet(BaseModel):
         task_sf_emb = torch.cat(task_parts, dim=-1)
 
         gate_input = torch.cat([dnn_input.detach(), domain_emb], dim=-1)
-        dnn_input = self.feature_gate(gate_input) * dnn_input
+        dnn_input = self.epnet(gate_input) * dnn_input
 
         task_logits = []
-        for block in self.ppn_blocks:
-            task_logits.append(block(o_ep=dnn_input, o_prior=task_sf_emb))
+        for block in self.ppnet_blocks:
+             task_logits.append(block(o_ep=dnn_input, o_prior=task_sf_emb))
 
         y = torch.cat(task_logits, dim=1)
         return self.prediction_layer(y)

nextrec/models/multi_task/ple.py

@@ -3,9 +3,8 @@ Date: create on 09/11/2025
 Checkpoint: edit on 23/12/2025
 Author: Yang Zhou,zyaztec@gmail.com
 Reference:
-[1] Tang H, Liu J, Zhao M, et al. Progressive layered extraction (PLE): A novel
-multi-task learning (MTL) model for personalized recommendations[C]//RecSys. 2020: 269-278.
-(https://dl.acm.org/doi/10.1145/3383313.3412236)
+- [1] Tang H, Liu J, Zhao M, Gong X. Progressive Layered Extraction (PLE): A Novel Multi-Task Learning (MTL) Model for Personalized Recommendations. In: Proceedings of the 14th ACM Conference on Recommender Systems (RecSys ’20), 2020, pp. 269–278.
+URL: https://dl.acm.org/doi/10.1145/3383313.3412236
 
 Progressive Layered Extraction (PLE) advances multi-task learning by stacking CGC
 (Customized Gate Control) blocks that mix shared and task-specific experts. Each
@@ -67,18 +66,21 @@ class CGCLayer(nn.Module):
         nums_task: int,
         num_shared_experts: int,
         num_specific_experts: int,
-        shared_expert_params: dict,
-        specific_expert_params: dict | list[dict],
+        shared_expert_mlp_params: dict,
+        specific_expert_mlp_params: list[dict],
     ):
         super().__init__()
         if nums_task < 1:
             raise ValueError("nums_task must be >= 1")
 
-        specific_params_list = self.normalize_specific_params(
-            specific_expert_params, nums_task
-        )
+        if len(specific_expert_mlp_params) != nums_task:
+            raise ValueError(
+                "Length of specific_expert_mlp_params "
+                f"({len(specific_expert_mlp_params)}) must match number of tasks ({nums_task})."
+            )
+        specific_params_list = [params.copy() for params in specific_expert_mlp_params]
 
-        self.output_dim = get_mlp_output_dim(shared_expert_params, input_dim)
+        self.output_dim = get_mlp_output_dim(shared_expert_mlp_params, input_dim)
         specific_dims = [
             get_mlp_output_dim(params, input_dim) for params in specific_params_list
         ]
@@ -94,7 +96,7 @@ class CGCLayer(nn.Module):
                 MLP(
                     input_dim=input_dim,
                     output_dim=None,
-                    **shared_expert_params,
+                    **shared_expert_mlp_params,
                 )
                 for _ in range(num_shared_experts)
             ]
@@ -166,18 +168,6 @@ class CGCLayer(nn.Module):
 
         return new_task_fea, new_shared
 
-    @staticmethod
-    def normalize_specific_params(
-        params: dict | list[dict], nums_task: int
-    ) -> list[dict]:
-        if isinstance(params, list):
-            if len(params) != nums_task:
-                raise ValueError(
-                    f"Length of specific_expert_params ({len(params)}) must match nums_task ({nums_task})."
-                )
-            return [p.copy() for p in params]
-        return [params.copy() for _ in range(nums_task)]
-
 
 class PLE(BaseModel):
     """
@@ -205,12 +195,12 @@ class PLE(BaseModel):
         dense_features: list[DenseFeature] | None = None,
         sparse_features: list[SparseFeature] | None = None,
         sequence_features: list[SequenceFeature] | None = None,
-        shared_expert_params: dict | None = None,
-        specific_expert_params: dict | list[dict] | None = None,
+        shared_expert_mlp_params: dict | None = None,
+        specific_expert_mlp_params: list[dict] | None = None,
         num_shared_experts: int = 2,
         num_specific_experts: int = 2,
         num_levels: int = 2,
-        tower_params_list: list[dict] | None = None,
+        tower_mlp_params_list: list[dict] | None = None,
         target: list[str] | None = None,
         task: str | list[str] | None = None,
         **kwargs,
@@ -218,24 +208,19 @@ class PLE(BaseModel):
 
         self.nums_task = len(target) if target is not None else 1
 
-        resolved_task = task
-        if resolved_task is None:
-            resolved_task = self.default_task
-        elif isinstance(resolved_task, str):
-            resolved_task = [resolved_task] * self.nums_task
-        elif len(resolved_task) == 1 and self.nums_task > 1:
-            resolved_task = resolved_task * self.nums_task
-        elif len(resolved_task) != self.nums_task:
+        shared_expert_mlp_params = shared_expert_mlp_params or {}
+        if specific_expert_mlp_params is None:
             raise ValueError(
-                f"Length of task ({len(resolved_task)}) must match number of targets ({self.nums_task})."
+                "specific_expert_mlp_params must be a list of dicts, one per task."
             )
+        tower_mlp_params_list = tower_mlp_params_list or []
 
         super(PLE, self).__init__(
             dense_features=dense_features,
             sparse_features=sparse_features,
             sequence_features=sequence_features,
             target=target,
-            task=resolved_task,
+            task=task,
             **kwargs,
         )
 
@@ -245,9 +230,10 @@ class PLE(BaseModel):
         self.num_specific_experts = num_specific_experts
         self.num_levels = num_levels
 
-        if len(tower_params_list) != self.nums_task:
+        if len(tower_mlp_params_list) != self.nums_task:
             raise ValueError(
-                f"Number of tower params ({len(tower_params_list)}) must match number of tasks ({self.nums_task})"
+                "Number of tower mlp params "
+                f"({len(tower_mlp_params_list)}) must match number of tasks ({self.nums_task})"
             )
         # Embedding layer
         self.embedding = EmbeddingLayer(features=self.all_features)
@@ -260,10 +246,10 @@ class PLE(BaseModel):
 
         # Get expert output dimension
         if (
-            "hidden_dims" in shared_expert_params
-            and len(shared_expert_params["hidden_dims"]) > 0
+            "hidden_dims" in shared_expert_mlp_params
+            and len(shared_expert_mlp_params["hidden_dims"]) > 0
         ):
-            expert_output_dim = shared_expert_params["hidden_dims"][-1]
+            expert_output_dim = shared_expert_mlp_params["hidden_dims"][-1]
         else:
             expert_output_dim = input_dim
 
@@ -276,8 +262,8 @@ class PLE(BaseModel):
                 nums_task=self.nums_task,
                 num_shared_experts=num_shared_experts,
                 num_specific_experts=num_specific_experts,
-                shared_expert_params=shared_expert_params,
-                specific_expert_params=specific_expert_params,
+                shared_expert_mlp_params=shared_expert_mlp_params,
+                specific_expert_mlp_params=specific_expert_mlp_params,
             )
             self.cgc_layers.append(cgc_layer)
             expert_output_dim = cgc_layer.output_dim
@@ -285,8 +271,8 @@ class PLE(BaseModel):
 
         # Task-specific towers
         self.towers = nn.ModuleList()
-        for tower_params in tower_params_list:
-            tower = MLP(input_dim=expert_output_dim, output_dim=1, **tower_params)
+        for tower_mlp_params in tower_mlp_params_list:
+            tower = MLP(input_dim=expert_output_dim, output_dim=1, **tower_mlp_params)
             self.towers.append(tower)
         self.prediction_layer = TaskHead(
             task_type=self.task, task_dims=[1] * self.nums_task

nextrec/models/multi_task/poso.py

@@ -3,7 +3,8 @@ Date: create on 28/11/2025
 Checkpoint: edit on 23/12/2025
 Author: Yang Zhou,zyaztec@gmail.com
 Reference:
-[1] Wang et al. "POSO: Personalized Cold Start Modules for Large-scale Recommender Systems", 2021.
+- [1] Dai S, Lin H, Zhao Z, Lin J, Wu H, Wang Z, Yang S, Liu J. POSO: Personalized Cold Start Modules for Large-scale Recommender Systems. arXiv preprint arXiv:2108.04690, 2021.
+URL: https://arxiv.org/abs/2108.04690
 
 POSO (Personalized cOld-start mOdules) augments backbone recommenders by injecting a
 personalized cold-start vector `pc` that gates hidden units layer by layer. Each fully
@@ -49,6 +50,7 @@ from nextrec.basic.layers import MLP, EmbeddingLayer
 from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
 from nextrec.utils.model import select_features
+from nextrec.utils.types import TaskTypeName
 
 
 class POSOGate(nn.Module):
@@ -306,9 +308,9 @@ class POSO(BaseModel):
         pc_dense_features: list[str] | None,
         pc_sparse_features: list[str] | None,
         pc_sequence_features: list[str] | None,
-        tower_params_list: list[dict],
+        tower_mlp_params_list: list[dict],
         target: list[str] | None = None,
-        task: str | list[str] = "binary",
+        task: TaskTypeName | list[TaskTypeName] | None = None,
         architecture: Literal["mlp", "mmoe"] = "mlp",
         # POSO gating defaults
         gate_hidden_dim: int = 32,
@@ -327,22 +329,10 @@ class POSO(BaseModel):
     ):
         self.nums_task = len(target)
 
-        # Normalize task to match nums_task
-        resolved_task = task
-        if resolved_task is None:
-            resolved_task = self.default_task
-        elif isinstance(resolved_task, str):
-            resolved_task = [resolved_task] * self.nums_task
-        elif len(resolved_task) == 1 and self.nums_task > 1:
-            resolved_task = resolved_task * self.nums_task
-        elif len(resolved_task) != self.nums_task:
+        if len(tower_mlp_params_list) != self.nums_task:
             raise ValueError(
-                f"Length of task ({len(resolved_task)}) must match number of targets ({self.nums_task})."
-            )
-
-        if len(tower_params_list) != self.nums_task:
-            raise ValueError(
-                f"Number of tower params ({len(tower_params_list)}) must match number of tasks ({self.nums_task})"
+                "Number of tower mlp params "
+                f"({len(tower_mlp_params_list)}) must match number of tasks ({self.nums_task})"
             )
 
         super().__init__(
@@ -350,7 +340,7 @@ class POSO(BaseModel):
             sparse_features=sparse_features,
             sequence_features=sequence_features,
             target=target,
-            task=resolved_task,
+            task=task,
             **kwargs,
         )
 
@@ -415,11 +405,12 @@ class POSO(BaseModel):
         if self.architecture == "mlp":
             self.towers = nn.ModuleList()
             self.tower_heads = nn.ModuleList()
-            for tower_params in tower_params_list:
+            for tower_params in tower_mlp_params_list:
                 hidden_dims = tower_params.get("hidden_dims")
                 if not hidden_dims:
                     raise ValueError(
-                        "tower_params must include a non-empty 'hidden_dims' list for POSO-MLP towers."
+                        "tower_mlp_params_list must include a non-empty 'hidden_dims' "
+                        "list for POSO-MLP towers."
                     )
                 dropout = tower_params.get("dropout", 0.0)
                 tower = POSOMLP(
@@ -463,7 +454,7 @@ class POSO(BaseModel):
                         output_dim=1,
                         **tower_params,
                     )
-                    for tower_params in tower_params_list
+                    for tower_params in tower_mlp_params_list
                 ]
             )
             self.tower_heads = None

nextrec/models/multi_task/share_bottom.py

@@ -2,9 +2,6 @@
 Date: create on 09/11/2025
 Checkpoint: edit on 23/12/2025
 Author: Yang Zhou,zyaztec@gmail.com
-Reference:
-[1] Caruana R. Multitask learning[J]. Machine Learning, 1997, 28: 41-75.
-(https://link.springer.com/article/10.1023/A:1007379606734)
 
 Shared-Bottom is the classic hard-parameter-sharing baseline for multi-task learning.
 All tasks share a common bottom network to learn general representations, and each
@@ -65,8 +62,8 @@ class ShareBottom(BaseModel):
         dense_features: list[DenseFeature],
         sparse_features: list[SparseFeature],
         sequence_features: list[SequenceFeature],
-        bottom_params: dict,
-        tower_params_list: list[dict],
+        bottom_mlp_params: dict,
+        tower_mlp_params_list: list[dict],
         target: list[str],
         task: str | list[str] | None = None,
         **kwargs,
@@ -74,32 +71,21 @@ class ShareBottom(BaseModel):
 
         self.nums_task = len(target)
 
-        resolved_task = task
-        if resolved_task is None:
-            resolved_task = self.default_task
-        elif isinstance(resolved_task, str):
-            resolved_task = [resolved_task] * self.nums_task
-        elif len(resolved_task) == 1 and self.nums_task > 1:
-            resolved_task = resolved_task * self.nums_task
-        elif len(resolved_task) != self.nums_task:
-            raise ValueError(
-                f"Length of task ({len(resolved_task)}) must match number of targets ({self.nums_task})."
-            )
-
         super(ShareBottom, self).__init__(
             dense_features=dense_features,
             sparse_features=sparse_features,
             sequence_features=sequence_features,
             target=target,
-            task=resolved_task,
+            task=task,
             **kwargs,
         )
 
         # Number of tasks
         self.nums_task = len(target)
-        if len(tower_params_list) != self.nums_task:
+        if len(tower_mlp_params_list) != self.nums_task:
             raise ValueError(
-                f"Number of tower params ({len(tower_params_list)}) must match number of tasks ({self.nums_task})"
+                "Number of tower mlp params "
+                f"({len(tower_mlp_params_list)}) must match number of tasks ({self.nums_task})"
             )
         # Embedding layer
         self.embedding = EmbeddingLayer(features=self.all_features)
@@ -110,19 +96,22 @@ class ShareBottom(BaseModel):
         # input_dim = emb_dim_total + dense_input_dim
 
         # Shared bottom network
-        self.bottom = MLP(input_dim=input_dim, output_dim=None, **bottom_params)
+        self.bottom = MLP(input_dim=input_dim, output_dim=None, **bottom_mlp_params)
         self.grad_norm_shared_modules = ["embedding", "bottom"]
 
         # Get bottom output dimension
-        if "hidden_dims" in bottom_params and len(bottom_params["hidden_dims"]) > 0:
-            bottom_output_dim = bottom_params["hidden_dims"][-1]
+        if (
+            "hidden_dims" in bottom_mlp_params
+            and len(bottom_mlp_params["hidden_dims"]) > 0
+        ):
+            bottom_output_dim = bottom_mlp_params["hidden_dims"][-1]
         else:
             bottom_output_dim = input_dim
 
         # Task-specific towers
         self.towers = nn.ModuleList()
-        for tower_params in tower_params_list:
-            tower = MLP(input_dim=bottom_output_dim, output_dim=1, **tower_params)
+        for tower_mlp_params in tower_mlp_params_list:
+            tower = MLP(input_dim=bottom_output_dim, output_dim=1, **tower_mlp_params)
             self.towers.append(tower)
         self.prediction_layer = TaskHead(
             task_type=self.task, task_dims=[1] * self.nums_task

nextrec/models/ranking/afm.py

@@ -3,8 +3,8 @@ Date: create on 09/11/2025
 Checkpoint: edit on 23/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 Reference:
-[1] Xiao J, Ye H, He X, et al. Attentional factorization machines: Learning the weight of
-feature interactions via attention networks[C]//IJCAI. 2017: 3119-3125.
+- [1] Xiao J, Ye H, He X, et al. Attentional Factorization Machines: Learning the Weight of Feature Interactions via Attention Networks
+URL: https://arxiv.org/abs/1708.04617
 
 Attentional Factorization Machine (AFM) builds on FM by learning an importance
 weight for every second-order interaction instead of treating all pairs equally.

nextrec/models/ranking/autoint.py

@@ -3,10 +3,8 @@ Date: create on 09/11/2025
 Checkpoint: edit on 23/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 Reference:
-[1] Song W, Shi C, Xiao Z, et al. Autoint: Automatic feature interaction learning via
-self-attentive neural networks[C]//Proceedings of the 28th ACM international conference
-on information and knowledge management. 2019: 1161-1170.
-(https://arxiv.org/abs/1810.11921)
+- [1] Song W, Shi C, Xiao Z, et al. AutoInt: Automatic feature interaction learning via self-attentive neural networks. In: Proceedings of the 28th ACM International Conference on Information and Knowledge Management (CIKM ’19), 2019, pp. 1161–1170.
+URL: https://arxiv.org/abs/1810.11921
 
 AutoInt is a CTR prediction model that leverages multi-head self-attention
 to automatically learn high-order feature interactions in an explicit and

nextrec/models/ranking/dcn.py

@@ -3,9 +3,8 @@ Date: create on 09/11/2025
 Checkpoint: edit on 23/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 Reference:
-[1] Wang R, Fu B, Fu G, et al. Deep & cross network for ad click predictions[C]
-//Proceedings of the ADKDD'17. 2017: 1-7.
-(https://arxiv.org/abs/1708.05123)
+- [1] Wang R, Fu B, Fu G, et al. Deep & cross network for ad click predictions[C] //Proceedings of the ADKDD'17. 2017: 1-7.
+URL: https://arxiv.org/abs/1708.05123
 
 Deep & Cross Network (DCN) mixes explicit polynomial feature crosses with a deep
 MLP branch to capture both low-order and high-order interactions for CTR-style

nextrec/models/ranking/dcn_v2.py

@@ -3,9 +3,8 @@ Date: create on 09/11/2025
 Checkpoint: edit on 23/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 Reference:
-[1] R. Wang et al. DCN V2: Improved Deep & Cross Network and Practical Lessons for
-Web-scale Learning to Rank Systems. KDD 2021.
-(https://arxiv.org/abs/2008.13535)
+- [1] R. Wang et al. DCN V2: Improved Deep & Cross Network and Practical Lessons for Web-scale Learning to Rank Systems. KDD 2021.
+URL: https://arxiv.org/abs/2008.13535
 
 DCN v2 enhances the original Deep & Cross Network by replacing the scalar cross
 weights with vector-wise (matrix) parameters and a Mixture-of-Low-Rank-Experts