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

nextrec/loss/pairwise.py

@@ -2,7 +2,7 @@
 Pairwise loss functions for learning-to-rank and matching tasks.
 
 Date: create on 27/10/2025
-Checkpoint: edit on 29/11/2025
+Checkpoint: edit on 29/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 """
 
@@ -18,7 +18,7 @@ class BPRLoss(nn.Module):
     Bayesian Personalized Ranking loss with support for multiple negatives.
     """
 
-    def __init__(self, reduction: str = "mean"):
+    def __init__(self, reduction: Literal["mean", "sum", "none"] = "mean"):
         super().__init__()
         self.reduction = reduction
 
@@ -42,7 +42,9 @@ class HingeLoss(nn.Module):
     Hinge loss for pairwise ranking.
     """
 
-    def __init__(self, margin: float = 1.0, reduction: str = "mean"):
+    def __init__(
+        self, margin: float = 1.0, reduction: Literal["mean", "sum", "none"] = "mean"
+    ):
         super().__init__()
         self.margin = margin
         self.reduction = reduction
@@ -69,7 +71,7 @@ class TripletLoss(nn.Module):
     def __init__(
         self,
         margin: float = 1.0,
-        reduction: str = "mean",
+        reduction: Literal["mean", "sum", "none"] = "mean",
         distance: Literal["euclidean", "cosine"] = "euclidean",
     ):
         super().__init__()

nextrec/loss/pointwise.py

@@ -2,11 +2,11 @@
 Pointwise loss functions, including imbalance-aware variants.
 
 Date: create on 27/10/2025
-Checkpoint: edit on 29/11/2025
+Checkpoint: edit on 29/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 """
 
-from typing import Optional, Sequence
+from typing import Optional, Sequence, Literal
 
 import torch
 import torch.nn as nn
@@ -18,7 +18,9 @@ class CosineContrastiveLoss(nn.Module):
     Contrastive loss using cosine similarity for positive/negative pairs.
     """
 
-    def __init__(self, margin: float = 0.5, reduction: str = "mean"):
+    def __init__(
+        self, margin: float = 0.5, reduction: Literal["mean", "sum", "none"] = "mean"
+    ):
         super().__init__()
         self.margin = margin
         self.reduction = reduction
@@ -50,7 +52,7 @@ class WeightedBCELoss(nn.Module):
     def __init__(
         self,
         pos_weight: float | torch.Tensor | None = None,
-        reduction: str = "mean",
+        reduction: Literal["mean", "sum", "none"] = "mean",
         logits: bool = False,
         auto_balance: bool = False,
     ):
@@ -110,7 +112,7 @@ class FocalLoss(nn.Module):
         self,
         gamma: float = 2.0,
         alpha: Optional[float | Sequence[float] | torch.Tensor] = None,
-        reduction: str = "mean",
+        reduction: Literal["mean", "sum", "none"] = "mean",
         logits: bool = False,
     ):
         super().__init__()
@@ -187,7 +189,7 @@ class ClassBalancedFocalLoss(nn.Module):
         class_counts: Sequence[int] | torch.Tensor,
         beta: float = 0.9999,
         gamma: float = 2.0,
-        reduction: str = "mean",
+        reduction: Literal["mean", "sum", "none"] = "mean",
     ):
         super().__init__()
         self.gamma = gamma

nextrec/models/multi_task/esmm.py

@@ -42,12 +42,12 @@ CVR 预测 P(conversion|click)，二者相乘得到 CTCVR 并在曝光标签上
 """
 
 import torch
-import torch.nn as nn
 
 from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
 from nextrec.basic.layers import MLP, EmbeddingLayer
 from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
+from nextrec.utils.types import TaskTypeName
 
 
 class ESMM(BaseModel):
@@ -77,23 +77,12 @@ class ESMM(BaseModel):
         sequence_features: list[SequenceFeature],
         ctr_params: dict,
         cvr_params: dict,
+        task: TaskTypeName | list[TaskTypeName] | None = None,
         target: list[str] | None = None,  # Note: ctcvr = ctr * cvr
-        task: 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,
     ):
 
         target = target or ["ctr", "ctcvr"]
-        optimizer_params = optimizer_params or {}
-        if loss is None:
-            loss = "bce"
 
         if len(target) != 2:
             raise ValueError(
@@ -120,15 +109,9 @@ class ESMM(BaseModel):
             sequence_features=sequence_features,
             target=target,
             task=resolved_task,  # Both CTR and CTCVR are binary classification
-            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
-
         self.embedding = EmbeddingLayer(features=self.all_features)
         input_dim = self.embedding.input_dim
 
@@ -138,17 +121,11 @@ class ESMM(BaseModel):
         # CVR tower
         self.cvr_tower = MLP(input_dim=input_dim, output_layer=True, **cvr_params)
         self.grad_norm_shared_modules = ["embedding"]
-        self.prediction_layer = TaskHead(task_type=self.default_task, task_dims=[1, 1])
+        self.prediction_layer = TaskHead(task_type=self.task, task_dims=[1, 1])
         # Register regularization weights
         self.register_regularization_weights(
             embedding_attr="embedding", include_modules=["ctr_tower", "cvr_tower"]
         )
-        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/multi_task/mmoe.py

@@ -82,14 +82,6 @@ class MMOE(BaseModel):
         tower_params_list: list[dict] | None = None,
         target: list[str] | str | None = None,
         task: str | list[str] = "binary",
-        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,
     ):
 
@@ -98,9 +90,7 @@ class MMOE(BaseModel):
         sequence_features = sequence_features or []
         expert_params = expert_params or {}
         tower_params_list = tower_params_list or []
-        optimizer_params = optimizer_params or {}
-        if loss is None:
-            loss = "bce"
+
         if target is None:
             target = []
         elif isinstance(target, str):
@@ -126,15 +116,9 @@ class MMOE(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
-
         # Number of tasks and experts
         self.nums_task = len(target)
         self.num_experts = num_experts
@@ -172,18 +156,12 @@ class MMOE(BaseModel):
             tower = MLP(input_dim=expert_output_dim, output_layer=True, **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=["experts", "gates", "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/ple.py

@@ -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})"
@@ -302,18 +286,12 @@ class PLE(BaseModel):
             tower = MLP(input_dim=expert_output_dim, output_layer=True, **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
@@ -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"
         )
@@ -488,7 +470,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 +481,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:
@@ -142,18 +125,12 @@ class ShareBottom(BaseModel):
             tower = MLP(input_dim=bottom_output_dim, output_layer=True, **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)