nextrec 0.3.1__py3-none-any.whl → 0.3.3__py3-none-any.whl

nextrec/models/generative/hstu.py

@@ -0,0 +1,399 @@
+"""
+[Info: this version is not released yet, i need to more research on source code and paper]
+Date: create on 01/12/2025
+Checkpoint: edit on 01/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.
+
+Hierarchical Sequential Transduction Unit (HSTU) is the core encoder behind 
+Meta’s Generative Recommenders. It replaces softmax attention with lightweight 
+pointwise activations, enabling extremely deep stacks on long behavior sequences.
+
+In each HSTU layer:
+  (1) Tokens are projected into four streams U, V, Q, K via a shared feed-forward block
+  (2) Softmax-free interactions combine QK^T with Relative Attention Bias (RAB) to encode distance
+  (3) Aggregated context is modulated by U-gating and mapped back through an output projection
+
+Stacking layers yields an efficient causal encoder for next-item 
+generation. With a tied-embedding LM head, HSTU forms 
+a full generative recommendation model.
+
+Key Advantages:
+- Softmax-free attention scales better on deep/long sequences
+- RAB captures temporal structure without extra attention heads
+- Causal masking and padding-aware normalization fit real logs
+- Weight tying reduces parameters and stabilizes training
+- Serves as a drop-in backbone for generative recommendation
+
+HSTU（层次化序列转导单元）是 Meta 生成式推荐的核心编码器，
+用点式激活替代 softmax 注意力，可在长序列上轻松堆叠深层结构。
+
+单层 HSTU 的主要步骤：
+  (1) 将输入一次性映射到 U、V、Q、K 四条通路
+  (2) 利用不含 softmax 的 QK^T 结合相对位置偏置（RAB）建模距离信息
+  (3) 用 U 对聚合上下文进行门控，再映射回输出空间
+
+多层堆叠后，可得到高效的因果编码器；与绑权 LM 头配合即可完成 next-item 预测。
+
+主要优势：
+- 摆脱 softmax，在长序列、深层模型上更易扩展
+- 相对位置偏置稳健刻画时序结构
+- 因果 mask 与 padding 感知归一化贴合真实日志
+- 绑权输出头降低参数量并提升训练稳定性
+- 直接作为生成式推荐的骨干网络
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from nextrec.basic.model import BaseModel
+from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
+
+
+def _relative_position_bucket(
+    relative_position: torch.Tensor,
+    num_buckets: int = 32,
+    max_distance: int = 128,
+) -> torch.Tensor:
+    """
+    map the relative position (i-j) to a bucket in [0, num_buckets).
+    """
+    # only need the negative part for causal attention
+    n = -relative_position
+    n = torch.clamp(n, min=0)
+
+    # when the distance is small, keep it exact
+    max_exact = num_buckets // 2
+    is_small = n < max_exact
+
+    # when the distance is too far, do log scaling
+    large_val = max_exact + ((torch.log(n.float() / max_exact + 1e-6) / math.log(max_distance / max_exact)) * (num_buckets - max_exact)).long()
+    large_val = torch.clamp(large_val, max=num_buckets - 1)
+
+    buckets = torch.where(is_small, n.long(), large_val)
+    return buckets
+
+
+class RelativePositionBias(nn.Module):
+    """
+    Compute relative position bias (RAB) for HSTU attention.
+    The input is the sequence length T, output is [1, num_heads, seq_len, seq_len].
+    """
+
+    def __init__(
+        self,
+        num_heads: int,
+        num_buckets: int = 32,
+        max_distance: int = 128,
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        self.num_buckets = num_buckets
+        self.max_distance = max_distance
+        self.embedding = nn.Embedding(num_buckets, num_heads)
+
+    def forward(self, seq_len: int, device: torch.device) -> torch.Tensor:
+        # positions: [T]
+        ctx = torch.arange(seq_len, device=device)[:, None]
+        mem = torch.arange(seq_len, device=device)[None, :]
+        rel_pos = mem - ctx  # a matrix to describe all relative positions for each [i,j] pair, shape = [seq_len, seq_len]
+        buckets = _relative_position_bucket(rel_pos, num_buckets=self.num_buckets, max_distance=self.max_distance,)  # map to buckets
+        values = self.embedding(buckets) # embedding vector for each [i,j] pair, shape = [seq_len, seq_len, embedding_dim=num_heads]
+        return values.permute(2, 0, 1).unsqueeze(0) # [1, num_heads, seq_len, seq_len]
+
+class HSTUPointwiseAttention(nn.Module):
+    """
+    Pointwise aggregation attention that implements HSTU without softmax:
+        1) [U, V, Q, K] = split( φ1(f1(X)) ), U: gate, V: value, Q: query, K: key
+        2) AV = φ2(QK^T + rab) V / N, av is attention-weighted value
+        3) Y  = f2( Norm(AV) ⊙ U ), y is output
+    φ1, φ2 use SiLU; Norm uses LayerNorm.
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        dropout: float = 0.1,
+        alpha: float | None = None
+    ):
+        super().__init__()
+        if d_model % num_heads != 0:
+            raise ValueError(f"[HSTUPointwiseAttention Error] d_model({d_model}) % num_heads({num_heads}) != 0")
+
+        self.d_model = d_model
+        self.num_heads = num_heads
+        self.d_head = d_model // num_heads
+        self.alpha = alpha if alpha is not None else (self.d_head ** -0.5)
+        # project input to 4 * d_model for U, V, Q, K
+        self.in_proj = nn.Linear(d_model, 4 * d_model, bias=True)
+        # project output back to d_model
+        self.out_proj = nn.Linear(d_model, d_model, bias=True)
+        self.dropout = nn.Dropout(dropout)
+        self.norm = nn.LayerNorm(d_model)
+
+    def _reshape_heads(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        [B, T, D] -> [B, H, T, d_head]
+        """
+        B, T, D = x.shape
+        return x.view(B, T, self.num_heads, self.d_head).transpose(1, 2)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,      # [T, T] with 0 or -inf
+        key_padding_mask: Optional[torch.Tensor] = None,  # [B, T], True = pad
+        rab: Optional[torch.Tensor] = None,            # [1, H, T, T] or None
+    ) -> torch.Tensor:
+        B, T, D = x.shape
+
+        # Eq.(1): [U, V, Q, K] = split( φ1(f1(X)) )
+        h = F.silu(self.in_proj(x))  # [B, T, 4D]
+        U, V, Q, K = h.chunk(4, dim=-1)  # each [B, T, D]
+
+        Qh = self._reshape_heads(Q)  # [B, H, T, d_head]
+        Kh = self._reshape_heads(K)  # [B, H, T, d_head]
+        Vh = self._reshape_heads(V)  # [B, H, T, d_head]
+        Uh = self._reshape_heads(U)  # [B, H, T, d_head]
+
+        # attention logits: QK^T (without 1/sqrt(d) and softmax)
+        logits = torch.matmul(Qh, Kh.transpose(-2, -1)) * self.alpha  # [B, H, T, T]
+
+        # add relative position bias (rab^p), and future extensible rab^t
+        if rab is not None:
+            # rab: [1, H, T, T] or [B, H, T, T]
+            logits = logits + rab
+
+        # construct an "allowed" mask to calculate N
+        # 1 indicates that the (query i, key j) pair is a valid attention pair; 0 indicates it is masked out
+        allowed = torch.ones_like(logits, dtype=torch.float)  # [B, H, T, T]
+
+        # causal mask: attn_mask is usually an upper triangular matrix of -inf with shape [T, T]
+        if attn_mask is not None:
+            allowed = allowed * (attn_mask.view(1, 1, T, T) == 0).float()
+            logits = logits + attn_mask.view(1, 1, T, T)
+
+        # padding mask: key_padding_mask is usually [B, T], True = pad
+        if key_padding_mask is not None:
+            # valid: 1 for non-pad, 0 for pad
+            valid = (~key_padding_mask).float()          # [B, T]
+            valid = valid.view(B, 1, 1, T)               # [B, 1, 1, T]
+            allowed = allowed * valid
+            logits = logits.masked_fill(valid == 0, float("-inf"))
+
+        # Eq.(2): A(X)V(X) = φ2(QK^T + rab) V(X) / N
+        attn = F.silu(logits)  # [B, H, T, T]
+        denom = allowed.sum(dim=-1, keepdim=True)  # [B, H, T, 1]
+        denom = denom.clamp(min=1.0)
+
+        attn = attn / denom  # [B, H, T, T]
+        AV = torch.matmul(attn, Vh)  # [B, H, T, d_head]
+        AV = AV.transpose(1, 2).contiguous().view(B, T, D) # reshape back to [B, T, D]
+        U_flat = Uh.transpose(1, 2).contiguous().view(B, T, D)
+        y = self.out_proj(self.dropout(self.norm(AV) * U_flat))  # [B, T, D]
+        return y
+
+
+class HSTULayer(nn.Module):
+    """
+    HSTUPointwiseAttention + Residual Connection
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        dropout: float = 0.1,
+        use_rab_pos: bool = True,
+        rab_num_buckets: int = 32,
+        rab_max_distance: int = 128,
+    ):
+        super().__init__()
+        self.attn = HSTUPointwiseAttention(d_model=d_model, num_heads=num_heads, dropout=dropout)
+        self.dropout = nn.Dropout(dropout)
+        self.use_rab_pos = use_rab_pos
+        self.rel_pos_bias = (RelativePositionBias(num_heads=num_heads, num_buckets=rab_num_buckets, max_distance=rab_max_distance) if use_rab_pos else None)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,
+        key_padding_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """
+        x: [B, T, D]
+        """
+        B, T, D = x.shape
+        device = x.device
+        rab = None
+        if self.use_rab_pos:
+            rab = self.rel_pos_bias(seq_len=T, device=device) # [1, H, T, T]
+        out = self.attn(x=x, attn_mask=attn_mask, key_padding_mask=key_padding_mask, rab=rab)
+        return x + self.dropout(out)
+
+
+class HSTU(BaseModel):
+    """
+    HSTU encoder for next-item prediction in a causal, generative setup.
+    Pipeline:
+      1) Embed tokens + positions from the behavior history
+      2) Apply stacked HSTU layers with causal mask and optional RAB
+      3) Use the last valid position to produce next-item logits via tied LM head
+    """
+
+    @property
+    def model_name(self) -> str:
+        return "HSTU"
+
+    @property
+    def task_type(self) -> str:
+        return "multiclass"
+
+    def __init__(
+        self,
+        sequence_features: list[SequenceFeature],
+        dense_features: Optional[list[DenseFeature]] = None,
+        sparse_features: Optional[list[SparseFeature]] = None,
+        d_model: Optional[int] = None,
+        num_heads: int = 8,
+        num_layers: int = 4,
+        max_seq_len: int = 200,
+        dropout: float = 0.1,
+        # RAB settings
+        use_rab_pos: bool = True,
+        rab_num_buckets: int = 32,
+        rab_max_distance: int = 128,
+
+        tie_embeddings: bool = True,
+        target: Optional[list[str] | str] = None,
+        optimizer: str = "adam",
+        optimizer_params: Optional[dict] = None,
+        scheduler: Optional[str] = None,
+        scheduler_params: Optional[dict] = None,
+        loss_params: Optional[dict] = None,
+        embedding_l1_reg: float = 0.0,
+        dense_l1_reg: float = 0.0,
+        embedding_l2_reg: float = 0.0,
+        dense_l2_reg: float = 0.0,
+        device: str = "cpu",
+        **kwargs,
+    ):
+        if not sequence_features:
+            raise ValueError("[HSTU Error] HSTU requires at least one SequenceFeature (user behavior history).")
+
+        # demo version: use the first SequenceFeature as the main sequence
+        self.history_feature = sequence_features[0]
+
+        hidden_dim = d_model or max(int(getattr(self.history_feature, "embedding_dim", 0) or 0), 32)
+        # Make hidden_dim divisible by num_heads
+        if hidden_dim % num_heads != 0:
+            hidden_dim = num_heads * math.ceil(hidden_dim / num_heads)
+
+        self.padding_idx = self.history_feature.padding_idx if self.history_feature.padding_idx is not None else 0
+        self.vocab_size = self.history_feature.vocab_size
+        self.max_seq_len = max_seq_len
+
+        super().__init__(
+            dense_features=dense_features,
+            sparse_features=sparse_features,
+            sequence_features=sequence_features,
+            target=target,
+            task=self.task_type,
+            device=device,
+            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,
+        )
+
+        # token & position embedding (paper usually includes pos embedding / RAB in encoder)
+        self.token_embedding = nn.Embedding(
+            num_embeddings=self.vocab_size,
+            embedding_dim=hidden_dim,
+            padding_idx=self.padding_idx,
+        )
+        self.position_embedding = nn.Embedding(max_seq_len, hidden_dim)
+        self.input_dropout = nn.Dropout(dropout)
+
+        # HSTU layers
+        self.layers = nn.ModuleList([HSTULayer(d_model=hidden_dim, num_heads=num_heads, dropout=dropout, use_rab_pos=use_rab_pos, 
+                                               rab_num_buckets=rab_num_buckets, rab_max_distance=rab_max_distance) for _ in range(num_layers)])
+
+        self.final_norm = nn.LayerNorm(hidden_dim)
+        self.lm_head = nn.Linear(hidden_dim, self.vocab_size, bias=False)
+        if tie_embeddings:
+            self.lm_head.weight = self.token_embedding.weight
+
+        # causal mask buffer
+        self.register_buffer("causal_mask", torch.empty(0), persistent=False)
+        self.ignore_index = self.padding_idx if self.padding_idx is not None else -100
+
+        optimizer_params = optimizer_params or {}
+        scheduler_params = scheduler_params or {}
+        loss_params = loss_params or {}
+        loss_params.setdefault("ignore_index", self.ignore_index)
+
+        self.compile(optimizer=optimizer, optimizer_params=optimizer_params, scheduler=scheduler, scheduler_params=scheduler_params, loss="crossentropy", loss_params=loss_params)
+        self.register_regularization_weights(embedding_attr="token_embedding", include_modules=["layers", "lm_head"])
+
+    def _build_causal_mask(self, seq_len: int, device: torch.device) -> torch.Tensor:
+        """
+        build causal mask of shape [T, T]: upper triangle is -inf, others are 0.
+        This will be added to the logits to simulate causal structure.
+        """
+        if self.causal_mask.numel() == 0 or self.causal_mask.size(0) < seq_len:
+            mask = torch.full((seq_len, seq_len), float("-inf"), device=device)
+            mask = torch.triu(mask, diagonal=1) 
+            self.causal_mask = mask
+        return self.causal_mask[:seq_len, :seq_len]
+
+    def _trim_sequence(self, seq: torch.Tensor) -> torch.Tensor:
+        if seq.size(1) <= self.max_seq_len:
+            return seq
+        return seq[:, -self.max_seq_len :]
+
+    def forward(self, x: dict[str, torch.Tensor]) -> torch.Tensor:
+        seq = x[self.history_feature.name].long()  # [B, T_raw]
+        seq = self._trim_sequence(seq)            # [B, T]
+
+        B, T = seq.shape
+        device = seq.device
+        # position ids: [B, T]
+        pos_ids = torch.arange(T, device=device).unsqueeze(0).expand(B, -1)
+        token_emb = self.token_embedding(seq)         # [B, T, D]
+        pos_emb = self.position_embedding(pos_ids)    # [B, T, D]
+        hidden_states = self.input_dropout(token_emb + pos_emb)
+
+        # padding mask：True = pad
+        padding_mask = seq.eq(self.padding_idx)       # [B, T]
+        attn_mask = self._build_causal_mask(seq_len=T, device=device)  # [T, T]
+
+        for layer in self.layers:
+            hidden_states = layer(x=hidden_states, attn_mask=attn_mask, key_padding_mask=padding_mask)
+        hidden_states = self.final_norm(hidden_states)  # [B, T, D]
+
+        valid_lengths = (~padding_mask).sum(dim=1)  # [B]
+        last_index = (valid_lengths - 1).clamp(min=0)
+        last_hidden = hidden_states[torch.arange(B, device=device), last_index]  # [B, D]
+
+        logits = self.lm_head(last_hidden)  # [B, vocab_size]
+        return logits
+
+    def compute_loss(self, y_pred, y_true):
+        """
+        y_true: [B] or [B, 1], the id of the next item.
+        """
+        if y_true is None:
+            raise ValueError("[HSTU-compute_loss] Training requires y_true (next item id).")
+        labels = y_true.view(-1).long()
+        return self.loss_fn[0](y_pred, labels)

nextrec/models/match/dssm.py

@@ -143,11 +143,11 @@ class DSSM(BaseMatchModel):
                 activation=dnn_activation
             )
         
-        self._register_regularization_weights(
+        self.register_regularization_weights(
             embedding_attr='user_embedding',
             include_modules=['user_dnn']
         )
-        self._register_regularization_weights(
+        self.register_regularization_weights(
             embedding_attr='item_embedding',
             include_modules=['item_dnn']
         )

nextrec/models/match/dssm_v2.py

@@ -134,11 +134,11 @@ class DSSM_v2(BaseMatchModel):
                 activation=dnn_activation
             )
         
-        self._register_regularization_weights(
+        self.register_regularization_weights(
             embedding_attr='user_embedding',
             include_modules=['user_dnn']
         )
-        self._register_regularization_weights(
+        self.register_regularization_weights(
             embedding_attr='item_embedding',
             include_modules=['item_dnn']
         )

nextrec/models/match/mind.py

@@ -258,11 +258,11 @@ class MIND(BaseMatchModel):
             else:
                 self.item_dnn = None
         
-        self._register_regularization_weights(
+        self.register_regularization_weights(
             embedding_attr='user_embedding',
             include_modules=['capsule_network']
         )
-        self._register_regularization_weights(
+        self.register_regularization_weights(
             embedding_attr='item_embedding',
             include_modules=['item_dnn'] if self.item_dnn else []
         )

nextrec/models/match/sdm.py

@@ -176,11 +176,11 @@ class SDM(BaseMatchModel):
             else:
                 self.item_dnn = None
         
-        self._register_regularization_weights(
+        self.register_regularization_weights(
             embedding_attr='user_embedding',
             include_modules=['rnn', 'user_dnn']
         )
-        self._register_regularization_weights(
+        self.register_regularization_weights(
             embedding_attr='item_embedding',
             include_modules=['item_dnn'] if self.item_dnn else []
         )

nextrec/models/match/youtube_dnn.py

@@ -140,11 +140,11 @@ class YoutubeDNN(BaseMatchModel):
                 activation=dnn_activation
             )
         
-        self._register_regularization_weights(
+        self.register_regularization_weights(
             embedding_attr='user_embedding',
             include_modules=['user_dnn']
         )
-        self._register_regularization_weights(
+        self.register_regularization_weights(
             embedding_attr='item_embedding',
             include_modules=['item_dnn']
         )

nextrec/models/multi_task/esmm.py

@@ -128,7 +128,7 @@ class ESMM(BaseModel):
         self.cvr_tower = MLP(input_dim=input_dim, output_layer=True, **cvr_params)
         self.prediction_layer = PredictionLayer(task_type=self.task_type, task_dims=[1, 1])
         # Register regularization weights
-        self._register_regularization_weights(embedding_attr='embedding', include_modules=['ctr_tower', 'cvr_tower'])
+        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):

nextrec/models/multi_task/mmoe.py

@@ -146,7 +146,7 @@ class MMOE(BaseModel):
             self.towers.append(tower)
         self.prediction_layer = PredictionLayer(task_type=self.task_type, task_dims=[1] * self.num_tasks)
         # Register regularization weights
-        self._register_regularization_weights(embedding_attr='embedding', include_modules=['experts', 'gates', 'towers'])
+        self.register_regularization_weights(embedding_attr='embedding', include_modules=['experts', 'gates', 'towers'])
         self.compile(optimizer=optimizer, optimizer_params=optimizer_params, loss=loss, loss_params=loss_params,)
 
     def forward(self, x):

nextrec/models/multi_task/ple.py

@@ -249,7 +249,7 @@ class PLE(BaseModel):
             self.towers.append(tower)
         self.prediction_layer = PredictionLayer(task_type=self.task_type, task_dims=[1] * self.num_tasks)
         # Register regularization weights
-        self._register_regularization_weights(embedding_attr='embedding', include_modules=['cgc_layers', 'towers'])
+        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):

nextrec/models/multi_task/poso.py

@@ -389,7 +389,7 @@ class POSO(BaseModel):
             self.tower_heads = None
         self.prediction_layer = PredictionLayer(task_type=self.task_type, task_dims=[1] * self.num_tasks,)
         include_modules = ["towers", "tower_heads"] if self.architecture == "mlp" else ["mmoe", "towers"]
-        self._register_regularization_weights(embedding_attr="embedding", include_modules=include_modules)
+        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):

nextrec/models/multi_task/share_bottom.py

@@ -122,7 +122,7 @@ class ShareBottom(BaseModel):
             self.towers.append(tower)
         self.prediction_layer = PredictionLayer(task_type=self.task_type, task_dims=[1] * self.num_tasks)
         # Register regularization weights
-        self._register_regularization_weights(embedding_attr='embedding', include_modules=['bottom', 'towers'])
+        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):

nextrec/models/ranking/afm.py

@@ -81,7 +81,7 @@ class AFM(BaseModel):
         self.prediction_layer = PredictionLayer(task_type=self.task_type)
 
         # Register regularization weights
-        self._register_regularization_weights(
+        self.register_regularization_weights(
             embedding_attr='embedding',
             include_modules=['linear', 'attention_linear', 'attention_p', 'output_projection']
         )

nextrec/models/ranking/autoint.py

@@ -150,7 +150,7 @@ class AutoInt(BaseModel):
         self.prediction_layer = PredictionLayer(task_type=self.task_type)
 
         # Register regularization weights
-        self._register_regularization_weights(
+        self.register_regularization_weights(
             embedding_attr='embedding',
             include_modules=['projection_layers', 'attention_layers', 'fc']
         )

nextrec/models/ranking/dcn.py

@@ -109,7 +109,7 @@ class DCN(BaseModel):
         self.prediction_layer = PredictionLayer(task_type=self.task_type)
 
         # Register regularization weights
-        self._register_regularization_weights(
+        self.register_regularization_weights(
             embedding_attr='embedding',
             include_modules=['cross_network', 'mlp', 'final_layer']
         )

nextrec/models/ranking/deepfm.py

@@ -107,7 +107,7 @@ class DeepFM(BaseModel):
         self.prediction_layer = PredictionLayer(task_type=self.task_type)
 
         # Register regularization weights
-        self._register_regularization_weights(embedding_attr='embedding', include_modules=['linear', 'mlp'])
+        self.register_regularization_weights(embedding_attr='embedding', include_modules=['linear', 'mlp'])
         self.compile(optimizer=optimizer, optimizer_params=optimizer_params, loss=loss, loss_params=loss_params)
 
     def forward(self, x):

nextrec/models/ranking/dien.py

@@ -237,7 +237,7 @@ class DIEN(BaseModel):
         self.mlp = MLP(input_dim=mlp_input_dim, **mlp_params)
         self.prediction_layer = PredictionLayer(task_type=self.task_type)
         # Register regularization weights
-        self._register_regularization_weights(embedding_attr='embedding', include_modules=['interest_extractor', 'interest_evolution', 'attention_layer', 'mlp', 'candidate_proj'])
+        self.register_regularization_weights(embedding_attr='embedding', include_modules=['interest_extractor', 'interest_evolution', 'attention_layer', 'mlp', 'candidate_proj'])
         self.compile(optimizer=optimizer, optimizer_params=optimizer_params, loss=loss, loss_params=loss_params)
 
     def forward(self, x):

nextrec/models/ranking/din.py

@@ -108,7 +108,7 @@ class DIN(BaseModel):
         self.prediction_layer = PredictionLayer(task_type=self.task_type)
 
         # Register regularization weights
-        self._register_regularization_weights(
+        self.register_regularization_weights(
             embedding_attr='embedding',
             include_modules=['attention', 'mlp', 'candidate_attention_proj']
         )

nextrec/models/ranking/fibinet.py

@@ -104,7 +104,7 @@ class FiBiNET(BaseModel):
         self.prediction_layer = PredictionLayer(task_type=self.task_type)
 
         # Register regularization weights
-        self._register_regularization_weights(
+        self.register_regularization_weights(
             embedding_attr='embedding',
             include_modules=['linear', 'senet', 'bilinear_standard', 'bilinear_senet', 'mlp']
         )

nextrec/models/ranking/fm.py

@@ -69,7 +69,7 @@ class FM(BaseModel):
         self.prediction_layer = PredictionLayer(task_type=self.task_type)
 
         # Register regularization weights
-        self._register_regularization_weights(
+        self.register_regularization_weights(
             embedding_attr='embedding',
             include_modules=['linear']
         )

nextrec/models/ranking/masknet.py

@@ -234,10 +234,10 @@ class MaskNet(BaseModel):
         self.prediction_layer = PredictionLayer(task_type=self.task_type)
 
         if self.model_type == "serial":
-            self._register_regularization_weights(embedding_attr="embedding", include_modules=["mask_blocks", "output_layer"],)
+            self.register_regularization_weights(embedding_attr="embedding", include_modules=["mask_blocks", "output_layer"],)
         # serial
         else:
-            self._register_regularization_weights(embedding_attr="embedding", include_modules=["mask_blocks", "final_mlp"])
+            self.register_regularization_weights(embedding_attr="embedding", include_modules=["mask_blocks", "final_mlp"])
         self.compile(optimizer=optimizer, optimizer_params=optimizer_params, loss=loss, loss_params=loss_params)
 
     def forward(self, x: dict[str, torch.Tensor]) -> torch.Tensor:

nextrec/models/ranking/pnn.py

@@ -91,7 +91,7 @@ class PNN(BaseModel):
         modules = ['mlp']
         if self.product_type == "outer":
             modules.append('kernel')
-        self._register_regularization_weights(
+        self.register_regularization_weights(
             embedding_attr='embedding',
             include_modules=modules
         )

nextrec/models/ranking/widedeep.py

@@ -111,7 +111,7 @@ class WideDeep(BaseModel):
         self.mlp = MLP(input_dim=input_dim, **mlp_params)
         self.prediction_layer = PredictionLayer(task_type=self.task_type)
         # Register regularization weights
-        self._register_regularization_weights(embedding_attr='embedding', include_modules=['linear', 'mlp'])
+        self.register_regularization_weights(embedding_attr='embedding', include_modules=['linear', 'mlp'])
         self.compile(optimizer=optimizer, optimizer_params=optimizer_params, loss=loss, loss_params=loss_params)
 
     def forward(self, x):

nextrec/models/ranking/xdeepfm.py

@@ -121,7 +121,7 @@ class xDeepFM(BaseModel):
         self.prediction_layer = PredictionLayer(task_type=self.task_type)
 
         # Register regularization weights
-        self._register_regularization_weights(
+        self.register_regularization_weights(
             embedding_attr='embedding',
             include_modules=['linear', 'cin', 'mlp']
         )