nextrec 0.4.33__py3-none-any.whl → 0.5.0__py3-none-any.whl

nextrec/__version__.py

@@ -1 +1 @@
-__version__ = "0.4.33"
+__version__ = "0.5.0"

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,18 @@ 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, affine=False)
 
     def forward(self, x):
-        # x shape: (batch_size, emb_size) or (batch_size, seq_len, emb_size)
-        original_shape = x.shape
-
-        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
+        # keep original shape for reshaping back after batch norm
+        orig_shape = x.shape  # x: [N, L, emb_size] or [N, emb_size]
+        x2 = x.reshape(-1, orig_shape[-1])  # x2:[N*L, emb_size] or [N, emb_size]
+        x_norm = self.bn(x2)
+        p = torch.sigmoid(x_norm).reshape(orig_shape)
+        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 25/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,
@@ -81,10 +79,12 @@ class PredictionLayer(nn.Module):
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         if x.dim() == 1:
             x = x.unsqueeze(0)  # (1 * total_dim)
-        if x.shape[-1] != self.total_dim:
-            raise ValueError(
-                f"[PredictionLayer Error]: Input last dimension ({x.shape[-1]}) does not match expected total dimension ({self.total_dim})."
-            )
+        if not torch.onnx.is_in_onnx_export():
+            if x.shape[-1] != self.total_dim:
+                raise ValueError(
+                    f"[PredictionLayer Error]: Input last dimension ({x.shape[-1]}) does not match expected total dimension ({self.total_dim})."
+                )
+
         logits = x if self.bias is None else x + self.bias
         outputs = []
         for task_type, (start, end) in zip(self.task_types, self.task_slices):
@@ -92,10 +92,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(
@@ -219,7 +218,7 @@ class EmbeddingLayer(nn.Module):
 
             elif isinstance(feature, SequenceFeature):
                 seq_input = x[feature.name].long()
-                if feature.max_len is not None and seq_input.size(1) > feature.max_len:
+                if feature.max_len is not None:
                     seq_input = seq_input[:, -feature.max_len :]
 
                 embed = self.embed_dict[feature.embedding_name]
@@ -282,10 +281,11 @@ class EmbeddingLayer(nn.Module):
             value = value.view(value.size(0), -1)  # [B, input_dim]
         input_dim = feature.input_dim
         assert_input_dim = self.dense_input_dims.get(feature.name, input_dim)
-        if value.shape[1] != assert_input_dim:
-            raise ValueError(
-                f"[EmbeddingLayer Error]:Dense feature '{feature.name}' expects {assert_input_dim} inputs but got {value.shape[1]}."
-            )
+        if not torch.onnx.is_in_onnx_export():
+            if value.shape[1] != assert_input_dim:
+                raise ValueError(
+                    f"[EmbeddingLayer Error]:Dense feature '{feature.name}' expects {assert_input_dim} inputs but got {value.shape[1]}."
+                )
         if not feature.use_projection:
             return value
         dense_layer = self.dense_transforms[feature.name]
@@ -331,29 +331,10 @@ class InputMask(nn.Module):
         feature: SequenceFeature,
         seq_tensor: torch.Tensor | None = None,
     ):
-        if seq_tensor is not None:
-            values = seq_tensor
-        else:
-            values = x[feature.name]
-        values = values.long()
+        values = seq_tensor if seq_tensor is not None else x[feature.name]
+        values = values.long().view(values.size(0), -1)
         padding_idx = feature.padding_idx if feature.padding_idx is not None else 0
-        mask = values != padding_idx
-
-        if mask.dim() == 1:
-            # [B] -> [B, 1, 1]
-            mask = mask.unsqueeze(1).unsqueeze(2)
-        elif mask.dim() == 2:
-            # [B, L] -> [B, 1, L]
-            mask = mask.unsqueeze(1)
-        elif mask.dim() == 3:
-            # [B, 1, L]
-            # [B, L, 1]  -> [B, L] -> [B, 1, L]
-            if mask.size(1) != 1 and mask.size(2) == 1:
-                mask = mask.squeeze(-1).unsqueeze(1)
-        else:
-            raise ValueError(
-                f"InputMask only supports 1D/2D/3D tensors, got shape {values.shape}"
-            )
+        mask = (values != padding_idx).unsqueeze(1)
         return mask.float()
 
 
@@ -897,30 +878,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__()
@@ -954,39 +912,22 @@ class AttentionPoolingLayer(nn.Module):
             output: [batch_size, embedding_dim] - attention pooled representation
         """
         batch_size, sequence_length, embedding_dim = keys.shape
-        assert query.shape == (
-            batch_size,
-            embedding_dim,
-        ), f"query shape {query.shape} != ({batch_size}, {embedding_dim})"
-        if mask is None and keys_length is not None:
-            # keys_length: (batch_size,)
-            device = keys.device
-            seq_range = torch.arange(sequence_length, device=device).unsqueeze(
-                0
-            )  # (1, sequence_length)
-            mask = (seq_range < keys_length.unsqueeze(1)).unsqueeze(-1).float()
-        if mask is not None:
-            if mask.dim() == 2:
-                # (B, L)
-                mask = mask.unsqueeze(-1)
-            elif (
-                mask.dim() == 3
-                and mask.shape[1] == 1
-                and mask.shape[2] == sequence_length
-            ):
-                # (B, 1, L) -> (B, L, 1)
-                mask = mask.transpose(1, 2)
-            elif (
-                mask.dim() == 3
-                and mask.shape[1] == sequence_length
-                and mask.shape[2] == 1
-            ):
-                pass
+        if mask is None:
+            if keys_length is None:
+                mask = torch.ones(
+                    (batch_size, sequence_length), device=keys.device, dtype=keys.dtype
+                )
             else:
+                device = keys.device
+                seq_range = torch.arange(sequence_length, device=device).unsqueeze(0)
+                mask = (seq_range < keys_length.unsqueeze(1)).to(keys.dtype)
+        else:
+            mask = mask.to(keys.dtype).reshape(batch_size, -1)
+            if mask.shape[1] != sequence_length:
                 raise ValueError(
                     f"[AttentionPoolingLayer Error]: Unsupported mask shape: {mask.shape}"
                 )
-            mask = mask.to(keys.dtype)
+        mask = mask.unsqueeze(-1)
         # Expand query to (B, L, D)
         query_expanded = query.unsqueeze(1).expand(-1, sequence_length, -1)
         # [query, key, query-key, query*key] -> (B, L, 4D)
@@ -1026,34 +967,3 @@ class RMSNorm(torch.nn.Module):
         variance = torch.mean(x**2, dim=-1, keepdim=True)
         x_normalized = x * torch.rsqrt(variance + self.eps)
         return self.weight * x_normalized
-
-
-class DomainBatchNorm(nn.Module):
-    """Domain-specific BatchNorm (applied per-domain with a shared interface)."""
-
-    def __init__(self, num_features: int, num_domains: int):
-        super().__init__()
-        if num_domains < 1:
-            raise ValueError("num_domains must be >= 1")
-        self.bns = nn.ModuleList(
-            [nn.BatchNorm1d(num_features) for _ in range(num_domains)]
-        )
-
-    def forward(self, x: torch.Tensor, domain_mask: torch.Tensor) -> torch.Tensor:
-        if x.dim() != 2:
-            raise ValueError("DomainBatchNorm expects 2D inputs [B, D].")
-        output = x.clone()
-        if domain_mask.dim() == 1:
-            domain_ids = domain_mask.long()
-            for idx, bn in enumerate(self.bns):
-                mask = domain_ids == idx
-                if mask.any():
-                    output[mask] = bn(x[mask])
-            return output
-        if domain_mask.dim() != 2:
-            raise ValueError("domain_mask must be 1D indices or 2D one-hot mask.")
-        for idx, bn in enumerate(self.bns):
-            mask = domain_mask[:, idx] > 0
-            if mask.any():
-                output[mask] = bn(x[mask])
-        return output

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