nextrec 0.4.31__py3-none-any.whl → 0.4.32__py3-none-any.whl

nextrec/__version__.py

@@ -1 +1 @@
-__version__ = "0.4.31"
+__version__ = "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/cli.py

@@ -682,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/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/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/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/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/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/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/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/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/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/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/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/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/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,
     ):
 

nextrec/models/multi_task/poso.py

@@ -1,6 +1,6 @@
 """
 Date: create on 28/11/2025
-Checkpoint: edit on 23/12/2025
+Checkpoint: edit on 01/14/2026
 Author: Yang Zhou,zyaztec@gmail.com
 Reference:
 - [1] Dai S, Lin H, Zhao Z, Lin J, Wu H, Wang Z, Yang S, Liu J. POSO: Personalized Cold Start Modules for Large-scale Recommender Systems. arXiv preprint arXiv:2108.04690, 2021.
@@ -50,7 +50,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 select_features
-from nextrec.utils.types import TaskTypeName
+from nextrec.utils.types import TaskTypeInput
 
 
 class POSOGate(nn.Module):
@@ -310,7 +310,7 @@ class POSO(BaseModel):
         pc_sequence_features: list[str] | None,
         tower_mlp_params_list: list[dict],
         target: list[str] | None = None,
-        task: TaskTypeName | list[TaskTypeName] | None = None,
+        task: TaskTypeInput | list[TaskTypeInput] | None = None,
         architecture: Literal["mlp", "mmoe"] = "mlp",
         # POSO gating defaults
         gate_hidden_dim: int = 32,

nextrec/models/multi_task/share_bottom.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
 
 Shared-Bottom is the classic hard-parameter-sharing baseline for multi-task learning.
@@ -43,7 +43,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 TaskTypeInput
 
 class ShareBottom(BaseModel):
     @property
@@ -65,7 +65,7 @@ class ShareBottom(BaseModel):
         bottom_mlp_params: dict,
         tower_mlp_params_list: list[dict],
         target: list[str],
-        task: str | list[str] | None = None,
+        task: TaskTypeInput | list[TaskTypeInput] | None = None,
         **kwargs,
     ):
 

nextrec/models/ranking/afm.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] Xiao J, Ye H, He X, et al. Attentional Factorization Machines: Learning the Weight of Feature Interactions via Attention Networks
@@ -43,6 +43,7 @@ from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
 from nextrec.basic.layers import EmbeddingLayer, InputMask
 from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
+from nextrec.utils.types import TaskTypeInput
 
 
 class AFM(BaseModel):
@@ -60,7 +61,7 @@ class AFM(BaseModel):
         sparse_features: list[SparseFeature] | None = None,
         sequence_features: list[SequenceFeature] | None = None,
         target: str | list[str] | None = None,
-        task: str | list[str] | None = None,
+        task: TaskTypeInput | list[TaskTypeInput] | None = None,
         attention_dim: int = 32,
         attention_dropout: float = 0.0,
         **kwargs,

nextrec/models/ranking/autoint.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] Song W, Shi C, Xiao Z, et al. AutoInt: Automatic feature interaction learning via self-attentive neural networks. In: Proceedings of the 28th ACM International Conference on Information and Knowledge Management (CIKM ’19), 2019, pp. 1161–1170.
@@ -59,6 +59,7 @@ from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
 from nextrec.basic.layers import EmbeddingLayer, MultiHeadSelfAttention
 from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
+from nextrec.utils.types import TaskTypeInput
 
 
 class AutoInt(BaseModel):
@@ -76,7 +77,7 @@ class AutoInt(BaseModel):
         sparse_features: list[SparseFeature],
         sequence_features: list[SequenceFeature],
         target: str | list[str] | None = None,
-        task: str | list[str] | None = None,
+        task: TaskTypeInput | list[TaskTypeInput] | None = None,
         att_layer_num: int = 3,
         att_embedding_dim: int = 8,
         att_head_num: int = 2,

nextrec/models/ranking/dcn.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] Wang R, Fu B, Fu G, et al. Deep & cross network for ad click predictions[C] //Proceedings of the ADKDD'17. 2017: 1-7.
@@ -56,6 +56,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 TaskTypeInput
 
 
 class CrossNetwork(nn.Module):
@@ -94,7 +95,7 @@ class DCN(BaseModel):
         sparse_features: list[SparseFeature] | None = None,
         sequence_features: list[SequenceFeature] | None = None,
         target: str | list[str] | None = None,
-        task: str | list[str] | None = None,
+        task: TaskTypeInput | list[TaskTypeInput] | None = None,
         cross_num: int = 3,
         mlp_params: dict | None = None,
         **kwargs,