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

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,31 +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):
-    """DataProcessor for data preprocessing including numeric, sparse, sequence features and target processing.
-    
-    Examples:
-        >>> processor = DataProcessor()
-        >>> processor.add_numeric_feature('age', scaler='standard')
-        >>> processor.add_sparse_feature('user_id', encode_method='hash', hash_size=10000)
-        >>> processor.add_sequence_feature('item_history', encode_method='label', max_len=50, pad_value=0)
-        >>> processor.add_target('label', target_type='binary')
-        >>> 
-        >>> # Fit and transform data
-        >>> processor.fit(train_df)
-        >>> processed_data = processor.transform(test_df)  # Returns dict of numpy arrays
-        >>> 
-        >>> # Save and load processor
-        >>> processor.save('processor.pkl')
-        >>> loaded_processor = DataProcessor.load('processor.pkl')
-        >>> 
-        >>> # Get vocabulary sizes for embedding layers
-        >>> vocab_sizes = processor.get_vocab_sizes()
-    """
+class DataProcessor(FeatureSet):
     def __init__(self):
         self.numeric_features: Dict[str, Dict[str, Any]] = {}
         self.sparse_features: Dict[str, Dict[str, Any]] = {}
@@ -70,7 +51,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,12 +110,12 @@ 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:
+    def hash_string(self, s: str, hash_size: int) -> int:
         return int(hashlib.md5(str(s).encode()).hexdigest(), 16) % hash_size
         
-    def _process_numeric_feature_fit(self, data: pd.Series, config: Dict[str, Any]):
+    def process_numeric_feature_fit(self, data: pd.Series, config: Dict[str, Any]):
         name = str(data.name)
         scaler_type = config['scaler']
         fill_na = config['fill_na']
@@ -163,7 +144,7 @@ class DataProcessor(FeatureSpecMixin):
             scaler.fit(values)
             self.scalers[name] = scaler
             
-    def _process_numeric_feature_transform(self, data: pd.Series, config: Dict[str, Any]) -> np.ndarray:
+    def process_numeric_feature_transform(self, data: pd.Series, config: Dict[str, Any]) -> np.ndarray:
         logger = logging.getLogger()
         name = str(data.name)
         scaler_type = config['scaler']
@@ -183,7 +164,7 @@ class DataProcessor(FeatureSpecMixin):
                 result = scaler.transform(values.reshape(-1, 1)).ravel()
         return result
         
-    def _process_sparse_feature_fit(self, data: pd.Series, config: Dict[str, Any]):
+    def process_sparse_feature_fit(self, data: pd.Series, config: Dict[str, Any]):
         name = str(data.name)
         encode_method = config['encode_method']
         fill_na = config['fill_na'] # <UNK>
@@ -196,7 +177,7 @@ class DataProcessor(FeatureSpecMixin):
         elif encode_method == 'hash':
             config['vocab_size'] = config['hash_size']
             
-    def _process_sparse_feature_transform(self, data: pd.Series, config: Dict[str, Any]) -> np.ndarray:
+    def process_sparse_feature_transform(self, data: pd.Series, config: Dict[str, Any]) -> np.ndarray:
         name = str(data.name)
         encode_method = config['encode_method']
         fill_na = config['fill_na']
@@ -214,11 +195,11 @@ class DataProcessor(FeatureSpecMixin):
             return encoded.to_numpy()
         if encode_method == 'hash':
             hash_size = config['hash_size']
-            hash_fn = self._hash_string
+            hash_fn = self.hash_string
             return np.fromiter((hash_fn(v, hash_size) for v in sparse_series.to_numpy()), dtype=np.int64, count=sparse_series.size,)
         return np.array([], dtype=np.int64)
             
-    def _process_sequence_feature_fit(self, data: pd.Series, config: Dict[str, Any]):
+    def process_sequence_feature_fit(self, data: pd.Series, config: Dict[str, Any]):
         name = str(data.name)
         encode_method = config['encode_method']
         separator = config['separator']
@@ -251,7 +232,7 @@ class DataProcessor(FeatureSpecMixin):
         elif encode_method == 'hash':
             config['vocab_size'] = config['hash_size']
             
-    def _process_sequence_feature_transform(self, data: pd.Series, config: Dict[str, Any]) -> np.ndarray:
+    def process_sequence_feature_transform(self, data: pd.Series, config: Dict[str, Any]) -> np.ndarray:
         """Optimized sequence transform with preallocation and cached vocab map."""
         name = str(data.name)
         encode_method = config['encode_method']
@@ -275,7 +256,7 @@ class DataProcessor(FeatureSpecMixin):
                 config['_class_to_idx'] = class_to_idx
         else:
             class_to_idx = None  # type: ignore
-        hash_fn = self._hash_string
+        hash_fn = self.hash_string
         hash_size = config.get('hash_size')
         for i, seq in enumerate(arr):
             # normalize sequence to a list of strings
@@ -300,11 +281,7 @@ class DataProcessor(FeatureSpecMixin):
             elif encode_method == 'hash':
                 if hash_size is None:
                     raise ValueError("hash_size must be set for hash encoding")
-                encoded = [
-                    hash_fn(str(token), hash_size)
-                    for token in tokens
-                    if str(token).strip()
-                ]
+                encoded = [hash_fn(str(token), hash_size) for token in tokens if str(token).strip()]
             else:
                 encoded = []
             if not encoded:
@@ -314,7 +291,7 @@ class DataProcessor(FeatureSpecMixin):
             output[i, : len(encoded)] = encoded
         return output
         
-    def _process_target_fit(self, data: pd.Series, config: Dict[str, Any]):
+    def process_target_fit(self, data: pd.Series, config: Dict[str, Any]):
         name = str(data.name)
         target_type = config['target_type']
         label_map = config.get('label_map')
@@ -333,7 +310,7 @@ class DataProcessor(FeatureSpecMixin):
                 config['label_map'] = label_map    
             self.target_encoders[name] = label_map
             
-    def _process_target_transform(self, data: pd.Series, config: Dict[str, Any]) -> np.ndarray:
+    def process_target_transform(self, data: pd.Series, config: Dict[str, Any]) -> np.ndarray:
         logger = logging.getLogger()
         name = str(data.name)
         target_type = config.get('target_type')
@@ -354,13 +331,13 @@ class DataProcessor(FeatureSpecMixin):
                     result.append(0)
             return np.array(result, dtype=np.int64 if target_type == 'multiclass' else np.float32)
     
-    def _load_dataframe_from_path(self, path: str) -> pd.DataFrame:
+    def load_dataframe_from_path(self, path: str) -> pd.DataFrame:
         """Load all data from a file or directory path into a single DataFrame."""
         file_paths, file_type = resolve_file_paths(path)
         frames = load_dataframes(file_paths, file_type)
         return pd.concat(frames, ignore_index=True) if len(frames) > 1 else frames[0]
 
-    def _extract_sequence_tokens(self, value: Any, separator: str) -> list[str]:
+    def extract_sequence_tokens(self, value: Any, separator: str) -> list[str]:
         """Extract sequence tokens from a single value."""
         if value is None:
             return []
@@ -373,7 +350,7 @@ class DataProcessor(FeatureSpecMixin):
             return [str(v) for v in value]
         return [str(value)]
 
-    def _fit_from_path(self, path: str, chunk_size: int) -> 'DataProcessor':
+    def fit_from_path(self, path: str, chunk_size: int) -> 'DataProcessor':
         """Fit processor statistics by streaming files to reduce memory usage."""
         logger = logging.getLogger()
         logger.info(colorize("Fitting DataProcessor (streaming path mode)...", color="cyan", bold=True))
@@ -432,7 +409,7 @@ class DataProcessor(FeatureSpecMixin):
                     series = chunk[name]
                     tokens = []
                     for val in series:
-                        tokens.extend(self._extract_sequence_tokens(val, separator))
+                        tokens.extend(self.extract_sequence_tokens(val, separator))
                     seq_vocab[name].update(tokens)
 
                 # target features
@@ -547,7 +524,7 @@ class DataProcessor(FeatureSpecMixin):
         logger.info(colorize("DataProcessor fitted successfully (streaming path mode)", color="green", bold=True))
         return self
 
-    def _transform_in_memory(
+    def transform_in_memory(
         self,
         data: Union[pd.DataFrame, Dict[str, Any]],
         return_dict: bool,
@@ -580,7 +557,7 @@ class DataProcessor(FeatureSpecMixin):
                 continue
             # Convert to Series for processing
             series_data = pd.Series(data_dict[name], name=name)
-            processed = self._process_numeric_feature_transform(series_data, config)
+            processed = self.process_numeric_feature_transform(series_data, config)
             result_dict[name] = processed
 
         # process sparse features
@@ -589,7 +566,7 @@ class DataProcessor(FeatureSpecMixin):
                 logger.warning(f"Sparse feature {name} not found in data")
                 continue
             series_data = pd.Series(data_dict[name], name=name)
-            processed = self._process_sparse_feature_transform(series_data, config)
+            processed = self.process_sparse_feature_transform(series_data, config)
             result_dict[name] = processed
 
         # process sequence features
@@ -598,7 +575,7 @@ class DataProcessor(FeatureSpecMixin):
                 logger.warning(f"Sequence feature {name} not found in data")
                 continue
             series_data = pd.Series(data_dict[name], name=name)
-            processed = self._process_sequence_feature_transform(series_data, config)
+            processed = self.process_sequence_feature_transform(series_data, config)
             result_dict[name] = processed
 
         # process target features
@@ -607,10 +584,10 @@ class DataProcessor(FeatureSpecMixin):
                 logger.warning(f"Target {name} not found in data")
                 continue
             series_data = pd.Series(data_dict[name], name=name)
-            processed = self._process_target_transform(series_data, config)
+            processed = self.process_target_transform(series_data, config)
             result_dict[name] = processed
 
-        def _dict_to_dataframe(result: Dict[str, np.ndarray]) -> pd.DataFrame:
+        def dict_to_dataframe(result: Dict[str, np.ndarray]) -> pd.DataFrame:
             # Convert all arrays to Series/lists at once to avoid fragmentation
             columns_dict = {}
             for key, value in result.items():
@@ -628,7 +605,7 @@ class DataProcessor(FeatureSpecMixin):
             effective_format = save_format or "parquet"
         result_df = None
         if (not return_dict) or persist:
-            result_df = _dict_to_dataframe(result_dict)
+            result_df = dict_to_dataframe(result_dict)
         if persist:
             if output_path is None:
                 raise ValueError("output_path must be provided when persisting transformed data.")
@@ -648,7 +625,7 @@ class DataProcessor(FeatureSpecMixin):
         assert result_df is not None, "DataFrame is None after transform"
         return result_df
 
-    def _transform_path(
+    def transform_path(
         self,
         input_path: str,
         output_path: Optional[str],
@@ -668,13 +645,7 @@ class DataProcessor(FeatureSpecMixin):
         saved_paths = []
         for file_path in tqdm.tqdm(file_paths, desc="Transforming files", unit="file"):
             df = read_table(file_path, file_type)
-            transformed_df = self._transform_in_memory(
-                df,
-                return_dict=False,
-                persist=False,
-                save_format=None,
-                output_path=None,
-            )
+            transformed_df = self.transform_in_memory(df, return_dict=False, persist=False, save_format=None, output_path=None)
             assert isinstance(transformed_df, pd.DataFrame), "Expected DataFrame when return_dict=False"
             source_path = Path(file_path)
             target_file = output_root / f"{source_path.stem}.{target_format}"
@@ -694,9 +665,9 @@ class DataProcessor(FeatureSpecMixin):
             uses_robust = any(cfg.get("scaler") == "robust" for cfg in self.numeric_features.values())
             if uses_robust:
                 logger.warning("Robust scaler requires full data; loading all files into memory. Consider smaller chunk_size or different scaler if memory is limited.")
-                data = self._load_dataframe_from_path(path_str)
+                data = self.load_dataframe_from_path(path_str)
             else:
-                return self._fit_from_path(path_str, chunk_size)
+                return self.fit_from_path(path_str, chunk_size)
         if isinstance(data, dict):
             data = pd.DataFrame(data)
         logger.info(colorize("Fitting DataProcessor...", color="cyan", bold=True))
@@ -704,22 +675,22 @@ class DataProcessor(FeatureSpecMixin):
             if name not in data.columns:
                 logger.warning(f"Numeric feature {name} not found in data")
                 continue
-            self._process_numeric_feature_fit(data[name], config)
+            self.process_numeric_feature_fit(data[name], config)
         for name, config in self.sparse_features.items():
             if name not in data.columns:
                 logger.warning(f"Sparse feature {name} not found in data")
                 continue
-            self._process_sparse_feature_fit(data[name], config)
+            self.process_sparse_feature_fit(data[name], config)
         for name, config in self.sequence_features.items():
             if name not in data.columns:
                 logger.warning(f"Sequence feature {name} not found in data")
                 continue
-            self._process_sequence_feature_fit(data[name], config)
+            self.process_sequence_feature_fit(data[name], config)
         for name, config in self.target_features.items():
             if name not in data.columns:
                 logger.warning(f"Target {name} not found in data")
                 continue
-            self._process_target_fit(data[name], config)
+            self.process_target_fit(data[name], config)
         self.is_fitted = True
         return self
         
@@ -735,14 +706,8 @@ class DataProcessor(FeatureSpecMixin):
         if isinstance(data, (str, os.PathLike)):
             if return_dict:
                 raise ValueError("Path transform writes files only; set return_dict=False when passing a path.")
-            return self._transform_path(str(data), output_path, save_format)
-        return self._transform_in_memory(
-            data=data,
-            return_dict=return_dict,
-            persist=output_path is not None,
-            save_format=save_format,
-            output_path=output_path,
-        )
+            return self.transform_path(str(data), output_path, save_format)
+        return self.transform_in_memory(data=data, return_dict=return_dict, persist=output_path is not None, save_format=save_format, output_path=output_path)
             
     def fit_transform(
         self, 

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
         )

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']
         )

nextrec/utils/__init__.py

@@ -1,7 +1,7 @@
 from .optimizer import get_optimizer, get_scheduler
 from .initializer import get_initializer
 from .embedding import get_auto_embedding_dim
-from .common import resolve_device
+from .common import resolve_device, to_tensor
 from . import optimizer, initializer, embedding, common
 
 __all__ = [
@@ -10,6 +10,7 @@ __all__ = [
     'get_initializer',
     'get_auto_embedding_dim',
     'resolve_device',
+    'to_tensor',
     'optimizer',
     'initializer',
     'embedding',