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

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)
 
@@ -127,13 +139,29 @@ class RecDataLoader(FeatureSet):
                          shuffle: bool = True,
                          load_full: bool = True,
                          chunk_size: int = 10000,
-                         num_workers: int = 0) -> DataLoader:
+                         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, num_workers=num_workers)
         elif isinstance(data, (dict, pd.DataFrame)):
-            return self.create_from_memory(data=data, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers)
+            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)}")
     
@@ -141,7 +169,9 @@ class RecDataLoader(FeatureSet):
                            data: dict | pd.DataFrame,
                            batch_size: int,
                            shuffle: bool,
-                           num_workers: int = 0) -> DataLoader:
+                           num_workers: int = 0,
+                           sampler=None) -> DataLoader:
+
         raw_data = data
 
         if self.processor is not None:
@@ -152,7 +182,7 @@ 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, num_workers=num_workers)
+        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,

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/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)