nextrec 0.4.8__py3-none-any.whl → 0.4.10__py3-none-any.whl

nextrec/cli.py

@@ -18,10 +18,10 @@ Checkpoint: edit on 18/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 """
 
-import sys
 import argparse
 import logging
 import pickle
+import sys
 import time
 from pathlib import Path
 from typing import Any, Dict, List
@@ -29,6 +29,7 @@ from typing import Any, Dict, List
 import pandas as pd
 
 from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
+from nextrec.basic.loggers import colorize, format_kv, setup_logger
 from nextrec.data.data_utils import split_dict_random
 from nextrec.data.dataloader import RecDataLoader
 from nextrec.data.preprocessor import DataProcessor
@@ -39,22 +40,29 @@ from nextrec.utils.config import (
     resolve_path,
     select_features,
 )
-from nextrec.utils.feature import normalize_to_list
-from nextrec.utils.file import (
+from nextrec.utils.console import get_nextrec_version
+from nextrec.utils.data import (
     iter_file_chunks,
     read_table,
     read_yaml,
     resolve_file_paths,
 )
-from nextrec.utils.cli_utils import (
-    get_nextrec_version,
-    log_startup_info,
-)
-from nextrec.basic.loggers import setup_logger
+from nextrec.utils.feature import normalize_to_list
 
 logger = logging.getLogger(__name__)
 
 
+def log_cli_section(title: str) -> None:
+    logger.info("")
+    logger.info(colorize(f"[{title}]", color="bright_blue", bold=True))
+    logger.info(colorize("-" * 80, color="bright_blue"))
+
+
+def log_kv_lines(items: list[tuple[str, Any]]) -> None:
+    for label, value in items:
+        logger.info(format_kv(label, value))
+
+
 def train_model(train_config_path: str) -> None:
     """
     Train a NextRec model using the provided configuration file.
@@ -77,8 +85,17 @@ def train_model(train_config_path: str) -> None:
     artifact_root = Path(session_cfg.get("artifact_root", "nextrec_logs"))
     session_dir = artifact_root / session_id
     setup_logger(session_id=session_id)
-    logger.info(
-        f"[NextRec CLI] Training start | version={get_nextrec_version()} | session_id={session_id} | artifacts={session_dir.resolve()}"
+
+    log_cli_section("CLI")
+    log_kv_lines(
+        [
+            ("Mode", "train"),
+            ("Version", get_nextrec_version()),
+            ("Session ID", session_id),
+            ("Artifacts", session_dir.resolve()),
+            ("Config", config_file.resolve()),
+            ("Command", " ".join(sys.argv)),
+        ]
     )
 
     processor_path = session_dir / "processor.pkl"
@@ -105,11 +122,53 @@ def train_model(train_config_path: str) -> None:
         cfg.get("model_config", "model_config.yaml"), config_dir
     )
 
+    log_cli_section("Config")
+    log_kv_lines(
+        [
+            ("Train config", config_file.resolve()),
+            ("Feature config", feature_cfg_path),
+            ("Model config", model_cfg_path),
+        ]
+    )
+
     feature_cfg = read_yaml(feature_cfg_path)
     model_cfg = read_yaml(model_cfg_path)
 
+    # Extract id_column from data config for GAUC metrics
+    id_column = data_cfg.get("id_column") or data_cfg.get("user_id_column")
+    id_columns = [id_column] if id_column else []
+
+    log_cli_section("Data")
+    log_kv_lines(
+        [
+            ("Data path", data_path),
+            ("Format", data_cfg.get("format", "auto")),
+            ("Streaming", streaming),
+            ("Target", target),
+            ("ID column", id_column or "(not set)"),
+        ]
+    )
+    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"):
+        logger.info(
+            format_kv(
+                "Validation path",
+                resolve_path(
+                    data_cfg.get("val_path") or data_cfg.get("valid_path"), config_dir
+                ),
+            )
+        )
+
     if streaming:
         file_paths, file_type = resolve_file_paths(str(data_path))
+        log_kv_lines(
+            [
+                ("File type", file_type),
+                ("Files", len(file_paths)),
+                ("Chunk size", dataloader_chunk_size),
+            ]
+        )
         first_file = file_paths[0]
         first_chunk_size = max(1, min(dataloader_chunk_size, 1000))
         chunk_iter = iter_file_chunks(first_file, file_type, first_chunk_size)
@@ -121,14 +180,12 @@ def train_model(train_config_path: str) -> None:
 
     else:
         df = read_table(data_path, data_cfg.get("format"))
+        logger.info(format_kv("Rows", len(df)))
+        logger.info(format_kv("Columns", len(df.columns)))
         df_columns = list(df.columns)
 
     dense_names, sparse_names, sequence_names = select_features(feature_cfg, df_columns)
 
-    # Extract id_column from data config for GAUC metrics
-    id_column = data_cfg.get("id_column") or data_cfg.get("user_id_column")
-    id_columns = [id_column] if id_column else []
-
     used_columns = dense_names + sparse_names + sequence_names + target + id_columns
 
     # keep order but drop duplicates
@@ -144,6 +201,17 @@ def train_model(train_config_path: str) -> None:
         processor, feature_cfg, dense_names, sparse_names, sequence_names
     )
 
+    log_cli_section("Features")
+    log_kv_lines(
+        [
+            ("Dense features", len(dense_names)),
+            ("Sparse features", len(sparse_names)),
+            ("Sequence features", len(sequence_names)),
+            ("Targets", len(target)),
+            ("Used columns", len(unique_used_columns)),
+        ]
+    )
+
     if streaming:
         processor.fit(str(data_path), chunk_size=dataloader_chunk_size)
         processed = None
@@ -247,7 +315,7 @@ def train_model(train_config_path: str) -> None:
             data=train_stream_source,
             batch_size=dataloader_cfg.get("train_batch_size", 512),
             shuffle=dataloader_cfg.get("train_shuffle", True),
-            load_full=False,
+            streaming=True,
             chunk_size=dataloader_chunk_size,
             num_workers=dataloader_cfg.get("num_workers", 0),
         )
@@ -258,7 +326,7 @@ def train_model(train_config_path: str) -> None:
                 data=str(val_data_resolved),
                 batch_size=dataloader_cfg.get("valid_batch_size", 512),
                 shuffle=dataloader_cfg.get("valid_shuffle", False),
-                load_full=False,
+                streaming=True,
                 chunk_size=dataloader_chunk_size,
                 num_workers=dataloader_cfg.get("num_workers", 0),
             )
@@ -267,7 +335,7 @@ def train_model(train_config_path: str) -> None:
                 data=streaming_valid_files,
                 batch_size=dataloader_cfg.get("valid_batch_size", 512),
                 shuffle=dataloader_cfg.get("valid_shuffle", False),
-                load_full=False,
+                streaming=True,
                 chunk_size=dataloader_chunk_size,
                 num_workers=dataloader_cfg.get("num_workers", 0),
             )
@@ -298,6 +366,15 @@ def train_model(train_config_path: str) -> None:
         device,
     )
 
+    log_cli_section("Model")
+    log_kv_lines(
+        [
+            ("Model", model.__class__.__name__),
+            ("Device", device),
+            ("Session ID", session_id),
+        ]
+    )
+
     model.compile(
         optimizer=train_cfg.get("optimizer", "adam"),
         optimizer_params=train_cfg.get("optimizer_params", {}),
@@ -328,13 +405,30 @@ def predict_model(predict_config_path: str) -> None:
     config_dir = config_file.resolve().parent
     cfg = read_yaml(config_file)
 
-    session_cfg = cfg.get("session", {}) or {}
-    session_id = session_cfg.get("id", "masknet_tutorial")
-    artifact_root = Path(session_cfg.get("artifact_root", "nextrec_logs"))
-    session_dir = Path(cfg.get("checkpoint_path") or (artifact_root / session_id))
+    # Checkpoint path is the primary configuration
+    if "checkpoint_path" not in cfg:
+        session_cfg = cfg.get("session", {}) or {}
+        session_id = session_cfg.get("id", "nextrec_session")
+        artifact_root = Path(session_cfg.get("artifact_root", "nextrec_logs"))
+        session_dir = artifact_root / session_id
+    else:
+        session_dir = Path(cfg["checkpoint_path"])
+        # Auto-infer session_id from checkpoint directory name
+        session_cfg = cfg.get("session", {}) or {}
+        session_id = session_cfg.get("id") or session_dir.name
+    
     setup_logger(session_id=session_id)
-    logger.info(
-        f"[NextRec CLI] Predict start | version={get_nextrec_version()} | session_id={session_id} | checkpoint={session_dir.resolve()}"
+
+    log_cli_section("CLI")
+    log_kv_lines(
+        [
+            ("Mode", "predict"),
+            ("Version", get_nextrec_version()),
+            ("Session ID", session_id),
+            ("Checkpoint", session_dir.resolve()),
+            ("Config", config_file.resolve()),
+            ("Command", " ".join(sys.argv)),
+        ]
     )
 
     processor_path = Path(session_dir / "processor.pkl")
@@ -342,24 +436,38 @@ def predict_model(predict_config_path: str) -> None:
         processor_path = session_dir / "processor" / "processor.pkl"
 
     predict_cfg = cfg.get("predict", {}) or {}
-    model_cfg_path = resolve_path(
-        cfg.get("model_config", "model_config.yaml"), config_dir
-    )
-    # feature_cfg_path = resolve_path(
-    #     cfg.get("feature_config", "feature_config.yaml"), config_dir
-    # )
+    
+    # Auto-find model_config in checkpoint directory if not specified
+    if "model_config" in cfg:
+        model_cfg_path = resolve_path(cfg["model_config"], config_dir)
+    else:
+        # Try to find model_config.yaml in checkpoint directory
+        auto_model_cfg = session_dir / "model_config.yaml"
+        if auto_model_cfg.exists():
+            model_cfg_path = auto_model_cfg
+        else:
+            # Fallback to config directory
+            model_cfg_path = resolve_path("model_config.yaml", config_dir)
 
     model_cfg = read_yaml(model_cfg_path)
-    # feature_cfg = read_yaml(feature_cfg_path)
     model_cfg.setdefault("session_id", session_id)
     model_cfg.setdefault("params", {})
 
+    log_cli_section("Config")
+    log_kv_lines(
+        [
+            ("Predict config", config_file.resolve()),
+            ("Model config", model_cfg_path),
+            ("Processor", processor_path),
+        ]
+    )
+
     processor = DataProcessor.load(processor_path)
 
     # Load checkpoint and ensure required parameters are passed
     checkpoint_base = Path(session_dir)
     if checkpoint_base.is_dir():
-        candidates = sorted(checkpoint_base.glob("*.model"))
+        candidates = sorted(checkpoint_base.glob("*.pt"))
         if not candidates:
             raise FileNotFoundError(
                 f"[NextRec CLI Error]: Unable to find model checkpoint: {checkpoint_base}"
@@ -368,7 +476,7 @@ def predict_model(predict_config_path: str) -> None:
         config_dir_for_features = checkpoint_base
     else:
         model_file = (
-            checkpoint_base.with_suffix(".model")
+            checkpoint_base.with_suffix(".pt")
             if checkpoint_base.suffix == ""
             else checkpoint_base
         )
@@ -418,40 +526,78 @@ def predict_model(predict_config_path: str) -> None:
         id_columns = [predict_cfg["id_column"]]
         model.id_columns = id_columns
 
+    effective_id_columns = id_columns or model.id_columns
+    log_cli_section("Features")
+    log_kv_lines(
+        [
+            ("Dense features", len(dense_features)),
+            ("Sparse features", len(sparse_features)),
+            ("Sequence features", len(sequence_features)),
+            ("Targets", len(target_cols)),
+            ("ID columns", len(effective_id_columns)),
+        ]
+    )
+
+    log_cli_section("Model")
+    log_kv_lines(
+        [
+            ("Model", model.__class__.__name__),
+            ("Checkpoint", model_file),
+            ("Device", predict_cfg.get("device", "cpu")),
+        ]
+    )
+
     rec_dataloader = RecDataLoader(
         dense_features=model.dense_features,
         sparse_features=model.sparse_features,
         sequence_features=model.sequence_features,
         target=None,
-        id_columns=id_columns or model.id_columns,
+        id_columns=effective_id_columns,
         processor=processor,
     )
 
     data_path = resolve_path(predict_cfg["data_path"], config_dir)
     batch_size = predict_cfg.get("batch_size", 512)
 
+    log_cli_section("Data")
+    log_kv_lines(
+        [
+            ("Data path", data_path),
+            ("Format", predict_cfg.get("source_data_format", predict_cfg.get("data_format", "auto"))),
+            ("Batch size", batch_size),
+            ("Chunk size", predict_cfg.get("chunk_size", 20000)),
+            ("Streaming", predict_cfg.get("streaming", True)),
+        ]
+    )
+    logger.info("")
     pred_loader = rec_dataloader.create_dataloader(
         data=str(data_path),
         batch_size=batch_size,
         shuffle=False,
-        load_full=predict_cfg.get("load_full", False),
+        streaming=predict_cfg.get("streaming", True),
         chunk_size=predict_cfg.get("chunk_size", 20000),
     )
 
-    output_path = resolve_path(predict_cfg["output_path"], config_dir)
-    output_path.parent.mkdir(parents=True, exist_ok=True)
+    # Build output path: {checkpoint_path}/predictions/{name}.{save_data_format}
+    save_format = predict_cfg.get("save_data_format", predict_cfg.get("save_format", "csv"))
+    pred_name = predict_cfg.get("name", "pred")
+    # Pass filename with extension to let model.predict handle path resolution
+    save_path = f"{pred_name}.{save_format}"
 
     start = time.time()
-    model.predict(
+    logger.info("")
+    result = model.predict(
         data=pred_loader,
         batch_size=batch_size,
         include_ids=bool(id_columns),
         return_dataframe=False,
-        save_path=output_path,
-        save_format=predict_cfg.get("save_format", "csv"),
+        save_path=save_path,
+        save_format=save_format,
         num_workers=predict_cfg.get("num_workers", 0),
     )
     duration = time.time() - start
+    # When return_dataframe=False, result is the actual file path
+    output_path = result if isinstance(result, Path) else checkpoint_base / "predictions" / save_path
     logger.info(f"Prediction completed, results saved to: {output_path}")
     logger.info(f"Total time: {duration:.2f} seconds")
 
@@ -495,8 +641,6 @@ Examples:
     parser.add_argument("--predict_config", help="Prediction configuration file path")
     args = parser.parse_args()
 
-    logger.info(get_nextrec_version())
-
     if not args.mode:
         parser.error("[NextRec CLI Error] --mode is required (train|predict)")
 
@@ -504,13 +648,11 @@ Examples:
         config_path = args.train_config
         if not config_path:
             parser.error("[NextRec CLI Error] train mode requires --train_config")
-        log_startup_info(logger, mode="train", config_path=config_path)
         train_model(config_path)
     else:
         config_path = args.predict_config
         if not config_path:
             parser.error("[NextRec CLI Error] predict mode requires --predict_config")
-        log_startup_info(logger, mode="predict", config_path=config_path)
         predict_model(config_path)
 
 

nextrec/data/__init__.py

@@ -1,29 +1,26 @@
-from nextrec.data.batch_utils import collate_fn, batch_to_dict, stack_section
+from nextrec.basic.features import FeatureSet
+from nextrec.data import data_utils
+from nextrec.data.batch_utils import batch_to_dict, collate_fn, stack_section
 from nextrec.data.data_processing import (
-    get_column_data,
-    split_dict_random,
     build_eval_candidates,
+    get_column_data,
     get_user_ids,
+    split_dict_random,
 )
-
-from nextrec.utils.file import (
-    resolve_file_paths,
-    iter_file_chunks,
-    read_table,
-    load_dataframes,
-    default_output_dir,
-)
-
 from nextrec.data.dataloader import (
-    TensorDictDataset,
     FileDataset,
     RecDataLoader,
+    TensorDictDataset,
     build_tensors_from_data,
 )
-
 from nextrec.data.preprocessor import DataProcessor
-from nextrec.basic.features import FeatureSet
-from nextrec.data import data_utils
+from nextrec.utils.data import (
+    default_output_dir,
+    iter_file_chunks,
+    load_dataframes,
+    read_table,
+    resolve_file_paths,
+)
 
 __all__ = [
     # Batch utilities

nextrec/data/batch_utils.py

@@ -5,10 +5,11 @@ Date: create on 03/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 """
 
-import torch
-import numpy as np
 from typing import Any, Mapping
 
+import numpy as np
+import torch
+
 
 def stack_section(batch: list[dict], section: str):
     entries = [item.get(section) for item in batch if item.get(section) is not None]

nextrec/data/data_processing.py

@@ -2,13 +2,16 @@
 Data processing utilities for NextRec
 
 Date: create on 03/12/2025
+Checkpoint: edit on 19/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 """
 
-import torch
+import hashlib
+from typing import Any
+
 import numpy as np
 import pandas as pd
-from typing import Any
+import torch
 
 
 def get_column_data(data: dict | pd.DataFrame, name: str):
@@ -166,3 +169,8 @@ def get_user_ids(
             return arr.reshape(arr.shape[0])
 
     return None
+
+
+def hash_md5_mod(value: str, hash_size: int) -> int:
+    digest = hashlib.md5(value.encode("utf-8")).digest()
+    return int.from_bytes(digest, byteorder="big", signed=False) % hash_size

nextrec/data/data_utils.py

@@ -1,30 +1,25 @@
 """
-Data processing utilities for NextRec (Refactored)
-
-This module now re-exports functions from specialized submodules:
-- batch_utils: collate_fn, batch_to_dict
-- data_processing: get_column_data, split_dict_random, build_eval_candidates, get_user_ids
-- nextrec.utils.file_utils: resolve_file_paths, iter_file_chunks, read_table, load_dataframes, default_output_dir
+Data processing utilities for NextRec
 
 Date: create on 27/10/2025
-Last update: 03/12/2025 (refactored)
+Last update: 19/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 """
 
 # Import from new organized modules
-from nextrec.data.batch_utils import collate_fn, batch_to_dict, stack_section
+from nextrec.data.batch_utils import batch_to_dict, collate_fn, stack_section
 from nextrec.data.data_processing import (
-    get_column_data,
-    split_dict_random,
     build_eval_candidates,
+    get_column_data,
     get_user_ids,
+    split_dict_random,
 )
-from nextrec.utils.file import (
-    resolve_file_paths,
+from nextrec.utils.data import (
+    default_output_dir,
     iter_file_chunks,
-    read_table,
     load_dataframes,
-    default_output_dir,
+    read_table,
+    resolve_file_paths,
 )
 
 __all__ = [

nextrec/data/dataloader.py

@@ -2,33 +2,32 @@
 Dataloader definitions
 
 Date: create on 27/10/2025
-Checkpoint: edit on 02/12/2025
+Checkpoint: edit on 19/12/2025
 Author: Yang Zhou,zyaztec@gmail.com
 """
 
-import os
-import torch
 import logging
+import os
+from pathlib import Path
+from typing import cast
+
 import numpy as np
 import pandas as pd
 import pyarrow.parquet as pq
-
-from pathlib import Path
-from typing import cast
+import torch
+from torch.utils.data import DataLoader, Dataset, IterableDataset
 
 from nextrec.basic.features import (
     DenseFeature,
-    SparseFeature,
-    SequenceFeature,
     FeatureSet,
+    SequenceFeature,
+    SparseFeature,
 )
-from nextrec.data.preprocessor import DataProcessor
-from torch.utils.data import DataLoader, Dataset, IterableDataset
-
-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
+from nextrec.data.preprocessor import DataProcessor
+from nextrec.utils.data import read_table, resolve_file_paths
+from nextrec.utils.torch_utils import to_tensor
 
 
 class TensorDictDataset(Dataset):
@@ -103,9 +102,8 @@ class FileDataset(FeatureSet, IterableDataset):
         self.current_file_index = 0
         for file_path in self.file_paths:
             self.current_file_index += 1
-            if self.total_files == 1:
-                file_name = os.path.basename(file_path)
-                logging.info(f"Processing file: {file_name}")
+            # Don't log file processing here to avoid interrupting progress bars
+            # File information is already displayed in the CLI data section
             if self.file_type == "csv":
                 yield from self.read_csv_chunks(file_path)
             elif self.file_type == "parquet":
@@ -191,7 +189,7 @@ class RecDataLoader(FeatureSet):
         ),
         batch_size: int = 32,
         shuffle: bool = True,
-        load_full: bool = True,
+        streaming: bool = False,
         chunk_size: int = 10000,
         num_workers: int = 0,
         sampler=None,
@@ -203,7 +201,7 @@ class RecDataLoader(FeatureSet):
             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.
+            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.
             sampler: Optional sampler for DataLoader, only used for distributed training.
@@ -218,7 +216,7 @@ class RecDataLoader(FeatureSet):
                 path=data,
                 batch_size=batch_size,
                 shuffle=shuffle,
-                load_full=load_full,
+                streaming=streaming,
                 chunk_size=chunk_size,
                 num_workers=num_workers,
             )
@@ -231,7 +229,7 @@ class RecDataLoader(FeatureSet):
                 path=data,
                 batch_size=batch_size,
                 shuffle=shuffle,
-                load_full=load_full,
+                streaming=streaming,
                 chunk_size=chunk_size,
                 num_workers=num_workers,
             )
@@ -291,7 +289,7 @@ class RecDataLoader(FeatureSet):
         path: str | os.PathLike | list[str] | list[os.PathLike],
         batch_size: int,
         shuffle: bool,
-        load_full: bool,
+        streaming: bool,
         chunk_size: int = 10000,
         num_workers: int = 0,
     ) -> DataLoader:
@@ -312,8 +310,17 @@ class RecDataLoader(FeatureSet):
                     f"[RecDataLoader Error] Unsupported file extension in list: {suffix}"
                 )
             file_type = "csv" if suffix == ".csv" else "parquet"
+        if streaming:
+            return self.load_files_streaming(
+                file_paths,
+                file_type,
+                batch_size,
+                chunk_size,
+                shuffle,
+                num_workers=num_workers,
+            )
         # Load full data into memory
-        if load_full:
+        else:
             dfs = []
             total_bytes = 0
             for file_path in file_paths:
@@ -326,26 +333,17 @@ class RecDataLoader(FeatureSet):
                     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."
+                        f"[RecDataLoader Error] Out of memory while reading {file_path}. Consider using streaming=True."
                     ) from exc
             try:
                 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."
+                    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
             )
-        else:
-            return self.load_files_streaming(
-                file_paths,
-                file_type,
-                batch_size,
-                chunk_size,
-                shuffle,
-                num_workers=num_workers,
-            )
 
     def load_files_streaming(
         self,