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

nextrec/models/ranking/deepfm.py

@@ -3,9 +3,8 @@ Date: create on 27/10/2025
 Checkpoint: edit on 23/12/2025
 Author: Yang Zhou,zyaztec@gmail.com
 Reference:
-[1] Guo H, Tang R, Ye Y, et al. DeepFM: A factorization-machine based neural network
-for CTR prediction[J]. arXiv preprint arXiv:1703.04247, 2017.
-(https://arxiv.org/abs/1703.04247)
+- [1] Guo H, Tang R, Ye Y, et al. DeepFM: A factorization-machine based neural network for CTR prediction[J]. arXiv preprint arXiv:1703.04247, 2017.
+URL: https://arxiv.org/abs/1703.04247
 
 DeepFM combines a Factorization Machine (FM) for explicit second-order feature
 interactions with a deep MLP for high-order nonlinear patterns. Both parts share

nextrec/models/ranking/dien.py

@@ -3,9 +3,7 @@ Date: create on 09/11/2025
 Author: Yang Zhou, zyaztec@gmail.com
 Checkpoint: edit on 09/12/2025
 Reference:
-[1] Zhou G, Mou N, Fan Y, et al. Deep interest evolution network for click-through
-rate prediction[C] // Proceedings of the AAAI conference on artificial intelligence.
-2019, 33(01): 5941-5948. (https://arxiv.org/abs/1809.03672)
+- [1] Zhou G, Mou N, Fan Y, et al. Deep interest evolution network for click-through rate prediction[C] // Proceedings of the AAAI conference on artificial intelligence. 2019, 33(01): 5941-5948. (https://arxiv.org/abs/1809.03672)
 
 DIEN is a CTR prediction model that explicitly models how user interests evolve
 over time. It introduces a two-stage pipeline:
@@ -58,7 +56,6 @@ from nextrec.basic.layers import (
 )
 from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
-from nextrec.utils.types import ActivationName
 
 
 class AUGRU(nn.Module):
@@ -211,8 +208,7 @@ class DIEN(BaseModel):
         neg_behavior_feature_name: str | None = None,
         mlp_params: dict | None = None,
         gru_hidden_size: int = 64,
-        attention_hidden_units: list[int] | None = None,
-        attention_activation: ActivationName = "sigmoid",
+        attention_mlp_params: dict | None = None,
         use_negsampling: bool = False,
         aux_loss_weight: float = 1.0,
         **kwargs,
@@ -222,7 +218,9 @@ class DIEN(BaseModel):
         sparse_features = sparse_features or []
         sequence_features = sequence_features or []
         mlp_params = mlp_params or {}
-        attention_hidden_units = attention_hidden_units or [80, 40]
+        attention_mlp_params = attention_mlp_params or {}
+        attention_mlp_params.setdefault("hidden_dims", [80, 40])
+        attention_mlp_params.setdefault("activation", "sigmoid")
 
         super(DIEN, self).__init__(
             dense_features=dense_features,
@@ -285,8 +283,8 @@ class DIEN(BaseModel):
 
         self.attention_layer = AttentionPoolingLayer(
             embedding_dim=gru_hidden_size,
-            hidden_units=attention_hidden_units,
-            activation=attention_activation,
+            hidden_units=attention_mlp_params["hidden_dims"],
+            activation=attention_mlp_params["activation"],
             use_softmax=False,
         )
 

nextrec/models/ranking/din.py

@@ -3,10 +3,8 @@ Date: create on 09/11/2025
 Checkpoint: edit on 09/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 Reference:
-[1] Zhou G, Zhu X, Song C, et al. Deep interest network for click-through rate
-prediction[C] //Proceedings of the 24th ACM SIGKDD international conference on
-knowledge discovery & data mining. 2018: 1059-1068.
-(https://arxiv.org/abs/1706.06978)
+- [1] Zhou G, Zhu X, Song C, et al. Deep interest network for click-through rate prediction[C] //Proceedings of the 24th ACM SIGKDD international conference on knowledge discovery & data mining. 2018: 1059-1068.
+  URL: https://arxiv.org/abs/1706.06978
 
 Deep Interest Network (DIN) is a CTR model that builds a target-aware user
 representation by attending over the historical behavior sequence. Instead of
@@ -58,7 +56,6 @@ from nextrec.basic.layers import (
 )
 from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
-from nextrec.utils.types import ActivationName
 
 
 class DIN(BaseModel):
@@ -78,8 +75,7 @@ class DIN(BaseModel):
         behavior_feature_name: str | None = None,
         candidate_feature_name: str | None = None,
         mlp_params: dict | None = None,
-        attention_hidden_units: list[int] | None = None,
-        attention_activation: ActivationName = "dice",
+        attention_mlp_params: dict | None = None,
         attention_use_softmax: bool = True,
         **kwargs,
     ):
@@ -88,7 +84,9 @@ class DIN(BaseModel):
         sparse_features = sparse_features or []
         sequence_features = sequence_features or []
         mlp_params = mlp_params or {}
-        attention_hidden_units = attention_hidden_units or [80, 40]
+        attention_mlp_params = attention_mlp_params or {}
+        attention_mlp_params.setdefault("hidden_dims", [80, 40])
+        attention_mlp_params.setdefault("activation", "dice")
 
         super(DIN, self).__init__(
             dense_features=dense_features,
@@ -135,8 +133,8 @@ class DIN(BaseModel):
             )
         self.attention = AttentionPoolingLayer(
             embedding_dim=behavior_emb_dim,
-            hidden_units=attention_hidden_units,
-            activation=attention_activation,
+            hidden_units=attention_mlp_params["hidden_dims"],
+            activation=attention_mlp_params["activation"],
             use_softmax=attention_use_softmax,
         )
 

nextrec/models/ranking/eulernet.py

@@ -3,8 +3,7 @@ Date: create on 09/11/2025
 Checkpoint: edit on 23/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 Reference:
-[1] Zhao Z, Zhang H, Tang H, et al. EulerNet: Efficient and Effective Feature
-Interaction Modeling with Euler's Formula. (SIGIR 2021)
+- [1] Zhao Z, Zhang H, Tang H, et al. EulerNet: Efficient and Effective Feature Interaction Modeling with Euler's Formula. (SIGIR 2021)
 
 EulerNet models feature interactions in the complex domain using Euler's
 formula. Each field embedding is transformed into amplitude and phase,

nextrec/models/ranking/ffm.py

@@ -3,8 +3,7 @@ Date: create on 19/12/2025
 Checkpoint: edit on 23/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 Reference:
-[1] Juan Y, Zhuang Y, Chin W-S, et al. Field-aware Factorization Machines for CTR
-Prediction[C]//RecSys. 2016: 43-50.
+- [1] Juan Y, Zhuang Y, Chin W-S, et al. Field-aware Factorization Machines for CTR Prediction[C]//RecSys. 2016: 43-50.
 
 Field-aware Factorization Machines (FFM) extend FM by learning a distinct
 embedding of each feature for every target field. For a pair of fields (i, j),

nextrec/models/ranking/fibinet.py

@@ -3,9 +3,8 @@ Date: create on 09/11/2025
 Checkpoint: edit on 09/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 Reference:
-[1] Huang T, Zhang Z, Zhang B, et al. FiBiNET: Combining feature importance and bilinear
-feature interaction for click-through rate prediction[C]//RecSys. 2019: 169-177.
-(https://arxiv.org/abs/1905.09433)
+- [1] Huang T, Zhang Z, Zhang B, et al. FiBiNET: Combining feature importance and bilinear feature interaction for click-through rate prediction[C]//RecSys. 2019: 169-177.
+  URL: https://arxiv.org/abs/1905.09433
 
 FiBiNET (Feature Importance and Bilinear Interaction Network) is a CTR model that
 jointly learns which fields matter most and how they interact. It first uses SENET

nextrec/models/ranking/fm.py

@@ -3,7 +3,7 @@ Date: create on 09/11/2025
 Checkpoint: edit on 23/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 Reference:
-[1] Rendle S. Factorization machines[C]//ICDM. 2010: 995-1000.
+- [1] Rendle S. Factorization machines[C]//ICDM. 2010: 995-1000.
 
 Factorization Machines (FM) capture second-order feature interactions with
 linear complexity by factorizing the pairwise interaction matrix. Each field

nextrec/models/ranking/lr.py

@@ -3,7 +3,7 @@ Date: create on 09/11/2025
 Checkpoint: edit on 23/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 Reference:
-[1] Hosmer D W, Lemeshow S, Sturdivant R X. Applied Logistic Regression.
+- [1] Hosmer D W, Lemeshow S, Sturdivant R X. Applied Logistic Regression.
 
 Logistic Regression (LR) is a classic linear baseline for CTR/ranking tasks.
 It maps each feature (dense, sparse, or sequence) into a numeric vector and

nextrec/models/ranking/masknet.py

@@ -3,8 +3,7 @@ Date: create on 09/11/2025
 Checkpoint: edit on 23/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 Reference:
-[1] Wang Z, She Q, Zhang J. MaskNet: Introducing Feature-Wise
-Multiplication to CTR Ranking Models by Instance-Guided Mask.
+- [1] Wang Z, She Q, Zhang J. MaskNet: Introducing Feature-Wise Multiplication to CTR Ranking Models by Instance-Guided Mask.
 
 MaskNet is a CTR prediction model that introduces instance-guided,
 feature-wise multiplicative interactions into deep ranking networks.

nextrec/models/ranking/pnn.py

@@ -3,8 +3,7 @@ Date: create on 09/11/2025
 Checkpoint: edit on 23/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 Reference:
-[1] Qu Y, Cai H, Ren K, et al. Product-based neural networks for user response
-prediction[C]//ICDM. 2016: 1149-1154. (https://arxiv.org/abs/1611.00144)
+- [1] Qu Y, Cai H, Ren K, et al. Product-based neural networks for user response prediction[C]//ICDM. 2016: 1149-1154. (https://arxiv.org/abs/1611.00144)
 
 Product-based Neural Networks (PNN) are CTR prediction models that explicitly
 encode feature interactions by combining:

nextrec/models/ranking/widedeep.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] Cheng H T, Koc L, Harmsen J, et al. Wide & Deep learning for recommender systems[C]
-//Proceedings of the 1st Workshop on Deep Learning for Recommender Systems. 2016: 7-10.
-(https://arxiv.org/abs/1606.07792)
+- [1] Cheng H T, Koc L, Harmsen J, et al. Wide & Deep learning for recommender systems[C] //Proceedings of the 1st Workshop on Deep Learning for Recommender Systems. 2016: 7-10.
+  URL: https://arxiv.org/abs/1606.07792
 
 Wide & Deep blends a linear wide component (memorization of cross features) with a
 deep neural network (generalization) sharing the same feature space. The wide part

nextrec/models/ranking/xdeepfm.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] Lian J, Zhou X, Zhang F, et al. xdeepfm: Combining explicit and implicit feature interactions
-for recommender systems[C]//Proceedings of the 24th ACM SIGKDD international conference on
-knowledge discovery & data mining. 2018: 1754-1763.
-(https://arxiv.org/abs/1803.05170)
+- [1] Lian J, Zhou X, Zhang F, et al. xdeepfm: Combining explicit and implicit feature interactions for recommender systems[C]//Proceedings of the 24th ACM SIGKDD international conference on knowledge discovery & data mining. 2018: 1754-1763.
+  URL: https://arxiv.org/abs/1803.05170
 
 xDeepFM is a CTR prediction model that unifies explicit and implicit
 feature interaction learning. It extends DeepFM by adding the

nextrec/models/representation/rqvae.py

@@ -4,11 +4,11 @@ Residual Quantized Variational AutoEncoder (RQ-VAE) for Generative Recommendatio
 Date: created on 11/12/2025
 Checkpoint: edit on 13/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
-Source code reference:
-[1] Tencent-Advertising-Algorithm-Competition-2025-Baseline
+Source Code Reference:
+- [1] Tencent-Advertising-Algorithm-Competition-2025-Baseline
 Reference:
-[1] Lee et al. Autoregressive Image Generation using Residual Quantization. CVPR 2022.
-[2] Zeghidour et al. SoundStream: An End-to-End Neural Audio Codec. IEEE/ACM TASLP 2021.
+- [1] Lee et al. Autoregressive Image Generation using Residual Quantization. CVPR 2022.
+- [2] Zeghidour et al. SoundStream: An End-to-End Neural Audio Codec. IEEE/ACM TASLP 2021.
 
 RQ-VAE learns hierarchical discrete representations via residual quantization.
 It encodes continuous embeddings (e.g., item/user embeddings) into multi-level

nextrec/models/retrieval/dssm.py

@@ -3,8 +3,7 @@ Date: create on 09/11/2025
 Checkpoint: edit on 18/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 Reference:
-[1] Huang P S, He X, Gao J, et al. Learning deep structured semantic models for web search using clickthrough data[C]
-//Proceedings of the 22nd ACM international conference on Information & Knowledge Management. 2013: 2333-2338.
+- [1] Huang P S, He X, Gao J, et al. Learning deep structured semantic models for web search using clickthrough data[C] //Proceedings of the 22nd ACM international conference on Information & Knowledge Management. 2013: 2333-2338.
 """
 
 from typing import Literal
@@ -40,11 +39,9 @@ class DSSM(BaseMatchModel):
         item_dense_features: list[DenseFeature] | None = None,
         item_sparse_features: list[SparseFeature] | None = None,
         item_sequence_features: list[SequenceFeature] | None = None,
-        user_dnn_hidden_units: list[int] = [256, 128, 64],
-        item_dnn_hidden_units: list[int] = [256, 128, 64],
+        user_mlp_params: dict | None = None,
+        item_mlp_params: dict | None = None,
         embedding_dim: int = 64,
-        dnn_activation: str = "relu",
-        dnn_dropout: float = 0.0,
         training_mode: Literal["pointwise", "pairwise", "listwise"] = "pointwise",
         num_negative_samples: int = 4,
         temperature: float = 1.0,
@@ -75,8 +72,17 @@ class DSSM(BaseMatchModel):
         )
 
         self.embedding_dim = embedding_dim
-        self.user_dnn_hidden_units = user_dnn_hidden_units
-        self.item_dnn_hidden_units = item_dnn_hidden_units
+        user_mlp_params = user_mlp_params or {}
+        item_mlp_params = item_mlp_params or {}
+
+        user_mlp_params.setdefault("hidden_dims", [256, 128, 64])
+        item_mlp_params.setdefault("hidden_dims", [256, 128, 64])
+        user_mlp_params.setdefault("activation", "relu")
+        user_mlp_params.setdefault("dropout", 0.0)
+        item_mlp_params.setdefault("activation", "relu")
+        item_mlp_params.setdefault("dropout", 0.0)
+        user_mlp_params.setdefault("output_dim", embedding_dim)
+        item_mlp_params.setdefault("output_dim", embedding_dim)
 
         # User tower embedding layer
         user_features = []
@@ -99,15 +105,8 @@ class DSSM(BaseMatchModel):
             for feat in user_sequence_features or []:
                 user_input_dim += feat.embedding_dim
 
-            # User DNN
-            user_dnn_units = user_dnn_hidden_units + [embedding_dim]
-            self.user_dnn = MLP(
-                input_dim=user_input_dim,
-                hidden_dims=user_dnn_units,
-                output_dim=None,
-                dropout=dnn_dropout,
-                activation=dnn_activation,
-            )
+            # User MLP
+            self.user_dnn = MLP(input_dim=user_input_dim, **user_mlp_params)
 
         # Item tower embedding layer
         item_features = []
@@ -130,15 +129,8 @@ class DSSM(BaseMatchModel):
             for feat in item_sequence_features or []:
                 item_input_dim += feat.embedding_dim
 
-            # Item DNN
-            item_dnn_units = item_dnn_hidden_units + [embedding_dim]
-            self.item_dnn = MLP(
-                input_dim=item_input_dim,
-                hidden_dims=item_dnn_units,
-                output_dim=None,
-                dropout=dnn_dropout,
-                activation=dnn_activation,
-            )
+            # Item MLP
+            self.item_dnn = MLP(input_dim=item_input_dim, **item_mlp_params)
 
         self.register_regularization_weights(
             embedding_attr="user_embedding", include_modules=["user_dnn"]

nextrec/models/retrieval/dssm_v2.py

@@ -36,11 +36,9 @@ class DSSM_v2(BaseMatchModel):
         item_dense_features: list[DenseFeature] | None = None,
         item_sparse_features: list[SparseFeature] | None = None,
         item_sequence_features: list[SequenceFeature] | None = None,
-        user_dnn_hidden_units: list[int] = [256, 128, 64],
-        item_dnn_hidden_units: list[int] = [256, 128, 64],
+        user_mlp_params: dict | None = None,
+        item_mlp_params: dict | None = None,
         embedding_dim: int = 64,
-        dnn_activation: str = "relu",
-        dnn_dropout: float = 0.0,
         training_mode: Literal["pointwise", "pairwise", "listwise"] = "pairwise",
         num_negative_samples: int = 4,
         temperature: float = 1.0,
@@ -71,8 +69,17 @@ class DSSM_v2(BaseMatchModel):
         )
 
         self.embedding_dim = embedding_dim
-        self.user_dnn_hidden_units = user_dnn_hidden_units
-        self.item_dnn_hidden_units = item_dnn_hidden_units
+        user_mlp_params = user_mlp_params or {}
+        item_mlp_params = item_mlp_params or {}
+
+        user_mlp_params.setdefault("hidden_dims", [256, 128, 64])
+        item_mlp_params.setdefault("hidden_dims", [256, 128, 64])
+        user_mlp_params.setdefault("activation", "relu")
+        user_mlp_params.setdefault("dropout", 0.0)
+        item_mlp_params.setdefault("activation", "relu")
+        item_mlp_params.setdefault("dropout", 0.0)
+        user_mlp_params.setdefault("output_dim", embedding_dim)
+        item_mlp_params.setdefault("output_dim", embedding_dim)
 
         # User tower
         user_features = []
@@ -94,14 +101,7 @@ class DSSM_v2(BaseMatchModel):
             for feat in user_sequence_features or []:
                 user_input_dim += feat.embedding_dim
 
-            user_dnn_units = user_dnn_hidden_units + [embedding_dim]
-            self.user_dnn = MLP(
-                input_dim=user_input_dim,
-                hidden_dims=user_dnn_units,
-                output_dim=None,
-                dropout=dnn_dropout,
-                activation=dnn_activation,
-            )
+            self.user_dnn = MLP(input_dim=user_input_dim, **user_mlp_params)
 
         # Item tower
         item_features = []
@@ -123,14 +123,7 @@ class DSSM_v2(BaseMatchModel):
             for feat in item_sequence_features or []:
                 item_input_dim += feat.embedding_dim
 
-            item_dnn_units = item_dnn_hidden_units + [embedding_dim]
-            self.item_dnn = MLP(
-                input_dim=item_input_dim,
-                hidden_dims=item_dnn_units,
-                output_dim=None,
-                dropout=dnn_dropout,
-                activation=dnn_activation,
-            )
+            self.item_dnn = MLP(input_dim=item_input_dim, **item_mlp_params)
 
         self.register_regularization_weights(
             embedding_attr="user_embedding", include_modules=["user_dnn"]

nextrec/models/retrieval/mind.py

@@ -3,8 +3,7 @@ Date: create on 09/11/2025
 Checkpoint: edit on 18/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 Reference:
-[1] Li C, Liu Z, Wu M, et al. Multi-interest network with dynamic routing for recommendation at Tmall[C]
-//Proceedings of the 28th ACM international conference on information and knowledge management. 2019: 2615-2623.
+- [1] Li C, Liu Z, Wu M, et al. Multi-interest network with dynamic routing for recommendation at Tmall[C] //Proceedings of the 28th ACM international conference on information and knowledge management. 2019: 2615-2623.
 """
 
 from typing import Literal
@@ -195,9 +194,7 @@ class MIND(BaseMatchModel):
         capsule_bilinear_type: int = 2,
         routing_times: int = 3,
         relu_layer: bool = False,
-        item_dnn_hidden_units: list[int] = [256, 128],
-        dnn_activation: str = "relu",
-        dnn_dropout: float = 0.0,
+        item_mlp_params: dict | None = None,
         training_mode: Literal["pointwise", "pairwise", "listwise"] = "pointwise",
         num_negative_samples: int = 100,
         temperature: float = 1.0,
@@ -229,7 +226,11 @@ class MIND(BaseMatchModel):
 
         self.embedding_dim = embedding_dim
         self.num_interests = num_interests
-        self.item_dnn_hidden_units = item_dnn_hidden_units
+        item_mlp_params = item_mlp_params or {}
+        item_mlp_params.setdefault("hidden_dims", [256, 128])
+        item_mlp_params.setdefault("activation", "relu")
+        item_mlp_params.setdefault("dropout", 0.0)
+        item_mlp_params.setdefault("output_dim", embedding_dim)
 
         user_features = []
         if user_dense_features:
@@ -291,15 +292,8 @@ class MIND(BaseMatchModel):
                 item_input_dim += feat.embedding_dim
 
             # Item DNN
-            if len(item_dnn_hidden_units) > 0:
-                item_dnn_units = item_dnn_hidden_units + [embedding_dim]
-                self.item_dnn = MLP(
-                    input_dim=item_input_dim,
-                    hidden_dims=item_dnn_units,
-                    output_dim=None,
-                    dropout=dnn_dropout,
-                    activation=dnn_activation,
-                )
+            if len(item_mlp_params["hidden_dims"]) > 0:
+                self.item_dnn = MLP(input_dim=item_input_dim, **item_mlp_params)
             else:
                 self.item_dnn = None
 

nextrec/models/retrieval/sdm.py

@@ -3,8 +3,7 @@ Date: create on 09/11/2025
 Checkpoint: edit on 18/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 Reference:
-[1] Ying H, Zhuang F, Zhang F, et al. Sequential recommender system based on hierarchical attention networks[C]
-//IJCAI. 2018: 3926-3932.
+- [1] Ying H, Zhuang F, Zhang F, et al. Sequential recommender system based on hierarchical attention networks[C] //IJCAI. 2018: 3926-3932.
 """
 
 from typing import Literal
@@ -37,14 +36,11 @@ class SDM(BaseMatchModel):
         item_sequence_features: list[SequenceFeature] | None = None,
         embedding_dim: int = 64,
         rnn_type: Literal["GRU", "LSTM"] = "GRU",
-        rnn_hidden_size: int = 64,
-        rnn_num_layers: int = 1,
-        rnn_dropout: float = 0.0,
+        rnn_params: dict | None = None,
         use_short_term: bool = True,
         use_long_term: bool = True,
-        item_dnn_hidden_units: list[int] = [256, 128],
-        dnn_activation: str = "relu",
-        dnn_dropout: float = 0.0,
+        user_mlp_params: dict | None = None,
+        item_mlp_params: dict | None = None,
         training_mode: Literal["pointwise", "pairwise", "listwise"] = "pointwise",
         num_negative_samples: int = 4,
         temperature: float = 1.0,
@@ -76,10 +72,26 @@ class SDM(BaseMatchModel):
 
         self.embedding_dim = embedding_dim
         self.rnn_type = rnn_type
-        self.rnn_hidden_size = rnn_hidden_size
         self.use_short_term = use_short_term
         self.use_long_term = use_long_term
-        self.item_dnn_hidden_units = item_dnn_hidden_units
+        rnn_params = rnn_params or {}
+        user_mlp_params = user_mlp_params or {}
+        item_mlp_params = item_mlp_params or {}
+
+        rnn_params.setdefault("hidden_size", 64)
+        rnn_params.setdefault("num_layers", 1)
+        rnn_params.setdefault("dropout", 0.0)
+        self.rnn_hidden_size = rnn_params["hidden_size"]
+
+        user_mlp_params.setdefault("hidden_dims", [self.rnn_hidden_size * 2])
+        user_mlp_params.setdefault("activation", "relu")
+        user_mlp_params.setdefault("dropout", 0.0)
+        user_mlp_params.setdefault("output_dim", embedding_dim)
+
+        item_mlp_params.setdefault("hidden_dims", [256, 128])
+        item_mlp_params.setdefault("activation", "relu")
+        item_mlp_params.setdefault("dropout", 0.0)
+        item_mlp_params.setdefault("output_dim", embedding_dim)
 
         # User tower
         user_features = []
@@ -101,25 +113,29 @@ class SDM(BaseMatchModel):
             if rnn_type == "GRU":
                 self.rnn = nn.GRU(
                     input_size=seq_emb_dim,
-                    hidden_size=rnn_hidden_size,
-                    num_layers=rnn_num_layers,
+                    hidden_size=self.rnn_hidden_size,
+                    num_layers=rnn_params["num_layers"],
                     batch_first=True,
-                    dropout=rnn_dropout if rnn_num_layers > 1 else 0.0,
+                    dropout=(
+                        rnn_params["dropout"] if rnn_params["num_layers"] > 1 else 0.0
+                    ),
                 )
             elif rnn_type == "LSTM":
                 self.rnn = nn.LSTM(
                     input_size=seq_emb_dim,
-                    hidden_size=rnn_hidden_size,
-                    num_layers=rnn_num_layers,
+                    hidden_size=self.rnn_hidden_size,
+                    num_layers=rnn_params["num_layers"],
                     batch_first=True,
-                    dropout=rnn_dropout if rnn_num_layers > 1 else 0.0,
+                    dropout=(
+                        rnn_params["dropout"] if rnn_params["num_layers"] > 1 else 0.0
+                    ),
                 )
             else:
                 raise ValueError(f"Unknown RNN type: {rnn_type}")
 
             user_final_dim = 0
             if use_long_term:
-                user_final_dim += rnn_hidden_size
+                user_final_dim += self.rnn_hidden_size
             if use_short_term:
                 user_final_dim += seq_emb_dim
 
@@ -129,13 +145,7 @@ class SDM(BaseMatchModel):
                 user_final_dim += feat.embedding_dim
 
             # User DNN to final embedding
-            self.user_dnn = MLP(
-                input_dim=user_final_dim,
-                hidden_dims=[rnn_hidden_size * 2, embedding_dim],
-                output_dim=None,
-                dropout=dnn_dropout,
-                activation=dnn_activation,
-            )
+            self.user_dnn = MLP(input_dim=user_final_dim, **user_mlp_params)
 
         # Item tower
         item_features = []
@@ -158,15 +168,8 @@ class SDM(BaseMatchModel):
                 item_input_dim += feat.embedding_dim
 
             # Item DNN
-            if len(item_dnn_hidden_units) > 0:
-                item_dnn_units = item_dnn_hidden_units + [embedding_dim]
-                self.item_dnn = MLP(
-                    input_dim=item_input_dim,
-                    hidden_dims=item_dnn_units,
-                    output_dim=None,
-                    dropout=dnn_dropout,
-                    activation=dnn_activation,
-                )
+            if len(item_mlp_params["hidden_dims"]) > 0:
+                self.item_dnn = MLP(input_dim=item_input_dim, **item_mlp_params)
             else:
                 self.item_dnn = None
 

nextrec/models/retrieval/youtube_dnn.py

@@ -3,8 +3,7 @@ Date: create on 09/11/2025
 Checkpoint: edit on 18/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 Reference:
-[1] Covington P, Adams J, Sargin E. Deep neural networks for youtube recommendations[C]
-//Proceedings of the 10th ACM conference on recommender systems. 2016: 191-198.
+- [1] Covington P, Adams J, Sargin E. Deep neural networks for youtube recommendations[C] //Proceedings of the 10th ACM conference on recommender systems. 2016: 191-198.
 """
 
 from typing import Literal
@@ -40,11 +39,9 @@ class YoutubeDNN(BaseMatchModel):
         item_dense_features: list[DenseFeature] | None = None,
         item_sparse_features: list[SparseFeature] | None = None,
         item_sequence_features: list[SequenceFeature] | None = None,
-        user_dnn_hidden_units: list[int] = [256, 128, 64],
-        item_dnn_hidden_units: list[int] = [256, 128, 64],
+        user_mlp_params: dict | None = None,
+        item_mlp_params: dict | None = None,
         embedding_dim: int = 64,
-        dnn_activation: str = "relu",
-        dnn_dropout: float = 0.0,
         training_mode: Literal["pointwise", "pairwise", "listwise"] = "listwise",
         num_negative_samples: int = 100,
         temperature: float = 1.0,
@@ -75,8 +72,17 @@ class YoutubeDNN(BaseMatchModel):
         )
 
         self.embedding_dim = embedding_dim
-        self.user_dnn_hidden_units = user_dnn_hidden_units
-        self.item_dnn_hidden_units = item_dnn_hidden_units
+        user_mlp_params = user_mlp_params or {}
+        item_mlp_params = item_mlp_params or {}
+
+        user_mlp_params.setdefault("hidden_dims", [256, 128, 64])
+        item_mlp_params.setdefault("hidden_dims", [256, 128, 64])
+        user_mlp_params.setdefault("activation", "relu")
+        user_mlp_params.setdefault("dropout", 0.0)
+        item_mlp_params.setdefault("activation", "relu")
+        item_mlp_params.setdefault("dropout", 0.0)
+        user_mlp_params.setdefault("output_dim", embedding_dim)
+        item_mlp_params.setdefault("output_dim", embedding_dim)
 
         # User tower
         user_features = []
@@ -99,14 +105,7 @@ class YoutubeDNN(BaseMatchModel):
                 # Sequence features are pooled before entering the DNN
                 user_input_dim += feat.embedding_dim
 
-            user_dnn_units = user_dnn_hidden_units + [embedding_dim]
-            self.user_dnn = MLP(
-                input_dim=user_input_dim,
-                hidden_dims=user_dnn_units,
-                output_dim=None,
-                dropout=dnn_dropout,
-                activation=dnn_activation,
-            )
+            self.user_dnn = MLP(input_dim=user_input_dim, **user_mlp_params)
 
         # Item tower
         item_features = []
@@ -128,14 +127,7 @@ class YoutubeDNN(BaseMatchModel):
             for feat in item_sequence_features or []:
                 item_input_dim += feat.embedding_dim
 
-            item_dnn_units = item_dnn_hidden_units + [embedding_dim]
-            self.item_dnn = MLP(
-                input_dim=item_input_dim,
-                hidden_dims=item_dnn_units,
-                output_dim=None,
-                dropout=dnn_dropout,
-                activation=dnn_activation,
-            )
+            self.item_dnn = MLP(input_dim=item_input_dim, **item_mlp_params)
 
         self.register_regularization_weights(
             embedding_attr="user_embedding", include_modules=["user_dnn"]

nextrec/models/sequential/hstu.py

@@ -4,8 +4,8 @@ Date: create on 01/12/2025
 Checkpoint: edit on 11/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 Reference:
-[1] Meta AI. Generative Recommenders (HSTU encoder) — https://github.com/meta-recsys/generative-recommenders
-[2] Ma W, Li P, Chen C, et al. Actions speak louder than words: Trillion-parameter sequential transducers for generative recommendations. arXiv:2402.17152.
+- [1] Meta AI. Generative Recommenders (HSTU encoder) — https://github.com/meta-recsys/generative-recommenders
+- [2] Ma W, Li P, Chen C, et al. Actions speak louder than words: Trillion-parameter sequential transducers for generative recommendations. arXiv:2402.17152.
 
 Hierarchical Sequential Transduction Unit (HSTU) is the core encoder behind
 Meta’s Generative Recommenders. It replaces softmax attention with lightweight