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

nextrec/__version__.py

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

nextrec/basic/activation.py

@@ -9,10 +9,10 @@ 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):
     """
     Dice activation function from the paper:

nextrec/basic/heads.py

@@ -15,6 +15,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 from nextrec.basic.layers import PredictionLayer
+from nextrec.utils.types import TaskTypeName
 
 
 class TaskHead(nn.Module):
@@ -27,9 +28,7 @@ class TaskHead(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,

nextrec/basic/metrics.py

@@ -2,7 +2,7 @@
 Metrics computation and configuration for model evaluation.
 
 Date: create on 27/10/2025
-Checkpoint: edit on 20/12/2025
+Checkpoint: edit on 29/12/2025
 Author: Yang Zhou,zyaztec@gmail.com
 """
 
@@ -39,7 +39,6 @@ REGRESSION_METRICS = {"mse", "mae", "rmse", "r2", "mape", "msle"}
 TASK_DEFAULT_METRICS = {
     "binary": ["auc", "gauc", "ks", "logloss", "accuracy", "precision", "recall", "f1"],
     "regression": ["mse", "mae", "rmse", "r2", "mape"],
-    "multilabel": ["auc", "hamming_loss", "subset_accuracy", "micro_f1", "macro_f1"],
     "matching": ["auc", "gauc", "precision@10", "hitrate@10", "map@10", "cosine"]
     + [f"recall@{k}" for k in (5, 10, 20)]
     + [f"ndcg@{k}" for k in (5, 10, 20)]

nextrec/basic/model.py

@@ -2,13 +2,14 @@
 Base Model & Base Match Model Class
 
 Date: create on 27/10/2025
-Checkpoint: edit on 28/12/2025
+Checkpoint: edit on 29/12/2025
 Author: Yang Zhou,zyaztec@gmail.com
 """
 
 import getpass
 import logging
 import os
+import sys
 import pickle
 import socket
 from pathlib import Path
@@ -16,6 +17,16 @@ from typing import Any, Literal
 
 import numpy as np
 import pandas as pd
+
+try:
+    import swanlab  # type: ignore
+except ModuleNotFoundError:
+    swanlab = None
+try:
+    import wandb  # type: ignore
+except ModuleNotFoundError:
+    wandb = None
+
 import torch
 import torch.distributed as dist
 import torch.nn as nn
@@ -60,8 +71,8 @@ from nextrec.loss import (
     InfoNCELoss,
     SampledSoftmaxLoss,
     TripletLoss,
-    get_loss_fn,
 )
+from nextrec.utils.loss import get_loss_fn
 from nextrec.loss.grad_norm import get_grad_norm_shared_params
 from nextrec.utils.console import display_metrics_table, progress
 from nextrec.utils.torch_utils import (
@@ -74,8 +85,20 @@ from nextrec.utils.torch_utils import (
     to_tensor,
 )
 from nextrec.utils.config import safe_value
-from nextrec.utils.model import compute_ranking_loss
-from nextrec.utils.types import LossName, OptimizerName, SchedulerName
+from nextrec.utils.model import (
+    compute_ranking_loss,
+    get_loss_list,
+    resolve_loss_weights,
+    get_training_modes,
+)
+from nextrec.utils.types import (
+    LossName,
+    OptimizerName,
+    SchedulerName,
+    TrainingModeName,
+    TaskTypeName,
+    MetricsName,
+)
 
 
 class BaseModel(SummarySet, FeatureSet, nn.Module):
@@ -84,7 +107,7 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
         raise NotImplementedError
 
     @property
-    def default_task(self) -> str | list[str]:
+    def default_task(self) -> TaskTypeName | list[TaskTypeName]:
         raise NotImplementedError
 
     def __init__(
@@ -94,11 +117,8 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
         sequence_features: list[SequenceFeature] | None = None,
         target: list[str] | str | None = None,
         id_columns: list[str] | str | None = None,
-        task: str | list[str] | None = None,
-        training_mode: (
-            Literal["pointwise", "pairwise", "listwise"]
-            | list[Literal["pointwise", "pairwise", "listwise"]]
-        ) = "pointwise",
+        task: TaskTypeName | list[TaskTypeName] | None = None,
+        training_mode: TrainingModeName | list[TrainingModeName] = "pointwise",
         embedding_l1_reg: float = 0.0,
         dense_l1_reg: float = 0.0,
         embedding_l2_reg: float = 0.0,
@@ -136,6 +156,9 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
             world_size: Number of processes (defaults to env WORLD_SIZE).
             local_rank: Local rank for selecting CUDA device (defaults to env LOCAL_RANK).
             ddp_find_unused_parameters: Default False, set it True only when exist unused parameters in ddp model, in most cases should be False.
+
+        Note:
+            Optimizer, scheduler, and loss are configured via compile().
         """
         super(BaseModel, self).__init__()
 
@@ -168,25 +191,12 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
             dense_features, sparse_features, sequence_features, target, id_columns
         )
 
-        self.task = self.default_task if task is None else task
+        self.task = task or self.default_task
         self.nums_task = len(self.task) if isinstance(self.task, list) else 1
-        if isinstance(training_mode, list):
-            training_modes = list(training_mode)
-            if len(training_modes) != self.nums_task:
-                raise ValueError(
-                    "[BaseModel-init Error] training_mode list length must match number of tasks."
-                )
-        else:
-            training_modes = [training_mode] * self.nums_task
-        if any(
-            mode not in {"pointwise", "pairwise", "listwise"}
-            for mode in training_modes
-        ):
-            raise ValueError(
-                "[BaseModel-init Error] training_mode must be one of {'pointwise', 'pairwise', 'listwise'}."
-            )
-        self.training_modes = training_modes
-        self.training_mode = training_modes if self.nums_task > 1 else training_modes[0]
+        self.training_modes = get_training_modes(training_mode, self.nums_task)
+        self.training_mode = (
+            self.training_modes if self.nums_task > 1 else self.training_modes[0]
+        )
 
         self.embedding_l1_reg = embedding_l1_reg
         self.dense_l1_reg = dense_l1_reg
@@ -194,7 +204,9 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
         self.dense_l2_reg = dense_l2_reg
         self.regularization_weights = []
         self.embedding_params = []
-        self.loss_weight = None
+
+        self.ignore_label = None
+        self.compiled = False
 
         self.max_gradient_norm = 1.0
         self.logger_initialized = False
@@ -407,6 +419,7 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
         loss: LossName | nn.Module | list[LossName | nn.Module] | None = "bce",
         loss_params: dict | list[dict] | None = None,
         loss_weights: int | float | list[int | float] | dict | str | None = None,
+        ignore_label: int | float | None = -1,
     ):
         """
         Configure the model for training.
@@ -419,34 +432,17 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
             loss_params: Loss function parameters, or list for multi-task. e.g., {'weight': tensor([0.25, 0.75])}.
             loss_weights: Weights for each task loss, int/float for single-task or list for multi-task. e.g., 1.0, or [1.0, 0.5].
                 Use "grad_norm" or {"method": "grad_norm", ...} to enable GradNorm for multi-task loss balancing.
+            ignore_label: Label value to ignore when computing loss. Use this to skip gradients for unknown labels.
         """
+        self.ignore_label = ignore_label
         default_losses = {
             "pointwise": "bce",
             "pairwise": "bpr",
             "listwise": "listnet",
         }
-        effective_loss = loss
-        if effective_loss is None:
-            loss_list = [default_losses[mode] for mode in self.training_modes]
-        elif isinstance(effective_loss, list):
-            if not effective_loss:
-                loss_list = [default_losses[mode] for mode in self.training_modes]
-            else:
-                if len(effective_loss) != self.nums_task:
-                    raise ValueError(
-                        f"[BaseModel-compile Error] Number of loss functions ({len(effective_loss)}) must match number of tasks ({self.nums_task})."
-                    )
-                loss_list = list(effective_loss)
-        else:
-            loss_list = [effective_loss] * self.nums_task
-
-        for idx, mode in enumerate(self.training_modes):
-            if isinstance(loss_list[idx], str) and loss_list[idx] in {
-                "bce",
-                "binary_crossentropy",
-            }:
-                if mode in {"pairwise", "listwise"}:
-                    loss_list[idx] = default_losses[mode]
+        loss_list = get_loss_list(
+            loss, self.training_modes, self.nums_task, default_losses
+        )
         self.loss_params = loss_params or {}
         optimizer_params = optimizer_params or {}
         self.optimizer_name = (
@@ -510,36 +506,16 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
                 nums_task=self.nums_task, device=self.device, **grad_norm_params
             )
             self.loss_weights = None
-        elif loss_weights is None:
-            self.loss_weights = None
-        elif self.nums_task == 1:
-            if isinstance(loss_weights, (list, tuple)):
-                if len(loss_weights) != 1:
-                    raise ValueError(
-                        "[BaseModel-compile Error] loss_weights list must have exactly one element for single-task setup."
-                    )
-                loss_weights = loss_weights[0]
-            self.loss_weights = [float(loss_weights)]  # type: ignore
         else:
-            if isinstance(loss_weights, (int, float)):
-                weights = [float(loss_weights)] * self.nums_task
-            elif isinstance(loss_weights, (list, tuple)):
-                weights = [float(w) for w in loss_weights]
-                if len(weights) != self.nums_task:
-                    raise ValueError(
-                        f"[BaseModel-compile Error] Number of loss_weights ({len(weights)}) must match number of tasks ({self.nums_task})."
-                    )
-            else:
-                raise TypeError(
-                    f"[BaseModel-compile Error] loss_weights must be int, float, list or tuple, got {type(loss_weights)}"
-                )
-            self.loss_weights = weights
+            self.loss_weights = resolve_loss_weights(loss_weights, self.nums_task)
+        self.compiled = True
 
     def compute_loss(self, y_pred, y_true):
         if y_true is None:
             raise ValueError(
                 "[BaseModel-compute_loss Error] Ground truth labels (y_true) are required."
             )
+
         # single-task
         if self.nums_task == 1:
             if y_pred.dim() == 1:
@@ -547,13 +523,24 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
             if y_true.dim() == 1:
                 y_true = y_true.view(-1, 1)
             if y_pred.shape != y_true.shape:
-                raise ValueError(f"Shape mismatch: {y_pred.shape} vs {y_true.shape}")
-            loss_fn = self.loss_fn[0] if getattr(self, "loss_fn", None) else None
-            if loss_fn is None:
                 raise ValueError(
-                    "[BaseModel-compute_loss Error] Loss function is not configured. Call compile() first."
+                    f"[BaseModel-compute_loss Error] Shape mismatch: {y_pred.shape} vs {y_true.shape}"
                 )
+
+            loss_fn = self.loss_fn[0]
+
+            if self.ignore_label is not None:
+                valid_mask = y_true != self.ignore_label
+                if valid_mask.dim() > 1:
+                    valid_mask = valid_mask.all(dim=1)
+                if not torch.any(valid_mask):  # if no valid labels, return zero loss
+                    return y_pred.sum() * 0.0
+
+                y_pred = y_pred[valid_mask]
+                y_true = y_true[valid_mask]
+
             mode = self.training_modes[0]
+
             task_dim = (
                 self.task_dims[0] if hasattr(self, "task_dims") else y_pred.shape[1]  # type: ignore
             )
@@ -584,7 +571,19 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
         for i, (start, end) in enumerate(slices):  # type: ignore
             y_pred_i = y_pred[:, start:end]
             y_true_i = y_true[:, start:end]
+            # mask ignored labels
+            if self.ignore_label is not None:
+                valid_mask = y_true_i != self.ignore_label
+                if valid_mask.dim() > 1:
+                    valid_mask = valid_mask.all(dim=1)
+                if not torch.any(valid_mask):
+                    task_losses.append(y_pred_i.sum() * 0.0)
+                    continue
+                y_pred_i = y_pred_i[valid_mask]
+                y_true_i = y_true_i[valid_mask]
+
             mode = self.training_modes[i]
+
             if mode in {"pairwise", "listwise"}:
                 task_loss = compute_ranking_loss(
                     training_mode=mode,
@@ -594,7 +593,11 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
                 )
             else:
                 task_loss = self.loss_fn[i](y_pred_i, y_true_i)
+            # task_loss = normalize_task_loss(
+            #     task_loss, valid_count, total_count
+            # )  # normalize by valid samples to avoid loss scale issues
             task_losses.append(task_loss)
+
         if self.grad_norm is not None:
             if self.grad_norm_shared_params is None:
                 self.grad_norm_shared_params = get_grad_norm_shared_params(
@@ -651,7 +654,7 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
         train_data=None,
         valid_data=None,
         metrics: (
-            list[str] | dict[str, list[str]] | None
+            list[MetricsName] | dict[str, list[MetricsName]] | None
         ) = None,  # ['auc', 'logloss'] or {'target1': ['auc', 'logloss'], 'target2': ['mse']}
         epochs: int = 1,
         shuffle: bool = True,
@@ -665,6 +668,8 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
         use_tensorboard: bool = True,
         use_wandb: bool = False,
         use_swanlab: bool = False,
+        wandb_api: str | None = None,
+        swanlab_api: str | None = None,
         wandb_kwargs: dict | None = None,
         swanlab_kwargs: dict | None = None,
         auto_ddp_sampler: bool = True,
@@ -694,6 +699,8 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
             use_tensorboard: Enable tensorboard logging.
             use_wandb: Enable Weights & Biases logging.
             use_swanlab: Enable SwanLab logging.
+            wandb_api: W&B API key for non-tty login.
+            swanlab_api: SwanLab API key for non-tty login.
             wandb_kwargs: Optional kwargs for wandb.init(...).
             swanlab_kwargs: Optional kwargs for swanlab.init(...).
             auto_ddp_sampler: Attach DistributedSampler automatically when distributed, set False to when data is already sharded per rank.
@@ -711,6 +718,16 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
         )
         self.to(self.device)
 
+        if not self.compiled:
+            self.compile(
+                optimizer="adam",
+                optimizer_params={},
+                scheduler=None,
+                scheduler_params={},
+                loss=None,
+                loss_params={},
+            )
+
         if (
             self.distributed
             and dist.is_available()
@@ -785,6 +802,24 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
             }
             training_config: dict = safe_value(training_config)  # type: ignore
 
+        if self.is_main_process:
+            is_tty = sys.stdin.isatty() and sys.stdout.isatty()
+            if not is_tty:
+                if use_wandb and wandb_api:
+                    if wandb is None:
+                        logging.warning(
+                            "[BaseModel-fit] wandb not installed, skip wandb login."
+                        )
+                    else:
+                        wandb.login(key=wandb_api)
+                if use_swanlab and swanlab_api:
+                    if swanlab is None:
+                        logging.warning(
+                            "[BaseModel-fit] swanlab not installed, skip swanlab login."
+                        )
+                    else:
+                        swanlab.login(api_key=swanlab_api)
+
         self.training_logger = (
             TrainingLogger(
                 session=self.session,
@@ -2164,7 +2199,7 @@ class BaseMatchModel(BaseModel):
             scheduler_params: Parameters for the scheduler. e.g., {'step_size': 10, 'gamma': 0.1}.
             loss: Loss function(s) to use (name, instance, or list). e.g., 'bce'.
             loss_params: Parameters for the loss function(s). e.g., {'reduction': 'mean'}.
-            loss_weights: Weights for the loss function(s). e.g., 1.0 or [0.7, 0.3]. 
+            loss_weights: Weights for the loss function(s). e.g., 1.0 or [0.7, 0.3].
         """
         if self.training_mode not in self.support_training_modes:
             raise ValueError(

nextrec/basic/summary.py

@@ -1,5 +1,9 @@
 """
 Summary utilities for BaseModel.
+
+Date: create on 03/12/2025
+Checkpoint: edit on 29/12/2025
+Author: Yang Zhou,zyaztec@gmail.com
 """
 
 from __future__ import annotations
@@ -12,9 +16,39 @@ from torch.utils.data import DataLoader
 
 from nextrec.basic.loggers import colorize, format_kv
 from nextrec.data.data_processing import extract_label_arrays, get_data_length
+from nextrec.utils.types import TaskTypeName
 
 
 class SummarySet:
+    model_name: str
+    dense_features: list[Any]
+    sparse_features: list[Any]
+    sequence_features: list[Any]
+    task: TaskTypeName | list[TaskTypeName]
+    target_columns: list[str]
+    nums_task: int
+    metrics: Any
+    device: Any
+    optimizer_name: str
+    optimizer_params: dict[str, Any]
+    scheduler_name: str | None
+    scheduler_params: dict[str, Any]
+    loss_config: Any
+    loss_weights: Any
+    grad_norm: Any
+    embedding_l1_reg: float
+    embedding_l2_reg: float
+    dense_l1_reg: float
+    dense_l2_reg: float
+    early_stop_patience: int
+    max_gradient_norm: float | None
+    metrics_sample_limit: int | None
+    session_id: str | None
+    features_config_path: str
+    checkpoint_path: str
+    train_data_summary: dict[str, Any] | None
+    valid_data_summary: dict[str, Any] | None
+
     def build_data_summary(
         self, data: Any, data_loader: DataLoader | None, sample_key: str
     ):
@@ -305,7 +339,7 @@ class SummarySet:
                         lines = details.get("lines", [])
                         logger.info(f"{target_name}:")
                         for label, value in lines:
-                            logger.info(format_kv(label, value))
+                            logger.info(f"  {format_kv(label, value)}")
 
             if self.valid_data_summary:
                 if self.train_data_summary:
@@ -320,4 +354,4 @@ class SummarySet:
                         lines = details.get("lines", [])
                         logger.info(f"{target_name}:")
                         for label, value in lines:
-                            logger.info(format_kv(label, value))
+                            logger.info(f"  {format_kv(label, value)}")