nextrec 0.4.32__py3-none-any.whl → 0.4.34__py3-none-any.whl

nextrec/__version__.py

@@ -1 +1 @@
-__version__ = "0.4.32"
+__version__ = "0.4.34"

nextrec/basic/activation.py

@@ -1,8 +1,8 @@
 """
-Activation function definitions for NextRec models.
+Activation function definitions.
 
 Date: create on 27/10/2025
-Checkpoint: edit on 28/12/2025
+Checkpoint: edit on 20/01/2026
 Author: Yang Zhou, zyaztec@gmail.com
 """
 
@@ -22,26 +22,24 @@ class Dice(nn.Module):
     where p(x) = sigmoid((x - E[x]) / sqrt(Var[x] + epsilon))
     """
 
-    def __init__(self, emb_size: int, epsilon: float = 1e-9):
+    def __init__(self, emb_size: int, epsilon: float = 1e-3):
         super(Dice, self).__init__()
-        self.epsilon = epsilon
         self.alpha = nn.Parameter(torch.zeros(emb_size))
-        self.bn = nn.BatchNorm1d(emb_size)
+        self.bn = nn.BatchNorm1d(emb_size, eps=epsilon)
 
     def forward(self, x):
         # x shape: (batch_size, emb_size) or (batch_size, seq_len, emb_size)
-        original_shape = x.shape
+        if x.dim() == 2:  # (B, E)
+            x_norm = self.bn(x)
+            p = torch.sigmoid(x_norm)
+            return x * (self.alpha + (1 - self.alpha) * p)
 
-        if x.dim() == 3:
-            # For 3D input (batch_size, seq_len, emb_size), reshape to 2D
-            batch_size, seq_len, emb_size = x.shape
-            x = x.view(-1, emb_size)
-        x_norm = self.bn(x)
-        p = torch.sigmoid(x_norm)
-        output = p * x + (1 - p) * self.alpha * x
-        if len(original_shape) == 3:
-            output = output.view(original_shape)
-        return output
+        if x.dim() == 3:  # (B, T, E)
+            b, t, e = x.shape
+            x2 = x.reshape(-1, e)  # (B*T, E)
+            x_norm = self.bn(x2)
+            p = torch.sigmoid(x_norm).reshape(b, t, e)
+            return x * (self.alpha + (1 - self.alpha) * p)
 
 
 def activation_layer(

nextrec/basic/asserts.py

@@ -8,7 +8,7 @@ Author: Yang Zhou, zyaztec@gmail.com
 
 from __future__ import annotations
 
-from nextrec.utils.types import TaskTypeName, TrainingModeName
+from nextrec.utils.types import TaskTypeName
 
 
 def assert_task(
@@ -49,24 +49,3 @@ def assert_task(
         raise ValueError(
             f"{model_name} requires task length {nums_task}, got {len(task)}."
         )
-
-
-def assert_training_mode(
-    training_mode: TrainingModeName | list[TrainingModeName],
-    nums_task: int,
-    *,
-    model_name: str,
-) -> None:
-    valid_modes = {"pointwise", "pairwise", "listwise"}
-    if not isinstance(training_mode, list):
-        raise TypeError(
-            f"[{model_name}-init Error] training_mode must be a list with length {nums_task}."
-        )
-    if len(training_mode) != nums_task:
-        raise ValueError(
-            f"[{model_name}-init Error] training_mode list length must match number of tasks."
-        )
-    if any(mode not in valid_modes for mode in training_mode):
-        raise ValueError(
-            f"[{model_name}-init Error] training_mode must be one of {'pointwise', 'pairwise', 'listwise'}."
-        )

nextrec/basic/callback.py

@@ -2,7 +2,7 @@
 Callback System for Training Process
 
 Date: create on 27/10/2025
-Checkpoint: edit on 27/12/2025
+Checkpoint: edit on 21/01/2026
 Author: Yang Zhou, zyaztec@gmail.com
 """
 
@@ -69,7 +69,7 @@ class Callback:
 
 class CallbackList:
     """
-    Generates a list of callbacks
+    Generates a list of callbacks, used to manage and invoke multiple callbacks during training.
     """
 
     def __init__(self, callbacks: Optional[list[Callback]] = None):

nextrec/basic/features.py

@@ -8,10 +8,9 @@ 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 to_list
+from nextrec.utils.torch_utils import to_list
+from nextrec.utils.types import EmbeddingInitType, SequenceCombinerType
 
 
 class BaseFeature:
@@ -29,15 +28,7 @@ class EmbeddingFeature(BaseFeature):
         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_type: EmbeddingInitType = "normal",
         init_params: dict | None = None,
         l1_reg: float = 0.0,
         l2_reg: float = 0.0,
@@ -73,23 +64,9 @@ class SequenceFeature(EmbeddingFeature):
         max_len: int = 50,
         embedding_name: str = "",
         embedding_dim: int | None = None,
-        combiner: Literal[
-            "mean",
-            "sum",
-            "concat",
-            "dot_attention",
-            "self_attention",
-        ] = "mean",
+        combiner: SequenceCombinerType = "mean",
         padding_idx: int = 0,
-        init_type: Literal[
-            "normal",
-            "uniform",
-            "xavier_uniform",
-            "xavier_normal",
-            "kaiming_uniform",
-            "kaiming_normal",
-            "orthogonal",
-        ] = "normal",
+        init_type: EmbeddingInitType = "normal",
         init_params: dict | None = None,
         l1_reg: float = 0.0,
         l2_reg: float = 0.0,
@@ -143,15 +120,7 @@ class SparseFeature(EmbeddingFeature):
         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_type: EmbeddingInitType = "normal",
         init_params: dict | None = None,
         l1_reg: float = 0.0,
         l2_reg: float = 0.0,

nextrec/basic/heads.py

@@ -2,7 +2,7 @@
 Task head implementations for NextRec models.
 
 Date: create on 23/12/2025
-Checkpoint: edit on 27/12/2025
+Checkpoint: edit on 22/01/2026
 Author: Yang Zhou, zyaztec@gmail.com
 """
 
@@ -24,6 +24,12 @@ class TaskHead(nn.Module):
 
     This wraps PredictionLayer so models can depend on a "Head" abstraction
     without changing their existing forward signatures.
+
+    Args:
+        task_type: The type of task(s) this head is responsible for.
+        task_dims: The dimensionality of each task's output.
+        use_bias: Whether to include a bias term in the prediction layer.
+        return_logits: Whether to return raw logits or apply activation.
     """
 
     def __init__(
@@ -56,6 +62,12 @@ class RetrievalHead(nn.Module):
 
     It computes similarity for pointwise training/inference, and returns
     raw embeddings for in-batch negative sampling in pairwise/listwise modes.
+
+    Args:
+        similarity_metric: The metric used to compute similarity between embeddings.
+        temperature: Scaling factor for similarity scores.
+        training_mode: The training mode, which can be pointwise, pairwise, or listwise.
+        apply_sigmoid: Whether to apply sigmoid activation to the similarity scores in pointwise mode.
     """
 
     def __init__(

nextrec/basic/layers.py

@@ -2,7 +2,7 @@
 Layer implementations used across NextRec.
 
 Date: create on 27/10/2025
-Checkpoint: edit on 27/12/2025
+Checkpoint: edit on 22/01/2026
 Author: Yang Zhou, zyaztec@gmail.com
 """
 
@@ -20,15 +20,13 @@ import torch.nn.functional as F
 from nextrec.basic.activation import activation_layer
 from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
 from nextrec.utils.torch_utils import get_initializer
-from nextrec.utils.types import ActivationName
+from nextrec.utils.types import ActivationName, TaskTypeName
 
 
 class PredictionLayer(nn.Module):
     def __init__(
         self,
-        task_type: (
-            Literal["binary", "regression"] | list[Literal["binary", "regression"]]
-        ) = "binary",
+        task_type: TaskTypeName | list[TaskTypeName] = "binary",
         task_dims: int | list[int] | None = None,
         use_bias: bool = True,
         return_logits: bool = False,
@@ -92,10 +90,9 @@ class PredictionLayer(nn.Module):
             if self.return_logits:
                 outputs.append(task_logits)
                 continue
-            task = task_type.lower()
-            if task == "binary":
+            if task_type == "binary":
                 outputs.append(torch.sigmoid(task_logits))
-            elif task == "regression":
+            elif task_type == "regression":
                 outputs.append(task_logits)
             else:
                 raise ValueError(
@@ -897,30 +894,7 @@ class AttentionPoolingLayer(nn.Module):
         self,
         embedding_dim: int,
         hidden_units: list = [80, 40],
-        activation: Literal[
-            "dice",
-            "relu",
-            "relu6",
-            "elu",
-            "selu",
-            "leaky_relu",
-            "prelu",
-            "gelu",
-            "sigmoid",
-            "tanh",
-            "softplus",
-            "softsign",
-            "hardswish",
-            "mish",
-            "silu",
-            "swish",
-            "hardsigmoid",
-            "tanhshrink",
-            "softshrink",
-            "none",
-            "linear",
-            "identity",
-        ] = "sigmoid",
+        activation: ActivationName = "sigmoid",
         use_softmax: bool = False,
     ):
         super().__init__()
@@ -1029,7 +1003,9 @@ class RMSNorm(torch.nn.Module):
 
 
 class DomainBatchNorm(nn.Module):
-    """Domain-specific BatchNorm (applied per-domain with a shared interface)."""
+    """
+    Domain-specific BatchNorm (applied per-domain with a shared interface).
+    """
 
     def __init__(self, num_features: int, num_domains: int):
         super().__init__()

nextrec/basic/loggers.py

@@ -2,7 +2,7 @@
 NextRec Basic Loggers
 
 Date: create on 27/10/2025
-Checkpoint: edit on 01/01/2026
+Checkpoint: edit on 22/01/2026
 Author: Yang Zhou, zyaztec@gmail.com
 """
 
@@ -99,7 +99,8 @@ def format_kv(label: str, value: Any, width: int = 34, indent: int = 0) -> str:
 
 
 def setup_logger(session_id: str | os.PathLike | None = None):
-    """Set up a logger that logs to both console and a file with ANSI formatting.
+    """
+    Set up a logger that logs to both console and a file with ANSI formatting.
     Only console output has colors; file output is stripped of ANSI codes.
 
     Logs are stored under ``log/<experiment_id>/logs`` by default. A stable

nextrec/basic/metrics.py

@@ -23,7 +23,6 @@ from sklearn.metrics import (
 )
 from nextrec.utils.types import TaskTypeName, MetricsName
 
-
 TASK_DEFAULT_METRICS = {
     "binary": ["auc", "gauc", "ks", "logloss", "accuracy", "precision", "recall", "f1"],
     "regression": ["mse", "mae", "rmse", "r2", "mape"],
@@ -334,6 +333,60 @@ def compute_map_at_k(
     return float(np.mean(aps)) if aps else 0.0
 
 
+def compute_topk_counts(
+    y_true: np.ndarray, y_pred: np.ndarray, k_percent: int
+) -> tuple[int, int, int]:
+    """Compute Top-K% sample size, hits, and positives for binary labels."""
+    y_true = (y_true > 0).astype(int)
+    n = y_true.size
+    if n == 0:
+        return 0, 0, 0
+    if k_percent <= 0:
+        return 0, 0, int(y_true.sum())
+    if k_percent >= 100:
+        k_count = n
+    else:
+        k_count = int(np.ceil(n * (k_percent / 100.0)))
+        k_count = max(k_count, 1)
+    order = np.argsort(y_pred)[::-1]
+    topk = order[:k_count]
+    hits = int(y_true[topk].sum())
+    total_pos = int(y_true.sum())
+    return k_count, hits, total_pos
+
+
+def compute_topk_precision(
+    y_true: np.ndarray, y_pred: np.ndarray, k_percent: int
+) -> float:
+    """Compute Top-K% Precision."""
+    k_count, hits, _ = compute_topk_counts(y_true, y_pred, k_percent)
+    if k_count == 0:
+        return 0.0
+    return float(hits / k_count)
+
+
+def compute_topk_recall(
+    y_true: np.ndarray, y_pred: np.ndarray, k_percent: int
+) -> float:
+    """Compute Top-K% Recall."""
+    _, hits, total_pos = compute_topk_counts(y_true, y_pred, k_percent)
+    if total_pos == 0:
+        return 0.0
+    return float(hits / total_pos)
+
+
+def compute_lift_at_k(y_true: np.ndarray, y_pred: np.ndarray, k_percent: int) -> float:
+    """Compute Lift@K from Top-K% precision and overall rate."""
+    k_count, hits, total_pos = compute_topk_counts(y_true, y_pred, k_percent)
+    if k_count == 0:
+        return 0.0
+    base_rate = total_pos / float(y_true.size)
+    if base_rate == 0.0:
+        return 0.0
+    precision = hits / float(k_count)
+    return float(precision / base_rate)
+
+
 def compute_cosine_separation(y_true: np.ndarray, y_pred: np.ndarray) -> float:
     """Compute Cosine Separation."""
     y_true = (y_true > 0).astype(int)
@@ -399,11 +452,11 @@ def configure_metrics(
         if primary_task not in TASK_DEFAULT_METRICS:
             raise ValueError(f"Unsupported task type: {primary_task}")
         metrics_list = TASK_DEFAULT_METRICS[primary_task]
-    best_metrics_mode = getbest_metric_mode(metrics_list[0], primary_task)
+    best_metrics_mode = get_best_metric_mode(metrics_list[0], primary_task)
     return metrics_list, task_specific_metrics, best_metrics_mode
 
 
-def getbest_metric_mode(first_metric: MetricsName, primary_task: TaskTypeName) -> str:
+def get_best_metric_mode(first_metric: MetricsName, primary_task: TaskTypeName) -> str:
     """Determine if metric should be maximized or minimized."""
     # Metrics that should be maximized
     if first_metric in {
@@ -429,6 +482,9 @@ def getbest_metric_mode(first_metric: MetricsName, primary_task: TaskTypeName) -
         or first_metric.startswith("mrr@")
         or first_metric.startswith("ndcg@")
         or first_metric.startswith("map@")
+        or first_metric.startswith("topk_recall@")
+        or first_metric.startswith("topk_precision@")
+        or first_metric.startswith("lift@")
     ):
         return "max"
     # Cosine separation should be maximized
@@ -457,6 +513,15 @@ def compute_single_metric(
 
     y_p_binary = (y_pred > 0.5).astype(int)
     try:
+        if metric.startswith("topk_recall@"):
+            k_percent = int(metric.split("@")[1])
+            return compute_topk_recall(y_true, y_pred, k_percent)
+        if metric.startswith("topk_precision@"):
+            k_percent = int(metric.split("@")[1])
+            return compute_topk_precision(y_true, y_pred, k_percent)
+        if metric.startswith("lift@"):
+            k_percent = int(metric.split("@")[1])
+            return compute_lift_at_k(y_true, y_pred, k_percent)
         if metric.startswith("recall@"):
             k = int(metric.split("@")[1])
             return compute_recall_at_k(y_true, y_pred, user_ids, k)  # type: ignore
@@ -650,7 +715,23 @@ def evaluate_metrics(
                 allowed_metrics = metric_allowlist.get(task_type)
             for metric in metrics:
                 if allowed_metrics is not None and metric not in allowed_metrics:
-                    continue
+                    if metric.startswith(
+                        (
+                            "recall@",
+                            "precision@",
+                            "hitrate@",
+                            "hr@",
+                            "mrr@",
+                            "ndcg@",
+                            "map@",
+                            "topk_recall@",
+                            "topk_precision@",
+                            "lift@",
+                        )
+                    ):
+                        pass
+                    else:
+                        continue
                 y_true_task = y_true[:, task_idx]
                 y_pred_task = y_pred[:, task_idx]
                 task_user_ids = user_ids

nextrec/basic/model.py

@@ -2,7 +2,7 @@
 Base Model & Base Match Model Class
 
 Date: create on 27/10/2025
-Checkpoint: edit on 01/01/2026
+Checkpoint: edit on 22/01/2026
 Author: Yang Zhou,zyaztec@gmail.com
 """
 
@@ -155,9 +155,9 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
             session_id: Session id for logging. If None, a default id with timestamps will be created. e.g., 'session_tutorial'.
 
             distributed: Enable DistributedDataParallel flow, set True to enable distributed training.
-            rank: Global rank (defaults to env RANK).
-            world_size: Number of processes (defaults to env WORLD_SIZE).
-            local_rank: Local rank for selecting CUDA device (defaults to env LOCAL_RANK).
+            rank: Global rank (defaults to env RANK). e.g., 0 for the main process.
+            world_size: Number of processes (defaults to env WORLD_SIZE). e.g., 4 for a 4-process training.
+            local_rank: Local rank for selecting CUDA device (defaults to env LOCAL_RANK). e.g., 0 for the first GPU.
             ddp_find_unused_parameters: Default False, set it True only when exist unused parameters in ddp model, in most cases should be False.
 
         Note:
@@ -933,6 +933,13 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
 
         existing_callbacks = self.callbacks.callbacks
 
+        has_validation = valid_data is not None or valid_split is not None
+        checkpoint_monitor = monitor_metric
+        checkpoint_mode = self.best_metrics_mode
+        if not has_validation:
+            checkpoint_monitor = "loss"
+            checkpoint_mode = "min"
+
         if self.early_stop_patience > 0 and not any(
             isinstance(cb, EarlyStopper) for cb in existing_callbacks
         ):
@@ -946,6 +953,8 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
                 )
             )
 
+        has_validation = valid_data is not None or valid_split is not None
+
         if self.is_main_process and not any(
             isinstance(cb, CheckpointSaver) for cb in existing_callbacks
         ):
@@ -953,9 +962,9 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
                 CheckpointSaver(
                     best_path=self.best_path,
                     checkpoint_path=self.checkpoint_path,
-                    monitor=monitor_metric,
-                    mode=self.best_metrics_mode,
-                    save_best_only=True,
+                    monitor=checkpoint_monitor,
+                    mode=checkpoint_mode,
+                    save_best_only=has_validation,
                     verbose=1,
                 )
             )
@@ -1246,11 +1255,6 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
                         epoch_logs[f"val_{k}"] = v
             else:
                 epoch_logs = {**train_log_payload}
-                if self.is_main_process:
-                    self.save_model(
-                        self.checkpoint_path, add_timestamp=False, verbose=False
-                    )
-                    self.best_checkpoint_path = self.checkpoint_path
 
             # Call on_epoch_end for all callbacks (handles early stopping, checkpointing, lr scheduling)
             self.callbacks.on_epoch_end(epoch, epoch_logs)
@@ -1347,6 +1351,9 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
             nn.utils.clip_grad_norm_(params, self.max_gradient_norm)
             self.optimizer_fn.step()
             if self.grad_norm is not None:
+                # Synchronize GradNorm buffers across DDP ranks before stepping
+                if self.distributed and dist.is_available() and dist.is_initialized():
+                    self.grad_norm.sync()
                 self.grad_norm.step()
             accumulated_loss += loss.item()