nextrec 0.4.30__tar.gz → 0.4.32__tar.gz

{nextrec-0.4.30 → nextrec-0.4.32}/LICENSE

@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2025 NextRec
+Copyright (c) 2026 NextRec
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

{nextrec-0.4.30 → nextrec-0.4.32}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nextrec
-Version: 0.4.30
+Version: 0.4.32
 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
@@ -69,7 +69,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.30-orange.svg)
+![Version](https://img.shields.io/badge/Version-0.4.32-orange.svg)
 [![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/zerolovesea/NextRec)
 
 中文文档 | [English Version](README_en.md)
@@ -254,11 +254,11 @@ nextrec --mode=predict --predict_config=path/to/predict_config.yaml
 
 预测结果固定保存到 `{checkpoint_path}/predictions/{name}.{save_data_format}`。
 
-> 截止当前版本0.4.30，NextRec CLI支持单机训练，分布式训练相关功能尚在开发中。
+> 截止当前版本0.4.32，NextRec CLI支持单机训练，分布式训练相关功能尚在开发中。
 
 ## 兼容平台
 
-当前最新版本为0.4.30，所有模型和测试代码均已在以下平台通过验证，如果开发者在使用中遇到兼容问题，请在issue区提出错误报告及系统版本：
+当前最新版本为0.4.32，所有模型和测试代码均已在以下平台通过验证，如果开发者在使用中遇到兼容问题，请在issue区提出错误报告及系统版本：
 
 | 平台 | 配置 | 
 |------|------|
@@ -400,7 +400,7 @@ NextRec 的开发受到以下优秀项目的启发：
 @misc{nextrec,
     title = {NextRec},
     author = {Yang Zhou},
-    year = {2025},
+    year = {2026},
     publisher = {GitHub},
     journal = {GitHub repository},
     howpublished = {\url{https://github.com/zerolovesea/NextRec}},

{nextrec-0.4.30 → nextrec-0.4.32}/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.30-orange.svg)
+![Version](https://img.shields.io/badge/Version-0.4.32-orange.svg)
 [![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/zerolovesea/NextRec)
 
 中文文档 | [English Version](README_en.md)
@@ -193,11 +193,11 @@ nextrec --mode=predict --predict_config=path/to/predict_config.yaml
 
 预测结果固定保存到 `{checkpoint_path}/predictions/{name}.{save_data_format}`。
 
-> 截止当前版本0.4.30，NextRec CLI支持单机训练，分布式训练相关功能尚在开发中。
+> 截止当前版本0.4.32，NextRec CLI支持单机训练，分布式训练相关功能尚在开发中。
 
 ## 兼容平台
 
-当前最新版本为0.4.30，所有模型和测试代码均已在以下平台通过验证，如果开发者在使用中遇到兼容问题，请在issue区提出错误报告及系统版本：
+当前最新版本为0.4.32，所有模型和测试代码均已在以下平台通过验证，如果开发者在使用中遇到兼容问题，请在issue区提出错误报告及系统版本：
 
 | 平台 | 配置 | 
 |------|------|
@@ -339,7 +339,7 @@ NextRec 的开发受到以下优秀项目的启发：
 @misc{nextrec,
     title = {NextRec},
     author = {Yang Zhou},
-    year = {2025},
+    year = {2026},
     publisher = {GitHub},
     journal = {GitHub repository},
     howpublished = {\url{https://github.com/zerolovesea/NextRec}},

{nextrec-0.4.30 → nextrec-0.4.32}/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.30-orange.svg)
+![Version](https://img.shields.io/badge/Version-0.4.32-orange.svg)
 [![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/zerolovesea/NextRec)
 
 English | [中文文档](README.md)
@@ -196,11 +196,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.30, NextRec CLI supports single-machine training; distributed training features are currently under development.
+> As of version 0.4.32, NextRec CLI supports single-machine training; distributed training features are currently under development.
 
 ## Platform Compatibility
 
-The current version is 0.4.30. 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.4.32. 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 | 
 |----------|---------------|
@@ -343,7 +343,7 @@ If you use this framework in your research or work, please consider citing:
 @misc{nextrec,
     title = {NextRec},
     author = {Yang Zhou},
-    year = {2025},
+    year = {2026},
     publisher = {GitHub},
     journal = {GitHub repository},
     howpublished = {\url{https://github.com/zerolovesea/NextRec}},

{nextrec-0.4.30 → nextrec-0.4.32}/docs/rtd/conf.py

@@ -9,9 +9,9 @@ PROJECT_ROOT = Path(__file__).resolve().parents[2]
 sys.path.insert(0, str(PROJECT_ROOT / "nextrec"))
 
 project = "NextRec"
-copyright = "2025, Yang Zhou"
+copyright = "2026, Yang Zhou"
 author = "Yang Zhou"
-release = "0.4.30"
+release = "0.4.32"
 
 extensions = [
     "myst_parser",

nextrec-0.4.32/nextrec/__version__.py

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

{nextrec-0.4.30 → nextrec-0.4.32}/nextrec/basic/model.py

@@ -13,7 +13,7 @@ import sys
 import pickle
 import socket
 from pathlib import Path
-from typing import Any, Literal
+from typing import Any, Literal, cast, overload
 
 import numpy as np
 import pandas as pd
@@ -97,6 +97,7 @@ from nextrec.utils.types import (
     SchedulerName,
     TrainingModeName,
     TaskTypeName,
+    TaskTypeInput,
     MetricsName,
 )
 
@@ -119,7 +120,7 @@ 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: TaskTypeName | list[TaskTypeName] | None = None,
+        task: TaskTypeInput | list[TaskTypeInput] | None = None,
         training_mode: TrainingModeName | list[TrainingModeName] | None = None,
         embedding_l1_reg: float = 0.0,
         dense_l1_reg: float = 0.0,
@@ -193,7 +194,7 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
             dense_features, sparse_features, sequence_features, target, id_columns
         )
 
-        self.task = task or self.default_task
+        self.task = cast(TaskTypeName | list[TaskTypeName], task or self.default_task)
         self.nums_task = len(self.task) if isinstance(self.task, list) else 1
 
         training_mode = training_mode or "pointwise"
@@ -1623,6 +1624,49 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
         )
         return metrics_dict
 
+    @overload
+    def predict(
+        self,
+        data: str | dict | pd.DataFrame | DataLoader,
+        batch_size: int = 32,
+        save_path: str | os.PathLike | None = None,
+        save_format: str = "csv",
+        include_ids: bool | None = None,
+        id_columns: str | list[str] | None = None,
+        return_dataframe: Literal[True] = True,
+        stream_chunk_size: int = 10000,
+        num_workers: int = 0,
+    ) -> pd.DataFrame: ...
+
+    @overload
+    def predict(
+        self,
+        data: str | dict | pd.DataFrame | DataLoader,
+        batch_size: int = 32,
+        save_path: None = None,
+        save_format: str = "csv",
+        include_ids: bool | None = None,
+        id_columns: str | list[str] | None = None,
+        return_dataframe: Literal[False] = False,
+        stream_chunk_size: int = 10000,
+        num_workers: int = 0,
+    ) -> np.ndarray: ...
+
+    @overload
+    def predict(
+        self,
+        data: str | dict | pd.DataFrame | DataLoader,
+        batch_size: int = 32,
+        *,
+        save_path: str | os.PathLike,
+        save_format: str = "csv",
+        include_ids: bool | None = None,
+        id_columns: str | list[str] | None = None,
+        return_dataframe: Literal[False] = False,
+        stream_chunk_size: int = 10000,
+        num_workers: int = 0,
+    ) -> Path: ...
+
     def predict(
         self,
         data: str | dict | pd.DataFrame | DataLoader,
@@ -2225,7 +2269,7 @@ class BaseMatchModel(BaseModel):
         dense_l2_reg: float = 0.0,
         target: list[str] | str | None = "label",
         id_columns: list[str] | str | None = None,
-        task: str | list[str] | None = None,
+        task: TaskTypeInput | list[TaskTypeInput] | None = None,
         session_id: str | None = None,
         distributed: bool = False,
         rank: int | None = None,

{nextrec-0.4.30 → nextrec-0.4.32}/nextrec/cli.py

@@ -407,6 +407,8 @@ def train_model(train_config_path: str) -> None:
         use_swanlab=train_cfg.get("use_swanlab", False),
         wandb_api=train_cfg.get("wandb_api"),
         swanlab_api=train_cfg.get("swanlab_api"),
+        wandb_kwargs=train_cfg.get("wandb_kwargs"),
+        swanlab_kwargs=train_cfg.get("swanlab_kwargs"),
         log_interval=train_cfg.get("log_interval", 1),
         note=train_cfg.get("note"),
     )
@@ -680,16 +682,22 @@ Examples:
     if not args.mode:
         parser.error("[NextRec CLI Error] --mode is required (train|predict)")
 
-    if args.mode == "train":
-        config_path = args.train_config
-        if not config_path:
-            parser.error("[NextRec CLI Error] train mode requires --train_config")
-        train_model(config_path)
-    else:
-        config_path = args.predict_config
-        if not config_path:
-            parser.error("[NextRec CLI Error] predict mode requires --predict_config")
-        predict_model(config_path)
+    try:
+        if args.mode == "train":
+            config_path = args.train_config
+            if not config_path:
+                parser.error("[NextRec CLI Error] train mode requires --train_config")
+            train_model(config_path)
+        else:
+            config_path = args.predict_config
+            if not config_path:
+                parser.error(
+                    "[NextRec CLI Error] predict mode requires --predict_config"
+                )
+            predict_model(config_path)
+    except Exception:
+        logging.getLogger(__name__).exception("[NextRec CLI Error] Unhandled exception")
+        raise
 
 
 if __name__ == "__main__":

{nextrec-0.4.30 → nextrec-0.4.32}/nextrec/data/batch_utils.py

@@ -12,7 +12,7 @@ import torch
 
 
 def stack_section(batch: list[dict], section: Literal["features", "labels", "ids"]):
-    """   
+    """
     input example:
     batch = [
         {"features": {"f1": tensor1, "f2": tensor2}, "labels": {"label": tensor3}},
@@ -24,7 +24,7 @@ def stack_section(batch: list[dict], section: Literal["features", "labels", "ids
         "f1": torch.stack([tensor1, tensor4], dim=0),
         "f2": torch.stack([tensor2, tensor5], dim=0),
     }
-    
+
     """
     entries = [item.get(section) for item in batch if item.get(section) is not None]
     if not entries:

{nextrec-0.4.30 → nextrec-0.4.32}/nextrec/data/preprocessor.py

@@ -13,7 +13,7 @@ import logging
 import os
 import pickle
 from pathlib import Path
-from typing import Any, Dict, Literal, Optional, Union
+from typing import Any, Dict, Literal, Optional, Union, overload
 
 import numpy as np
 import pandas as pd
@@ -895,6 +895,28 @@ class DataProcessor(FeatureSet):
         )
         return self
 
+    @overload
+    def transform_in_memory(
+        self,
+        data: Union[pd.DataFrame, Dict[str, Any]],
+        return_dict: Literal[True],
+        persist: bool,
+        save_format: Optional[str],
+        output_path: Optional[str],
+        warn_missing: bool = True,
+    ) -> Dict[str, np.ndarray]: ...
+
+    @overload
+    def transform_in_memory(
+        self,
+        data: Union[pd.DataFrame, Dict[str, Any]],
+        return_dict: Literal[False],
+        persist: bool,
+        save_format: Optional[str],
+        output_path: Optional[str],
+        warn_missing: bool = True,
+    ) -> pd.DataFrame: ...
+
     def transform_in_memory(
         self,
         data: Union[pd.DataFrame, Dict[str, Any]],
@@ -1238,6 +1260,36 @@ class DataProcessor(FeatureSet):
         self.is_fitted = True
         return self
 
+    @overload
+    def transform(
+        self,
+        data: Union[pd.DataFrame, Dict[str, Any]],
+        return_dict: Literal[True] = True,
+        save_format: Optional[str] = None,
+        output_path: Optional[str] = None,
+        chunk_size: int = 200000,
+    ) -> Dict[str, np.ndarray]: ...
+
+    @overload
+    def transform(
+        self,
+        data: Union[pd.DataFrame, Dict[str, Any]],
+        return_dict: Literal[False] = False,
+        save_format: Optional[str] = None,
+        output_path: Optional[str] = None,
+        chunk_size: int = 200000,
+    ) -> pd.DataFrame: ...
+
+    @overload
+    def transform(
+        self,
+        data: str | os.PathLike,
+        return_dict: Literal[False] = False,
+        save_format: Optional[str] = None,
+        output_path: Optional[str] = None,
+        chunk_size: int = 200000,
+    ) -> list[str]: ...
+
     def transform(
         self,
         data: Union[pd.DataFrame, Dict[str, Any], str, os.PathLike],

{nextrec-0.4.30 → nextrec-0.4.32}/nextrec/models/multi_task/[pre]aitm.py

@@ -1,6 +1,6 @@
 """
 Date: create on 01/01/2026 - prerelease version: need to overwrite compute_loss later
-Checkpoint: edit on 01/01/2026
+Checkpoint: edit on 01/14/2026
 Author: Yang Zhou, zyaztec@gmail.com
 Reference:
 - [1] Xi D, Chen Z, Yan P, Zhang Y, Zhu Y, Zhuang F, Chen Y. Modeling the Sequential Dependence among Audience Multi-step Conversions with Multi-task Learning in Targeted Display Advertising. Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining (KDD ’21), 2021, pp. 3745–3755.
@@ -20,7 +20,7 @@ from nextrec.basic.layers import MLP, EmbeddingLayer
 from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
 from nextrec.utils.model import get_mlp_output_dim
-from nextrec.utils.types import TaskTypeName
+from nextrec.utils.types import TaskTypeInput
 
 
 class AITMTransfer(nn.Module):
@@ -76,7 +76,7 @@ class AITM(BaseModel):
         tower_mlp_params_list: list[dict] | None = None,
         calibrator_alpha: float = 0.1,
         target: list[str] | str | None = None,
-        task: list[TaskTypeName] | None = None,
+        task: list[TaskTypeInput] | None = None,
         **kwargs,
     ):
         dense_features = dense_features or []

{nextrec-0.4.30 → nextrec-0.4.32}/nextrec/models/multi_task/[pre]snr_trans.py

@@ -1,6 +1,6 @@
 """
 Date: create on 01/01/2026 - prerelease version: still need to align with the source paper
-Checkpoint: edit on 01/01/2026
+Checkpoint: edit on 01/14/2026
 Author: Yang Zhou, zyaztec@gmail.com
 Reference:
 - [1] Ma J, Zhao Z, Chen J, Li A, Hong L, Chi EH. SNR: Sub-Network Routing for Flexible Parameter Sharing in Multi-Task Learning in E-Commerce by Exploiting Task Relationships in the Label Space. Proceedings of the 33rd AAAI Conference on Artificial Intelligence (AAAI 2019), 2019, pp. 216-223.
@@ -22,7 +22,7 @@ from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
 from nextrec.basic.layers import EmbeddingLayer, MLP
 from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
-from nextrec.utils.types import TaskTypeName
+from nextrec.utils.types import TaskTypeInput, TaskTypeName
 
 
 class SNRTransGate(nn.Module):
@@ -101,7 +101,7 @@ class SNRTrans(BaseModel):
         num_experts: int = 4,
         tower_mlp_params_list: list[dict] | None = None,
         target: list[str] | str | None = None,
-        task: TaskTypeName | list[TaskTypeName] | None = None,
+        task: TaskTypeInput | list[TaskTypeInput] | None = None,
         **kwargs,
     ) -> None:
         dense_features = dense_features or []

{nextrec-0.4.30 → nextrec-0.4.32}/nextrec/models/multi_task/[pre]star.py

@@ -1,6 +1,6 @@
 """
 Date: create on 01/01/2026 - prerelease version: still need to align with the source paper
-Checkpoint: edit on 01/01/2026
+Checkpoint: edit on 01/14/2026
 Author: Yang Zhou, zyaztec@gmail.com
 Reference:
 - [1] Sheng XR, Zhao L, Zhou G, Ding X, Dai B, Luo Q, Yang S, Lv J, Zhang C, Deng H, Zhu X. One Model to Serve All: Star Topology Adaptive Recommender for Multi-Domain CTR Prediction. arXiv preprint arXiv:2101.11427, 2021.
@@ -22,7 +22,7 @@ from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
 from nextrec.basic.heads import TaskHead
 from nextrec.basic.layers import DomainBatchNorm, EmbeddingLayer
 from nextrec.basic.model import BaseModel
-from nextrec.utils.types import TaskTypeName
+from nextrec.utils.types import TaskTypeInput, TaskTypeName
 
 
 class SharedSpecificLinear(nn.Module):
@@ -73,7 +73,7 @@ class STAR(BaseModel):
         sparse_features: list[SparseFeature] | None = None,
         sequence_features: list[SequenceFeature] | None = None,
         target: list[str] | str | None = None,
-        task: TaskTypeName | list[TaskTypeName] | None = None,
+        task: TaskTypeInput | list[TaskTypeInput] | None = None,
         mlp_params: dict | None = None,
         use_shared: bool = True,
         **kwargs,

{nextrec-0.4.30 → nextrec-0.4.32}/nextrec/models/multi_task/apg.py

@@ -1,6 +1,6 @@
 """
 Date: create on 01/01/2026
-Checkpoint: edit on 01/01/2026
+Checkpoint: edit on 01/14/2026
 Author: Yang Zhou, zyaztec@gmail.com
 Reference:
 - [1] Yan B, Wang P, Zhang K, Li F, Deng H, Xu J, Zheng B. APG: Adaptive Parameter Generation Network for Click-Through Rate Prediction. Advances in Neural Information Processing Systems 35 (NeurIPS 2022), 2022.
@@ -20,7 +20,7 @@ from nextrec.basic.layers import EmbeddingLayer, MLP
 from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
 from nextrec.utils.model import select_features
-from nextrec.utils.types import ActivationName, TaskTypeName
+from nextrec.utils.types import ActivationName, TaskTypeInput, TaskTypeName
 
 
 class APGLayer(nn.Module):
@@ -233,7 +233,7 @@ class APG(BaseModel):
         sparse_features: list[SparseFeature] | None = None,
         sequence_features: list[SequenceFeature] | None = None,
         target: list[str] | str | None = None,
-        task: TaskTypeName | list[TaskTypeName] | None = None,
+        task: TaskTypeInput | list[TaskTypeInput] | None = None,
         mlp_params: dict | None = None,
         inner_activation: ActivationName | None = None,
         generate_activation: ActivationName | None = None,

{nextrec-0.4.30 → nextrec-0.4.32}/nextrec/models/multi_task/cross_stitch.py

@@ -1,6 +1,6 @@
 """
 Date: create on 01/01/2026
-Checkpoint: edit on 01/01/2026
+Checkpoint: edit on 01/14/2026
 Author: Yang Zhou, zyaztec@gmail.com
 Reference:
 - [1] Misra I, Shrivastava A, Gupta A, Hebert M. Cross-Stitch Networks for Multi-Task Learning. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR 2016), 2016, pp. 3994–4003.
@@ -21,7 +21,7 @@ from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
 from nextrec.basic.layers import EmbeddingLayer, MLP
 from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
-from nextrec.utils.types import TaskTypeName
+from nextrec.utils.types import TaskTypeInput, TaskTypeName
 
 
 class CrossStitchLayer(nn.Module):
@@ -76,7 +76,7 @@ class CrossStitch(BaseModel):
         sparse_features: list[SparseFeature] | None = None,
         sequence_features: list[SequenceFeature] | None = None,
         target: list[str] | str | None = None,
-        task: TaskTypeName | list[TaskTypeName] | None = None,
+        task: TaskTypeInput | list[TaskTypeInput] | None = None,
         shared_mlp_params: dict | None = None,
         task_mlp_params: dict | None = None,
         tower_mlp_params: dict | None = None,

{nextrec-0.4.30 → nextrec-0.4.32}/nextrec/models/multi_task/escm.py

@@ -1,6 +1,6 @@
 """
 Date: create on 01/01/2026
-Checkpoint: edit on 01/01/2026
+Checkpoint: edit on 01/14/2026
 Author: Yang Zhou, zyaztec@gmail.com
 Reference:
 - [1] Wang H, Chang T-W, Liu T, Huang J, Chen Z, Yu C, Li R, Chu W. ESCM²: Entire Space Counterfactual Multi-Task Model for Post-Click Conversion Rate Estimation. Proceedings of the 45th International ACM SIGIR Conference on Research and Development in Information Retrieval (SIGIR ’22), 2022:363–372.
@@ -23,7 +23,7 @@ from nextrec.basic.layers import EmbeddingLayer, MLP
 from nextrec.basic.model import BaseModel
 from nextrec.loss.grad_norm import get_grad_norm_shared_params
 from nextrec.utils.model import compute_ranking_loss
-from nextrec.utils.types import TaskTypeName
+from nextrec.utils.types import TaskTypeInput, TaskTypeName
 
 
 class ESCM(BaseModel):
@@ -52,7 +52,7 @@ class ESCM(BaseModel):
         imp_mlp_params: dict | None = None,
         use_dr: bool = False,
         target: list[str] | str | None = None,
-        task: TaskTypeName | list[TaskTypeName] | None = None,
+        task: TaskTypeInput | list[TaskTypeInput] | None = None,
         **kwargs,
     ) -> None:
         dense_features = dense_features or []

{nextrec-0.4.30 → nextrec-0.4.32}/nextrec/models/multi_task/esmm.py

@@ -1,6 +1,6 @@
 """
 Date: create on 09/11/2025
-Checkpoint: edit on 23/12/2025
+Checkpoint: edit on 01/14/2026
 Author: Yang Zhou,zyaztec@gmail.com
 Reference:
 - [1] Ma X, Zhao L, Huang G, Wang Z, Hu Z, Zhu X, Gai K. Entire Space Multi-Task Model: An Effective Approach for Estimating Post-Click Conversion Rate. In: Proceedings of the 41st International ACM SIGIR Conference on Research and Development in Information Retrieval (SIGIR ’18), 2018, pp. 1137–1140.
@@ -46,7 +46,7 @@ from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
 from nextrec.basic.layers import MLP, EmbeddingLayer
 from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
-from nextrec.utils.types import TaskTypeName
+from nextrec.utils.types import TaskTypeInput
 
 
 class ESMM(BaseModel):
@@ -76,7 +76,7 @@ class ESMM(BaseModel):
         sequence_features: list[SequenceFeature],
         ctr_mlp_params: dict,
         cvr_mlp_params: dict,
-        task: list[TaskTypeName] | None = None,
+        task: list[TaskTypeInput] | None = None,
         target: list[str] | None = None,  # Note: ctcvr = ctr * cvr
         **kwargs,
     ):

{nextrec-0.4.30 → nextrec-0.4.32}/nextrec/models/multi_task/hmoe.py

@@ -1,6 +1,6 @@
 """
 Date: create on 01/01/2026
-Checkpoint: edit on 01/01/2026
+Checkpoint: edit on 01/14/2026
 Author: Yang Zhou, zyaztec@gmail.com
 [1] Zhao Z, Liu Y, Jin R, Zhu X, He X. HMOE: Improving Multi-Scenario Learning to Rank in E-commerce by Exploiting Task Relationships in the Label Space. Proceedings of the 29th ACM International Conference on Information & Knowledge Management (CIKM ’20), 2020, pp. 2069–2078.
 URL: https://dl.acm.org/doi/10.1145/3340531.3412713
@@ -23,7 +23,7 @@ from nextrec.basic.layers import MLP, EmbeddingLayer
 from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
 from nextrec.utils.model import get_mlp_output_dim
-from nextrec.utils.types import TaskTypeName
+from nextrec.utils.types import TaskTypeInput, TaskTypeName
 
 
 class HMOE(BaseModel):
@@ -53,7 +53,7 @@ class HMOE(BaseModel):
         tower_mlp_params_list: list[dict] | None = None,
         task_weight_mlp_params: list[dict] | None = None,
         target: list[str] | str | None = None,
-        task: TaskTypeName | list[TaskTypeName] | None = None,
+        task: TaskTypeInput | list[TaskTypeInput] | None = None,
         **kwargs,
     ) -> None:
         dense_features = dense_features or []

{nextrec-0.4.30 → nextrec-0.4.32}/nextrec/models/multi_task/mmoe.py

@@ -1,6 +1,6 @@
 """
 Date: create on 09/11/2025
-Checkpoint: edit on 23/12/2025
+Checkpoint: edit on 01/14/2026
 Author: Yang Zhou,zyaztec@gmail.com
 Reference:
 - [1] Ma J, Zhao Z, Yi X, Chen J, Hong L, Chi E H. Modeling Task Relationships in Multi-task Learning with Multi-gate Mixture-of-Experts. In: Proceedings of the 24th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining (KDD ’18), 2018, pp. 1930–1939.
@@ -48,7 +48,7 @@ from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
 from nextrec.basic.layers import MLP, EmbeddingLayer
 from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
-from nextrec.utils.types import TaskTypeName
+from nextrec.utils.types import TaskTypeInput
 
 
 class MMOE(BaseModel):
@@ -81,7 +81,7 @@ class MMOE(BaseModel):
         num_experts: int = 3,
         tower_mlp_params_list: list[dict] | None = None,
         target: list[str] | str | None = None,
-        task: TaskTypeName | list[TaskTypeName] | None = None,
+        task: TaskTypeInput | list[TaskTypeInput] | None = None,
         **kwargs,
     ):
 

{nextrec-0.4.30 → nextrec-0.4.32}/nextrec/models/multi_task/pepnet.py

@@ -1,6 +1,6 @@
 """
 Date: create on 01/01/2026
-Checkpoint: edit on 01/01/2026
+Checkpoint: edit on 01/14/2026
 Author: Yang Zhou, zyaztec@gmail.com
 Reference:
 - [1] Chang J, Zhang C, Hui Y, Leng D, Niu Y, Song Y, Gai K. PEPNet: Parameter and Embedding Personalized Network for Infusing with Personalized Prior Information. In: Proceedings of the 29th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining (KDD ’23), 2023.
@@ -58,7 +58,7 @@ from nextrec.basic.layers import EmbeddingLayer, GateMLP
 from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
 from nextrec.utils.model import select_features
-from nextrec.utils.types import TaskTypeName
+from nextrec.utils.types import TaskTypeInput, TaskTypeName
 
 
 class PPNet(nn.Module):
@@ -184,7 +184,7 @@ class PEPNet(BaseModel):
         sparse_features: list[SparseFeature] | None = None,
         sequence_features: list[SequenceFeature] | None = None,
         target: list[str] | str | None = None,
-        task: TaskTypeName | list[TaskTypeName] | None = None,
+        task: TaskTypeInput | list[TaskTypeInput] | None = None,
         mlp_params: dict | None = None,
         feature_gate_mlp_params: dict | None = None,
         gate_mlp_params: dict | None = None,
@@ -334,7 +334,7 @@ class PEPNet(BaseModel):
 
         task_logits = []
         for block in self.ppnet_blocks:
-             task_logits.append(block(o_ep=dnn_input, o_prior=task_sf_emb))
+            task_logits.append(block(o_ep=dnn_input, o_prior=task_sf_emb))
 
         y = torch.cat(task_logits, dim=1)
         return self.prediction_layer(y)

{nextrec-0.4.30 → nextrec-0.4.32}/nextrec/models/multi_task/ple.py

@@ -1,6 +1,6 @@
 """
 Date: create on 09/11/2025
-Checkpoint: edit on 23/12/2025
+Checkpoint: edit on 01/14/2026
 Author: Yang Zhou,zyaztec@gmail.com
 Reference:
 - [1] Tang H, Liu J, Zhao M, Gong X. Progressive Layered Extraction (PLE): A Novel Multi-Task Learning (MTL) Model for Personalized Recommendations. In: Proceedings of the 14th ACM Conference on Recommender Systems (RecSys ’20), 2020, pp. 269–278.
@@ -52,7 +52,7 @@ from nextrec.basic.layers import MLP, EmbeddingLayer
 from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
 from nextrec.utils.model import get_mlp_output_dim
-
+from nextrec.utils.types import TaskTypeInput
 
 class CGCLayer(nn.Module):
     """
@@ -202,7 +202,7 @@ class PLE(BaseModel):
         num_levels: int = 2,
         tower_mlp_params_list: list[dict] | None = None,
         target: list[str] | None = None,
-        task: str | list[str] | None = None,
+        task: TaskTypeInput | list[TaskTypeInput] | None = None,
         **kwargs,
     ):