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

nextrec/__version__.py

@@ -1 +1 @@
-__version__ = "0.4.31"
+__version__ = "0.4.33"

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"
@@ -932,6 +933,13 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
 
         existing_callbacks = self.callbacks.callbacks
 
+        has_validation = valid_data is not None or valid_split is not None
+        checkpoint_monitor = monitor_metric
+        checkpoint_mode = self.best_metrics_mode
+        if not has_validation:
+            checkpoint_monitor = "loss"
+            checkpoint_mode = "min"
+
         if self.early_stop_patience > 0 and not any(
             isinstance(cb, EarlyStopper) for cb in existing_callbacks
         ):
@@ -945,6 +953,8 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
                 )
             )
 
+        has_validation = valid_data is not None or valid_split is not None
+
         if self.is_main_process and not any(
             isinstance(cb, CheckpointSaver) for cb in existing_callbacks
         ):
@@ -952,9 +962,9 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
                 CheckpointSaver(
                     best_path=self.best_path,
                     checkpoint_path=self.checkpoint_path,
-                    monitor=monitor_metric,
-                    mode=self.best_metrics_mode,
-                    save_best_only=True,
+                    monitor=checkpoint_monitor,
+                    mode=checkpoint_mode,
+                    save_best_only=has_validation,
                     verbose=1,
                 )
             )
@@ -1245,11 +1255,6 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
                         epoch_logs[f"val_{k}"] = v
             else:
                 epoch_logs = {**train_log_payload}
-                if self.is_main_process:
-                    self.save_model(
-                        self.checkpoint_path, add_timestamp=False, verbose=False
-                    )
-                    self.best_checkpoint_path = self.checkpoint_path
 
             # Call on_epoch_end for all callbacks (handles early stopping, checkpointing, lr scheduling)
             self.callbacks.on_epoch_end(epoch, epoch_logs)
@@ -1623,6 +1628,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 +2273,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/basic/summary.py

@@ -73,7 +73,8 @@ class SummarySet:
     def build_data_summary(
         self, data: Any, data_loader: DataLoader | None, sample_key: str
     ):
-        dataset = data_loader.dataset if data_loader else None
+        
+        dataset = data_loader.dataset if data_loader is not None else None
 
         train_size = get_data_length(dataset)
         if train_size is None:

nextrec/cli.py

@@ -152,16 +152,19 @@ def train_model(train_config_path: str) -> None:
     )
     if data_cfg.get("valid_ratio") is not None:
         logger.info(format_kv("Valid ratio", data_cfg.get("valid_ratio")))
-    if data_cfg.get("val_path") or data_cfg.get("valid_path"):
+    if data_cfg.get("valid_path"):
         logger.info(
             format_kv(
                 "Validation path",
                 resolve_path(
-                    data_cfg.get("val_path") or data_cfg.get("valid_path"), config_dir
+                    data_cfg.get("valid_path"), config_dir
                 ),
             )
         )
 
+    # Determine validation dataset path early for streaming split / fitting
+    val_data_path = data_cfg.get("valid_path")
+
     if streaming:
         file_paths, file_type = resolve_file_paths(str(data_path))
         log_kv_lines(
@@ -180,6 +183,34 @@ def train_model(train_config_path: str) -> None:
             raise ValueError(f"Data file is empty: {first_file}") from exc
         df_columns = list(first_chunk.columns)
 
+        # Decide training/validation file lists before fitting processor, to avoid
+        # leaking validation statistics into preprocessing (scalers/encoders).
+        streaming_train_files = file_paths
+        streaming_valid_ratio = data_cfg.get("valid_ratio")
+        if val_data_path:
+            streaming_valid_files = None
+        elif streaming_valid_ratio is not None:
+            ratio = float(streaming_valid_ratio)
+            if not (0 < ratio < 1):
+                raise ValueError(
+                    f"[NextRec CLI Error] Valid_ratio must be between 0 and 1, current value is {streaming_valid_ratio}"
+                )
+            total_files = len(file_paths)
+            if total_files < 2:
+                raise ValueError(
+                    "[NextRec CLI Error] Must provide valid_path or increase the number of data files. At least 2 files are required for streaming validation split."
+                )
+            val_count = max(1, int(round(total_files * ratio)))
+            if val_count >= total_files:
+                val_count = total_files - 1
+            streaming_valid_files = file_paths[-val_count:]
+            streaming_train_files = file_paths[:-val_count]
+            logger.info(
+                f"Split files for streaming training and validation using valid_ratio={ratio:.3f}: training {len(streaming_train_files)} files, validation {len(streaming_valid_files)} files"
+            )
+        else:
+            streaming_valid_files = None
+
     else:
         df = read_table(data_path, data_cfg.get("format"))
         logger.info(format_kv("Rows", len(df)))
@@ -215,7 +246,13 @@ def train_model(train_config_path: str) -> None:
     )
 
     if streaming:
-        processor.fit(str(data_path), chunk_size=dataloader_chunk_size)
+        if file_type is None:
+            raise ValueError("[NextRec CLI Error] Streaming mode requires a valid file_type")
+        processor.fit_from_files(
+            file_paths=streaming_train_files or file_paths,
+            file_type=file_type,
+            chunk_size=dataloader_chunk_size,
+        )
         processed = None
         df = None  # type: ignore[assignment]
     else:
@@ -232,34 +269,6 @@ def train_model(train_config_path: str) -> None:
         sequence_names,
     )
 
-    # Check if validation dataset path is specified
-    val_data_path = data_cfg.get("val_path") or data_cfg.get("valid_path")
-    if streaming:
-        if not file_paths:
-            file_paths, file_type = resolve_file_paths(str(data_path))
-        streaming_train_files = file_paths
-        streaming_valid_ratio = data_cfg.get("valid_ratio")
-        if val_data_path:
-            streaming_valid_files = None
-        elif streaming_valid_ratio is not None:
-            ratio = float(streaming_valid_ratio)
-            if not (0 < ratio < 1):
-                raise ValueError(
-                    f"[NextRec CLI Error] Valid_ratio must be between 0 and 1, current value is {streaming_valid_ratio}"
-                )
-            total_files = len(file_paths)
-            if total_files < 2:
-                raise ValueError(
-                    "[NextRec CLI Error] Must provide val_path or increase the number of data files. At least 2 files are required for streaming validation split."
-                )
-            val_count = max(1, int(round(total_files * ratio)))
-            if val_count >= total_files:
-                val_count = total_files - 1
-            streaming_valid_files = file_paths[-val_count:]
-            streaming_train_files = file_paths[:-val_count]
-            logger.info(
-                f"Split files for streaming training and validation using valid_ratio={ratio:.3f}: training {len(streaming_train_files)} files, validation {len(streaming_valid_files)} files"
-            )
     train_data: Dict[str, Any]
     valid_data: Dict[str, Any] | None
 
@@ -682,16 +691,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
@@ -566,35 +566,16 @@ class DataProcessor(FeatureSet):
             return [str(v) for v in value]
         return [str(value)]
 
-    def fit_from_path(self, path: str, chunk_size: int) -> "DataProcessor":
-        """
-        Fit processor statistics by streaming files to reduce memory usage.
-
-        Args:
-            path (str): File or directory path.
-            chunk_size (int): Number of rows per chunk.
-
-        Returns:
-            DataProcessor: Fitted DataProcessor instance.
-        """
+    def fit_from_file_paths(
+        self, file_paths: list[str], file_type: str, chunk_size: int
+    ) -> "DataProcessor":
         logger = logging.getLogger()
-        logger.info(
-            colorize(
-                "Fitting DataProcessor (streaming path mode)...",
-                color="cyan",
-                bold=True,
-            )
-        )
-        for config in self.sparse_features.values():
-            config.pop("_min_freq_logged", None)
-        for config in self.sequence_features.values():
-            config.pop("_min_freq_logged", None)
-        file_paths, file_type = resolve_file_paths(path)
+        if not file_paths:
+            raise ValueError("[DataProcessor Error] Empty file list for streaming fit")
         if not check_streaming_support(file_type):
             raise ValueError(
                 f"[DataProcessor Error] Format '{file_type}' does not support streaming. "
-                "fit_from_path only supports streaming formats (csv, parquet) to avoid high memory usage. "
-                "Use fit(dataframe) with in-memory data or convert the data format."
+                "Streaming fit only supports csv, parquet to avoid high memory usage."
             )
 
         numeric_acc = {}
@@ -636,6 +617,7 @@ class DataProcessor(FeatureSet):
         target_values: Dict[str, set[Any]] = {
             name: set() for name in self.target_features.keys()
         }
+
         missing_features = set()
         for file_path in file_paths:
             for chunk in iter_file_chunks(file_path, file_type, chunk_size):
@@ -702,10 +684,12 @@ class DataProcessor(FeatureSet):
                 for name in self.target_features.keys() & columns:
                     vals = chunk[name].dropna().tolist()
                     target_values[name].update(vals)
+
         if missing_features:
             logger.warning(
                 f"The following configured features were not found in provided files: {sorted(missing_features)}"
             )
+
         # finalize numeric scalers
         for name, config in self.numeric_features.items():
             acc = numeric_acc[name]
@@ -895,6 +879,91 @@ class DataProcessor(FeatureSet):
         )
         return self
 
+    def fit_from_files(
+        self, file_paths: list[str], file_type: str, chunk_size: int
+    ) -> "DataProcessor":
+        """Fit processor statistics by streaming an explicit list of files.
+
+        This is useful when you want to fit statistics on training files only (exclude
+        validation files) in streaming mode.
+        """
+        logger = logging.getLogger()
+        logger.info(
+            colorize(
+                "Fitting DataProcessor (streaming files mode)...",
+                color="cyan",
+                bold=True,
+            )
+        )
+        for config in self.sparse_features.values():
+            config.pop("_min_freq_logged", None)
+        for config in self.sequence_features.values():
+            config.pop("_min_freq_logged", None)
+        uses_robust = any(
+            cfg.get("scaler") == "robust" for cfg in self.numeric_features.values()
+        )
+        if uses_robust:
+            logger.warning(
+                "Robust scaler requires full data; loading provided files into memory. "
+                "Consider smaller chunk_size or different scaler if memory is limited."
+            )
+            frames = [read_table(p, file_type) for p in file_paths]
+            df = pd.concat(frames, ignore_index=True) if len(frames) > 1 else frames[0]
+            return self.fit(df)
+        return self.fit_from_file_paths(file_paths=file_paths, file_type=file_type, chunk_size=chunk_size)
+
+    def fit_from_path(self, path: str, chunk_size: int) -> "DataProcessor":
+        """
+        Fit processor statistics by streaming files to reduce memory usage.
+
+        Args:
+            path (str): File or directory path.
+            chunk_size (int): Number of rows per chunk.
+
+        Returns:
+            DataProcessor: Fitted DataProcessor instance.
+        """
+        logger = logging.getLogger()
+        logger.info(
+            colorize(
+                "Fitting DataProcessor (streaming path mode)...",
+                color="cyan",
+                bold=True,
+            )
+        )
+        for config in self.sparse_features.values():
+            config.pop("_min_freq_logged", None)
+        for config in self.sequence_features.values():
+            config.pop("_min_freq_logged", None)
+        file_paths, file_type = resolve_file_paths(path)
+        return self.fit_from_file_paths(
+            file_paths=file_paths,
+            file_type=file_type,
+            chunk_size=chunk_size,
+        )
+
+    @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 +1307,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 []