nextrec 0.4.20__py3-none-any.whl → 0.4.21__py3-none-any.whl

nextrec/__version__.py

@@ -1 +1 @@
-__version__ = "0.4.20"
+__version__ = "0.4.21"

nextrec/basic/activation.py

@@ -1,14 +1,17 @@
 """
-Activation function definitions
+Activation function definitions for NextRec models.
 
 Date: create on 27/10/2025
-Checkpoint: edit on 29/11/2025
+Checkpoint: edit on 28/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 """
 
 import torch
 import torch.nn as nn
 
+from typing import Literal
+
+from nextrec.utils.types import ActivationName
 
 class Dice(nn.Module):
     """
@@ -41,9 +44,11 @@ class Dice(nn.Module):
         return output
 
 
-def activation_layer(activation: str, emb_size: int | None = None):
+def activation_layer(
+    activation: ActivationName = "none",
+    emb_size: int | None = None,
+):
     """Create an activation layer based on the given activation name."""
-    activation = activation.lower()
     if activation == "dice":
         if emb_size is None:
             raise ValueError(

nextrec/basic/callback.py

@@ -2,7 +2,7 @@
 Callback System for Training Process
 
 Date: create on 27/10/2025
-Checkpoint: edit on 19/12/2025
+Checkpoint: edit on 27/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 """
 
@@ -61,16 +61,16 @@ class Callback:
         self.params = params
 
     def should_run(self) -> bool:
-        if not getattr(self, "run_on_main_process_only", False):
+        if not self.run_on_main_process_only:
             return True
-        model = getattr(self, "model", None)
-        if model is None:
-            return True
-        return bool(getattr(model, "is_main_process", True))
+        model = self.model
+        return bool(model.is_main_process)
 
 
 class CallbackList:
-    """Generates a list of callbacks"""
+    """
+    Generates a list of callbacks
+    """
 
     def __init__(self, callbacks: Optional[list[Callback]] = None):
         self.callbacks = callbacks or []
@@ -85,7 +85,8 @@ class CallbackList:
             getattr(callback, fn_name)(*args, **kwargs)
 
     def set_model(self, model):
-        self.call("set_model", model)
+        for callback in self.callbacks:
+            callback.set_model(model)
 
     def set_params(self, params: dict):
         self.call("set_params", params)
@@ -194,9 +195,8 @@ class EarlyStopper(Callback):
             self.wait += 1
             if self.wait >= self.patience:
                 self.stopped_epoch = epoch
-                if hasattr(self.model, "stop_training"):
-                    self.model.stop_training = True
-                if self.verbose > 0:
+                self.model.stop_training = True
+                if self.verbose == 1:
                     logging.info(
                         f"Early stopping triggered at epoch {epoch + 1}. "
                         f"Best {self.monitor}: {self.best_value:.6f} at epoch {self.best_epoch + 1}"
@@ -218,14 +218,15 @@ class EarlyStopper(Callback):
 
 
 class CheckpointSaver(Callback):
-    """Callback to save model checkpoints during training.
+    """
+    Callback to save model checkpoints during training.
 
     Args:
         save_path: Path to save checkpoints.
         monitor: Metric name to monitor for saving best model.
         mode: One of {'min', 'max'}.
         save_best_only: If True, only save when the model is considered the "best".
-        save_freq: Frequency of checkpoint saving ('epoch' or integer for every N epochs).
+        save_freq: Frequency of checkpoint saving (integer for every N epochs).
         verbose: Verbosity mode.
         run_on_main_process_only: Whether to run this callback only on the main process in DDP.
     """
@@ -237,7 +238,7 @@ class CheckpointSaver(Callback):
         monitor: str = "val_auc",
         mode: str = "max",
         save_best_only: bool = False,
-        save_freq: str | int = "epoch",
+        save_freq: int = 1,
         verbose: int = 1,
         run_on_main_process_only: bool = True,
     ):
@@ -272,7 +273,7 @@ class CheckpointSaver(Callback):
         logs = logs or {}
 
         should_save = False
-        if self.save_freq == "epoch":
+        if self.save_freq == 1:
             should_save = True
         elif isinstance(self.save_freq, int) and (epoch + 1) % self.save_freq == 0:
             should_save = True
@@ -306,12 +307,10 @@ class CheckpointSaver(Callback):
 
     def save_checkpoint(self, path: Path, epoch: int, logs: dict):
 
-        # Get the actual model (unwrap DDP if needed)
-        model_to_save = (
-            self.model.ddp_model.module
-            if getattr(self.model, "ddp_model", None) is not None
-            else self.model
-        )
+        if hasattr(self.model, "ddp_model") and self.model.ddp_model is not None:
+            model_to_save = self.model.ddp_model.module
+        else:
+            model_to_save = self.model
 
         # Save only state_dict to match BaseModel.save_model() format
         torch.save(model_to_save.state_dict(), path)
@@ -328,12 +327,13 @@ class CheckpointSaver(Callback):
             with open(config_path, "wb") as f:
                 pickle.dump(features_config, f)
 
-        if self.verbose > 1:
+        if self.verbose == 1:
             logging.info(f"Saved checkpoint to {path}")
 
 
 class LearningRateScheduler(Callback):
-    """Callback for learning rate scheduling.
+    """
+    Callback for learning rate scheduling.
 
     Args:
         scheduler: Learning rate scheduler instance or name.
@@ -346,73 +346,25 @@ class LearningRateScheduler(Callback):
         self.verbose = verbose
 
     def on_train_begin(self, logs: Optional[dict] = None):
-        if self.scheduler is None and hasattr(self.model, "scheduler_fn"):
+        if self.scheduler is None:
             self.scheduler = self.model.scheduler_fn
 
     def on_epoch_end(self, epoch: int, logs: Optional[dict] = None):
         if self.scheduler is not None:
-            # Get current lr before step
-            if hasattr(self.model, "optimizer_fn"):
-                old_lr = self.model.optimizer_fn.param_groups[0]["lr"]
-
-            # Step the scheduler
-            if hasattr(self.scheduler, "step"):
-                # Some schedulers need metrics
-                if logs is None:
-                    logs = {}
-                if "val_loss" in logs and hasattr(self.scheduler, "mode"):
-                    self.scheduler.step(logs["val_loss"])
-                else:
-                    self.scheduler.step()
+            old_lr = self.model.optimizer_fn.param_groups[0]["lr"]
+            if logs is None:
+                logs = {}
 
-            # Log new lr
-            if self.verbose > 0 and hasattr(self.model, "optimizer_fn"):
-                if getattr(self.model, "is_main_process", True):
-                    new_lr = self.model.optimizer_fn.param_groups[0]["lr"]
-                    if new_lr != old_lr:
-                        logging.info(
-                            f"Learning rate changed from {old_lr:.6e} to {new_lr:.6e}"
-                        )
+            # step for ReduceLROnPlateau
+            if "val_loss" in logs and hasattr(self.scheduler, "mode"):
+                self.scheduler.step(logs["val_loss"])
+            else:
+                self.scheduler.step()
 
-
-class MetricsLogger(Callback):
-    """Callback for logging training metrics.
-
-    Args:
-        log_freq: Frequency of logging ('epoch', 'batch', or integer for every N epochs/batches).
-        verbose: Verbosity mode.
-    """
-
-    def __init__(self, log_freq: str | int = "epoch", verbose: int = 1):
-        super().__init__()
-        self.run_on_main_process_only = True
-        self.log_freq = log_freq
-        self.verbose = verbose
-
-    def on_epoch_end(self, epoch: int, logs: Optional[dict] = None):
-        if self.verbose > 0 and (
-            self.log_freq == "epoch"
-            or (isinstance(self.log_freq, int) and (epoch + 1) % self.log_freq == 0)
-        ):
-            logs = logs or {}
-            metrics_str = " - ".join(
-                [
-                    f"{k}: {v:.6f}" if isinstance(v, float) else f"{k}: {v}"
-                    for k, v in logs.items()
-                ]
-            )
-            logging.info(f"Epoch {epoch + 1}: {metrics_str}")
-
-    def on_batch_end(self, batch: int, logs: Optional[dict] = None):
-        if self.verbose > 1 and (
-            self.log_freq == "batch"
-            or (isinstance(self.log_freq, int) and (batch + 1) % self.log_freq == 0)
-        ):
-            logs = logs or {}
-            metrics_str = " - ".join(
-                [
-                    f"{k}: {v:.6f}" if isinstance(v, float) else f"{k}: {v}"
-                    for k, v in logs.items()
-                ]
-            )
-            logging.info(f"Batch {batch}: {metrics_str}")
+            # Log new lr
+            if self.verbose == 1:
+                new_lr = self.model.optimizer_fn.param_groups[0]["lr"]
+                if new_lr != old_lr:
+                    logging.info(
+                        f"Learning rate changed from {old_lr:.6e} to {new_lr:.6e}"
+                    )

nextrec/basic/features.py

@@ -1,15 +1,17 @@
 """
-Feature definitions
+Feature definitions for NextRec models.
 
 Date: create on 27/10/2025
-Checkpoint: edit on 20/12/2025
+Checkpoint: edit on 27/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 """
 
 import torch
 
+from typing import Literal
+
 from nextrec.utils.embedding import get_auto_embedding_dim
-from nextrec.utils.feature import normalize_to_list
+from nextrec.utils.feature import to_list
 
 
 class BaseFeature:
@@ -25,12 +27,20 @@ class EmbeddingFeature(BaseFeature):
         name: str,
         vocab_size: int,
         embedding_name: str = "",
-        embedding_dim: int | None = 4,
-        padding_idx: int | None = None,
-        init_type: str = "normal",
+        embedding_dim: int | None = None,
+        padding_idx: int = 0,
+        init_type: Literal[
+            "normal",
+            "uniform",
+            "xavier_uniform",
+            "xavier_normal",
+            "kaiming_uniform",
+            "kaiming_normal",
+            "orthogonal",
+        ] = "normal",
         init_params: dict | None = None,
         l1_reg: float = 0.0,
-        l2_reg: float = 1e-5,
+        l2_reg: float = 0.0,
         trainable: bool = True,
         pretrained_weight: torch.Tensor | None = None,
         freeze_pretrained: bool = False,
@@ -55,23 +65,57 @@ class EmbeddingFeature(BaseFeature):
 
 
 class SequenceFeature(EmbeddingFeature):
+
     def __init__(
         self,
         name: str,
         vocab_size: int,
-        max_len: int = 20,
+        max_len: int = 50,
         embedding_name: str = "",
-        embedding_dim: int | None = 4,
-        combiner: str = "mean",
-        padding_idx: int | None = None,
-        init_type: str = "normal",
+        embedding_dim: int | None = None,
+        combiner: Literal[
+            "mean",
+            "sum",
+            "concat",
+            "dot_attention",
+            "self_attention",
+        ] = "mean",
+        padding_idx: int = 0,
+        init_type: Literal[
+            "normal",
+            "uniform",
+            "xavier_uniform",
+            "xavier_normal",
+            "kaiming_uniform",
+            "kaiming_normal",
+            "orthogonal",
+        ] = "normal",
         init_params: dict | None = None,
         l1_reg: float = 0.0,
-        l2_reg: float = 1e-5,
+        l2_reg: float = 0.0,
         trainable: bool = True,
         pretrained_weight: torch.Tensor | None = None,
         freeze_pretrained: bool = False,
     ):
+        """
+        Sequence feature for variable-length categorical id sequences.
+
+        Args:
+            name: Feature name used as input key.
+            vocab_size: Number of unique ids in the sequence vocabulary.
+            max_len: Maximum sequence length for padding/truncation.
+            embedding_name: Shared embedding table name. Defaults to ``name``.
+            embedding_dim: Embedding dimension. Set to ``None`` for auto sizing.
+            combiner: Pooling method for sequence embeddings, e.g. ``"mean"`` or ``"sum"``.
+            padding_idx: Index used for padding tokens.
+            init_type: Embedding initializer type.
+            init_params: Initializer parameters.
+            l1_reg: L1 regularization weight on embedding.
+            l2_reg: L2 regularization weight on embedding.
+            trainable: Whether the embedding is trainable. [TODO] This is for representation learning.
+            pretrained_weight: Optional pretrained embedding weights. [TODO] This is for representation learning.
+            freeze_pretrained: If True, keep pretrained weights frozen. [TODO] This is for representation learning.
+        """
         super().__init__(
             name=name,
             vocab_size=vocab_size,
@@ -91,28 +135,105 @@ class SequenceFeature(EmbeddingFeature):
 
 
 class SparseFeature(EmbeddingFeature):
-    pass
+
+    def __init__(
+        self,
+        name: str,
+        vocab_size: int,
+        embedding_name: str = "",
+        embedding_dim: int | None = None,
+        padding_idx: int = 0,
+        init_type: Literal[
+            "normal",
+            "uniform",
+            "xavier_uniform",
+            "xavier_normal",
+            "kaiming_uniform",
+            "kaiming_normal",
+            "orthogonal",
+        ] = "normal",
+        init_params: dict | None = None,
+        l1_reg: float = 0.0,
+        l2_reg: float = 0.0,
+        trainable: bool = True,
+        pretrained_weight: torch.Tensor | None = None,
+        freeze_pretrained: bool = False,
+    ):
+        """
+        Sparse feature for categorical ids.
+
+        Args:
+            name: Feature name used as input key.
+            vocab_size: Number of unique categorical ids.
+            embedding_name: Shared embedding table name. Defaults to ``name``.
+            embedding_dim: Embedding dimension. Set to ``None`` for auto sizing.
+            padding_idx: Index used for padding tokens.
+            init_type: Embedding initializer type.
+            init_params: Initializer parameters.
+            l1_reg: L1 regularization weight on embedding.
+            l2_reg: L2 regularization weight on embedding.
+            trainable: Whether the embedding is trainable.
+            pretrained_weight: Optional pretrained embedding weights.
+            freeze_pretrained: If True, keep pretrained weights frozen.
+        """
+        super().__init__(
+            name=name,
+            vocab_size=vocab_size,
+            embedding_name=embedding_name,
+            embedding_dim=embedding_dim,
+            padding_idx=padding_idx,
+            init_type=init_type,
+            init_params=init_params,
+            l1_reg=l1_reg,
+            l2_reg=l2_reg,
+            trainable=trainable,
+            pretrained_weight=pretrained_weight,
+            freeze_pretrained=freeze_pretrained,
+        )
 
 
 class DenseFeature(BaseFeature):
+
     def __init__(
         self,
         name: str,
-        embedding_dim: int | None = 1,
         input_dim: int = 1,
-        use_embedding: bool = False,
+        proj_dim: int | None = 0,
+        use_projection: bool = False,
+        trainable: bool = True,
+        pretrained_weight: torch.Tensor | None = None,
+        freeze_pretrained: bool = False,
     ):
+        """
+        Dense feature for continuous values.
+
+        Args:
+            name: Feature name used as input key.
+            input_dim: Input dimension for continuous values.
+            proj_dim: Projection dimension. If None or 0, no projection is applied.
+            use_projection: Whether to project inputs to higher dimension.
+            trainable: Whether the projection is trainable.
+            pretrained_weight: Optional pretrained projection weights.
+            freeze_pretrained: If True, keep pretrained weights frozen.
+        """
         self.name = name
-        self.input_dim = max(int(input_dim or 1), 1)
-        self.embedding_dim = self.input_dim if embedding_dim is None else embedding_dim
-        if use_embedding and self.embedding_dim == 0:
+        self.input_dim = max(int(input_dim), 1)
+        self.proj_dim = self.input_dim if proj_dim is None else proj_dim
+        if use_projection and self.proj_dim == 0:
             raise ValueError(
-                "[Features Error] DenseFeature: use_embedding=True is incompatible with embedding_dim=0"
+                "[Features Error] DenseFeature: use_projection=True is incompatible with proj_dim=0"
             )
-        if embedding_dim is not None and embedding_dim > 1:
-            self.use_embedding = True
+        if proj_dim is not None and proj_dim > 1:
+            self.use_projection = True
         else:
-            self.use_embedding = use_embedding  # user decides for dim <= 1
+            self.use_projection = use_projection
+        self.embedding_dim = (
+            self.input_dim if not self.use_projection else self.proj_dim
+        )  # for compatibility
+
+        self.trainable = trainable
+        self.pretrained_weight = pretrained_weight
+        self.freeze_pretrained = freeze_pretrained
 
 
 class FeatureSet:
@@ -123,7 +244,7 @@ class FeatureSet:
         sequence_features: list[SequenceFeature] | None = None,
         target: str | list[str] | None = None,
         id_columns: str | list[str] | None = None,
-    ) -> None:
+    ):
         self.dense_features = list(dense_features) if dense_features else []
         self.sparse_features = list(sparse_features) if sparse_features else []
         self.sequence_features = list(sequence_features) if sequence_features else []
@@ -132,13 +253,13 @@ class FeatureSet:
             self.dense_features + self.sparse_features + self.sequence_features
         )
         self.feature_names = [feat.name for feat in self.all_features]
-        self.target_columns = normalize_to_list(target)
-        self.id_columns = normalize_to_list(id_columns)
+        self.target_columns = to_list(target)
+        self.id_columns = to_list(id_columns)
 
     def set_target_id(
         self,
         target: str | list[str] | None = None,
         id_columns: str | list[str] | None = None,
     ) -> None:
-        self.target_columns = normalize_to_list(target)
-        self.id_columns = normalize_to_list(id_columns)
+        self.target_columns = to_list(target)
+        self.id_columns = to_list(id_columns)

nextrec/basic/heads.py

@@ -2,6 +2,7 @@
 Task head implementations for NextRec models.
 
 Date: create on 23/12/2025
+Checkpoint: edit on 27/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 """
 
@@ -26,7 +27,9 @@ class TaskHead(nn.Module):
 
     def __init__(
         self,
-        task_type: str | list[str] = "binary",
+        task_type: (
+            Literal["binary", "regression"] | list[Literal["binary", "regression"]]
+        ) = "binary",
         task_dims: int | list[int] | None = None,
         use_bias: bool = True,
         return_logits: bool = False,