nextrec 0.4.22__py3-none-any.whl → 0.4.24__py3-none-any.whl

nextrec/models/multi_task/ple.py

@@ -93,7 +93,7 @@ class CGCLayer(nn.Module):
             [
                 MLP(
                     input_dim=input_dim,
-                    output_layer=False,
+                    output_dim=None,
                     **shared_expert_params,
                 )
                 for _ in range(num_shared_experts)
@@ -105,7 +105,7 @@ class CGCLayer(nn.Module):
                 [
                     MLP(
                         input_dim=input_dim,
-                        output_layer=False,
+                        output_dim=None,
                         **params,
                     )
                     for _ in range(num_specific_experts)
@@ -202,29 +202,21 @@ class PLE(BaseModel):
 
     def __init__(
         self,
-        dense_features: list[DenseFeature],
-        sparse_features: list[SparseFeature],
-        sequence_features: list[SequenceFeature],
-        shared_expert_params: dict,
-        specific_expert_params: dict | list[dict],
-        num_shared_experts: int,
-        num_specific_experts: int,
-        num_levels: int,
-        tower_params_list: list[dict],
-        target: list[str],
+        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,
+        num_shared_experts: int = 2,
+        num_specific_experts: int = 2,
+        num_levels: int = 2,
+        tower_params_list: list[dict] | None = None,
+        target: list[str] | None = None,
         task: str | list[str] | None = None,
-        optimizer: str = "adam",
-        optimizer_params: dict | None = None,
-        loss: str | nn.Module | list[str | nn.Module] | None = "bce",
-        loss_params: dict | list[dict] | None = None,
-        embedding_l1_reg=0.0,
-        dense_l1_reg=0.0,
-        embedding_l2_reg=0.0,
-        dense_l2_reg=0.0,
         **kwargs,
     ):
 
-        self.nums_task = len(target)
+        self.nums_task = len(target) if target is not None else 1
 
         resolved_task = task
         if resolved_task is None:
@@ -244,23 +236,15 @@ class PLE(BaseModel):
             sequence_features=sequence_features,
             target=target,
             task=resolved_task,
-            embedding_l1_reg=embedding_l1_reg,
-            dense_l1_reg=dense_l1_reg,
-            embedding_l2_reg=embedding_l2_reg,
-            dense_l2_reg=dense_l2_reg,
             **kwargs,
         )
 
-        self.loss = loss
-        if self.loss is None:
-            self.loss = "bce"
         # Number of tasks, experts, and levels
         self.nums_task = len(target)
         self.num_shared_experts = num_shared_experts
         self.num_specific_experts = num_specific_experts
         self.num_levels = num_levels
-        if optimizer_params is None:
-            optimizer_params = {}
+
         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})"
@@ -275,8 +259,11 @@ class PLE(BaseModel):
         # input_dim = emb_dim_total + dense_input_dim
 
         # Get expert output dimension
-        if "dims" in shared_expert_params and len(shared_expert_params["dims"]) > 0:
-            expert_output_dim = shared_expert_params["dims"][-1]
+        if (
+            "hidden_dims" in shared_expert_params
+            and len(shared_expert_params["hidden_dims"]) > 0
+        ):
+            expert_output_dim = shared_expert_params["hidden_dims"][-1]
         else:
             expert_output_dim = input_dim
 
@@ -299,21 +286,15 @@ 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_layer=True, **tower_params)
+            tower = MLP(input_dim=expert_output_dim, output_dim=1, **tower_params)
             self.towers.append(tower)
         self.prediction_layer = TaskHead(
-            task_type=self.default_task, task_dims=[1] * self.nums_task
+            task_type=self.task, task_dims=[1] * self.nums_task
         )
         # Register regularization weights
         self.register_regularization_weights(
             embedding_attr="embedding", include_modules=["cgc_layers", "towers"]
         )
-        self.compile(
-            optimizer=optimizer,
-            optimizer_params=optimizer_params,
-            loss=self.loss,
-            loss_params=loss_params,
-        )
 
     def forward(self, x):
         # Get all embeddings and flatten

nextrec/models/multi_task/poso.py

@@ -42,7 +42,7 @@ from __future__ import annotations
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-
+from typing import Literal
 from nextrec.basic.activation import activation_layer
 from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
 from nextrec.basic.layers import MLP, EmbeddingLayer
@@ -127,7 +127,7 @@ class POSOMLP(nn.Module):
     """
     POSO-enhanced MLP that stacks multiple POSOFC layers.
 
-    dims: e.g., [256, 128, 64] means
+    hidden_dims: e.g., [256, 128, 64] means
         in_dim -> 256 -> 128 -> 64
     Each layer has its own gate g_l(pc) following Eq. (11).
     """
@@ -136,7 +136,7 @@ class POSOMLP(nn.Module):
         self,
         input_dim: int,
         pc_dim: int,
-        dims: list[int],
+        hidden_dims: list[int],
         gate_hidden_dim: int = 32,
         scale_factor: float = 2.0,
         activation: str = "relu",
@@ -147,7 +147,7 @@ class POSOMLP(nn.Module):
 
         layers = []
         in_dim = input_dim
-        for out_dim in dims:
+        for out_dim in hidden_dims:
             layers.append(
                 POSOFC(
                     in_dim=in_dim,
@@ -213,8 +213,8 @@ class POSOMMoE(nn.Module):
             [
                 MLP(
                     input_dim=input_dim,
-                    output_layer=False,
-                    dims=expert_hidden_dims,
+                    output_dim=None,
+                    hidden_dims=expert_hidden_dims,
                     activation=activation,
                     dropout=expert_dropout,
                 )
@@ -307,9 +307,9 @@ class POSO(BaseModel):
         pc_sparse_features: list[str] | None,
         pc_sequence_features: list[str] | None,
         tower_params_list: list[dict],
-        target: list[str],
+        target: list[str] | None = None,
         task: str | list[str] = "binary",
-        architecture: str = "mlp",
+        architecture: Literal["mlp", "mmoe"] = "mlp",
         # POSO gating defaults
         gate_hidden_dim: int = 32,
         gate_scale_factor: float = 2.0,
@@ -323,14 +323,6 @@ class POSO(BaseModel):
         expert_gate_hidden_dim: int = 32,
         expert_gate_scale_factor: float = 2.0,
         gate_use_softmax: bool = True,
-        optimizer: str = "adam",
-        optimizer_params: dict | None = None,
-        loss: str | nn.Module | list[str | nn.Module] | None = "bce",
-        loss_params: dict | list[dict] | None = None,
-        embedding_l1_reg=0.0,
-        dense_l1_reg=0.0,
-        embedding_l2_reg=0.0,
-        dense_l2_reg=0.0,
         **kwargs,
     ):
         self.nums_task = len(target)
@@ -359,10 +351,6 @@ class POSO(BaseModel):
             sequence_features=sequence_features,
             target=target,
             task=resolved_task,
-            embedding_l1_reg=embedding_l1_reg,
-            dense_l1_reg=dense_l1_reg,
-            embedding_l2_reg=embedding_l2_reg,
-            dense_l2_reg=dense_l2_reg,
             **kwargs,
         )
 
@@ -373,12 +361,6 @@ class POSO(BaseModel):
         self.pc_sparse_feature_names = list(pc_sparse_features or [])
         self.pc_sequence_feature_names = list(pc_sequence_features or [])
 
-        if loss is None:
-            self.loss = "bce"
-        self.loss = loss
-
-        optimizer_params = optimizer_params or {}
-
         self.main_dense_features = select_features(
             self.dense_features, self.main_dense_feature_names, "main_dense_features"
         )
@@ -434,16 +416,16 @@ class POSO(BaseModel):
             self.towers = nn.ModuleList()
             self.tower_heads = nn.ModuleList()
             for tower_params in tower_params_list:
-                dims = tower_params.get("dims")
-                if not dims:
+                hidden_dims = tower_params.get("hidden_dims")
+                if not hidden_dims:
                     raise ValueError(
-                        "tower_params must include a non-empty 'dims' list for POSO-MLP towers."
+                        "tower_params must include a non-empty 'hidden_dims' list for POSO-MLP towers."
                     )
                 dropout = tower_params.get("dropout", 0.0)
                 tower = POSOMLP(
                     input_dim=self.main_input_dim,
                     pc_dim=self.pc_input_dim,
-                    dims=dims,
+                    hidden_dims=hidden_dims,
                     gate_hidden_dim=tower_params.get(
                         "gate_hidden_dim", gate_hidden_dim
                     ),
@@ -453,7 +435,9 @@ class POSO(BaseModel):
                     dropout=dropout,
                 )
                 self.towers.append(tower)
-                tower_output_dim = dims[-1] if dims else self.main_input_dim
+                tower_output_dim = (
+                    hidden_dims[-1] if hidden_dims else self.main_input_dim
+                )
                 self.tower_heads.append(nn.Linear(tower_output_dim, 1))
         else:
             if expert_hidden_dims is None or not expert_hidden_dims:
@@ -476,7 +460,7 @@ class POSO(BaseModel):
                 [
                     MLP(
                         input_dim=self.mmoe.expert_output_dim,
-                        output_layer=True,
+                        output_dim=1,
                         **tower_params,
                     )
                     for tower_params in tower_params_list
@@ -488,7 +472,7 @@ class POSO(BaseModel):
         else:
             self.grad_norm_shared_modules = ["embedding", "mmoe"]
         self.prediction_layer = TaskHead(
-            task_type=self.default_task,
+            task_type=self.task,
             task_dims=[1] * self.nums_task,
         )
         include_modules = (
@@ -499,12 +483,6 @@ class POSO(BaseModel):
         self.register_regularization_weights(
             embedding_attr="embedding", include_modules=include_modules
         )
-        self.compile(
-            optimizer=optimizer,
-            optimizer_params=optimizer_params,
-            loss=loss,
-            loss_params=loss_params,
-        )
 
     def forward(self, x):
         # Embed main and PC features separately so PC can gate hidden units

nextrec/models/multi_task/share_bottom.py

@@ -69,19 +69,9 @@ class ShareBottom(BaseModel):
         tower_params_list: list[dict],
         target: list[str],
         task: str | list[str] | None = None,
-        optimizer: str = "adam",
-        optimizer_params: dict | None = None,
-        loss: str | nn.Module | list[str | nn.Module] | None = "bce",
-        loss_params: dict | list[dict] | None = None,
-        embedding_l1_reg=0.0,
-        dense_l1_reg=0.0,
-        embedding_l2_reg=0.0,
-        dense_l2_reg=0.0,
         **kwargs,
     ):
 
-        optimizer_params = optimizer_params or {}
-
         self.nums_task = len(target)
 
         resolved_task = task
@@ -102,16 +92,9 @@ class ShareBottom(BaseModel):
             sequence_features=sequence_features,
             target=target,
             task=resolved_task,
-            embedding_l1_reg=embedding_l1_reg,
-            dense_l1_reg=dense_l1_reg,
-            embedding_l2_reg=embedding_l2_reg,
-            dense_l2_reg=dense_l2_reg,
             **kwargs,
         )
 
-        self.loss = loss
-        if self.loss is None:
-            self.loss = "bce"
         # Number of tasks
         self.nums_task = len(target)
         if len(tower_params_list) != self.nums_task:
@@ -127,33 +110,27 @@ class ShareBottom(BaseModel):
         # input_dim = emb_dim_total + dense_input_dim
 
         # Shared bottom network
-        self.bottom = MLP(input_dim=input_dim, output_layer=False, **bottom_params)
+        self.bottom = MLP(input_dim=input_dim, output_dim=None, **bottom_params)
         self.grad_norm_shared_modules = ["embedding", "bottom"]
 
         # Get bottom output dimension
-        if "dims" in bottom_params and len(bottom_params["dims"]) > 0:
-            bottom_output_dim = bottom_params["dims"][-1]
+        if "hidden_dims" in bottom_params and len(bottom_params["hidden_dims"]) > 0:
+            bottom_output_dim = bottom_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_layer=True, **tower_params)
+            tower = MLP(input_dim=bottom_output_dim, output_dim=1, **tower_params)
             self.towers.append(tower)
         self.prediction_layer = TaskHead(
-            task_type=self.default_task, task_dims=[1] * self.nums_task
+            task_type=self.task, task_dims=[1] * self.nums_task
         )
         # Register regularization weights
         self.register_regularization_weights(
             embedding_attr="embedding", include_modules=["bottom", "towers"]
         )
-        self.compile(
-            optimizer=optimizer,
-            optimizer_params=optimizer_params,
-            loss=loss,
-            loss_params=loss_params,
-        )
 
     def forward(self, x):
         # Get all embeddings and flatten

nextrec/models/ranking/afm.py

@@ -61,36 +61,17 @@ class AFM(BaseModel):
         sequence_features: list[SequenceFeature] | None = None,
         attention_dim: int = 32,
         attention_dropout: float = 0.0,
-        target: list[str] | str | None = None,
-        task: str | list[str] | None = None,
-        optimizer: str = "adam",
-        optimizer_params: dict | None = None,
-        loss: str | nn.Module | None = "bce",
-        loss_params: dict | list[dict] | None = None,
-        embedding_l1_reg=0.0,
-        dense_l1_reg=0.0,
-        embedding_l2_reg=0.0,
-        dense_l2_reg=0.0,
         **kwargs,
     ):
 
         dense_features = dense_features or []
         sparse_features = sparse_features or []
         sequence_features = sequence_features or []
-        optimizer_params = optimizer_params or {}
-        if loss is None:
-            loss = "bce"
 
         super(AFM, self).__init__(
             dense_features=dense_features,
             sparse_features=sparse_features,
             sequence_features=sequence_features,
-            target=target,
-            task=task or self.default_task,
-            embedding_l1_reg=embedding_l1_reg,
-            dense_l1_reg=dense_l1_reg,
-            embedding_l2_reg=embedding_l2_reg,
-            dense_l2_reg=dense_l2_reg,
             **kwargs,
         )
 
@@ -139,8 +120,10 @@ class AFM(BaseModel):
         self.attention_linear = nn.Linear(self.embedding_dim, attention_dim)
         self.attention_p = nn.Linear(attention_dim, 1, bias=False)
         self.attention_dropout = nn.Dropout(attention_dropout)
+
         self.output_projection = nn.Linear(self.embedding_dim, 1, bias=False)
-        self.prediction_layer = TaskHead(task_type=self.default_task)
+
+        self.prediction_layer = TaskHead(task_type=self.task)
         self.input_mask = InputMask()
 
         # Register regularization weights
@@ -158,13 +141,6 @@ class AFM(BaseModel):
             emb.weight for emb in self.first_order_embeddings.values()
         )
 
-        self.compile(
-            optimizer=optimizer,
-            optimizer_params=optimizer_params,
-            loss=loss,
-            loss_params=loss_params,
-        )
-
     def forward(self, x):
         field_emb = self.embedding(
             x=x, features=self.fm_features, squeeze_dim=False

nextrec/models/ranking/autoint.py

@@ -82,16 +82,6 @@ class AutoInt(BaseModel):
         att_head_num: int = 2,
         att_dropout: float = 0.0,
         att_use_residual: bool = True,
-        target: list[str] | None = None,
-        task: str | list[str] | None = None,
-        optimizer: str = "adam",
-        optimizer_params: dict | None = None,
-        loss: str | nn.Module | None = "bce",
-        loss_params: dict | list[dict] | None = None,
-        embedding_l1_reg=0.0,
-        dense_l1_reg=0.0,
-        embedding_l2_reg=0.0,
-        dense_l2_reg=0.0,
         **kwargs,
     ):
 
@@ -99,21 +89,12 @@ class AutoInt(BaseModel):
             dense_features=dense_features,
             sparse_features=sparse_features,
             sequence_features=sequence_features,
-            target=target,
-            task=task or self.default_task,
-            embedding_l1_reg=embedding_l1_reg,
-            dense_l1_reg=dense_l1_reg,
-            embedding_l2_reg=embedding_l2_reg,
-            dense_l2_reg=dense_l2_reg,
             **kwargs,
         )
 
-        if target is None:
-            target = []
-        if optimizer_params is None:
-            optimizer_params = {}
-        if loss is None:
-            loss = "bce"
+        dense_features = dense_features or []
+        sparse_features = sparse_features or []
+        sequence_features = sequence_features or []
 
         self.att_layer_num = att_layer_num
         self.att_embedding_dim = att_embedding_dim
@@ -123,12 +104,7 @@ class AutoInt(BaseModel):
         # if you want to follow the paper strictly, set dense_features=[]
         # or modify the code accordingly
         self.interaction_features = dense_features + sparse_features + sequence_features
-
-        # All features for embedding
-        self.all_features = dense_features + sparse_features + sequence_features
-
-        # Embedding layer
-        self.embedding = EmbeddingLayer(features=self.all_features)
+        self.embedding = EmbeddingLayer(features=self.interaction_features)
 
         # Project embeddings to attention embedding dimension
         num_fields = len(self.interaction_features)
@@ -159,23 +135,14 @@ class AutoInt(BaseModel):
             ]
         )
 
-        # Final prediction layer
         self.fc = nn.Linear(num_fields * att_embedding_dim, 1)
-        self.prediction_layer = TaskHead(task_type=self.default_task)
+        self.prediction_layer = TaskHead(task_type=self.task)
 
-        # Register regularization weights
         self.register_regularization_weights(
             embedding_attr="embedding",
             include_modules=["projection_layers", "attention_layers", "fc"],
         )
 
-        self.compile(
-            optimizer=optimizer,
-            optimizer_params=optimizer_params,
-            loss=loss,
-            loss_params=loss_params,
-        )
-
     def forward(self, x):
         # Get embeddings field-by-field so mixed dimensions can be projected safely
         field_embeddings = []

nextrec/models/ranking/dcn.py

@@ -96,36 +96,18 @@ class DCN(BaseModel):
         sequence_features: list[SequenceFeature] | None = None,
         cross_num: int = 3,
         mlp_params: dict | None = None,
-        target: list[str] | str | None = None,
-        task: str | list[str] | None = None,
-        optimizer: str = "adam",
-        optimizer_params: dict | None = None,
-        loss: str | nn.Module | None = "bce",
-        loss_params: dict | list[dict] | None = None,
-        embedding_l1_reg=0.0,
-        dense_l1_reg=0.0,
-        embedding_l2_reg=0.0,
-        dense_l2_reg=0.0,
         **kwargs,
     ):
 
         dense_features = dense_features or []
         sparse_features = sparse_features or []
         sequence_features = sequence_features or []
-        optimizer_params = optimizer_params or {}
-        if loss is None:
-            loss = "bce"
+        mlp_params = mlp_params or {}
 
         super(DCN, self).__init__(
             dense_features=dense_features,
             sparse_features=sparse_features,
             sequence_features=sequence_features,
-            target=target,
-            task=task or self.default_task,
-            embedding_l1_reg=embedding_l1_reg,
-            dense_l1_reg=dense_l1_reg,
-            embedding_l2_reg=embedding_l2_reg,
-            dense_l2_reg=dense_l2_reg,
             **kwargs,
         )
 
@@ -173,13 +155,6 @@ class DCN(BaseModel):
             include_modules=["cross_network", "mlp", "final_layer"],
         )
 
-        self.compile(
-            optimizer=optimizer,
-            optimizer_params=optimizer_params,
-            loss=loss,
-            loss_params=loss_params,
-        )
-
     def forward(self, x):
         # Get all embeddings and flatten
         input_flat = self.embedding(x=x, features=self.all_features, squeeze_dim=True)

nextrec/models/ranking/dcn_v2.py

@@ -45,7 +45,7 @@ DCN v2 在原始 DCN 基础上，将标量交叉权重升级为向量/矩阵参
 
 import torch
 import torch.nn as nn
-
+from typing import Literal
 from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
 from nextrec.basic.layers import MLP, EmbeddingLayer
 from nextrec.basic.heads import TaskHead
@@ -193,53 +193,32 @@ class DCNv2(BaseModel):
         sparse_features: list[SparseFeature] | None = None,
         sequence_features: list[SequenceFeature] | None = None,
         cross_num: int = 3,
-        cross_type: str = "matrix",
-        architecture: str = "parallel",
+        cross_type: Literal["matrix", "mix", "low_rank"] = "matrix",
+        architecture: Literal["parallel", "stacked"] = "parallel",
         low_rank: int = 32,
         num_experts: int = 4,
         mlp_params: dict | None = None,
-        target: list[str] | str | None = None,
-        task: str | list[str] | None = None,
-        optimizer: str = "adam",
-        optimizer_params: dict | None = None,
-        loss: str | nn.Module | None = "bce",
-        loss_params: dict | list[dict] | None = None,
-        embedding_l1_reg=0.0,
-        dense_l1_reg=0.0,
-        embedding_l2_reg=0.0,
-        dense_l2_reg=0.0,
         **kwargs,
     ):
         dense_features = dense_features or []
         sparse_features = sparse_features or []
         sequence_features = sequence_features or []
-        optimizer_params = optimizer_params or {}
-        if loss is None:
-            loss = "bce"
+        mlp_params = mlp_params or {}
 
         super(DCNv2, self).__init__(
             dense_features=dense_features,
             sparse_features=sparse_features,
             sequence_features=sequence_features,
-            target=target,
-            task=task or self.default_task,
-            embedding_l1_reg=embedding_l1_reg,
-            dense_l1_reg=dense_l1_reg,
-            embedding_l2_reg=embedding_l2_reg,
-            dense_l2_reg=dense_l2_reg,
             **kwargs,
         )
 
-        self.all_features = dense_features + sparse_features + sequence_features
         self.embedding = EmbeddingLayer(features=self.all_features)
         input_dim = self.embedding.input_dim
 
-        architecture = architecture.lower()
         if architecture not in {"parallel", "stacked"}:
             raise ValueError("architecture must be 'parallel' or 'stacked'.")
         self.architecture = architecture
 
-        cross_type = cross_type.lower()
         if cross_type == "matrix":
             self.cross_network = CrossNetV2(input_dim=input_dim, num_layers=cross_num)
         elif cross_type in {"mix", "low_rank"}:
@@ -255,7 +234,7 @@ class DCNv2(BaseModel):
         if mlp_params is not None:
             self.use_dnn = True
             dnn_params = dict(mlp_params)
-            dnn_params.setdefault("output_layer", False)
+            dnn_params.setdefault("output_dim", None)
             self.mlp = MLP(input_dim=input_dim, **dnn_params)
             deep_dim = self.mlp.output_dim
             final_input_dim = (
@@ -271,20 +250,13 @@ class DCNv2(BaseModel):
             final_input_dim = input_dim
 
         self.final_layer = nn.Linear(final_input_dim, 1)
-        self.prediction_layer = TaskHead(task_type=self.default_task)
+        self.prediction_layer = TaskHead(task_type=self.task)
 
         self.register_regularization_weights(
             embedding_attr="embedding",
             include_modules=["cross_network", "mlp", "final_layer"],
         )
 
-        self.compile(
-            optimizer=optimizer,
-            optimizer_params=optimizer_params,
-            loss=loss,
-            loss_params=loss_params,
-        )
-
     def forward(self, x) -> torch.Tensor:
         input_flat = self.embedding(x=x, features=self.all_features, squeeze_dim=True)
         cross_out = self.cross_network(input_flat)