nextrec 0.4.34__tar.gz → 0.5.1__tar.gz

{nextrec-0.4.34 → nextrec-0.5.1}/.gitignore

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

{nextrec-0.4.34 → nextrec-0.5.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nextrec
-Version: 0.4.34
+Version: 0.5.1
 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.34-orange.svg)
+![Version](https://img.shields.io/badge/Version-0.5.1-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.34，NextRec CLI支持单机训练，分布式训练相关功能尚在开发中。
+> 截止当前版本0.5.1，NextRec CLI支持单机训练，分布式训练相关功能尚在开发中。
 
 ## 兼容平台
 
-当前最新版本为0.4.34，所有模型和测试代码均已在以下平台通过验证，如果开发者在使用中遇到兼容问题，请在issue区提出错误报告及系统版本：
+当前最新版本为0.5.1，所有模型和测试代码均已在以下平台通过验证，如果开发者在使用中遇到兼容问题，请在issue区提出错误报告及系统版本：
 
 | 平台 | 配置 | 
 |------|------|

{nextrec-0.4.34 → nextrec-0.5.1}/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.34-orange.svg)
+![Version](https://img.shields.io/badge/Version-0.5.1-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.34，NextRec CLI支持单机训练，分布式训练相关功能尚在开发中。
+> 截止当前版本0.5.1，NextRec CLI支持单机训练，分布式训练相关功能尚在开发中。
 
 ## 兼容平台
 
-当前最新版本为0.4.34，所有模型和测试代码均已在以下平台通过验证，如果开发者在使用中遇到兼容问题，请在issue区提出错误报告及系统版本：
+当前最新版本为0.5.1，所有模型和测试代码均已在以下平台通过验证，如果开发者在使用中遇到兼容问题，请在issue区提出错误报告及系统版本：
 
 | 平台 | 配置 | 
 |------|------|

{nextrec-0.4.34 → nextrec-0.5.1}/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.34-orange.svg)
+![Version](https://img.shields.io/badge/Version-0.5.1-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.34, NextRec CLI supports single-machine training; distributed training features are currently under development.
+> As of version 0.5.1, NextRec CLI supports single-machine training; distributed training features are currently under development.
 
 ## Platform Compatibility
 
-The current version is 0.4.34. 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.1. 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.34 → nextrec-0.5.1}/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.34"
+release = "0.5.1"
 
 extensions = [
     "myst_parser",

nextrec-0.5.1/nextrec/__version__.py

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

{nextrec-0.4.34 → nextrec-0.5.1}/nextrec/basic/activation.py

@@ -25,21 +25,15 @@ class Dice(nn.Module):
     def __init__(self, emb_size: int, epsilon: float = 1e-3):
         super(Dice, self).__init__()
         self.alpha = nn.Parameter(torch.zeros(emb_size))
-        self.bn = nn.BatchNorm1d(emb_size, eps=epsilon)
+        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)
-        if x.dim() == 2:  # (B, E)
-            x_norm = self.bn(x)
-            p = torch.sigmoid(x_norm)
-            return x * (self.alpha + (1 - self.alpha) * p)
-
-        if x.dim() == 3:  # (B, T, E)
-            b, t, e = x.shape
-            x2 = x.reshape(-1, e)  # (B*T, E)
-            x_norm = self.bn(x2)
-            p = torch.sigmoid(x_norm).reshape(b, t, e)
-            return x * (self.alpha + (1 - self.alpha) * p)
+        # 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.34 → nextrec-0.5.1}/nextrec/basic/layers.py

@@ -2,7 +2,7 @@
 Layer implementations used across NextRec.
 
 Date: create on 27/10/2025
-Checkpoint: edit on 22/01/2026
+Checkpoint: edit on 25/01/2026
 Author: Yang Zhou, zyaztec@gmail.com
 """
 
@@ -79,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):
@@ -216,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]
@@ -279,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]
@@ -328,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()
 
 
@@ -928,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)
@@ -1000,36 +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