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

nextrec/models/retrieval/youtube_dnn.py

@@ -50,12 +50,10 @@ class YoutubeDNN(BaseMatchModel):
         num_negative_samples: int = 100,
         temperature: float = 1.0,
         similarity_metric: Literal["dot", "cosine", "euclidean"] = "dot",
-        device: str = "cpu",
-        embedding_l1_reg: float = 0.0,
-        dense_l1_reg: float = 0.0,
-        embedding_l2_reg: float = 0.0,
-        dense_l2_reg: float = 0.0,
-        early_stop_patience: int = 20,
+        embedding_l1_reg=0.0,
+        dense_l1_reg=0.0,
+        embedding_l2_reg=0.0,
+        dense_l2_reg=0.0,
         optimizer: str | torch.optim.Optimizer = "adam",
         optimizer_params: dict | None = None,
         scheduler: (
@@ -81,12 +79,10 @@ class YoutubeDNN(BaseMatchModel):
             num_negative_samples=num_negative_samples,
             temperature=temperature,
             similarity_metric=similarity_metric,
-            device=device,
             embedding_l1_reg=embedding_l1_reg,
             dense_l1_reg=dense_l1_reg,
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
-            early_stop_patience=early_stop_patience,
             **kwargs,
         )
 
@@ -169,8 +165,6 @@ class YoutubeDNN(BaseMatchModel):
             loss_params=loss_params,
         )
 
-        self.to(device)
-
     def user_tower(self, user_input: dict) -> torch.Tensor:
         """
         User tower to encode historical behavior sequences and user features.

nextrec/models/sequential/hstu.py

@@ -332,7 +332,6 @@ class HSTU(BaseModel):
         dense_l1_reg: float = 0.0,
         embedding_l2_reg: float = 0.0,
         dense_l2_reg: float = 0.0,
-        device: str = "cpu",
         **kwargs,
     ):
         raise NotImplementedError(
@@ -348,7 +347,7 @@ class HSTU(BaseModel):
         )[0]
 
         self.hidden_dim = hidden_dim or max(
-            int(getattr(self.item_history_feature, "embedding_dim", 0) or 0), 32
+            int(self.item_history_feature.embedding_dim or 0), 32
         )
         # Make hidden_dim divisible by num_heads
         if self.hidden_dim % num_heads != 0:
@@ -368,7 +367,6 @@ class HSTU(BaseModel):
             sequence_features=sequence_features,
             target=target,
             task=task or self.default_task,
-            device=device,
             embedding_l1_reg=embedding_l1_reg,
             dense_l1_reg=dense_l1_reg,
             embedding_l2_reg=embedding_l2_reg,

nextrec/utils/__init__.py

@@ -6,7 +6,7 @@ Last update: 19/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 """
 
-from . import console, data, embedding, torch_utils
+from . import console, data, embedding, loss, torch_utils
 from .config import (
     build_feature_objects,
     build_model_instance,
@@ -14,6 +14,7 @@ from .config import (
     load_model_class,
     register_processor_features,
     resolve_path,
+    safe_value,
     select_features,
 )
 from .console import (
@@ -35,23 +36,20 @@ from .data import (
     resolve_file_paths,
 )
 from .embedding import get_auto_embedding_dim
-from .feature import normalize_to_list
+from .feature import to_list
 from .model import compute_pair_scores, get_mlp_output_dim, merge_features
+from .loss import normalize_task_loss
 from .torch_utils import (
     add_distributed_sampler,
-    concat_tensors,
-    configure_device,
+    get_device,
     gather_numpy,
-    get_device_info,
     get_initializer,
     get_optimizer,
     get_scheduler,
     init_process_group,
-    pad_sequence_tensors,
-    resolve_device,
-    stack_tensors,
     to_tensor,
 )
+from .types import LossName, OptimizerName, SchedulerName, ActivationName
 
 __all__ = [
     # Console utilities
@@ -67,17 +65,12 @@ __all__ = [
     # Embedding utilities
     "get_auto_embedding_dim",
     # Device utilities (torch utils)
-    "resolve_device",
-    "get_device_info",
-    "configure_device",
+    "get_device",
     "init_process_group",
     "gather_numpy",
     "add_distributed_sampler",
     # Tensor utilities
     "to_tensor",
-    "stack_tensors",
-    "concat_tensors",
-    "pad_sequence_tensors",
     # Data utilities
     "resolve_file_paths",
     "read_table",
@@ -89,10 +82,13 @@ __all__ = [
     "merge_features",
     "get_mlp_output_dim",
     "compute_pair_scores",
+    # Loss utilities
+    "normalize_task_loss",
     # Feature utilities
-    "normalize_to_list",
+    "to_list",
     # Config utilities
     "resolve_path",
+    "safe_value",
     "register_processor_features",
     "build_feature_objects",
     "extract_feature_groups",
@@ -108,5 +104,11 @@ __all__ = [
     "console",
     "data",
     "embedding",
+    "loss",
     "torch_utils",
+    # Type aliases
+    "OptimizerName",
+    "SchedulerName",
+    "LossName",
+    "ActivationName",
 ]

nextrec/utils/config.py

@@ -21,7 +21,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Tuple
 import pandas as pd
 import torch
 
-from nextrec.utils.feature import normalize_to_list
+from nextrec.utils.feature import to_list
 
 if TYPE_CHECKING:
     from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
@@ -52,6 +52,16 @@ def resolve_path(
     )
 
 
+def safe_value(value: Any):
+    if isinstance(value, (str, int, float, bool)) or value is None:
+        return value
+    if isinstance(value, dict):
+        return {str(k): safe_value(v) for k, v in value.items()}
+    if isinstance(value, (list, tuple)):
+        return [safe_value(v) for v in value]
+    return str(value)
+
+
 def select_features(
     feature_cfg: Dict[str, Any], df_columns: List[str]
 ) -> Tuple[List[str], List[str], List[str]]:
@@ -152,9 +162,9 @@ def build_feature_objects(
         dense_features.append(
             DenseFeature(
                 name=name,
-                embedding_dim=embed_cfg.get("embedding_dim"),
+                proj_dim=embed_cfg.get("proj_dim"),
                 input_dim=embed_cfg.get("input_dim", 1),
-                use_embedding=embed_cfg.get("use_embedding", False),
+                use_projection=embed_cfg.get("use_projection", False),
             )
         )
 
@@ -239,7 +249,7 @@ def extract_feature_groups(
     collected: List[str] = []
 
     for group_name, names in feature_groups.items():
-        name_list = normalize_to_list(names)
+        name_list = to_list(names)
         filtered = []
         missing_defined = [n for n in name_list if n not in defined]
         missing_cols = [n for n in name_list if n not in available_cols]
@@ -441,7 +451,7 @@ def build_model_instance(
             direct_features = binding.get("features") or binding.get("feature_names")
             if direct_features and (accepts(param_name) or accepts_var_kwargs):
                 init_kwargs[param_name] = _select(
-                    normalize_to_list(direct_features),
+                    to_list(direct_features),
                     feature_pool,
                     f"feature_bindings.{param_name}",
                 )

nextrec/utils/console.py

@@ -36,7 +36,7 @@ from rich.progress import (
 from rich.table import Table
 from rich.text import Text
 
-from nextrec.utils.feature import as_float, normalize_to_list
+from nextrec.utils.feature import as_float, to_list
 
 T = TypeVar("T")
 
@@ -283,7 +283,7 @@ def display_metrics_table(
     if not is_main_process:
         return
 
-    target_list = normalize_to_list(target_names)
+    target_list = to_list(target_names)
     task_order, grouped = group_metrics_by_task(metrics, target_names=target_names)
 
     if isinstance(loss, np.ndarray) and target_list:

nextrec/utils/feature.py

@@ -2,7 +2,7 @@
 Feature processing utilities for NextRec
 
 Date: create on 03/12/2025
-Checkpoint: edit on 19/12/2025
+Checkpoint: edit on 27/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 """
 
@@ -10,7 +10,7 @@ import numbers
 from typing import Any
 
 
-def normalize_to_list(value: str | list[str] | None) -> list[str]:
+def to_list(value: str | list[str] | None) -> list[str]:
     if value is None:
         return []
     if isinstance(value, str):

nextrec/{loss/loss_utils.py → utils/loss.py}

@@ -1,13 +1,13 @@
 """
 Loss utilities for NextRec.
 
-Date: create on 27/10/2025
-Checkpoint: edit on 19/12/2025
+Date: create on 28/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 """
 
-from typing import Literal
+from __future__ import annotations
 
+import torch
 import torch.nn as nn
 
 from nextrec.loss.listwise import (
@@ -19,39 +19,21 @@ from nextrec.loss.listwise import (
 )
 from nextrec.loss.pairwise import BPRLoss, HingeLoss, TripletLoss
 from nextrec.loss.pointwise import ClassBalancedFocalLoss, FocalLoss, WeightedBCELoss
+from nextrec.utils.types import LossName
 
-VALID_TASK_TYPES = [
-    "binary",
-    "multilabel",
-    "regression",
-]
-
-# Define all supported loss types
-LossType = Literal[
-    # Pointwise losses
-    "bce",
-    "binary_crossentropy",
-    "weighted_bce",
-    "focal",
-    "focal_loss",
-    "cb_focal",
-    "class_balanced_focal",
-    "crossentropy",
-    "ce",
-    "mse",
-    "mae",
-    # Pairwise ranking losses
-    "bpr",
-    "hinge",
-    "triplet",
-    # Listwise ranking losses
-    "sampled_softmax",
-    "softmax",
-    "infonce",
-    "listnet",
-    "listmle",
-    "approx_ndcg",
-]
+
+def normalize_task_loss(
+    task_loss,
+    valid_count,
+    total_count,
+    eps=1e-8,
+) -> torch.Tensor:
+    if not torch.is_tensor(valid_count):
+        valid_count = torch.tensor(float(valid_count), device=task_loss.device)
+    if not torch.is_tensor(total_count):
+        total_count = torch.tensor(float(total_count), device=task_loss.device)
+    scale = valid_count.to(task_loss.dtype) / (total_count.to(task_loss.dtype) + eps)
+    return task_loss * scale
 
 
 def build_cb_focal(kw):
@@ -60,7 +42,10 @@ def build_cb_focal(kw):
     return ClassBalancedFocalLoss(**kw)
 
 
-def get_loss_fn(loss=None, **kw) -> nn.Module:
+def get_loss_fn(
+    loss: LossName | None | nn.Module = None,
+    **kw,
+) -> nn.Module:
     """
     Get loss function by name or return the provided loss module.
 

nextrec/utils/torch_utils.py

@@ -3,12 +3,16 @@ PyTorch-related utilities for NextRec.
 
 This module groups device setup, distributed helpers, optimizers/schedulers,
 initialization, and tensor helpers.
+
+Date: create on 27/10/2025
+Checkpoint: edit on 27/12/2025
+Author: Yang Zhou, zyaztec@gmail.com
 """
 
 from __future__ import annotations
 
 import logging
-from typing import Any, Dict, Iterable, Set
+from typing import Any, Dict, Iterable, Literal
 
 import numpy as np
 import torch
@@ -18,26 +22,25 @@ from torch.utils.data import DataLoader, IterableDataset
 from torch.utils.data.distributed import DistributedSampler
 
 from nextrec.basic.loggers import colorize
-
-KNOWN_NONLINEARITIES: Set[str] = {
-    "linear",
-    "conv1d",
-    "conv2d",
-    "conv3d",
-    "conv_transpose1d",
-    "conv_transpose2d",
-    "conv_transpose3d",
-    "sigmoid",
-    "tanh",
-    "relu",
-    "leaky_relu",
-    "selu",
-    "gelu",
-}
+from nextrec.utils.types import OptimizerName, SchedulerName
 
 
 def resolve_nonlinearity(activation: str) -> str:
-    if activation in KNOWN_NONLINEARITIES:
+    if activation in [
+        "linear",
+        "conv1d",
+        "conv2d",
+        "conv3d",
+        "conv_transpose1d",
+        "conv_transpose2d",
+        "conv_transpose3d",
+        "sigmoid",
+        "tanh",
+        "relu",
+        "leaky_relu",
+        "selu",
+        "gelu",
+    ]:
         return activation
     return "linear"
 
@@ -53,8 +56,30 @@ def resolve_gain(activation: str, param: Dict[str, Any]) -> float:
 
 
 def get_initializer(
-    init_type: str = "normal",
-    activation: str = "linear",
+    init_type: Literal[
+        "xavier_uniform",
+        "xavier_normal",
+        "kaiming_uniform",
+        "kaiming_normal",
+        "orthogonal",
+        "normal",
+        "uniform",
+    ] = "normal",
+    activation: Literal[
+        "linear",
+        "conv1d",
+        "conv2d",
+        "conv3d",
+        "conv_transpose1d",
+        "conv_transpose2d",
+        "conv_transpose3d",
+        "sigmoid",
+        "tanh",
+        "relu",
+        "leaky_relu",
+        "selu",
+        "gelu",
+    ] = "linear",
     param: Dict[str, Any] | None = None,
 ):
     param = param or {}
@@ -89,47 +114,14 @@ def get_initializer(
     return initializer_fn
 
 
-def resolve_device() -> str:
-    if torch.cuda.is_available():
-        return "cuda"
-    if torch.backends.mps.is_available():
-        import platform
-
-        mac_ver = platform.mac_ver()[0]
-        try:
-            major, _ = (int(x) for x in mac_ver.split(".")[:2])
-        except Exception:
-            major, _ = 0, 0
-        if major >= 14:
-            return "mps"
-    return "cpu"
-
-
-def get_device_info() -> dict:
-    info = {
-        "cuda_available": torch.cuda.is_available(),
-        "cuda_device_count": (
-            torch.cuda.device_count() if torch.cuda.is_available() else 0
-        ),
-        "mps_available": torch.backends.mps.is_available(),
-        "current_device": resolve_device(),
-    }
-
-    if torch.cuda.is_available():
-        info["cuda_device_name"] = torch.cuda.get_device_name(0)
-        info["cuda_capability"] = torch.cuda.get_device_capability(0)
-
-    return info
-
-
-def configure_device(
+def get_device(
     distributed: bool, local_rank: int, base_device: torch.device | str = "cpu"
 ) -> torch.device:
     try:
         device = torch.device(base_device)
     except Exception:
         logging.warning(
-            "[configure_device Warning] Invalid base_device, falling back to CPU."
+            "[get_device Warning] Invalid base_device, falling back to CPU."
         )
         return torch.device("cpu")
 
@@ -158,7 +150,7 @@ def configure_device(
 
 
 def get_optimizer(
-    optimizer: str | torch.optim.Optimizer = "adam",
+    optimizer: OptimizerName | torch.optim.Optimizer = "adam",
     params: Iterable[torch.nn.Parameter] | None = None,
     **optimizer_params,
 ):
@@ -191,7 +183,7 @@ def get_optimizer(
 
 def get_scheduler(
     scheduler: (
-        str
+        SchedulerName
         | torch.optim.lr_scheduler._LRScheduler
         | torch.optim.lr_scheduler.LRScheduler
         | type[torch.optim.lr_scheduler._LRScheduler]
@@ -241,51 +233,6 @@ def to_tensor(
     return tensor
 
 
-def stack_tensors(tensors: list[torch.Tensor], dim: int = 0) -> torch.Tensor:
-    if not tensors:
-        raise ValueError("[Tensor Utils Error] Cannot stack empty list of tensors.")
-    return torch.stack(tensors, dim=dim)
-
-
-def concat_tensors(tensors: list[torch.Tensor], dim: int = 0) -> torch.Tensor:
-    if not tensors:
-        raise ValueError(
-            "[Tensor Utils Error] Cannot concatenate empty list of tensors."
-        )
-    return torch.cat(tensors, dim=dim)
-
-
-def pad_sequence_tensors(
-    tensors: list[torch.Tensor],
-    max_len: int | None = None,
-    padding_value: float = 0.0,
-    padding_side: str = "right",
-) -> torch.Tensor:
-    if not tensors:
-        raise ValueError("[Tensor Utils Error] Cannot pad empty list of tensors.")
-    if max_len is None:
-        max_len = max(t.size(0) for t in tensors)
-    batch_size = len(tensors)
-    padded = torch.full(
-        (batch_size, max_len),
-        padding_value,
-        dtype=tensors[0].dtype,
-        device=tensors[0].device,
-    )
-
-    for i, tensor in enumerate(tensors):
-        length = min(tensor.size(0), max_len)
-        if padding_side == "right":
-            padded[i, :length] = tensor[:length]
-        elif padding_side == "left":
-            padded[i, -length:] = tensor[:length]
-        else:
-            raise ValueError(
-                f"[Tensor Utils Error] padding_side must be 'right' or 'left', got {padding_side}"
-            )
-    return padded
-
-
 def init_process_group(
     distributed: bool, rank: int, world_size: int, device_id: int | None = None
 ) -> None:
@@ -350,7 +297,7 @@ def add_distributed_sampler(
     # return if already has DistributedSampler
     if isinstance(loader.sampler, DistributedSampler):
         return loader, loader.sampler
-    dataset = getattr(loader, "dataset", None)
+    dataset = loader.dataset
     if dataset is None:
         return loader, None
     if isinstance(dataset, IterableDataset):
@@ -379,25 +326,23 @@ def add_distributed_sampler(
         "collate_fn": loader.collate_fn,
         "drop_last": drop_last,
     }
-    if getattr(loader, "pin_memory", False):
+    if loader.pin_memory:
         loader_kwargs["pin_memory"] = True
-    pin_memory_device = getattr(loader, "pin_memory_device", None)
+    pin_memory_device = loader.pin_memory_device
     if pin_memory_device:
         loader_kwargs["pin_memory_device"] = pin_memory_device
-    timeout = getattr(loader, "timeout", None)
+    timeout = loader.timeout
     if timeout:
         loader_kwargs["timeout"] = timeout
-    worker_init_fn = getattr(loader, "worker_init_fn", None)
+    worker_init_fn = loader.worker_init_fn
     if worker_init_fn is not None:
         loader_kwargs["worker_init_fn"] = worker_init_fn
-    generator = getattr(loader, "generator", None)
+    generator = loader.generator
     if generator is not None:
         loader_kwargs["generator"] = generator
     if loader.num_workers > 0:
-        loader_kwargs["persistent_workers"] = getattr(
-            loader, "persistent_workers", False
-        )
-        prefetch_factor = getattr(loader, "prefetch_factor", None)
+        loader_kwargs["persistent_workers"] = loader.persistent_workers
+        prefetch_factor = loader.prefetch_factor
         if prefetch_factor is not None:
             loader_kwargs["prefetch_factor"] = prefetch_factor
     distributed_loader = DataLoader(dataset, **loader_kwargs)

nextrec/utils/types.py

@@ -0,0 +1,63 @@
+"""
+Shared type aliases for NextRec.
+
+Keep Literal-based public string options centralized to avoid drift.
+"""
+
+from typing import Literal
+
+OptimizerName = Literal["adam", "sgd", "adamw", "adagrad", "rmsprop"]
+
+SchedulerName = Literal["step", "cosine"]
+
+LossName = Literal[
+    "bce",
+    "binary_crossentropy",
+    "weighted_bce",
+    "focal",
+    "focal_loss",
+    "cb_focal",
+    "class_balanced_focal",
+    "crossentropy",
+    "ce",
+    "mse",
+    "mae",
+    "bpr",
+    "hinge",
+    "triplet",
+    "sampled_softmax",
+    "softmax",
+    "infonce",
+    "listnet",
+    "listmle",
+    "approx_ndcg",
+]
+
+ActivationName = Literal[
+    "dice",
+    "relu",
+    "relu6",
+    "elu",
+    "selu",
+    "leaky_relu",
+    "prelu",
+    "gelu",
+    "sigmoid",
+    "tanh",
+    "softplus",
+    "softsign",
+    "hardswish",
+    "mish",
+    "silu",
+    "swish",
+    "hardsigmoid",
+    "tanhshrink",
+    "softshrink",
+    "none",
+    "linear",
+    "identity",
+]
+
+TrainingModeName = Literal["pointwise", "pairwise", "listwise"]
+
+TaskTypeName = Literal["binary", "regression"]