nextrec 0.2.1__py3-none-any.whl → 0.2.2__py3-none-any.whl

nextrec/models/match/dssm_v2.py

@@ -44,6 +44,12 @@ class DSSM_v2(BaseMatchModel):
                  embedding_l2_reg: float = 0.0,
                  dense_l2_reg: float = 0.0,
                  early_stop_patience: int = 20,
+                 optimizer: str | torch.optim.Optimizer = "adam",
+                 optimizer_params: dict | None = None,
+                 scheduler: str | torch.optim.lr_scheduler._LRScheduler | type[torch.optim.lr_scheduler._LRScheduler] | None = None,
+                 scheduler_params: dict | None = None,
+                 loss: str | nn.Module | list[str | nn.Module] | None = "bce",
+                 loss_params: dict | list[dict] | None = None,
                  **kwargs):
         
         super(DSSM_v2, self).__init__(
@@ -137,6 +143,18 @@ class DSSM_v2(BaseMatchModel):
             include_modules=['item_dnn']
         )
         
+        if optimizer_params is None:
+            optimizer_params = {"lr": 1e-3, "weight_decay": 1e-5}
+
+        self.compile(
+            optimizer=optimizer,
+            optimizer_params=optimizer_params,
+            scheduler=scheduler,
+            scheduler_params=scheduler_params,
+            loss=loss,
+            loss_params=loss_params,
+        )
+
         self.to(device)
     
     def user_tower(self, user_input: dict) -> torch.Tensor:

nextrec/models/match/mind.py

@@ -41,7 +41,7 @@ class MIND(BaseMatchModel):
                  item_dnn_hidden_units: list[int] = [256, 128],
                  dnn_activation: str = 'relu',
                  dnn_dropout: float = 0.0,
-                 training_mode: Literal['pointwise', 'pairwise', 'listwise'] = 'listwise',
+                 training_mode: Literal['pointwise', 'pairwise', 'listwise'] = 'pointwise',
                  num_negative_samples: int = 100,
                  temperature: float = 1.0,
                  similarity_metric: Literal['dot', 'cosine', 'euclidean'] = 'dot',
@@ -51,6 +51,12 @@ class MIND(BaseMatchModel):
                  embedding_l2_reg: float = 0.0,
                  dense_l2_reg: float = 0.0,
                  early_stop_patience: int = 20,
+                 optimizer: str | torch.optim.Optimizer = "adam",
+                 optimizer_params: dict | None = None,
+                 scheduler: str | torch.optim.lr_scheduler._LRScheduler | type[torch.optim.lr_scheduler._LRScheduler] | None = None,
+                 scheduler_params: dict | None = None,
+                 loss: str | nn.Module | list[str | nn.Module] | None = "bce",
+                 loss_params: dict | list[dict] | None = None,
                  **kwargs):
         
         super(MIND, self).__init__(
@@ -152,6 +158,15 @@ class MIND(BaseMatchModel):
             include_modules=['item_dnn'] if self.item_dnn else []
         )
         
+        self.compile(
+            optimizer=optimizer,
+            optimizer_params=optimizer_params,
+            scheduler=scheduler,
+            scheduler_params=scheduler_params,
+            loss=loss,
+            loss_params=loss_params,
+        )
+
         self.to(device)
     
     def user_tower(self, user_input: dict) -> torch.Tensor:

nextrec/models/match/sdm.py

@@ -52,6 +52,12 @@ class SDM(BaseMatchModel):
                  embedding_l2_reg: float = 0.0,
                  dense_l2_reg: float = 0.0,
                  early_stop_patience: int = 20,
+                 optimizer: str | torch.optim.Optimizer = "adam",
+                 optimizer_params: dict | None = None,
+                 scheduler: str | torch.optim.lr_scheduler._LRScheduler | type[torch.optim.lr_scheduler._LRScheduler] | None = None,
+                 scheduler_params: dict | None = None,
+                 loss: str | nn.Module | list[str | nn.Module] | None = "bce",
+                 loss_params: dict | list[dict] | None = None,
                  **kwargs):
         
         super(SDM, self).__init__(
@@ -179,6 +185,15 @@ class SDM(BaseMatchModel):
             include_modules=['item_dnn'] if self.item_dnn else []
         )
         
+        self.compile(
+            optimizer=optimizer,
+            optimizer_params=optimizer_params,
+            scheduler=scheduler,
+            scheduler_params=scheduler_params,
+            loss=loss,
+            loss_params=loss_params,
+        )
+
         self.to(device)
     
     def user_tower(self, user_input: dict) -> torch.Tensor:

nextrec/models/match/youtube_dnn.py

@@ -17,11 +17,10 @@ from nextrec.basic.layers import MLP, EmbeddingLayer, AveragePooling
 
 class YoutubeDNN(BaseMatchModel):
     """
-    YouTube Deep Neural Network for Recommendations
-    
-    用户塔：历史行为序列 + 用户特征 -> 用户embedding
-    物品塔：物品特征 -> 物品embedding
-    训练：sampled softmax loss (listwise)
+    YouTube Deep Neural Network for Recommendations.
+    User tower: behavior sequence + user features -> user embedding.
+    Item tower: item features -> item embedding.
+    Training usually uses listwise / sampled softmax style objectives.
     """
     
     @property
@@ -50,6 +49,12 @@ class YoutubeDNN(BaseMatchModel):
                  embedding_l2_reg: float = 0.0,
                  dense_l2_reg: float = 0.0,
                  early_stop_patience: int = 20,
+                 optimizer: str | torch.optim.Optimizer = "adam",
+                 optimizer_params: dict | None = None,
+                 scheduler: str | torch.optim.lr_scheduler._LRScheduler | type[torch.optim.lr_scheduler._LRScheduler] | None = None,
+                 scheduler_params: dict | None = None,
+                 loss: str | nn.Module | list[str | nn.Module] | None = "bce",
+                 loss_params: dict | list[dict] | None = None,
                  **kwargs):
         
         super(YoutubeDNN, self).__init__(
@@ -94,7 +99,7 @@ class YoutubeDNN(BaseMatchModel):
             for feat in user_sparse_features or []:
                 user_input_dim += feat.embedding_dim
             for feat in user_sequence_features or []:
-                # 序列特征通过平均池化聚合
+                # Sequence features are pooled before entering the DNN
                 user_input_dim += feat.embedding_dim
             
             user_dnn_units = user_dnn_hidden_units + [embedding_dim]
@@ -144,12 +149,20 @@ class YoutubeDNN(BaseMatchModel):
             include_modules=['item_dnn']
         )
         
+        self.compile(
+            optimizer=optimizer,
+            optimizer_params=optimizer_params,
+            scheduler=scheduler,
+            scheduler_params=scheduler_params,
+            loss=loss,
+            loss_params=loss_params,
+        )
+
         self.to(device)
     
     def user_tower(self, user_input: dict) -> torch.Tensor:
         """
-        User tower
-        处理用户历史行为序列和其他用户特征
+        User tower to encode historical behavior sequences and user features.
         """
         all_user_features = self.user_dense_features + self.user_sparse_features + self.user_sequence_features
         user_emb = self.user_embedding(user_input, all_user_features, squeeze_dim=True)

nextrec/models/multi_task/esmm.py

@@ -1,9 +1,7 @@
 """
 Date: create on 09/11/2025
-Author:
-    Yang Zhou,zyaztec@gmail.com
-Reference:
-    [1] Ma X, Zhao L, Huang G, et al. Entire space multi-task model: An effective approach for estimating post-click conversion rate[C]//SIGIR. 2018: 1137-1140.
+Author: Yang Zhou,zyaztec@gmail.com
+Reference: [1] Ma X, Zhao L, Huang G, et al. Entire space multi-task model: An effective approach for estimating post-click conversion rate[C]//SIGIR. 2018: 1137-1140.
 """
 
 import torch
@@ -46,6 +44,7 @@ class ESMM(BaseModel):
                  optimizer: str = "adam",
                  optimizer_params: dict = {},
                  loss: str | nn.Module | list[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,
@@ -106,7 +105,8 @@ class ESMM(BaseModel):
         self.compile(
             optimizer=optimizer,
             optimizer_params=optimizer_params,
-            loss=loss
+            loss=loss,
+            loss_params=loss_params,
         )
 
     def forward(self, x):

nextrec/models/multi_task/mmoe.py

@@ -1,9 +1,7 @@
 """
 Date: create on 09/11/2025
-Author:
-    Yang Zhou,zyaztec@gmail.com
-Reference:
-    [1] Ma J, Zhao Z, Yi X, et al. Modeling task relationships in multi-task learning with multi-gate mixture-of-experts[C]//KDD. 2018: 1930-1939.
+Author: Yang Zhou,zyaztec@gmail.com
+Reference: [1] Ma J, Zhao Z, Yi X, et al. Modeling task relationships in multi-task learning with multi-gate mixture-of-experts[C]//KDD. 2018: 1930-1939.
 """
 
 import torch
@@ -44,6 +42,7 @@ class MMOE(BaseModel):
                  optimizer: str = "adam",
                  optimizer_params: dict = {},
                  loss: str | nn.Module | list[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,
@@ -128,7 +127,8 @@ class MMOE(BaseModel):
         self.compile(
             optimizer=optimizer,
             optimizer_params=optimizer_params,
-            loss=loss
+            loss=loss,
+            loss_params=loss_params,
         )
 
     def forward(self, x):

nextrec/models/multi_task/ple.py

@@ -1,9 +1,7 @@
 """
 Date: create on 09/11/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.
+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.
 """
 
 import torch
@@ -47,6 +45,7 @@ class PLE(BaseModel):
                  optimizer: str = "adam",
                  optimizer_params: dict = {},
                  loss: str | nn.Module | list[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,
@@ -166,7 +165,8 @@ class PLE(BaseModel):
         self.compile(
             optimizer=optimizer,
             optimizer_params=optimizer_params,
-            loss=loss
+            loss=loss,
+            loss_params=loss_params,
         )
 
     def forward(self, x):

nextrec/models/multi_task/share_bottom.py

@@ -1,9 +1,7 @@
 """
 Date: create on 09/11/2025
-Author:
-    Yang Zhou,zyaztec@gmail.com
-Reference:
-    [1] Caruana R. Multitask learning[J]. Machine learning, 1997, 28: 41-75.
+Author: Yang Zhou,zyaztec@gmail.com
+Reference: [1] Caruana R. Multitask learning[J]. Machine learning, 1997, 28: 41-75.
 """
 
 import torch
@@ -35,6 +33,7 @@ class ShareBottom(BaseModel):
                  optimizer: str = "adam",
                  optimizer_params: dict = {},
                  loss: str | nn.Module | list[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,
@@ -105,7 +104,8 @@ class ShareBottom(BaseModel):
         self.compile(
             optimizer=optimizer,
             optimizer_params=optimizer_params,
-            loss=loss
+            loss=loss,
+            loss_params=loss_params,
         )
 
     def forward(self, x):

nextrec/models/ranking/__init__.py

@@ -1,3 +1,7 @@
+from .fm import FM
+from .afm import AFM
+from .masknet import MaskNet
+from .pnn import PNN
 from .deepfm import DeepFM
 from .autoint import AutoInt
 from .widedeep import WideDeep
@@ -14,4 +18,8 @@ __all__ = [
     'DCN',
     'DIN',
     'DIEN',
+    'FM',
+    'AFM',
+    'MaskNet',
+    'PNN',
 ]

nextrec/models/ranking/afm.py

@@ -34,6 +34,7 @@ class AFM(BaseModel):
                  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,
@@ -88,7 +89,8 @@ class AFM(BaseModel):
         self.compile(
             optimizer=optimizer,
             optimizer_params=optimizer_params,
-            loss=loss
+            loss=loss,
+            loss_params=loss_params,
         )
 
     def forward(self, x):

nextrec/models/ranking/autoint.py

@@ -39,6 +39,7 @@ class AutoInt(BaseModel):
                  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,
@@ -113,7 +114,8 @@ class AutoInt(BaseModel):
         self.compile(
             optimizer=optimizer,
             optimizer_params=optimizer_params,
-            loss=loss
+            loss=loss,
+            loss_params=loss_params,
         )
 
     def forward(self, x):

nextrec/models/ranking/dcn.py

@@ -35,6 +35,7 @@ class DCN(BaseModel):
                  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,
@@ -97,7 +98,8 @@ class DCN(BaseModel):
         self.compile(
             optimizer=optimizer,
             optimizer_params=optimizer_params,
-            loss=loss
+            loss=loss,
+            loss_params=loss_params,
         )
 
     def forward(self, x):

nextrec/models/ranking/deepfm.py

@@ -31,6 +31,7 @@ class DeepFM(BaseModel):
                  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,
@@ -79,7 +80,8 @@ class DeepFM(BaseModel):
         self.compile(
             optimizer=optimizer,
             optimizer_params=optimizer_params,
-            loss=loss
+            loss=loss,
+            loss_params=loss_params,
         )
 
     def forward(self, x):

nextrec/models/ranking/dien.py

@@ -38,6 +38,7 @@ class DIEN(BaseModel):
                  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,
@@ -128,7 +129,8 @@ class DIEN(BaseModel):
         self.compile(
             optimizer=optimizer,
             optimizer_params=optimizer_params,
-            loss=loss
+            loss=loss,
+            loss_params=loss_params,
         )
 
     def forward(self, x):

nextrec/models/ranking/din.py

@@ -37,6 +37,7 @@ class DIN(BaseModel):
                  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,
@@ -115,7 +116,8 @@ class DIN(BaseModel):
         self.compile(
             optimizer=optimizer,
             optimizer_params=optimizer_params,
-            loss=loss
+            loss=loss,
+            loss_params=loss_params,
         )
 
     def forward(self, x):

nextrec/models/ranking/fibinet.py

@@ -42,6 +42,7 @@ class FiBiNET(BaseModel):
                  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,
@@ -111,7 +112,8 @@ class FiBiNET(BaseModel):
         self.compile(
             optimizer=optimizer,
             optimizer_params=optimizer_params,
-            loss=loss
+            loss=loss,
+            loss_params=loss_params,
         )
 
     def forward(self, x):

nextrec/models/ranking/fm.py

@@ -30,6 +30,7 @@ class FM(BaseModel):
                  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,
@@ -76,7 +77,8 @@ class FM(BaseModel):
         self.compile(
             optimizer=optimizer,
             optimizer_params=optimizer_params,
-            loss=loss
+            loss=loss,
+            loss_params=loss_params,
         )
 
     def forward(self, x):

nextrec/models/ranking/masknet.py

@@ -36,6 +36,7 @@ class MaskNet(BaseModel):
                  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,
@@ -98,7 +99,8 @@ class MaskNet(BaseModel):
         self.compile(
             optimizer=optimizer,
             optimizer_params=optimizer_params,
-            loss=loss
+            loss=loss,
+            loss_params=loss_params,
         )
 
     def forward(self, x):

nextrec/models/ranking/pnn.py

@@ -34,6 +34,7 @@ class PNN(BaseModel):
                  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,
@@ -98,7 +99,8 @@ class PNN(BaseModel):
         self.compile(
             optimizer=optimizer,
             optimizer_params=optimizer_params,
-            loss=loss
+            loss=loss,
+            loss_params=loss_params,
         )
 
     def forward(self, x):

nextrec/models/ranking/widedeep.py

@@ -34,6 +34,7 @@ class WideDeep(BaseModel):
                  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,
@@ -88,7 +89,8 @@ class WideDeep(BaseModel):
         self.compile(
             optimizer=optimizer,
             optimizer_params=optimizer_params,
-            loss=loss
+            loss=loss,
+            loss_params=loss_params,
         )
 
     def forward(self, x):

nextrec/models/ranking/xdeepfm.py

@@ -37,6 +37,7 @@ class xDeepFM(BaseModel):
                  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,
@@ -95,7 +96,8 @@ class xDeepFM(BaseModel):
         self.compile(
             optimizer=optimizer,
             optimizer_params=optimizer_params,
-            loss=loss
+            loss=loss,
+            loss_params=loss_params,
         )
 
     def forward(self, x):

nextrec/utils/__init__.py

@@ -1,12 +1,12 @@
-from .optimizer import get_optimizer_fn, get_scheduler_fn
-from .initializer import get_initializer_fn
+from .optimizer import get_optimizer, get_scheduler
+from .initializer import get_initializer
 from .embedding import get_auto_embedding_dim
 from . import optimizer, initializer, embedding
 
 __all__ = [
-    'get_optimizer_fn',
-    'get_scheduler_fn',
-    'get_initializer_fn',
+    'get_optimizer',
+    'get_scheduler',
+    'get_initializer',
     'get_auto_embedding_dim',
     'optimizer',
     'initializer',

nextrec/utils/initializer.py

@@ -8,14 +8,14 @@ Author: Yang Zhou, zyaztec@gmail.com
 import torch.nn as nn
 
 
-def get_initializer_fn(init_type='normal', activation='linear', param=None):
+def get_initializer(init_type='normal', activation='linear', param=None):
     """
     Get parameter initialization function.
         
     Examples:
-        >>> init_fn = get_initializer_fn('xavier_uniform', 'relu')
+        >>> init_fn = get_initializer('xavier_uniform', 'relu')
         >>> init_fn(tensor)
-        >>> init_fn = get_initializer_fn('normal', param={'mean': 0.0, 'std': 0.01})
+        >>> init_fn = get_initializer('normal', param={'mean': 0.0, 'std': 0.01})
     """
     param = param or {}
 

nextrec/utils/optimizer.py

@@ -9,7 +9,7 @@ import torch
 from typing import Iterable
 
 
-def get_optimizer_fn(
+def get_optimizer(
     optimizer: str = "adam",
     params: Iterable[torch.nn.Parameter] | None = None,
     **optimizer_params
@@ -18,8 +18,8 @@ def get_optimizer_fn(
     Get optimizer function based on optimizer name or instance.
         
     Examples:
-        >>> optimizer = get_optimizer_fn("adam", model.parameters(), lr=1e-3)
-        >>> optimizer = get_optimizer_fn("sgd", model.parameters(), lr=0.01, momentum=0.9)
+        >>> optimizer = get_optimizer("adam", model.parameters(), lr=1e-3)
+        >>> optimizer = get_optimizer("sgd", model.parameters(), lr=0.01, momentum=0.9)
     """
     if params is None:
         raise ValueError("params cannot be None. Please provide model parameters.")
@@ -51,13 +51,13 @@ def get_optimizer_fn(
     return optimizer_fn
 
 
-def get_scheduler_fn(scheduler, optimizer, **scheduler_params):
+def get_scheduler(scheduler, optimizer, **scheduler_params):
     """
     Get learning rate scheduler function.
     
     Examples:
-        >>> scheduler = get_scheduler_fn("step", optimizer, step_size=10, gamma=0.1)
-        >>> scheduler = get_scheduler_fn("cosine", optimizer, T_max=100)
+        >>> scheduler = get_scheduler("step", optimizer, step_size=10, gamma=0.1)
+        >>> scheduler = get_scheduler("cosine", optimizer, T_max=100)
     """
     if isinstance(scheduler, str):
         if scheduler == "step":

{nextrec-0.2.1.dist-info → nextrec-0.2.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nextrec
-Version: 0.2.1
+Version: 0.2.2
 Summary: A comprehensive recommendation library with match, ranking, and multi-task learning models
 Project-URL: Homepage, https://github.com/zerolovesea/NextRec
 Project-URL: Repository, https://github.com/zerolovesea/NextRec
@@ -61,7 +61,7 @@ Description-Content-Type: text/markdown
 ![Python](https://img.shields.io/badge/Python-3.10+-blue.svg)
 ![PyTorch](https://img.shields.io/badge/PyTorch-1.10+-ee4c2c.svg)
 ![License](https://img.shields.io/badge/License-Apache%202.0-green.svg)
-![Version](https://img.shields.io/badge/Version-0.2.1-orange.svg)
+![Version](https://img.shields.io/badge/Version-0.2.2-orange.svg)
 
 English | [中文版](README_zh.md)