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

nextrec/__version__.py

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

nextrec/basic/layers.py

@@ -20,6 +20,7 @@ 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
 
 
 class PredictionLayer(nn.Module):
@@ -590,71 +591,48 @@ class MLP(nn.Module):
     def __init__(
         self,
         input_dim: int,
-        output_layer: bool = True,
-        dims: list[int] | None = None,
+        hidden_dims: list[int] | None = None,
+        output_dim: int | None = 1,
         dropout: float = 0.0,
-        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",
-        ] = "relu",
-        use_norm: bool = True,
-        norm_type: Literal["batch_norm", "layer_norm"] = "layer_norm",
+        activation: ActivationName = "relu",
+        norm_type: Literal["batch_norm", "layer_norm", "none"] = "none",
+        output_activation: ActivationName = "none",
     ):
         """
         Multi-Layer Perceptron (MLP) module.
 
         Args:
             input_dim: Dimension of the input features.
-            output_layer: Whether to include the final output layer. If False, the MLP will output the last hidden layer, else it will output a single value.
-            dims: List of hidden layer dimensions. If None, no hidden layers are added.
+            output_dim: Output dimension of the final layer. If None, no output layer is added.
+            hidden_dims: List of hidden layer dimensions. If None, no hidden layers are added.
             dropout: Dropout rate between layers.
             activation: Activation function to use between layers.
-            use_norm: Whether to use normalization layers.
-            norm_type: Type of normalization to use ("batch_norm" or "layer_norm").
+            norm_type: Type of normalization to use ("batch_norm", "layer_norm", or "none").
+            output_activation: Activation function applied after the output layer.
         """
         super().__init__()
-        if dims is None:
-            dims = []
+        hidden_dims = hidden_dims or []
         layers = []
         current_dim = input_dim
-        for i_dim in dims:
+        for i_dim in hidden_dims:
             layers.append(nn.Linear(current_dim, i_dim))
-            if use_norm:
-                if norm_type == "batch_norm":
-                    # **IMPORTANT** be careful when using BatchNorm1d in distributed training, nextrec does not support sync batch norm now
-                    layers.append(nn.BatchNorm1d(i_dim))
-                elif norm_type == "layer_norm":
-                    layers.append(nn.LayerNorm(i_dim))
-                else:
-                    raise ValueError(f"Unsupported norm_type: {norm_type}")
+            if norm_type == "batch_norm":
+                # **IMPORTANT** be careful when using BatchNorm1d in distributed training, nextrec does not support sync batch norm now
+                layers.append(nn.BatchNorm1d(i_dim))
+            elif norm_type == "layer_norm":
+                layers.append(nn.LayerNorm(i_dim))
+            elif norm_type != "none":
+                raise ValueError(f"Unsupported norm_type: {norm_type}")
 
             layers.append(activation_layer(activation))
             layers.append(nn.Dropout(p=dropout))
             current_dim = i_dim
         # output layer
-        if output_layer:
-            layers.append(nn.Linear(current_dim, 1))
-            self.output_dim = 1
+        if output_dim is not None:
+            layers.append(nn.Linear(current_dim, output_dim))
+            if output_activation != "none":
+                layers.append(activation_layer(output_activation))
+            self.output_dim = output_dim
         else:
             self.output_dim = current_dim
         self.mlp = nn.Sequential(*layers)
@@ -663,6 +641,47 @@ class MLP(nn.Module):
         return self.mlp(x)
 
 
+class GateMLP(nn.Module):
+    """
+    Lightweight gate network: sigmoid MLP scaled by a constant factor.
+
+    Args:
+        input_dim: Dimension of the input features.
+        hidden_dim: Dimension of the hidden layer. If None, defaults to output_dim.
+        output_dim: Output dimension of the gate.
+        activation: Activation function to use in the hidden layer.
+        dropout: Dropout rate between layers.
+        use_bn: Whether to use batch normalization.
+        scale_factor: Scaling factor applied to the sigmoid output.
+    """
+
+    def __init__(
+        self,
+        input_dim: int,
+        hidden_dim: int | None,
+        output_dim: int,
+        activation: ActivationName = "relu",
+        dropout: float = 0.0,
+        use_bn: bool = False,
+        scale_factor: float = 2.0,
+    ) -> None:
+        super().__init__()
+        hidden_dim = output_dim if hidden_dim is None else hidden_dim
+        self.gate = MLP(
+            input_dim=input_dim,
+            hidden_dims=[hidden_dim],
+            output_dim=output_dim,
+            activation=activation,
+            dropout=dropout,
+            norm_type="batch_norm" if use_bn else "none",
+            output_activation="sigmoid",
+        )
+        self.scale_factor = scale_factor
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        return self.gate(inputs) * self.scale_factor
+
+
 class FM(nn.Module):
     def __init__(self, reduce_sum: bool = True):
         super().__init__()
@@ -1007,3 +1026,34 @@ 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/metrics.py

@@ -2,7 +2,7 @@
 Metrics computation and configuration for model evaluation.
 
 Date: create on 27/10/2025
-Checkpoint: edit on 29/12/2025
+Checkpoint: edit on 30/12/2025
 Author: Yang Zhou,zyaztec@gmail.com
 """
 
@@ -21,21 +21,9 @@ from sklearn.metrics import (
     recall_score,
     roc_auc_score,
 )
+from nextrec.utils.types import TaskTypeName, MetricsName
+
 
-CLASSIFICATION_METRICS = {
-    "auc",
-    "gauc",
-    "ks",
-    "logloss",
-    "accuracy",
-    "acc",
-    "precision",
-    "recall",
-    "f1",
-    "micro_f1",
-    "macro_f1",
-}
-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"],
@@ -58,7 +46,7 @@ def check_user_id(*metric_sources: Any) -> bool:
             stack.extend(item.values())
             continue
         if isinstance(item, str):
-            metric_names.add(item.lower())
+            metric_names.add(item)
             continue
         try:
             stack.extend(item)
@@ -361,9 +349,9 @@ def compute_cosine_separation(y_true: np.ndarray, y_pred: np.ndarray) -> float:
 
 
 def configure_metrics(
-    task: str | list[str],  # 'binary' or ['binary', 'regression']
+    task: TaskTypeName | list[TaskTypeName],  # 'binary' or ['binary', 'regression']
     metrics: (
-        list[str] | dict[str, list[str]] | None
+        list[MetricsName] | dict[str, list[MetricsName]] | None
     ),  # ['auc', 'logloss'] or {'task1': ['auc'], 'task2': ['mse']}
     target_names: list[str],  # ['target1', 'target2']
 ) -> tuple[list[str], dict[str, list[str]] | None, str]:
@@ -383,13 +371,12 @@ def configure_metrics(
                     f"[Metrics Warning] Task {task_name} not found in targets {target_names}, skipping its metrics"
                 )
                 continue
-            lowered = [m.lower() for m in task_metrics]
-            task_specific_metrics[task_name] = lowered
-            for metric in lowered:
+            task_specific_metrics[task_name] = task_metrics
+            for metric in task_metrics:
                 if metric not in metrics_list:
                     metrics_list.append(metric)
     elif metrics:
-        metrics_list = [m.lower() for m in metrics]
+        metrics_list = [m for m in metrics]
     else:
         # No user provided metrics, derive per task type
         if nums_task > 1 and isinstance(task, list):
@@ -416,11 +403,10 @@ def configure_metrics(
     return metrics_list, task_specific_metrics, best_metrics_mode
 
 
-def getbest_metric_mode(first_metric: str, primary_task: str) -> str:
+def getbest_metric_mode(first_metric: MetricsName, primary_task: TaskTypeName) -> str:
     """Determine if metric should be maximized or minimized."""
-    first_metric_lower = first_metric.lower()
     # Metrics that should be maximized
-    if first_metric_lower in {
+    if first_metric in {
         "auc",
         "gauc",
         "ks",
@@ -436,20 +422,20 @@ def getbest_metric_mode(first_metric: str, primary_task: str) -> str:
         return "max"
     # Ranking metrics that should be maximized (with @K suffix)
     if (
-        first_metric_lower.startswith("recall@")
-        or first_metric_lower.startswith("precision@")
-        or first_metric_lower.startswith("hitrate@")
-        or first_metric_lower.startswith("hr@")
-        or first_metric_lower.startswith("mrr@")
-        or first_metric_lower.startswith("ndcg@")
-        or first_metric_lower.startswith("map@")
+        first_metric.startswith("recall@")
+        or first_metric.startswith("precision@")
+        or first_metric.startswith("hitrate@")
+        or first_metric.startswith("hr@")
+        or first_metric.startswith("mrr@")
+        or first_metric.startswith("ndcg@")
+        or first_metric.startswith("map@")
     ):
         return "max"
     # Cosine separation should be maximized
-    if first_metric_lower == "cosine":
+    if first_metric == "cosine":
         return "max"
     # Metrics that should be minimized
-    if first_metric_lower in {"logloss", "mse", "mae", "rmse", "mape", "msle"}:
+    if first_metric in {"logloss", "mse", "mae", "rmse", "mape", "msle"}:
         return "min"
     # Default based on task type
     if primary_task == "regression":
@@ -458,7 +444,7 @@ def getbest_metric_mode(first_metric: str, primary_task: str) -> str:
 
 
 def compute_single_metric(
-    metric: str,
+    metric: MetricsName,
     y_true: np.ndarray,
     y_pred: np.ndarray,
     task_type: str,
@@ -466,30 +452,32 @@ def compute_single_metric(
 ) -> float:
     """Compute a single metric given true and predicted values."""
 
+    if y_true.size == 0:
+        return 0.0
+
     y_p_binary = (y_pred > 0.5).astype(int)
-    metric_lower = metric.lower()
     try:
-        if metric_lower.startswith("recall@"):
-            k = int(metric_lower.split("@")[1])
+        if metric.startswith("recall@"):
+            k = int(metric.split("@")[1])
             return compute_recall_at_k(y_true, y_pred, user_ids, k)  # type: ignore
-        if metric_lower.startswith("precision@"):
-            k = int(metric_lower.split("@")[1])
+        if metric.startswith("precision@"):
+            k = int(metric.split("@")[1])
             return compute_precision_at_k(y_true, y_pred, user_ids, k)  # type: ignore
-        if metric_lower.startswith("hitrate@") or metric_lower.startswith("hr@"):
-            k_str = metric_lower.split("@")[1]
+        if metric.startswith("hitrate@") or metric.startswith("hr@"):
+            k_str = metric.split("@")[1]
             k = int(k_str)
             return compute_hitrate_at_k(y_true, y_pred, user_ids, k)  # type: ignore
-        if metric_lower.startswith("mrr@"):
-            k = int(metric_lower.split("@")[1])
+        if metric.startswith("mrr@"):
+            k = int(metric.split("@")[1])
             return compute_mrr_at_k(y_true, y_pred, user_ids, k)  # type: ignore
-        if metric_lower.startswith("ndcg@"):
-            k = int(metric_lower.split("@")[1])
+        if metric.startswith("ndcg@"):
+            k = int(metric.split("@")[1])
             return compute_ndcg_at_k(y_true, y_pred, user_ids, k)  # type: ignore
-        if metric_lower.startswith("map@"):
-            k = int(metric_lower.split("@")[1])
+        if metric.startswith("map@"):
+            k = int(metric.split("@")[1])
             return compute_map_at_k(y_true, y_pred, user_ids, k)  # type: ignore
         # cosine for matching task
-        if metric_lower == "cosine":
+        if metric == "cosine":
             return compute_cosine_separation(y_true, y_pred)
         if metric == "auc":
             value = float(
@@ -570,15 +558,31 @@ def compute_single_metric(
 def evaluate_metrics(
     y_true: np.ndarray | None,
     y_pred: np.ndarray | None,
-    metrics: list[str],  # example: ['auc', 'logloss']
-    task: str | list[str],  # example: 'binary' or ['binary', 'regression']
-    target_names: list[str],  # example: ['target1', 'target2']
-    task_specific_metrics: (
-        dict[str, list[str]] | None
-    ) = None,  # example: {'target1': ['auc', 'logloss'], 'target2': ['mse']}
-    user_ids: np.ndarray | None = None,  # example: User IDs for GAUC computation
-) -> dict:  # {'auc': 0.75, 'logloss': 0.45, 'mse_target2': 3.2}
-    """Evaluate specified metrics for given true and predicted values."""
+    metrics: list[MetricsName],
+    task: TaskTypeName | list[TaskTypeName],
+    target_names: list[str],
+    task_specific_metrics: dict[str, list[MetricsName]] | None = None,
+    user_ids: np.ndarray | None = None,
+    ignore_label: int | float | None = None,
+) -> dict:
+    """
+    Evaluate specified metrics for given true and predicted values.
+    Supports single-task and multi-task evaluation.
+    Handles optional ignore_label to exclude certain samples.
+
+    Args:
+        y_true: Ground truth labels.
+        y_pred: Predicted values.
+        metrics: List of metric names to compute.
+        task: Task type(s) - 'binary', 'regression', etc.
+        target_names: Names of target variables. e.g., ['target1', 'target2']
+        task_specific_metrics: Optional dict mapping target names to specific metrics. e.g., {'target1': ['auc', 'logloss'], 'target2': ['mse']}
+        user_ids: Optional user IDs for GAUC and ranking metrics. e.g., User IDs for GAUC computation
+        ignore_label: Optional label value to ignore during evaluation.
+
+    Returns: Dictionary of computed metric values. {'auc': 0.75, 'logloss': 0.45, 'mse_target2': 3.2}
+
+    """
 
     result = {}
     if y_true is None or y_pred is None:
@@ -588,70 +592,81 @@ def evaluate_metrics(
     nums_task = len(task) if isinstance(task, list) else 1
     # Single task evaluation
     if nums_task == 1:
+        if ignore_label is not None:
+            valid_mask = y_true != ignore_label
+            if np.any(valid_mask):
+                y_true = y_true[valid_mask]
+                y_pred = y_pred[valid_mask]
+                if user_ids is not None:
+                    user_ids = user_ids[valid_mask]
+            else:
+                return result
         for metric in metrics:
-            metric_lower = metric.lower()
             value = compute_single_metric(
-                metric_lower, y_true, y_pred, primary_task, user_ids
+                metric, y_true, y_pred, primary_task, user_ids
             )
-            result[metric_lower] = value
+            result[metric] = value
     # Multi-task evaluation
     else:
-        for metric in metrics:
-            metric_lower = metric.lower()
-            for task_idx in range(nums_task):
-                # Check if metric should be computed for given task
-                should_compute = True
-                if task_specific_metrics is not None and task_idx < len(target_names):
-                    task_name = target_names[task_idx]
-                    should_compute = metric_lower in task_specific_metrics.get(
-                        task_name, []
-                    )
-                else:
-                    # Get task type for specific index
-                    if isinstance(task, list) and task_idx < len(task):
-                        task_type = task[task_idx]
-                    elif isinstance(task, str):
-                        task_type = task
-                    else:
-                        task_type = "binary"
-                    if task_type in ["binary", "multilabel"]:
-                        should_compute = metric_lower in {
-                            "auc",
-                            "gauc",
-                            "ks",
-                            "logloss",
-                            "accuracy",
-                            "acc",
-                            "precision",
-                            "recall",
-                            "f1",
-                            "micro_f1",
-                            "macro_f1",
-                        }
-                    elif task_type == "regression":
-                        should_compute = metric_lower in {
-                            "mse",
-                            "mae",
-                            "rmse",
-                            "r2",
-                            "mape",
-                            "msle",
-                        }
-                if not should_compute:
+        task_types = []
+        for task_idx in range(nums_task):
+            if isinstance(task, list) and task_idx < len(task):
+                task_types.append(task[task_idx])
+            elif isinstance(task, str):
+                task_types.append(task)
+            else:
+                task_types.append("binary")
+        metric_allowlist = {
+            "binary": {
+                "auc",
+                "gauc",
+                "ks",
+                "logloss",
+                "accuracy",
+                "acc",
+                "precision",
+                "recall",
+                "f1",
+                "micro_f1",
+                "macro_f1",
+            },
+            "regression": {
+                "mse",
+                "mae",
+                "rmse",
+                "r2",
+                "mape",
+                "msle",
+            },
+        }
+        for task_idx in range(nums_task):
+            task_type = task_types[task_idx]
+            target_name = target_names[task_idx]
+            if task_specific_metrics is not None and task_idx < len(target_names):
+                allowed_metrics = {
+                    m for m in task_specific_metrics.get(target_name, [])
+                }
+            else:
+                allowed_metrics = metric_allowlist.get(task_type)
+            for metric in metrics:
+                if allowed_metrics is not None and metric not in allowed_metrics:
                     continue
-                target_name = target_names[task_idx]
-                # Get task type for specific index
-                if isinstance(task, list) and task_idx < len(task):
-                    task_type = task[task_idx]
-                elif isinstance(task, str):
-                    task_type = task
-                else:
-                    task_type = "binary"
                 y_true_task = y_true[:, task_idx]
                 y_pred_task = y_pred[:, task_idx]
+                task_user_ids = user_ids
+                if ignore_label is not None:
+                    valid_mask = y_true_task != ignore_label
+                    if np.any(valid_mask):
+                        y_true_task = y_true_task[valid_mask]
+                        y_pred_task = y_pred_task[valid_mask]
+                        if task_user_ids is not None:
+                            task_user_ids = task_user_ids[valid_mask]
+                    else:
+                        result[f"{metric}_{target_name}"] = 0.0
+                        continue
                 # Compute metric
                 value = compute_single_metric(
-                    metric_lower, y_true_task, y_pred_task, task_type, user_ids
+                    metric, y_true_task, y_pred_task, task_type, task_user_ids
                 )
-                result[f"{metric_lower}_{target_name}"] = value
+                result[f"{metric}_{target_name}"] = value
     return result

nextrec/basic/model.py

@@ -2,7 +2,7 @@
 Base Model & Base Match Model Class
 
 Date: create on 27/10/2025
-Checkpoint: edit on 29/12/2025
+Checkpoint: edit on 30/12/2025
 Author: Yang Zhou,zyaztec@gmail.com
 """
 
@@ -646,6 +646,8 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
             sampler=sampler,
             collate_fn=collate_fn,
             num_workers=num_workers,
+            pin_memory=self.device.type == "cuda",
+            persistent_workers=num_workers > 0,
         )
         return (loader, dataset) if return_dataset else loader
 
@@ -1119,16 +1121,17 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
                     train_log_payload, step=epoch + 1, split="train"
                 )
             if valid_loader is not None:
-                self.callbacks.on_validation_begin()
-                val_metrics = self.evaluate(
-                    valid_loader,
-                    user_ids=valid_user_ids if self.needs_user_ids else None,
-                    num_workers=num_workers,
-                )
-                should_log_valid = (epoch + 1) % log_interval == 0 or (
+                should_eval_valid = (epoch + 1) % log_interval == 0 or (
                     epoch + 1
                 ) == epochs
-                if should_log_valid:
+                val_metrics = None
+                if should_eval_valid:
+                    self.callbacks.on_validation_begin()
+                    val_metrics = self.evaluate(
+                        valid_loader,
+                        user_ids=valid_user_ids if self.needs_user_ids else None,
+                        num_workers=num_workers,
+                    )
                     display_metrics_table(
                         epoch=epoch + 1,
                         epochs=epochs,
@@ -1142,23 +1145,24 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
                         is_main_process=self.is_main_process,
                         colorize=lambda s: colorize("  " + s, color="cyan"),
                     )
-                self.callbacks.on_validation_end()
-                if should_log_valid and val_metrics and self.training_logger:
-                    self.training_logger.log_metrics(
-                        val_metrics, step=epoch + 1, split="valid"
-                    )
+                    self.callbacks.on_validation_end()
+                    if val_metrics and self.training_logger:
+                        self.training_logger.log_metrics(
+                            val_metrics, step=epoch + 1, split="valid"
+                        )
                 if not val_metrics:
-                    if self.is_main_process:
+                    if should_eval_valid and self.is_main_process:
                         logging.info(
                             colorize(
                                 "Warning: No validation metrics computed. Skipping validation for this epoch.",
                                 color="yellow",
                             )
                         )
-                    continue
-                epoch_logs = {**train_log_payload}
-                for k, v in val_metrics.items():
-                    epoch_logs[f"val_{k}"] = v
+                    epoch_logs = {**train_log_payload}
+                else:
+                    epoch_logs = {**train_log_payload}
+                    for k, v in val_metrics.items():
+                        epoch_logs[f"val_{k}"] = v
             else:
                 epoch_logs = {**train_log_payload}
                 if self.is_main_process:
@@ -1340,6 +1344,7 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
                 target_names=self.target_columns,
                 task_specific_metrics=self.task_specific_metrics,
                 user_ids=combined_user_ids,
+                ignore_label=self.ignore_label,
             )
             return avg_loss, metrics_dict
         return avg_loss
@@ -1387,6 +1392,8 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
                 sampler=valid_sampler,
                 collate_fn=collate_fn,
                 num_workers=num_workers,
+                pin_memory=self.device.type == "cuda",
+                persistent_workers=num_workers > 0,
             )
         valid_user_ids = None
         if needs_user_ids:
@@ -1532,6 +1539,7 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
             target_names=self.target_columns,
             task_specific_metrics=self.task_specific_metrics,
             user_ids=final_user_ids,
+            ignore_label=self.ignore_label,
         )
         return metrics_dict
 

nextrec/data/dataloader.py

@@ -282,6 +282,8 @@ class RecDataLoader(FeatureSet):
             sampler=sampler,
             collate_fn=collate_fn,
             num_workers=num_workers,
+            pin_memory=torch.cuda.is_available(),
+            persistent_workers=num_workers > 0,
         )
 
     def create_from_path(