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

nextrec/data/__init__.py

@@ -18,7 +18,7 @@ from nextrec.data.data_utils import (
     read_table,
     load_dataframes,
 )
-from nextrec.basic.features import FeatureSpecMixin
+from nextrec.basic.features import FeatureSet
 from nextrec.data import data_utils
 from nextrec.data.dataloader import (
     TensorDictDataset,
@@ -38,7 +38,7 @@ __all__ = [
     'iter_file_chunks',
     'read_table',
     'load_dataframes',
-    'FeatureSpecMixin',
+    'FeatureSet',
     'data_utils',
     'TensorDictDataset',
     'FileDataset',

nextrec/data/data_utils.py

@@ -5,8 +5,9 @@ import numpy as np
 import pandas as pd
 import pyarrow.parquet as pq
 from pathlib import Path
+from typing import Any, Mapping, Sequence
 
-def _stack_section(batch: list[dict], section: str):
+def stack_section(batch: list[dict], section: str):
     """Stack one section of the batch (features/labels/ids)."""
     entries = [item.get(section) for item in batch if item.get(section) is not None]
     if not entries:
@@ -39,9 +40,9 @@ def collate_fn(batch):
                 "ids": first.get("ids"),
             }
         return {
-            "features": _stack_section(batch, "features") or {},
-            "labels": _stack_section(batch, "labels"),
-            "ids": _stack_section(batch, "ids"),
+            "features": stack_section(batch, "features") or {},
+            "labels": stack_section(batch, "labels"),
+            "ids": stack_section(batch, "ids"),
         }
 
     # Fallback: stack tuples/lists of tensors
@@ -190,3 +191,78 @@ def build_eval_candidates(
     eval_df = eval_df.merge(user_features, on=user_col, how='left')
     eval_df = eval_df.merge(item_features, on=item_col, how='left')
     return eval_df
+
+def batch_to_dict(batch_data: Any, include_ids: bool = True) -> dict:
+    """Standardize a dataloader batch into a dict of features, labels, and ids."""
+    if not (isinstance(batch_data, Mapping) and "features" in batch_data):
+        raise TypeError(
+            "[BaseModel-batch_to_dict Error] Batch data must be a dict with 'features' produced by the current DataLoader."
+        )
+    return {
+        "features": batch_data.get("features", {}),
+        "labels": batch_data.get("labels"),
+        "ids": batch_data.get("ids") if include_ids else None,
+    }
+
+
+# def get_user_ids(
+#     data: dict | pd.DataFrame | None, user_id_column: str = "user_id"
+# ) -> np.ndarray | None:
+#     """Extract user IDs from a dataset dict or DataFrame."""
+#     if data is None:
+#         return None
+#     if isinstance(data, pd.DataFrame) and user_id_column in data.columns:
+#         return np.asarray(data[user_id_column].values)
+#     if isinstance(data, dict) and user_id_column in data:
+#         return np.asarray(data[user_id_column])
+#     return None
+
+
+# def get_user_ids_from_batch(
+#     batch_dict: Mapping[str, Any], id_columns: Sequence[str] | None = None
+# ) -> np.ndarray | None:
+#     """Extract the prioritized user id column from a batch dict."""
+#     ids_container = batch_dict.get("ids") if isinstance(batch_dict, Mapping) else None
+#     if not ids_container:
+#         return None
+
+#     batch_user_id = None
+#     if id_columns:
+#         for id_name in id_columns:
+#             if id_name in ids_container:
+#                 batch_user_id = ids_container[id_name]
+#                 break
+#     if batch_user_id is None:
+#         batch_user_id = next(iter(ids_container.values()), None)
+#     if batch_user_id is None:
+#         return None
+
+#     if isinstance(batch_user_id, torch.Tensor):
+#         ids_np = batch_user_id.detach().cpu().numpy()
+#     else:
+#         ids_np = np.asarray(batch_user_id)
+#     if ids_np.ndim == 0:
+#         ids_np = ids_np.reshape(1)
+#     return ids_np.reshape(ids_np.shape[0])
+
+
+def get_user_ids(data, id_columns: list[str] | str | None = None) -> np.ndarray | None:
+    id_columns = id_columns if isinstance(id_columns, list) else [id_columns] if isinstance(id_columns, str) else []
+    if not id_columns:
+        return None
+
+    main_id = id_columns[0]  
+    if isinstance(data, pd.DataFrame) and main_id in data.columns:
+        arr = np.asarray(data[main_id].values)
+        return arr.reshape(arr.shape[0])
+    if isinstance(data, dict):
+        ids_container = data.get("ids")
+        if isinstance(ids_container, dict) and main_id in ids_container:
+            val = ids_container[main_id]
+            val = val.detach().cpu().numpy() if isinstance(val, torch.Tensor) else np.asarray(val)
+            return val.reshape(val.shape[0])
+        if main_id in data:
+            arr = np.asarray(data[main_id])
+            return arr.reshape(arr.shape[0])
+
+    return None

nextrec/data/dataloader.py

@@ -2,11 +2,10 @@
 Dataloader definitions
 
 Date: create on 27/10/2025
-Checkpoint: edit on 29/11/2025
+Checkpoint: edit on 02/12/2025
 Author: Yang Zhou,zyaztec@gmail.com
 """
 import os
-import tqdm
 import torch
 import logging
 import numpy as np
@@ -18,15 +17,11 @@ from typing import cast
 
 from torch.utils.data import DataLoader, Dataset, IterableDataset
 from nextrec.data.preprocessor import DataProcessor
-from nextrec.basic.features import DenseFeature, SparseFeature, SequenceFeature, FeatureSpecMixin
+from nextrec.basic.features import DenseFeature, SparseFeature, SequenceFeature, FeatureSet
 
 from nextrec.basic.loggers import colorize
-from nextrec.data import (
-    get_column_data,
-    collate_fn,
-    resolve_file_paths,
-    read_table,
-)
+from nextrec.data import get_column_data, collate_fn, resolve_file_paths, read_table
+from nextrec.utils import to_tensor
 
 class TensorDictDataset(Dataset):
     """Dataset returning sample-level dicts matching the unified batch schema."""
@@ -52,7 +47,7 @@ class TensorDictDataset(Dataset):
         sample_ids = {name: tensor[idx] for name, tensor in self.ids.items()} if self.ids else None
         return {"features": sample_features, "labels": sample_labels, "ids": sample_ids}
 
-class FileDataset(FeatureSpecMixin, IterableDataset):
+class FileDataset(FeatureSet, IterableDataset):
     def __init__(self, 
                  file_paths: list[str],                      # file paths to read, containing CSV or Parquet files
                  dense_features: list[DenseFeature],         # dense feature definitions
@@ -67,44 +62,37 @@ class FileDataset(FeatureSpecMixin, IterableDataset):
         self.chunk_size = chunk_size
         self.file_type = file_type
         self.processor = processor
-        self._set_feature_config(dense_features, sparse_features, sequence_features, target_columns, id_columns)
+        self.set_all_features(dense_features, sparse_features, sequence_features, target_columns, id_columns)
         self.current_file_index = 0
         self.total_files = len(file_paths)
     
     def __iter__(self):
         self.current_file_index = 0
-        self._file_pbar = None
-        if self.total_files > 1:
-            self._file_pbar = tqdm.tqdm(total=self.total_files, desc="Files", unit="file", position=0, leave=True, bar_format='{desc}: {percentage:3.0f}%|{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}]')
         for file_path in self.file_paths:
             self.current_file_index += 1
-            if self._file_pbar is not None:
-                self._file_pbar.update(1)
-            elif self.total_files == 1:
+            if self.total_files == 1:
                 file_name = os.path.basename(file_path)
                 logging.info(f"Processing file: {file_name}")
             if self.file_type == 'csv':
-                yield from self._read_csv_chunks(file_path)
+                yield from self.read_csv_chunks(file_path)
             elif self.file_type == 'parquet':
-                yield from self._read_parquet_chunks(file_path)
-        if self._file_pbar is not None:
-            self._file_pbar.close()
+                yield from self.read_parquet_chunks(file_path)
     
-    def _read_csv_chunks(self, file_path: str):
+    def read_csv_chunks(self, file_path: str):
         chunk_iterator = pd.read_csv(file_path, chunksize=self.chunk_size)
         for chunk in chunk_iterator:
-            tensors = self._dataframe_to_tensors(chunk)
+            tensors = self.dataframeto_tensors(chunk)
             yield tensors
     
-    def _read_parquet_chunks(self, file_path: str):
+    def read_parquet_chunks(self, file_path: str):
         parquet_file = pq.ParquetFile(file_path)
         for batch in parquet_file.iter_batches(batch_size=self.chunk_size):
             chunk = batch.to_pandas()            
-            tensors = self._dataframe_to_tensors(chunk)
+            tensors = self.dataframeto_tensors(chunk)
             yield tensors
             del chunk
                 
-    def _dataframe_to_tensors(self, df: pd.DataFrame) -> dict | None:
+    def dataframeto_tensors(self, df: pd.DataFrame) -> dict | None:
         if self.processor is not None:
             if not self.processor.is_fitted:
                 raise ValueError("[DataLoader Error] DataProcessor must be fitted before using in streaming mode")
@@ -120,7 +108,7 @@ class FileDataset(FeatureSpecMixin, IterableDataset):
         return batch
 
 
-class RecDataLoader(FeatureSpecMixin):
+class RecDataLoader(FeatureSet):
     def __init__(self,
                  dense_features: list[DenseFeature] | None = None,
                  sparse_features: list[SparseFeature] | None = None,
@@ -129,7 +117,7 @@ class RecDataLoader(FeatureSpecMixin):
                  id_columns: str | list[str] | None = None,
                  processor: DataProcessor | None = None):
         self.processor = processor
-        self._set_feature_config(dense_features, sparse_features, sequence_features, target, id_columns)
+        self.set_all_features(dense_features, sparse_features, sequence_features, target, id_columns)
 
     def create_dataloader(self,
                          data: dict | pd.DataFrame | str | DataLoader,
@@ -140,13 +128,13 @@ class RecDataLoader(FeatureSpecMixin):
         if isinstance(data, DataLoader):
             return data
         elif isinstance(data, (str, os.PathLike)):
-            return self._create_from_path(path=data, batch_size=batch_size, shuffle=shuffle, load_full=load_full, chunk_size=chunk_size)
+            return self.create_from_path(path=data, batch_size=batch_size, shuffle=shuffle, load_full=load_full, chunk_size=chunk_size)
         elif isinstance(data, (dict, pd.DataFrame)):
-            return self._create_from_memory(data=data, batch_size=batch_size, shuffle=shuffle)
+            return self.create_from_memory(data=data, batch_size=batch_size, shuffle=shuffle)
         else:
             raise ValueError(f"[RecDataLoader Error] Unsupported data type: {type(data)}")
     
-    def _create_from_memory(self, 
+    def create_from_memory(self, 
                            data: dict | pd.DataFrame,
                            batch_size: int,
                            shuffle: bool) -> DataLoader:
@@ -162,7 +150,7 @@ class RecDataLoader(FeatureSpecMixin):
         dataset = TensorDictDataset(tensors)
         return DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, collate_fn=collate_fn)
     
-    def _create_from_path(self,
+    def create_from_path(self,
                          path: str,
                          batch_size: int,
                          shuffle: bool,
@@ -179,7 +167,6 @@ class RecDataLoader(FeatureSpecMixin):
                 except OSError:
                     pass
                 try:
-                    df = read_table(file_path, file_type)
                     dfs.append(df)
                 except MemoryError as exc:
                     raise MemoryError(f"[RecDataLoader Error] Out of memory while reading {file_path}. Consider using load_full=False with streaming.") from exc
@@ -187,11 +174,11 @@ class RecDataLoader(FeatureSpecMixin):
                 combined_df = pd.concat(dfs, ignore_index=True)
             except MemoryError as exc:
                 raise MemoryError(f"[RecDataLoader Error] Out of memory while concatenating loaded data (approx {total_bytes / (1024**3):.2f} GB). Use load_full=False to stream or reduce chunk_size.") from exc
-            return self._create_from_memory(combined_df, batch_size, shuffle,)
+            return self.create_from_memory(combined_df, batch_size, shuffle,)
         else:
-            return self._load_files_streaming(file_paths, file_type, batch_size, chunk_size, shuffle)
+            return self.load_files_streaming(file_paths, file_type, batch_size, chunk_size, shuffle)
 
-    def _load_files_streaming(self,
+    def load_files_streaming(self,
                              file_paths: list[str],
                              file_type: str,
                              batch_size: int,
@@ -201,20 +188,10 @@ class RecDataLoader(FeatureSpecMixin):
             logging.warning("[RecDataLoader Warning] Shuffle is ignored in streaming mode (IterableDataset).")
         if batch_size != 1:
             logging.warning("[RecDataLoader Warning] Streaming mode enforces batch_size=1; tune chunk_size to control memory/throughput.")
-        dataset = FileDataset(
-            file_paths=file_paths,
-            dense_features=self.dense_features,
-            sparse_features=self.sparse_features,
-            sequence_features=self.sequence_features,
-            target_columns=self.target_columns,
-            id_columns=self.id_columns,
-            chunk_size=chunk_size,
-            file_type=file_type,
-            processor=self.processor
-        )
+        dataset = FileDataset(file_paths=file_paths, dense_features=self.dense_features, sparse_features=self.sparse_features, sequence_features=self.sequence_features, target_columns=self.target_columns, id_columns=self.id_columns, chunk_size=chunk_size, file_type=file_type, processor=self.processor)
         return DataLoader(dataset, batch_size=1, collate_fn=collate_fn)
 
-def _normalize_sequence_column(column, feature: SequenceFeature) -> np.ndarray:
+def normalize_sequence_column(column, feature: SequenceFeature) -> np.ndarray:
     if isinstance(column, pd.Series):
         column = column.tolist()
     if isinstance(column, (list, tuple)):
@@ -250,25 +227,27 @@ def _normalize_sequence_column(column, feature: SequenceFeature) -> np.ndarray:
         column = column.reshape(-1, 1)
     return np.asarray(column, dtype=np.int64)
 
-
-def build_tensors_from_data(  # noqa: C901
+def build_tensors_from_data( 
     data: dict | pd.DataFrame,
     raw_data: dict | pd.DataFrame,
     features: list,
     target_columns: list[str],
     id_columns: list[str]
 ) -> dict | None:
-    feature_tensors: dict[str, torch.Tensor] = {}
+    feature_tensors = {}
     for feature in features:
         column = get_column_data(data, feature.name)
         if column is None:
             raise ValueError(f"[RecDataLoader Error] Feature column '{feature.name}' not found in data")
-        if isinstance(feature, SequenceFeature):
-            tensor = torch.from_numpy(_normalize_sequence_column(column, feature))
+        if isinstance(feature, SequenceFeature): # sequence feature will do padding/truncation again to avoid the case when input data is not preprocessed
+            arr = normalize_sequence_column(column, feature)
+            tensor = to_tensor(arr, dtype=torch.long)
         elif isinstance(feature, DenseFeature):
-            tensor = torch.from_numpy(np.asarray(column, dtype=np.float32))
+            arr = np.asarray(column, dtype=np.float32)
+            tensor = to_tensor(arr, dtype=torch.float32)
         else:
-            tensor = torch.from_numpy(np.asarray(column, dtype=np.int64))
+            arr = np.asarray(column, dtype=np.int64)
+            tensor = to_tensor(arr, dtype=torch.long)
         feature_tensors[feature.name] = tensor
     label_tensors = None
     if target_columns:
@@ -277,7 +256,7 @@ def build_tensors_from_data(  # noqa: C901
             column = get_column_data(data, target_name)
             if column is None:
                 continue
-            label_tensor = torch.from_numpy(np.asarray(column, dtype=np.float32))
+            label_tensor = to_tensor(np.asarray(column, dtype=np.float32), dtype=torch.float32)
             if label_tensor.dim() == 2 and label_tensor.shape[0] == 1 and label_tensor.shape[1] > 1:
                 label_tensor = label_tensor.t()
             if label_tensor.shape[1:] == (1,):
@@ -298,7 +277,7 @@ def build_tensors_from_data(  # noqa: C901
                 id_arr = np.asarray(column, dtype=np.int64)
             except Exception as exc:
                 raise TypeError( f"[RecDataLoader Error] ID column '{id_col}' must contain numeric values. Received dtype={np.asarray(column).dtype}, error: {exc}") from exc
-            id_tensors[id_col] = torch.from_numpy(id_arr)
+            id_tensors[id_col] = to_tensor(id_arr, dtype=torch.long)
     if not feature_tensors:
         return None
     return {"features": feature_tensors, "labels": label_tensors, "ids": id_tensors}

nextrec/data/preprocessor.py

@@ -2,6 +2,7 @@
 DataProcessor for data preprocessing including numeric, sparse, sequence features and target processing.
 
 Date: create on 13/11/2025
+Checkpoint: edit on 02/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 """
 from __future__ import annotations
@@ -32,11 +33,11 @@ from nextrec.data.data_utils import (
     default_output_dir,
 )
 from nextrec.basic.session import resolve_save_path
-from nextrec.basic.features import FeatureSpecMixin
+from nextrec.basic.features import FeatureSet
 from nextrec.__version__ import __version__
 
 
-class DataProcessor(FeatureSpecMixin):
+class DataProcessor(FeatureSet):
     """DataProcessor for data preprocessing including numeric, sparse, sequence features and target processing.
     
     Examples:
@@ -70,7 +71,7 @@ class DataProcessor(FeatureSpecMixin):
         self.scalers: Dict[str, Any] = {}
         self.label_encoders: Dict[str, LabelEncoder] = {}
         self.target_encoders: Dict[str, Dict[str, int]] = {}
-        self._set_target_id_config([], [])
+        self.set_target_id([], [])
 
     def add_numeric_feature(
         self, 
@@ -129,7 +130,7 @@ class DataProcessor(FeatureSpecMixin):
             'target_type': target_type,
             'label_map': label_map
         }
-        self._set_target_id_config(list(self.target_features.keys()), [])
+        self.set_target_id(list(self.target_features.keys()), [])
         
     def _hash_string(self, s: str, hash_size: int) -> int:
         return int(hashlib.md5(str(s).encode()).hexdigest(), 16) % hash_size

nextrec/models/generative/hstu.py

@@ -344,7 +344,7 @@ class HSTU(BaseModel):
         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"])
+        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:
         """

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
         )