nextrec 0.3.5__py3-none-any.whl → 0.4.1__py3-none-any.whl

nextrec/basic/session.py

@@ -22,6 +22,7 @@ class Session:
 
     experiment_id: str
     root: Path
+    log_basename: str  # The base name for log files, without path separators
 
     @property
     def logs_dir(self) -> Path:
@@ -60,7 +61,6 @@ class Session:
         return path
 
 def create_session(experiment_id: str | Path | None = None) -> Session:
-    """Create a :class:`Session` instance with prepared directories."""
 
     if experiment_id is not None and str(experiment_id).strip():
         exp_id = str(experiment_id).strip()
@@ -68,6 +68,8 @@ def create_session(experiment_id: str | Path | None = None) -> Session:
         # Use local time for session naming
         exp_id = "nextrec_session_" + datetime.now().strftime("%Y%m%d")
 
+    log_basename = Path(exp_id).name if exp_id else exp_id
+
     if (
         os.getenv("PYTEST_CURRENT_TEST")
         or os.getenv("PYTEST_RUNNING")
@@ -82,7 +84,7 @@ def create_session(experiment_id: str | Path | None = None) -> Session:
     session_path.mkdir(parents=True, exist_ok=True)
     root = session_path.resolve()
 
-    return Session(experiment_id=exp_id, root=root)
+    return Session(experiment_id=exp_id, root=root, log_basename=log_basename)
 
 def resolve_save_path(
     path: str | os.PathLike | Path | None,

nextrec/data/__init__.py

@@ -1,22 +1,4 @@
-"""
-Data utilities package for NextRec
-
-This package provides data processing and manipulation utilities organized by category:
-- batch_utils: Batch collation and processing
-- data_processing: Data manipulation and user ID extraction
-- data_utils: Legacy module (re-exports from specialized modules)
-- dataloader: Dataset and DataLoader implementations
-- preprocessor: Data preprocessing pipeline
-
-Date: create on 13/11/2025
-Last update: 03/12/2025 (refactored)
-Author: Yang Zhou, zyaztec@gmail.com
-"""
-
-# Batch utilities
 from nextrec.data.batch_utils import collate_fn, batch_to_dict, stack_section
-
-# Data processing utilities
 from nextrec.data.data_processing import (
     get_column_data,
     split_dict_random,
@@ -24,7 +6,6 @@ from nextrec.data.data_processing import (
     get_user_ids,
 )
 
-# File utilities (from utils package)
 from nextrec.utils.file import (
     resolve_file_paths,
     iter_file_chunks,
@@ -33,7 +14,6 @@ from nextrec.utils.file import (
     default_output_dir,
 )
 
-# DataLoader components
 from nextrec.data.dataloader import (
     TensorDictDataset,
     FileDataset,
@@ -41,13 +21,8 @@ from nextrec.data.dataloader import (
     build_tensors_from_data,
 )
 
-# Preprocessor
 from nextrec.data.preprocessor import DataProcessor
-
-# Feature definitions
 from nextrec.basic.features import FeatureSet
-
-# Legacy module (for backward compatibility)
 from nextrec.data import data_utils
 
 __all__ = [

nextrec/data/data_processing.py

@@ -11,7 +11,10 @@ import pandas as pd
 from typing import Any, Mapping
 
 
-def get_column_data(data: dict | pd.DataFrame, name: str):
+def get_column_data(
+        data: dict | pd.DataFrame, 
+        name: str):
+    
     if isinstance(data, dict):
         return data[name] if name in data else None
     elif isinstance(data, pd.DataFrame):
@@ -24,10 +27,11 @@ def get_column_data(data: dict | pd.DataFrame, name: str):
         raise KeyError(f"Unsupported data type for extracting column {name}")
 
 def split_dict_random(
-    data_dict: dict,
-    test_size: float = 0.2,
-    random_state: int | None = None
-):
+        data_dict: dict,
+        test_size: float = 0.2,
+        random_state: int | None = None
+        ):
+    
     lengths = [len(v) for v in data_dict.values()]
     if len(set(lengths)) != 1:
         raise ValueError(f"Length mismatch: {lengths}")
@@ -51,18 +55,27 @@ def split_dict_random(
     test_dict = {k: take(v, test_idx) for k, v in data_dict.items()}
     return train_dict, test_dict
 
+def split_data(
+        df: pd.DataFrame, 
+        test_size: float = 0.2
+        ) -> tuple[pd.DataFrame, pd.DataFrame]:
+    
+    split_idx = int(len(df) * (1 - test_size))
+    train_df = df.iloc[:split_idx].reset_index(drop=True)
+    valid_df = df.iloc[split_idx:].reset_index(drop=True)
+    return train_df, valid_df
 
 def build_eval_candidates(
-    df_all: pd.DataFrame,
-    user_col: str,
-    item_col: str,
-    label_col: str,
-    user_features: pd.DataFrame,
-    item_features: pd.DataFrame,
-    num_pos_per_user: int = 5,
-    num_neg_per_pos: int = 50,
-    random_seed: int = 2025,
-) -> pd.DataFrame:
+        df_all: pd.DataFrame,
+        user_col: str,
+        item_col: str,
+        label_col: str,
+        user_features: pd.DataFrame,
+        item_features: pd.DataFrame,
+        num_pos_per_user: int = 5,
+        num_neg_per_pos: int = 50,
+        random_seed: int = 2025,
+        ) -> pd.DataFrame:
     """
     Build evaluation candidates with positive and negative samples for each user.
     
@@ -111,11 +124,10 @@ def build_eval_candidates(
     eval_df = eval_df.merge(item_features, on=item_col, how='left')
     return eval_df
 
-
 def get_user_ids(
-    data: Any,
-    id_columns: list[str] | str | None = None
-) -> np.ndarray | None:
+        data: Any,
+        id_columns: list[str] | str | None = None
+        ) -> np.ndarray | None:
     """
     Extract user IDs from various data structures.
     

nextrec/data/dataloader.py

@@ -15,15 +15,15 @@ import pyarrow.parquet as pq
 from pathlib import Path
 from typing import cast
 
-from torch.utils.data import DataLoader, Dataset, IterableDataset
-from nextrec.data.preprocessor import DataProcessor
+from nextrec.basic.loggers import colorize
 from nextrec.basic.features import DenseFeature, SparseFeature, SequenceFeature, FeatureSet
+from nextrec.data.preprocessor import DataProcessor
+from torch.utils.data import DataLoader, Dataset, IterableDataset
 
-from nextrec.basic.loggers import colorize
-from nextrec.data.data_processing import get_column_data
-from nextrec.data.batch_utils import collate_fn
-from nextrec.utils.file import resolve_file_paths, read_table
 from nextrec.utils.tensor import to_tensor
+from nextrec.utils.file import resolve_file_paths, read_table
+from nextrec.data.batch_utils import collate_fn
+from nextrec.data.data_processing import get_column_data
 
 class TensorDictDataset(Dataset):
     """Dataset returning sample-level dicts matching the unified batch schema."""
@@ -118,6 +118,18 @@ class RecDataLoader(FeatureSet):
                  target: list[str] | None | str = None,
                  id_columns: str | list[str] | None = None,
                  processor: DataProcessor | None = None):
+        """
+        RecDataLoader is a unified dataloader for supporting in-memory and streaming data.
+        Basemodel will accept RecDataLoader to create dataloaders for training/evaluation/prediction.
+
+        Args:  
+            dense_features: list of DenseFeature definitions
+            sparse_features: list of SparseFeature definitions
+            sequence_features: list of SequenceFeature definitions
+            target: target column name(s), e.g. 'label' or ['ctr', 'ctcvr']
+            id_columns: id column name(s) to carry through (not used for model inputs), e.g. 'user_id' or ['user_id', 'item_id']
+            processor: an instance of DataProcessor, if provided, will be used to transform data before creating tensors.
+        """
         self.processor = processor
         self.set_all_features(dense_features, sparse_features, sequence_features, target, id_columns)
 
@@ -126,20 +138,40 @@ class RecDataLoader(FeatureSet):
                          batch_size: int = 32,
                          shuffle: bool = True,
                          load_full: bool = True,
-                         chunk_size: int = 10000) -> DataLoader:
+                         chunk_size: int = 10000,
+                         num_workers: int = 0,
+                         sampler = None) -> DataLoader:
+        """
+        Create a DataLoader from various data sources.
+
+        Args:
+            data: Data source, can be a dict, pd.DataFrame, file path (str), or existing DataLoader.
+            batch_size: Batch size for DataLoader.
+            shuffle: Whether to shuffle the data (ignored in streaming mode).
+            load_full: If True, load full data into memory; if False, use streaming mode for large files.
+            chunk_size: Chunk size for streaming mode (number of rows per chunk).
+            num_workers: Number of worker processes for data loading.
+            sampler: Optional sampler for DataLoader, only used for distributed training.
+        Returns:
+            DataLoader instance.
+        """
+
         if isinstance(data, DataLoader):
             return data
         elif isinstance(data, (str, os.PathLike)):
-            return self.create_from_path(path=data, batch_size=batch_size, shuffle=shuffle, load_full=load_full, chunk_size=chunk_size)
+            return self.create_from_path(path=data, batch_size=batch_size, shuffle=shuffle, load_full=load_full, chunk_size=chunk_size, num_workers=num_workers)
         elif isinstance(data, (dict, pd.DataFrame)):
-            return self.create_from_memory(data=data, batch_size=batch_size, shuffle=shuffle)
+            return self.create_from_memory(data=data, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers, sampler=sampler)
         else:
             raise ValueError(f"[RecDataLoader Error] Unsupported data type: {type(data)}")
     
     def create_from_memory(self, 
                            data: dict | pd.DataFrame,
                            batch_size: int,
-                           shuffle: bool) -> DataLoader:
+                           shuffle: bool,
+                           num_workers: int = 0,
+                           sampler=None) -> DataLoader:
+
         raw_data = data
 
         if self.processor is not None:
@@ -150,14 +182,15 @@ class RecDataLoader(FeatureSet):
         if tensors is None:
             raise ValueError("[RecDataLoader Error] No valid tensors could be built from the provided data.")
         dataset = TensorDictDataset(tensors)
-        return DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, collate_fn=collate_fn)
+        return DataLoader(dataset, batch_size=batch_size, shuffle=False if sampler is not None else shuffle, sampler=sampler, collate_fn=collate_fn, num_workers=num_workers)
     
     def create_from_path(self,
                          path: str,
                          batch_size: int,
                          shuffle: bool,
                          load_full: bool,
-                         chunk_size: int = 10000) -> DataLoader:
+                         chunk_size: int = 10000,
+                         num_workers: int = 0) -> DataLoader:
         file_paths, file_type = resolve_file_paths(str(Path(path)))
         # Load full data into memory
         if load_full:
@@ -169,6 +202,7 @@ class RecDataLoader(FeatureSet):
                 except OSError:
                     pass
                 try:
+                    df = read_table(file_path, file_type=file_type)
                     dfs.append(df)
                 except MemoryError as exc:
                     raise MemoryError(f"[RecDataLoader Error] Out of memory while reading {file_path}. Consider using load_full=False with streaming.") from exc
@@ -176,22 +210,23 @@ class RecDataLoader(FeatureSet):
                 combined_df = pd.concat(dfs, ignore_index=True)
             except MemoryError as exc:
                 raise MemoryError(f"[RecDataLoader Error] Out of memory while concatenating loaded data (approx {total_bytes / (1024**3):.2f} GB). Use load_full=False to stream or reduce chunk_size.") from exc
-            return self.create_from_memory(combined_df, batch_size, shuffle,)
+            return self.create_from_memory(combined_df, batch_size, shuffle, num_workers=num_workers)
         else:
-            return self.load_files_streaming(file_paths, file_type, batch_size, chunk_size, shuffle)
+            return self.load_files_streaming(file_paths, file_type, batch_size, chunk_size, shuffle, num_workers=num_workers)
 
     def load_files_streaming(self,
                              file_paths: list[str],
                              file_type: str,
                              batch_size: int,
                              chunk_size: int,
-                             shuffle: bool) -> DataLoader:
+                             shuffle: bool,
+                             num_workers: int = 0) -> DataLoader:
         if shuffle:
             logging.info("[RecDataLoader Info] Shuffle is ignored in streaming mode (IterableDataset).")
         if batch_size != 1:
             logging.info("[RecDataLoader Info] Streaming mode enforces batch_size=1; tune chunk_size to control memory/throughput.")
         dataset = FileDataset(file_paths=file_paths, dense_features=self.dense_features, sparse_features=self.sparse_features, sequence_features=self.sequence_features, target_columns=self.target_columns, id_columns=self.id_columns, chunk_size=chunk_size, file_type=file_type, processor=self.processor)
-        return DataLoader(dataset, batch_size=1, collate_fn=collate_fn)
+        return DataLoader(dataset, batch_size=1, collate_fn=collate_fn, num_workers=num_workers)
 
 def normalize_sequence_column(column, feature: SequenceFeature) -> np.ndarray:
     if isinstance(column, pd.Series):

nextrec/models/generative/__init__.py

@@ -1,5 +0,0 @@
-from .hstu import HSTU
-
-__all__ = [
-    "HSTU",
-]

nextrec/models/generative/hstu.py

@@ -255,7 +255,7 @@ class HSTU(BaseModel):
         return "HSTU"
 
     @property
-    def task_type(self) -> str:
+    def default_task(self) -> str:
         return "multiclass"
 
     def __init__(
@@ -275,6 +275,7 @@ class HSTU(BaseModel):
 
         tie_embeddings: bool = True,
         target: Optional[list[str] | str] = None,
+        task: str | list[str] | None = None,
         optimizer: str = "adam",
         optimizer_params: Optional[dict] = None,
         scheduler: Optional[str] = None,
@@ -307,7 +308,7 @@ class HSTU(BaseModel):
             sparse_features=sparse_features,
             sequence_features=sequence_features,
             target=target,
-            task=self.task_type,
+            task=task or self.default_task,
             device=device,
             embedding_l1_reg=embedding_l1_reg,
             dense_l1_reg=dense_l1_reg,

nextrec/models/match/__init__.py

@@ -1,13 +0,0 @@
-from .dssm import DSSM
-from .dssm_v2 import DSSM_v2
-from .youtube_dnn import YoutubeDNN
-from .mind import MIND
-from .sdm import SDM
-
-__all__ = [
-    'DSSM',
-    'DSSM_v2',
-    'YoutubeDNN',
-    'MIND',
-    'SDM',
-]

nextrec/models/match/dssm.py

@@ -73,7 +73,6 @@ class DSSM(BaseMatchModel):
             dense_l1_reg=dense_l1_reg,
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
-            early_stop_patience=early_stop_patience,
             **kwargs
         )
         

nextrec/models/match/dssm_v2.py

@@ -68,7 +68,6 @@ class DSSM_v2(BaseMatchModel):
             dense_l1_reg=dense_l1_reg,
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
-            early_stop_patience=early_stop_patience,
             **kwargs
         )
         

nextrec/models/match/mind.py

@@ -184,7 +184,6 @@ class MIND(BaseMatchModel):
             dense_l1_reg=dense_l1_reg,
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
-            early_stop_patience=early_stop_patience,
             **kwargs
         )
         

nextrec/models/match/sdm.py

@@ -76,7 +76,6 @@ class SDM(BaseMatchModel):
             dense_l1_reg=dense_l1_reg,
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
-            early_stop_patience=early_stop_patience,
             **kwargs
         )
         

nextrec/models/match/youtube_dnn.py

@@ -73,7 +73,6 @@ class YoutubeDNN(BaseMatchModel):
             dense_l1_reg=dense_l1_reg,
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
-            early_stop_patience=early_stop_patience,
             **kwargs
         )
         

nextrec/models/multi_task/esmm.py

@@ -64,10 +64,9 @@ class ESMM(BaseModel):
     @property
     def model_name(self):
         return "ESMM"
-
+    
     @property
-    def task_type(self):
-        # ESMM has fixed task types: CTR (binary) and CVR (binary)
+    def default_task(self):
         return ['binary', 'binary']
     
     def __init__(self,
@@ -77,7 +76,7 @@ class ESMM(BaseModel):
                  ctr_params: dict,
                  cvr_params: dict,
                  target: list[str] = ['ctr', 'ctcvr'],  # Note: ctcvr = ctr * cvr
-                 task: list[str] = ['binary', 'binary'],
+                 task: list[str] | None = None,
                  optimizer: str = "adam",
                  optimizer_params: dict = {},
                  loss: str | nn.Module | list[str | nn.Module] | None = "bce",
@@ -98,13 +97,12 @@ class ESMM(BaseModel):
             sparse_features=sparse_features,
             sequence_features=sequence_features,
             target=target,
-            task=task,  # Both CTR and CTCVR are binary classification
+            task=task or self.default_task,  # Both CTR and CTCVR are binary classification
             device=device,
             embedding_l1_reg=embedding_l1_reg,
             dense_l1_reg=dense_l1_reg,
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
-            early_stop_patience=20,
             **kwargs
         )
 
@@ -126,7 +124,7 @@ class ESMM(BaseModel):
         
         # CVR tower
         self.cvr_tower = MLP(input_dim=input_dim, output_layer=True, **cvr_params)
-        self.prediction_layer = PredictionLayer(task_type=self.task_type, task_dims=[1, 1])
+        self.prediction_layer = PredictionLayer(task_type=self.default_task, task_dims=[1, 1])
         # Register regularization weights
         self.register_regularization_weights(embedding_attr='embedding', include_modules=['ctr_tower', 'cvr_tower'])
         self.compile(optimizer=optimizer, optimizer_params=optimizer_params, loss=loss, loss_params=loss_params)

nextrec/models/multi_task/mmoe.py

@@ -65,8 +65,11 @@ class MMOE(BaseModel):
         return "MMOE"
 
     @property
-    def task_type(self):
-        return self.task if isinstance(self.task, list) else [self.task]
+    def default_task(self):
+        num_tasks = getattr(self, "num_tasks", None)
+        if num_tasks is not None and num_tasks > 0:
+            return ['binary'] * num_tasks
+        return ['binary']
     
     def __init__(self,
                  dense_features: list[DenseFeature]=[],
@@ -76,7 +79,7 @@ class MMOE(BaseModel):
                  num_experts: int=3,
                  tower_params_list: list[dict]=[],
                  target: list[str]=[],
-                 task: str | list[str] = 'binary',
+                 task: str | list[str] | None = None,
                  optimizer: str = "adam",
                  optimizer_params: dict = {},
                  loss: str | nn.Module | list[str | nn.Module] | None = "bce",
@@ -88,18 +91,19 @@ class MMOE(BaseModel):
                  dense_l2_reg=1e-4,
                  **kwargs):
         
+        self.num_tasks = len(target)
+
         super(MMOE, self).__init__(
             dense_features=dense_features,
             sparse_features=sparse_features,
             sequence_features=sequence_features,
             target=target,
-            task=task,
+            task=task or self.default_task,
             device=device,
             embedding_l1_reg=embedding_l1_reg,
             dense_l1_reg=dense_l1_reg,
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
-            early_stop_patience=20,
             **kwargs
         )
 
@@ -144,7 +148,7 @@ class MMOE(BaseModel):
         for tower_params in tower_params_list:
             tower = MLP(input_dim=expert_output_dim, output_layer=True, **tower_params)
             self.towers.append(tower)
-        self.prediction_layer = PredictionLayer(task_type=self.task_type, task_dims=[1] * self.num_tasks)
+        self.prediction_layer = PredictionLayer(task_type=self.default_task, task_dims=[1] * self.num_tasks)
         # Register regularization weights
         self.register_regularization_weights(embedding_attr='embedding', include_modules=['experts', 'gates', 'towers'])
         self.compile(optimizer=optimizer, optimizer_params=optimizer_params, loss=loss, loss_params=loss_params,)

nextrec/models/multi_task/ple.py

@@ -159,8 +159,11 @@ class PLE(BaseModel):
         return "PLE"
 
     @property
-    def task_type(self):
-        return self.task if isinstance(self.task, list) else [self.task]
+    def default_task(self):
+        num_tasks = getattr(self, "num_tasks", None)
+        if num_tasks is not None and num_tasks > 0:
+            return ['binary'] * num_tasks
+        return ['binary']
     
     def __init__(self,
                  dense_features: list[DenseFeature],
@@ -173,7 +176,7 @@ class PLE(BaseModel):
                  num_levels: int,
                  tower_params_list: list[dict],
                  target: list[str],
-                 task: str | list[str] = 'binary',
+                 task: str | list[str] | None = None,
                  optimizer: str = "adam",
                  optimizer_params: dict | None = None,
                  loss: str | nn.Module | list[str | nn.Module] | None = "bce",
@@ -185,18 +188,19 @@ class PLE(BaseModel):
                  dense_l2_reg=1e-4,
                  **kwargs):
         
+        self.num_tasks = len(target)
+
         super(PLE, self).__init__(
             dense_features=dense_features,
             sparse_features=sparse_features,
             sequence_features=sequence_features,
             target=target,
-            task=task,
+            task=task or self.default_task,
             device=device,
             embedding_l1_reg=embedding_l1_reg,
             dense_l1_reg=dense_l1_reg,
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
-            early_stop_patience=20,
             **kwargs
         )
 
@@ -247,7 +251,7 @@ class PLE(BaseModel):
         for tower_params in tower_params_list:
             tower = MLP(input_dim=expert_output_dim, output_layer=True, **tower_params)
             self.towers.append(tower)
-        self.prediction_layer = PredictionLayer(task_type=self.task_type, task_dims=[1] * self.num_tasks)
+        self.prediction_layer = PredictionLayer(task_type=self.default_task, task_dims=[1] * self.num_tasks)
         # Register regularization weights
         self.register_regularization_weights(embedding_attr='embedding', include_modules=['cgc_layers', 'towers'])
         self.compile(optimizer=optimizer, optimizer_params=optimizer_params, loss=self.loss, loss_params=loss_params)

nextrec/models/multi_task/poso.py

@@ -261,8 +261,11 @@ class POSO(BaseModel):
         return "POSO"
 
     @property
-    def task_type(self) -> list[str]:
-        return self.task if isinstance(self.task, list) else [self.task]
+    def default_task(self) -> list[str]:
+        num_tasks = getattr(self, "num_tasks", None)
+        if num_tasks is not None and num_tasks > 0:
+            return ["binary"] * num_tasks
+        return ["binary"]
 
     def __init__(
         self,
@@ -274,7 +277,7 @@ class POSO(BaseModel):
         pc_sequence_features: list[SequenceFeature] | None,
         tower_params_list: list[dict],
         target: list[str],
-        task: str | list[str] = "binary",
+        task: str | list[str] | None = None,
         architecture: str = "mlp",
         # POSO gating defaults
         gate_hidden_dim: int = 32,
@@ -307,6 +310,7 @@ class POSO(BaseModel):
         self.pc_dense_features = list(pc_dense_features or [])
         self.pc_sparse_features = list(pc_sparse_features or [])
         self.pc_sequence_features = list(pc_sequence_features or [])
+        self.num_tasks = len(target)
 
         if not self.pc_dense_features and not self.pc_sparse_features and not self.pc_sequence_features:
             raise ValueError("POSO requires at least one PC feature for personalization.")
@@ -320,13 +324,12 @@ class POSO(BaseModel):
             sparse_features=sparse_features,
             sequence_features=sequence_features,
             target=target,
-            task=task,
+            task=task or self.default_task,
             device=device,
             embedding_l1_reg=embedding_l1_reg,
             dense_l1_reg=dense_l1_reg,
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
-            early_stop_patience=20,
             **kwargs,
         )
 
@@ -387,7 +390,7 @@ class POSO(BaseModel):
             )
             self.towers = nn.ModuleList([MLP(input_dim=self.mmoe.expert_output_dim, output_layer=True, **tower_params,) for tower_params in tower_params_list])
             self.tower_heads = None
-        self.prediction_layer = PredictionLayer(task_type=self.task_type, task_dims=[1] * self.num_tasks,)
+        self.prediction_layer = PredictionLayer(task_type=self.default_task, task_dims=[1] * self.num_tasks,)
         include_modules = ["towers", "tower_heads"] if self.architecture == "mlp" else ["mmoe", "towers"]
         self.register_regularization_weights(embedding_attr="embedding", include_modules=include_modules)
         self.compile(optimizer=optimizer, optimizer_params=optimizer_params, loss=loss, loss_params=loss_params)

nextrec/models/multi_task/share_bottom.py

@@ -53,9 +53,11 @@ class ShareBottom(BaseModel):
         return "ShareBottom"
 
     @property
-    def task_type(self):
-        # Multi-task model, return list of task types
-        return self.task if isinstance(self.task, list) else [self.task]
+    def default_task(self):
+        num_tasks = getattr(self, "num_tasks", None)
+        if num_tasks is not None and num_tasks > 0:
+            return ['binary'] * num_tasks
+        return ['binary']
     
     def __init__(self,
                  dense_features: list[DenseFeature],
@@ -64,7 +66,7 @@ class ShareBottom(BaseModel):
                  bottom_params: dict,
                  tower_params_list: list[dict],
                  target: list[str],
-                 task: str | list[str] = 'binary',
+                 task: str | list[str] | None = None,
                  optimizer: str = "adam",
                  optimizer_params: dict = {},
                  loss: str | nn.Module | list[str | nn.Module] | None = "bce",
@@ -76,18 +78,19 @@ class ShareBottom(BaseModel):
                  dense_l2_reg=1e-4,
                  **kwargs):
         
+        self.num_tasks = len(target)
+
         super(ShareBottom, self).__init__(
             dense_features=dense_features,
             sparse_features=sparse_features,
             sequence_features=sequence_features,
             target=target,
-            task=task,
+            task=task or self.default_task,
             device=device,
             embedding_l1_reg=embedding_l1_reg,
             dense_l1_reg=dense_l1_reg,
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
-            early_stop_patience=20,
             **kwargs
         )
 
@@ -120,7 +123,7 @@ class ShareBottom(BaseModel):
         for tower_params in tower_params_list:
             tower = MLP(input_dim=bottom_output_dim, output_layer=True, **tower_params)
             self.towers.append(tower)
-        self.prediction_layer = PredictionLayer(task_type=self.task_type, task_dims=[1] * self.num_tasks)
+        self.prediction_layer = PredictionLayer(task_type=self.default_task, task_dims=[1] * self.num_tasks)
         # Register regularization weights
         self.register_regularization_weights(embedding_attr='embedding', include_modules=['bottom', 'towers'])
         self.compile(optimizer=optimizer, optimizer_params=optimizer_params, loss=loss, loss_params=loss_params)