nextrec 0.4.33__tar.gz → 0.5.0__tar.gz

{nextrec-0.4.33 → nextrec-0.5.0}/.gitignore

@@ -128,3 +128,6 @@ pypirc.template
 
 # Sphinx build
 docs/rtd/_build/
+
+*.onnx
+artifacts/

{nextrec-0.4.33 → nextrec-0.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nextrec
-Version: 0.4.33
+Version: 0.5.0
 Summary: A comprehensive recommendation library with match, ranking, and multi-task learning models
 Project-URL: Homepage, https://github.com/zerolovesea/NextRec
 Project-URL: Repository, https://github.com/zerolovesea/NextRec
@@ -24,10 +24,14 @@ Requires-Dist: numpy<2.0,>=1.21; sys_platform == 'linux' and python_version < '3
 Requires-Dist: numpy<3.0,>=1.26; sys_platform == 'linux' and python_version >= '3.12'
 Requires-Dist: numpy>=1.23.0; sys_platform == 'win32'
 Requires-Dist: numpy>=1.24.0; sys_platform == 'darwin'
+Requires-Dist: onnx>=1.16.0
+Requires-Dist: onnxruntime>=1.18.0
+Requires-Dist: onnxscript>=0.1.1
 Requires-Dist: pandas<2.0,>=1.5; sys_platform == 'linux' and python_version < '3.12'
 Requires-Dist: pandas<2.3.0,>=2.1.0; sys_platform == 'win32'
 Requires-Dist: pandas>=2.0.0; sys_platform == 'darwin'
 Requires-Dist: pandas>=2.1.0; sys_platform == 'linux' and python_version >= '3.12'
+Requires-Dist: polars>=0.20.0
 Requires-Dist: pyarrow<13.0.0,>=10.0.0; sys_platform == 'linux' and python_version < '3.12'
 Requires-Dist: pyarrow<15.0.0,>=12.0.0; sys_platform == 'win32'
 Requires-Dist: pyarrow>=12.0.0; sys_platform == 'darwin'
@@ -69,7 +73,7 @@ Description-Content-Type: text/markdown
 ![Python](https://img.shields.io/badge/Python-3.10+-blue.svg)
 ![PyTorch](https://img.shields.io/badge/PyTorch-1.10+-ee4c2c.svg)
 ![License](https://img.shields.io/badge/License-Apache%202.0-green.svg)
-![Version](https://img.shields.io/badge/Version-0.4.33-orange.svg)
+![Version](https://img.shields.io/badge/Version-0.5.0-orange.svg)
 [![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/zerolovesea/NextRec)
 
 中文文档 | [English Version](README_en.md)
@@ -102,6 +106,7 @@ NextRec是一个基于PyTorch的现代推荐系统框架，旨在为研究工程
 - **高效训练与评估**：内置多种优化器、学习率调度、早停、模型检查点与详细的日志管理，开箱即用。
 
 ## NextRec近期进展
+- **28/01/2026** 在v0.4.39中加入了对onnx导出和加载的支持，并大大加速了数据预处理速度（最高9x加速）
 - **01/01/2026** 新年好，在v0.4.27中加入了多个多目标模型的支持：[APG](nextrec/models/multi_task/apg.py), [ESCM](nextrec/models/multi_task/escm.py), [HMoE](nextrec/models/multi_task/hmoe.py), [Cross Stitch](nextrec/models/multi_task/cross_stitch.py)
 - **28/12/2025** 在v0.4.21中加入了对SwanLab和Wandb的支持，通过model的`fit`方法进行配置：`use_swanlab=True, swanlab_kwargs={"project": "NextRec","name":"tutorial_movielens_deepfm"},`
 - **21/12/2025** 在v0.4.16中加入了对[GradNorm](/nextrec/loss/grad_norm.py)的支持，通过compile的`loss_weight='grad_norm'`进行配置
@@ -136,6 +141,7 @@ pip install nextrec # or pip install -e .
 - [example_multitask.py](/tutorials/example_multitask.py) - 电商数据集上的ESMM多任务学习训练示例
 - [movielen_match_dssm.py](/tutorials/movielen_match_dssm.py) - 基于movielen 100k数据集训练的 DSSM 召回模型示例
 
+- [example_onnx.py](/tutorials/example_onnx.py) - 使用NextRec训练和导出onnx模型
 - [example_distributed_training.py](/tutorials/distributed/example_distributed_training.py) - 使用NextRec进行单机多卡训练的代码示例
 
 - [run_all_ranking_models.py](/tutorials/run_all_ranking_models.py) - 快速校验所有排序模型的可用性
@@ -254,11 +260,11 @@ nextrec --mode=predict --predict_config=path/to/predict_config.yaml
 
 预测结果固定保存到 `{checkpoint_path}/predictions/{name}.{save_data_format}`。
 
-> 截止当前版本0.4.33，NextRec CLI支持单机训练，分布式训练相关功能尚在开发中。
+> 截止当前版本0.5.0，NextRec CLI支持单机训练，分布式训练相关功能尚在开发中。
 
 ## 兼容平台
 
-当前最新版本为0.4.33，所有模型和测试代码均已在以下平台通过验证，如果开发者在使用中遇到兼容问题，请在issue区提出错误报告及系统版本：
+当前最新版本为0.5.0，所有模型和测试代码均已在以下平台通过验证，如果开发者在使用中遇到兼容问题，请在issue区提出错误报告及系统版本：
 
 | 平台 | 配置 | 
 |------|------|

{nextrec-0.4.33 → nextrec-0.5.0}/README.md

@@ -8,7 +8,7 @@
 ![Python](https://img.shields.io/badge/Python-3.10+-blue.svg)
 ![PyTorch](https://img.shields.io/badge/PyTorch-1.10+-ee4c2c.svg)
 ![License](https://img.shields.io/badge/License-Apache%202.0-green.svg)
-![Version](https://img.shields.io/badge/Version-0.4.33-orange.svg)
+![Version](https://img.shields.io/badge/Version-0.5.0-orange.svg)
 [![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/zerolovesea/NextRec)
 
 中文文档 | [English Version](README_en.md)
@@ -41,6 +41,7 @@ NextRec是一个基于PyTorch的现代推荐系统框架，旨在为研究工程
 - **高效训练与评估**：内置多种优化器、学习率调度、早停、模型检查点与详细的日志管理，开箱即用。
 
 ## NextRec近期进展
+- **28/01/2026** 在v0.4.39中加入了对onnx导出和加载的支持，并大大加速了数据预处理速度（最高9x加速）
 - **01/01/2026** 新年好，在v0.4.27中加入了多个多目标模型的支持：[APG](nextrec/models/multi_task/apg.py), [ESCM](nextrec/models/multi_task/escm.py), [HMoE](nextrec/models/multi_task/hmoe.py), [Cross Stitch](nextrec/models/multi_task/cross_stitch.py)
 - **28/12/2025** 在v0.4.21中加入了对SwanLab和Wandb的支持，通过model的`fit`方法进行配置：`use_swanlab=True, swanlab_kwargs={"project": "NextRec","name":"tutorial_movielens_deepfm"},`
 - **21/12/2025** 在v0.4.16中加入了对[GradNorm](/nextrec/loss/grad_norm.py)的支持，通过compile的`loss_weight='grad_norm'`进行配置
@@ -75,6 +76,7 @@ pip install nextrec # or pip install -e .
 - [example_multitask.py](/tutorials/example_multitask.py) - 电商数据集上的ESMM多任务学习训练示例
 - [movielen_match_dssm.py](/tutorials/movielen_match_dssm.py) - 基于movielen 100k数据集训练的 DSSM 召回模型示例
 
+- [example_onnx.py](/tutorials/example_onnx.py) - 使用NextRec训练和导出onnx模型
 - [example_distributed_training.py](/tutorials/distributed/example_distributed_training.py) - 使用NextRec进行单机多卡训练的代码示例
 
 - [run_all_ranking_models.py](/tutorials/run_all_ranking_models.py) - 快速校验所有排序模型的可用性
@@ -193,11 +195,11 @@ nextrec --mode=predict --predict_config=path/to/predict_config.yaml
 
 预测结果固定保存到 `{checkpoint_path}/predictions/{name}.{save_data_format}`。
 
-> 截止当前版本0.4.33，NextRec CLI支持单机训练，分布式训练相关功能尚在开发中。
+> 截止当前版本0.5.0，NextRec CLI支持单机训练，分布式训练相关功能尚在开发中。
 
 ## 兼容平台
 
-当前最新版本为0.4.33，所有模型和测试代码均已在以下平台通过验证，如果开发者在使用中遇到兼容问题，请在issue区提出错误报告及系统版本：
+当前最新版本为0.5.0，所有模型和测试代码均已在以下平台通过验证，如果开发者在使用中遇到兼容问题，请在issue区提出错误报告及系统版本：
 
 | 平台 | 配置 | 
 |------|------|

{nextrec-0.4.33 → nextrec-0.5.0}/README_en.md

@@ -8,7 +8,7 @@
 ![Python](https://img.shields.io/badge/Python-3.10+-blue.svg)
 ![PyTorch](https://img.shields.io/badge/PyTorch-1.10+-ee4c2c.svg)
 ![License](https://img.shields.io/badge/License-Apache%202.0-green.svg)
-![Version](https://img.shields.io/badge/Version-0.4.33-orange.svg)
+![Version](https://img.shields.io/badge/Version-0.5.0-orange.svg)
 [![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/zerolovesea/NextRec)
 
 English | [中文文档](README.md)
@@ -44,6 +44,7 @@ NextRec is a modern recommendation framework built on PyTorch, delivering a unif
 
 ## NextRec Progress
 
+- **28/01/2026** Added support for ONNX export and loading in v0.4.39, and significantly accelerated data preprocessing speed (up to 9x speedup)
 - **01/01/2026** Happy New Year! In v0.4.27, added support for multiple multi-task models: [APG](/nextrec/models/multi_task/apg.py), [ESCM](/nextrec/models/multi_task/escm.py), [HMoE](/nextrec/models/multi_task/hmoe.py), [Cross Stitch](/nextrec/models/multi_task/cross_stitch.py)
 - **28/12/2025** Added support for SwanLab and Weights & Biases in v0.4.21, configurable via the model `fit` method: `use_swanlab=True, swanlab_kwargs={"project": "NextRec","name":"tutorial_movielens_deepfm"},`
 - **21/12/2025** Added support for [GradNorm](/nextrec/loss/grad_norm.py) in v0.4.16, configurable via `loss_weight='grad_norm'` in the compile method
@@ -79,6 +80,7 @@ See `tutorials/` for examples covering ranking, retrieval, multi-task learning,
 - [example_multitask.py](/tutorials/example_multitask.py) — ESMM multi-task learning training on e-commerce dataset
 - [movielen_match_dssm.py](/tutorials/movielen_match_dssm.py) — DSSM retrieval model training on MovieLens 100k dataset
 
+- [example_onnx.py](/tutorials/example_onnx.py) — Train and export models to ONNX format with NextRec
 - [example_distributed_training.py](/tutorials/distributed/example_distributed_training.py) — Single-machine multi-GPU training with NextRec
 
 - [run_all_ranking_models.py](/tutorials/run_all_ranking_models.py) — Quickly validate availability of all ranking models
@@ -196,11 +198,11 @@ nextrec --mode=predict --predict_config=path/to/predict_config.yaml
 
 Prediction outputs are saved under `{checkpoint_path}/predictions/{name}.{save_data_format}`.
 
-> As of version 0.4.33, NextRec CLI supports single-machine training; distributed training features are currently under development.
+> As of version 0.5.0, NextRec CLI supports single-machine training; distributed training features are currently under development.
 
 ## Platform Compatibility
 
-The current version is 0.4.33. All models and test code have been validated on the following platforms. If you encounter compatibility issues, please report them in the issue tracker with your system version:
+The current version is 0.5.0. All models and test code have been validated on the following platforms. If you encounter compatibility issues, please report them in the issue tracker with your system version:
 
 | Platform | Configuration | 
 |----------|---------------|

{nextrec-0.4.33 → nextrec-0.5.0}/docs/rtd/conf.py

@@ -11,7 +11,7 @@ sys.path.insert(0, str(PROJECT_ROOT / "nextrec"))
 project = "NextRec"
 copyright = "2026, Yang Zhou"
 author = "Yang Zhou"
-release = "0.4.33"
+release = "0.5.0"
 
 extensions = [
     "myst_parser",

{nextrec-0.4.33 → nextrec-0.5.0}/docs/rtd/nextrec.utils.rst

@@ -28,14 +28,6 @@ nextrec.utils.data module
    :undoc-members:
    :show-inheritance:
 
-nextrec.utils.feature module
-----------------------------
-
-.. automodule:: nextrec.utils.feature
-   :members:
-   :undoc-members:
-   :show-inheritance:
-
 nextrec.utils.model module
 --------------------------
 

nextrec-0.5.0/nextrec/__version__.py

@@ -0,0 +1 @@
+__version__ = "0.5.0"

{nextrec-0.4.33 → nextrec-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-0.4.33 → nextrec-0.5.0}/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-0.4.33 → nextrec-0.5.0}/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-0.4.33 → nextrec-0.5.0}/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-0.4.33 → nextrec-0.5.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-0.4.33 → nextrec-0.5.0}/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-0.4.33 → nextrec-0.5.0}/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