nextrec 0.4.33__py3-none-any.whl → 0.5.0__py3-none-any.whl

nextrec/utils/onnx_utils.py

@@ -0,0 +1,252 @@
+"""
+ONNX utilities for NextRec.
+
+Date: create on 25/01/2026
+Author: Yang Zhou, zyaztec@gmail.com
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Iterable, Sequence
+
+import torch
+import numpy as np
+import onnxruntime as ort
+
+from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
+from nextrec.utils.torch_utils import to_numpy
+
+
+class OnnxModelWrapper(torch.nn.Module):
+    """Wrap a NextRec model to accept positional ONNX inputs."""
+
+    def __init__(self, model: torch.nn.Module, feature_names: Sequence[str]):
+        super().__init__()
+        self.model = model
+        self.feature_names = list(feature_names)
+
+    def forward(self, *inputs: torch.Tensor):
+        if len(inputs) != len(self.feature_names):
+            raise ValueError(
+                "[OnnxWrapper Error] Number of inputs does not match feature names."
+            )
+        x = {name: tensor for name, tensor in zip(self.feature_names, inputs)}
+        output = self.model(x)
+        if isinstance(output, list):
+            return tuple(output)
+        return output
+
+
+def create_dummy_inputs(
+    features: Iterable[object],
+    batch_size: int,
+    device: torch.device,
+    default_seq_len: int = 10,
+    seq_len_map: dict[str, int] | None = None,
+) -> list[torch.Tensor]:
+    tensors: list[torch.Tensor] = []
+    for feature in features:
+        if isinstance(feature, DenseFeature):
+            input_dim = max(int(feature.input_dim), 1)
+            tensors.append(
+                torch.zeros((batch_size, input_dim), dtype=torch.float32, device=device)
+            )
+        elif isinstance(feature, SequenceFeature):
+            seq_len = None
+            if seq_len_map:
+                seq_len = seq_len_map.get(feature.name)
+            if seq_len is None:
+                seq_len = (
+                    feature.max_len if feature.max_len is not None else default_seq_len
+                )
+            seq_len = max(int(seq_len), 1)
+            tensors.append(
+                torch.zeros((batch_size, seq_len), dtype=torch.long, device=device)
+            )
+        else:
+            tensors.append(torch.zeros((batch_size,), dtype=torch.long, device=device))
+    return tensors
+
+
+def normalize_dense(feature: DenseFeature, array: object) -> np.ndarray:
+    arr = np.asarray(array, dtype=np.float32)
+    if arr.ndim == 1:
+        arr = arr.reshape(-1, 1)
+    else:
+        arr = arr.reshape(arr.shape[0], -1)
+    expected = max(int(feature.input_dim), 1)
+    if arr.shape[1] != expected:
+        raise ValueError(
+            f"[ONNX Input Error] Dense feature '{feature.name}' expects {expected} dims but got {arr.shape[1]}."
+        )
+    return arr
+
+
+def normalize_sparse(feature: SparseFeature, array: object) -> np.ndarray:
+    arr = np.asarray(array, dtype=np.int64)
+    if arr.ndim == 2 and arr.shape[1] == 1:
+        arr = arr.reshape(-1)
+    elif arr.ndim != 1:
+        arr = arr.reshape(arr.shape[0], -1)
+        if arr.shape[1] != 1:
+            raise ValueError(
+                f"[ONNX Input Error] Sparse feature '{feature.name}' expects 1 dim but got {arr.shape}."
+            )
+        arr = arr.reshape(-1)
+    return arr
+
+
+def normalize_sequence(feature: SequenceFeature, array: object) -> np.ndarray:
+    arr = np.asarray(array, dtype=np.int64)
+    if arr.ndim == 1:
+        arr = arr.reshape(1, -1)
+    elif arr.ndim > 2:
+        arr = arr.reshape(arr.shape[0], -1)
+    max_len = feature.max_len if feature.max_len is not None else arr.shape[1]
+    max_len = max(int(max_len), 1)
+    if arr.shape[1] > max_len:
+        arr = arr[:, :max_len]
+    elif arr.shape[1] < max_len:
+        pad_value = feature.padding_idx if feature.padding_idx is not None else 0
+        pad_width = max_len - arr.shape[1]
+        arr = np.pad(arr, ((0, 0), (0, pad_width)), constant_values=pad_value)
+    return arr
+
+
+def build_onnx_input_feed(
+    features: Iterable[object],
+    feature_batch: dict[str, object],
+    input_names: Sequence[str] | None = None,
+) -> dict[str, np.ndarray]:
+    feed: dict[str, np.ndarray] = {}
+    for feature in features:
+        if input_names is not None and feature.name not in input_names:
+            continue
+        if feature.name not in feature_batch:
+            raise KeyError(
+                f"[ONNX Input Error] Feature '{feature.name}' missing from batch data."
+            )
+        value = to_numpy(feature_batch[feature.name])
+        if isinstance(feature, DenseFeature):
+            value = normalize_dense(feature, value)
+        elif isinstance(feature, SequenceFeature):
+            value = normalize_sequence(feature, value)
+        else:
+            value = normalize_sparse(feature, value)
+        feed[feature.name] = value
+    return feed
+
+
+def pad_tensor(
+    value: torch.Tensor, pad_rows: int, pad_value: int | float
+) -> torch.Tensor:
+    if pad_rows <= 0:
+        return value
+    pad_shape = (pad_rows, *value.shape[1:])
+    pad = torch.full(pad_shape, pad_value, dtype=value.dtype, device=value.device)
+    return torch.cat([value, pad], dim=0)
+
+
+def pad_array(value: np.ndarray, pad_rows: int, pad_value: int | float) -> np.ndarray:
+    if pad_rows <= 0:
+        return value
+    pad_shape = (pad_rows, *value.shape[1:])
+    pad = np.full(pad_shape, pad_value, dtype=value.dtype)
+    return np.concatenate([value, pad], axis=0)
+
+
+def pad_onnx_inputs(
+    features: Iterable[object],
+    feature_batch: dict[str, object],
+    target_batch: int,
+) -> tuple[dict[str, object], int]:
+    if target_batch <= 0:
+        return feature_batch, 0
+    padded: dict[str, object] = {}
+    orig_batch = None
+    for feature in features:
+        if feature.name not in feature_batch:
+            continue
+        value = feature_batch[feature.name]
+        if isinstance(value, torch.Tensor):
+            batch = value.shape[0] if value.dim() > 0 else 1
+        else:
+            arr = np.asarray(value)
+            batch = arr.shape[0] if arr.ndim > 0 else 1
+        if orig_batch is None:
+            orig_batch = int(batch)
+        pad_rows = max(int(target_batch) - int(batch), 0)
+        if isinstance(feature, DenseFeature):
+            pad_value: int | float = 0.0
+        elif isinstance(feature, SequenceFeature):
+            pad_value = feature.padding_idx if feature.padding_idx is not None else 0
+        else:
+            pad_value = 0
+        if isinstance(value, torch.Tensor):
+            padded[feature.name] = pad_tensor(value, pad_rows, pad_value)
+        else:
+            padded[feature.name] = pad_array(np.asarray(value), pad_rows, pad_value)
+    if orig_batch is None:
+        orig_batch = 0
+    return padded, orig_batch
+
+
+def pad_id_batch(
+    id_batch: dict[str, object],
+    target_batch: int,
+    pad_value: int | float = 0,
+) -> tuple[dict[str, object], int]:
+    if target_batch <= 0:
+        return id_batch, 0
+    padded: dict[str, object] = {}
+    orig_batch = None
+    for name, value in id_batch.items():
+        if isinstance(value, torch.Tensor):
+            batch = value.shape[0] if value.dim() > 0 else 1
+        else:
+            arr = np.asarray(value)
+            batch = arr.shape[0] if arr.ndim > 0 else 1
+        if orig_batch is None:
+            orig_batch = int(batch)
+        pad_rows = max(int(target_batch) - int(batch), 0)
+        if isinstance(value, torch.Tensor):
+            padded[name] = pad_tensor(value, pad_rows, pad_value)
+        else:
+            padded[name] = pad_array(np.asarray(value), pad_rows, pad_value)
+    if orig_batch is None:
+        orig_batch = 0
+    return padded, orig_batch
+
+
+def merge_onnx_outputs(outputs: Sequence[np.ndarray]) -> np.ndarray:
+    if not outputs:
+        raise ValueError("[ONNX Output Error] Empty ONNX output list.")
+    if len(outputs) == 1:
+        return outputs[0]
+    normalized: list[np.ndarray] = []
+    batch = outputs[0].shape[0] if outputs[0].ndim > 0 else None
+    for out in outputs:
+        arr = np.asarray(out)
+        if arr.ndim == 0:
+            arr = arr.reshape(1, 1)
+        elif arr.ndim == 1:
+            arr = arr.reshape(-1, 1)
+        if batch is not None and arr.shape[0] != batch:
+            raise ValueError(
+                "[ONNX Output Error] Output batch size mismatch across ONNX outputs."
+            )
+        normalized.append(arr)
+    return np.concatenate(normalized, axis=1)
+
+
+def load_onnx_session(
+    onnx_path: str | Path,
+):
+    available = ort.get_available_providers()
+    preferred = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+    selected = [p for p in preferred if p in available]
+    if not selected:
+        selected = available
+
+    return ort.InferenceSession(str(onnx_path), providers=selected)

nextrec/utils/torch_utils.py

@@ -5,14 +5,15 @@ 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
+Checkpoint: edit on 22/01/2026
 Author: Yang Zhou, zyaztec@gmail.com
 """
 
 from __future__ import annotations
 
 import logging
-from typing import Any, Dict, Iterable, Literal
+import numbers
+from typing import Any, Dict, Iterable
 
 import numpy as np
 import torch
@@ -22,7 +23,55 @@ from torch.utils.data import DataLoader, IterableDataset
 from torch.utils.data.distributed import DistributedSampler
 
 from nextrec.basic.loggers import colorize
-from nextrec.utils.types import OptimizerName, SchedulerName
+from nextrec.utils.types import (
+    EmbeddingInitType,
+    InitializerActivationType,
+    OptimizerName,
+    SchedulerName,
+)
+
+
+def to_list(value: str | list[str] | None) -> list[str]:
+    if value is None:
+        return []
+    if isinstance(value, str):
+        return [value]
+    return list(value)
+
+
+def as_float(value: Any) -> float | None:
+    if isinstance(value, numbers.Number):
+        return float(value)
+    if hasattr(value, "item"):
+        try:
+            return float(value.item())
+        except Exception:
+            return None
+    return None
+
+
+def to_numpy(values: Any) -> np.ndarray:
+    if isinstance(values, torch.Tensor):
+        return values.detach().cpu().numpy()
+    return np.asarray(values)
+
+
+def to_tensor(
+    value: Any, dtype: torch.dtype, device: torch.device | str | None = None
+) -> torch.Tensor:
+    if value is None:
+        raise ValueError("[Tensor Utils Error] Cannot convert None to tensor.")
+    tensor = value if isinstance(value, torch.Tensor) else torch.as_tensor(value)
+    if tensor.dtype != dtype:
+        tensor = tensor.to(dtype=dtype)
+
+    if device is not None:
+        target_device = (
+            device if isinstance(device, torch.device) else torch.device(device)
+        )
+        if tensor.device != target_device:
+            tensor = tensor.to(target_device)
+    return tensor
 
 
 def resolve_nonlinearity(activation: str) -> str:
@@ -56,30 +105,8 @@ def resolve_gain(activation: str, param: Dict[str, Any]) -> float:
 
 
 def get_initializer(
-    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",
+    init_type: EmbeddingInitType = "normal",
+    activation: InitializerActivationType = "linear",
     param: Dict[str, Any] | None = None,
 ):
     param = param or {}
@@ -108,7 +135,7 @@ def get_initializer(
         elif init_type == "uniform":
             nn.init.uniform_(tensor, a=param.get("a", -0.05), b=param.get("b", 0.05))
         else:
-            raise ValueError(f"Unknown init_type: {init_type}")
+            raise ValueError(f"[Initializer Error] Unknown init_type: {init_type}")
         return tensor
 
     return initializer_fn
@@ -172,12 +199,14 @@ def get_optimizer(
         elif opt_name == "rmsprop":
             opt_class = torch.optim.RMSprop
         else:
-            raise NotImplementedError(f"Unsupported optimizer: {optimizer}")
+            raise NotImplementedError(
+                f"[Optimizer Error] Unsupported optimizer: {optimizer}"
+            )
         optimizer_fn = opt_class(params=params, **optimizer_params)
     elif isinstance(optimizer, torch.optim.Optimizer):
         optimizer_fn = optimizer
     else:
-        raise TypeError(f"Invalid optimizer type: {type(optimizer)}")
+        raise TypeError(f"[Optimizer Error] Invalid optimizer type: {type(optimizer)}")
     return optimizer_fn
 
 
@@ -203,7 +232,9 @@ def get_scheduler(
                 optimizer, **scheduler_params
             )
         else:
-            raise NotImplementedError(f"Unsupported scheduler: {scheduler}")
+            raise NotImplementedError(
+                f"[Scheduler Error] Unsupported scheduler: {scheduler}"
+            )
     elif isinstance(scheduler, type) and issubclass(
         scheduler,
         (torch.optim.lr_scheduler._LRScheduler, torch.optim.lr_scheduler.LRScheduler),
@@ -215,35 +246,11 @@ def get_scheduler(
     ):
         scheduler_fn = scheduler
     else:
-        raise TypeError(f"Invalid scheduler type: {type(scheduler)}")
+        raise TypeError(f"[Scheduler Error] Invalid scheduler type: {type(scheduler)}")
 
     return scheduler_fn
 
 
-def to_numpy(values: Any) -> np.ndarray:
-    if isinstance(values, torch.Tensor):
-        return values.detach().cpu().numpy()
-    return np.asarray(values)
-
-
-def to_tensor(
-    value: Any, dtype: torch.dtype, device: torch.device | str | None = None
-) -> torch.Tensor:
-    if value is None:
-        raise ValueError("[Tensor Utils Error] Cannot convert None to tensor.")
-    tensor = value if isinstance(value, torch.Tensor) else torch.as_tensor(value)
-    if tensor.dtype != dtype:
-        tensor = tensor.to(dtype=dtype)
-
-    if device is not None:
-        target_device = (
-            device if isinstance(device, torch.device) else torch.device(device)
-        )
-        if tensor.device != target_device:
-            tensor = tensor.to(target_device)
-    return tensor
-
-
 def init_process_group(
     distributed: bool, rank: int, world_size: int, device_id: int | None = None
 ) -> None:

nextrec/utils/types.py

@@ -64,6 +64,40 @@ TaskTypeName = Literal["binary", "regression"]
 
 TaskTypeInput = TaskTypeName | str
 
+EmbeddingInitType = Literal[
+    "normal",
+    "uniform",
+    "xavier_uniform",
+    "xavier_normal",
+    "kaiming_uniform",
+    "kaiming_normal",
+    "orthogonal",
+]
+
+SequenceCombinerType = Literal[
+    "mean",
+    "sum",
+    "concat",
+    "dot_attention",
+    "self_attention",
+]
+
+InitializerActivationType = Literal[
+    "linear",
+    "conv1d",
+    "conv2d",
+    "conv3d",
+    "conv_transpose1d",
+    "conv_transpose2d",
+    "conv_transpose3d",
+    "sigmoid",
+    "tanh",
+    "relu",
+    "leaky_relu",
+    "selu",
+    "gelu",
+]
+
 MetricsName = Literal[
     "auc",
     "gauc",
@@ -97,4 +131,13 @@ MetricsName = Literal[
     "mrr@5",
     "mrr@10",
     "mrr@20",
+    "topk_recall@5",
+    "topk_recall@10",
+    "topk_recall@20",
+    "topk_precision@5",
+    "topk_precision@10",
+    "topk_precision@20",
+    "lift@5",
+    "lift@10",
+    "lift@20",
 ]

{nextrec-0.4.33.dist-info → nextrec-0.5.0.dist-info}/METADATA

@@ -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区提出错误报告及系统版本：
 
 | 平台 | 配置 | 
 |------|------|