nextrec 0.4.24__py3-none-any.whl → 0.4.27__py3-none-any.whl

nextrec/basic/summary.py

@@ -48,6 +48,27 @@ class SummarySet:
     checkpoint_path: str
     train_data_summary: dict[str, Any] | None
     valid_data_summary: dict[str, Any] | None
+    note: str | None
+
+    def collect_dataloader_summary(self, data_loader: DataLoader | None):
+        if data_loader is None:
+            return None
+
+        summary = {
+            "batch_size": data_loader.batch_size,
+            "num_workers": data_loader.num_workers,
+            "pin_memory": data_loader.pin_memory,
+            "persistent_workers": data_loader.persistent_workers,
+        }
+        prefetch_factor = getattr(data_loader, "prefetch_factor", None)
+        if prefetch_factor is not None:
+            summary["prefetch_factor"] = prefetch_factor
+
+        sampler = getattr(data_loader, "sampler", None)
+        if sampler is not None:
+            summary["sampler"] = sampler.__class__.__name__
+
+        return summary or None
 
     def build_data_summary(
         self, data: Any, data_loader: DataLoader | None, sample_key: str
@@ -66,6 +87,10 @@ class SummarySet:
         if train_size is not None:
             summary[sample_key] = int(train_size)
 
+        dataloader_summary = self.collect_dataloader_summary(data_loader)
+        if dataloader_summary:
+            summary["dataloader"] = dataloader_summary
+
         if labels:
             task_types = list(self.task) if isinstance(self.task, list) else [self.task]
             if len(task_types) != len(self.target_columns):
@@ -321,6 +346,7 @@ class SummarySet:
             logger.info(f"  Session ID:            {self.session_id}")
             logger.info(f"  Features Config Path:  {self.features_config_path}")
             logger.info(f"  Latest Checkpoint:     {self.checkpoint_path}")
+            logger.info(f"  Note:                  {self.note}")
 
         if "Data Summary" in selected_sections and (
             self.train_data_summary or self.valid_data_summary
@@ -341,6 +367,22 @@ class SummarySet:
                         for label, value in lines:
                             logger.info(f"  {format_kv(label, value)}")
 
+                dataloader_info = self.train_data_summary.get("dataloader")
+                if isinstance(dataloader_info, dict):
+                    logger.info("Train DataLoader:")
+                    for key in (
+                        "batch_size",
+                        "num_workers",
+                        "pin_memory",
+                        "persistent_workers",
+                        "sampler",
+                    ):
+                        if key in dataloader_info:
+                            label = key.replace("_", " ").title()
+                            logger.info(
+                                format_kv(label, dataloader_info[key], indent=2)
+                            )
+
             if self.valid_data_summary:
                 if self.train_data_summary:
                     logger.info("")
@@ -355,3 +397,19 @@ class SummarySet:
                         logger.info(f"{target_name}:")
                         for label, value in lines:
                             logger.info(f"  {format_kv(label, value)}")
+
+                dataloader_info = self.valid_data_summary.get("dataloader")
+                if isinstance(dataloader_info, dict):
+                    logger.info("Valid DataLoader:")
+                    for key in (
+                        "batch_size",
+                        "num_workers",
+                        "pin_memory",
+                        "persistent_workers",
+                        "sampler",
+                    ):
+                        if key in dataloader_info:
+                            label = key.replace("_", " ").title()
+                            logger.info(
+                                format_kv(label, dataloader_info[key], indent=2)
+                            )

nextrec/cli.py

@@ -320,6 +320,7 @@ def train_model(train_config_path: str) -> None:
             streaming=True,
             chunk_size=dataloader_chunk_size,
             num_workers=dataloader_cfg.get("num_workers", 0),
+            prefetch_factor=dataloader_cfg.get("prefetch_factor"),
         )
         valid_loader = None
         if val_data_path:
@@ -331,6 +332,7 @@ def train_model(train_config_path: str) -> None:
                 streaming=True,
                 chunk_size=dataloader_chunk_size,
                 num_workers=dataloader_cfg.get("num_workers", 0),
+                prefetch_factor=dataloader_cfg.get("prefetch_factor"),
             )
         elif streaming_valid_files:
             valid_loader = dataloader.create_dataloader(
@@ -340,6 +342,7 @@ def train_model(train_config_path: str) -> None:
                 streaming=True,
                 chunk_size=dataloader_chunk_size,
                 num_workers=dataloader_cfg.get("num_workers", 0),
+                prefetch_factor=dataloader_cfg.get("prefetch_factor"),
             )
     else:
         train_loader = dataloader.create_dataloader(
@@ -347,12 +350,14 @@ def train_model(train_config_path: str) -> None:
             batch_size=dataloader_cfg.get("train_batch_size", 512),
             shuffle=dataloader_cfg.get("train_shuffle", True),
             num_workers=dataloader_cfg.get("num_workers", 0),
+            prefetch_factor=dataloader_cfg.get("prefetch_factor"),
         )
         valid_loader = dataloader.create_dataloader(
             data=valid_data,
             batch_size=dataloader_cfg.get("valid_batch_size", 512),
             shuffle=dataloader_cfg.get("valid_shuffle", False),
             num_workers=dataloader_cfg.get("num_workers", 0),
+            prefetch_factor=dataloader_cfg.get("prefetch_factor"),
         )
 
     model_cfg.setdefault("session_id", session_id)
@@ -383,6 +388,7 @@ def train_model(train_config_path: str) -> None:
         loss=train_cfg.get("loss", "focal"),
         loss_params=train_cfg.get("loss_params", {}),
         loss_weights=train_cfg.get("loss_weights"),
+        ignore_label=train_cfg.get("ignore_label", -1),
     )
 
     model.fit(
@@ -397,6 +403,12 @@ def train_model(train_config_path: str) -> None:
         num_workers=dataloader_cfg.get("num_workers", 0),
         user_id_column=id_column,
         use_tensorboard=False,
+        use_wandb=train_cfg.get("use_wandb", False),
+        use_swanlab=train_cfg.get("use_swanlab", False),
+        wandb_api=train_cfg.get("wandb_api"),
+        swanlab_api=train_cfg.get("swanlab_api"),
+        log_interval=train_cfg.get("log_interval", 1),
+        note=train_cfg.get("note"),
     )
 
 
@@ -583,6 +595,7 @@ def predict_model(predict_config_path: str) -> None:
         shuffle=False,
         streaming=predict_cfg.get("streaming", True),
         chunk_size=predict_cfg.get("chunk_size", 20000),
+        prefetch_factor=predict_cfg.get("prefetch_factor"),
     )
 
     save_format = predict_cfg.get(

nextrec/data/data_processing.py

@@ -13,6 +13,8 @@ import numpy as np
 import pandas as pd
 import torch
 
+from nextrec.utils.torch_utils import to_numpy
+
 
 def get_column_data(data: dict | pd.DataFrame, name: str):
 
@@ -23,15 +25,7 @@ def get_column_data(data: dict | pd.DataFrame, name: str):
             return None
         return data[name].values
     else:
-        if hasattr(data, name):
-            return getattr(data, name)
-        raise KeyError(f"Unsupported data type for extracting column {name}")
-
-
-def to_numpy(values: Any) -> np.ndarray:
-    if isinstance(values, torch.Tensor):
-        return values.detach().cpu().numpy()
-    return np.asarray(values)
+        raise KeyError(f"Only dict or DataFrame supported, got {type(data)}")
 
 
 def get_data_length(data: Any) -> int | None:

nextrec/data/dataloader.py

@@ -194,6 +194,7 @@ class RecDataLoader(FeatureSet):
         streaming: bool = False,
         chunk_size: int = 10000,
         num_workers: int = 0,
+        prefetch_factor: int | None = None,
         sampler=None,
     ) -> DataLoader:
         """
@@ -206,6 +207,7 @@ class RecDataLoader(FeatureSet):
             streaming: If True, use streaming mode for large files; if False, load full data into memory.
             chunk_size: Chunk size for streaming mode (number of rows per chunk).
             num_workers: Number of worker processes for data loading.
+            prefetch_factor: Number of batches loaded in advance by each worker.
             sampler: Optional sampler for DataLoader, only used for distributed training.
         Returns:
             DataLoader instance.
@@ -234,6 +236,7 @@ class RecDataLoader(FeatureSet):
                 streaming=streaming,
                 chunk_size=chunk_size,
                 num_workers=num_workers,
+                prefetch_factor=prefetch_factor,
             )
 
         if isinstance(data, (dict, pd.DataFrame)):
@@ -242,6 +245,7 @@ class RecDataLoader(FeatureSet):
                 batch_size=batch_size,
                 shuffle=shuffle,
                 num_workers=num_workers,
+                prefetch_factor=prefetch_factor,
                 sampler=sampler,
             )
 
@@ -253,6 +257,7 @@ class RecDataLoader(FeatureSet):
         batch_size: int,
         shuffle: bool,
         num_workers: int = 0,
+        prefetch_factor: int | None = None,
         sampler=None,
     ) -> DataLoader:
         raw_data = data
@@ -275,6 +280,9 @@ class RecDataLoader(FeatureSet):
                 "[RecDataLoader Error] No valid tensors could be built from the provided data."
             )
         dataset = TensorDictDataset(tensors)
+        loader_kwargs = {}
+        if num_workers > 0 and prefetch_factor is not None:
+            loader_kwargs["prefetch_factor"] = prefetch_factor
         return DataLoader(
             dataset,
             batch_size=batch_size,
@@ -284,6 +292,7 @@ class RecDataLoader(FeatureSet):
             num_workers=num_workers,
             pin_memory=torch.cuda.is_available(),
             persistent_workers=num_workers > 0,
+            **loader_kwargs,
         )
 
     def create_from_path(
@@ -294,6 +303,7 @@ class RecDataLoader(FeatureSet):
         streaming: bool,
         chunk_size: int = 10000,
         num_workers: int = 0,
+        prefetch_factor: int | None = None,
     ) -> DataLoader:
         if isinstance(path, (str, os.PathLike)):
             file_paths, file_type = resolve_file_paths(str(Path(path)))
@@ -327,6 +337,7 @@ class RecDataLoader(FeatureSet):
                 chunk_size,
                 shuffle,
                 num_workers=num_workers,
+                prefetch_factor=prefetch_factor,
             )
 
         dfs = []
@@ -350,7 +361,11 @@ class RecDataLoader(FeatureSet):
                 f"[RecDataLoader Error] Out of memory while concatenating loaded data (approx {total_bytes / (1024**3):.2f} GB). Use streaming=True or reduce chunk_size."
             ) from exc
         return self.create_from_memory(
-            combined_df, batch_size, shuffle, num_workers=num_workers
+            combined_df,
+            batch_size,
+            shuffle,
+            num_workers=num_workers,
+            prefetch_factor=prefetch_factor,
         )
 
     def load_files_streaming(
@@ -361,6 +376,7 @@ class RecDataLoader(FeatureSet):
         chunk_size: int,
         shuffle: bool,
         num_workers: int = 0,
+        prefetch_factor: int | None = None,
     ) -> DataLoader:
         if not check_streaming_support(file_type):
             raise ValueError(
@@ -393,8 +409,15 @@ class RecDataLoader(FeatureSet):
             file_type=file_type,
             processor=self.processor,
         )
+        loader_kwargs = {}
+        if num_workers > 0 and prefetch_factor is not None:
+            loader_kwargs["prefetch_factor"] = prefetch_factor
         return DataLoader(
-            dataset, batch_size=1, collate_fn=collate_fn, num_workers=num_workers
+            dataset,
+            batch_size=1,
+            collate_fn=collate_fn,
+            num_workers=num_workers,
+            **loader_kwargs,
         )