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

nextrec/basic/summary.py

@@ -8,6 +8,7 @@ Author: Yang Zhou,zyaztec@gmail.com
 
 from __future__ import annotations
 
+import inspect
 import logging
 from typing import Any, Literal
 
@@ -34,6 +35,7 @@ class SummarySet:
     scheduler_name: str | None
     scheduler_params: dict[str, Any]
     loss_config: Any
+    loss_params: Any
     loss_weights: Any
     grad_norm: Any
     embedding_l1_reg: float
@@ -73,7 +75,7 @@ class SummarySet:
     def build_data_summary(
         self, data: Any, data_loader: DataLoader | None, sample_key: str
     ):
-        
+
         dataset = data_loader.dataset if data_loader is not None else None
 
         train_size = get_data_length(dataset)
@@ -325,6 +327,73 @@ class SummarySet:
 
             if hasattr(self, "loss_config"):
                 logger.info(f"Loss Function:           {self.loss_config}")
+
+            loss_params_summary: list[str] = []
+            loss_fn = getattr(self, "loss_fn", None)
+            if loss_fn is not None:
+                loss_modules = (
+                    list(loss_fn) if isinstance(loss_fn, (list, tuple)) else [loss_fn]
+                )
+                loss_config = getattr(self, "loss_config", None)
+                if isinstance(loss_config, list):
+                    loss_names = loss_config
+                elif loss_config is not None:
+                    loss_names = [loss_config] * len(loss_modules)
+                else:
+                    loss_names = [None] * len(loss_modules)
+
+                loss_params = getattr(self, "loss_params", None)
+                if isinstance(loss_params, list):
+                    explicit_params = loss_params
+                elif isinstance(loss_params, dict):
+                    explicit_params = [loss_params] * len(loss_modules)
+                else:
+                    explicit_params = [None] * len(loss_modules)
+
+                for idx, loss_module in enumerate(loss_modules):
+                    params: dict[str, Any] = {}
+                    explicit = (
+                        explicit_params[idx] if idx < len(explicit_params) else None
+                    )
+                    if explicit:
+                        params.update(explicit)
+                    try:
+                        signature = inspect.signature(loss_module.__class__.__init__)
+                    except (TypeError, ValueError):
+                        signature = None
+                    if signature is not None:
+                        for name, param in signature.parameters.items():
+                            if name == "self" or name.startswith("_"):
+                                continue
+                            if hasattr(loss_module, name):
+                                value = getattr(loss_module, name)
+                                if callable(value):
+                                    continue
+                                params.setdefault(name, value)
+                            elif (
+                                param.default is not inspect._empty
+                                and param.default is not None
+                            ):
+                                params.setdefault(name, param.default)
+                    if not params:
+                        continue
+
+                    loss_name = loss_names[idx] if idx < len(loss_names) else None
+                    if len(loss_modules) > 1:
+                        header = f"  [{idx}]"
+                        if loss_name is not None:
+                            header = f"{header} {loss_name}"
+                        loss_params_summary.append(header)
+                        indent = "    "
+                    else:
+                        indent = "  "
+                    for key, value in params.items():
+                        loss_params_summary.append(f"{indent}{key:25s}: {value}")
+
+            if loss_params_summary:
+                logger.info("Loss Params:")
+                for line in loss_params_summary:
+                    logger.info(line)
             if hasattr(self, "loss_weights"):
                 logger.info(f"Loss Weights:            {self.loss_weights}")
             if hasattr(self, "grad_norm"):
@@ -355,53 +424,30 @@ class SummarySet:
             logger.info("")
             logger.info(colorize("Data Summary", color="cyan", bold=True))
             logger.info(colorize("-" * 80, color="cyan"))
-            if self.train_data_summary:
-                train_samples = self.train_data_summary.get("train_samples")
-                if train_samples is not None:
-                    logger.info(format_kv("Train Samples", f"{train_samples:,}"))
-
-                label_distributions = self.train_data_summary.get("label_distributions")
-                if isinstance(label_distributions, dict):
-                    for target_name, details in label_distributions.items():
-                        lines = details.get("lines", [])
-                        logger.info(f"{target_name}:")
-                        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:
+            for label, data_summary in (
+                ("Train", self.train_data_summary),
+                ("Valid", self.valid_data_summary),
+            ):
+                if not data_summary:
+                    continue
+                if label == "Valid" and self.train_data_summary:
                     logger.info("")
-                valid_samples = self.valid_data_summary.get("valid_samples")
-                if valid_samples is not None:
-                    logger.info(format_kv("Valid Samples", f"{valid_samples:,}"))
+                sample_key = "train_samples" if label == "Train" else "valid_samples"
+                samples = data_summary.get(sample_key)
+                if samples is not None:
+                    logger.info(format_kv(f"{label} Samples", f"{samples:,}"))
 
-                label_distributions = self.valid_data_summary.get("label_distributions")
+                label_distributions = data_summary.get("label_distributions")
                 if isinstance(label_distributions, dict):
                     for target_name, details in label_distributions.items():
                         lines = details.get("lines", [])
                         logger.info(f"{target_name}:")
-                        for label, value in lines:
-                            logger.info(f"  {format_kv(label, value)}")
+                        for line_label, value in lines:
+                            logger.info(f"  {format_kv(line_label, value)}")
 
-                dataloader_info = self.valid_data_summary.get("dataloader")
+                dataloader_info = data_summary.get("dataloader")
                 if isinstance(dataloader_info, dict):
-                    logger.info("Valid DataLoader:")
+                    logger.info(f"{label} DataLoader:")
                     for key in (
                         "batch_size",
                         "num_workers",
@@ -410,7 +456,7 @@ class SummarySet:
                         "sampler",
                     ):
                         if key in dataloader_info:
-                            label = key.replace("_", " ").title()
+                            field_label = key.replace("_", " ").title()
                             logger.info(
-                                format_kv(label, dataloader_info[key], indent=2)
+                                format_kv(field_label, dataloader_info[key], indent=2)
                             )

nextrec/cli.py

@@ -48,7 +48,7 @@ from nextrec.utils.data import (
     read_yaml,
     resolve_file_paths,
 )
-from nextrec.utils.feature import to_list
+from nextrec.utils.torch_utils import to_list
 
 logger = logging.getLogger(__name__)
 
@@ -156,9 +156,7 @@ def train_model(train_config_path: str) -> None:
         logger.info(
             format_kv(
                 "Validation path",
-                resolve_path(
-                    data_cfg.get("valid_path"), config_dir
-                ),
+                resolve_path(data_cfg.get("valid_path"), config_dir),
             )
         )
 
@@ -247,7 +245,9 @@ def train_model(train_config_path: str) -> None:
 
     if streaming:
         if file_type is None:
-            raise ValueError("[NextRec CLI Error] Streaming mode requires a valid file_type")
+            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,
@@ -422,6 +422,49 @@ def train_model(train_config_path: str) -> None:
         note=train_cfg.get("note"),
     )
 
+    export_cfg = train_cfg.get("export_onnx")
+    if export_cfg is None:
+        export_cfg = cfg.get("export_onnx")
+    export_enabled = False
+    export_options: dict[str, Any] = {}
+    if isinstance(export_cfg, bool):
+        export_enabled = export_cfg
+    elif isinstance(export_cfg, dict):
+        export_options = export_cfg
+        export_enabled = bool(export_cfg.get("enable", False))
+
+    if export_enabled:
+        log_cli_section("ONNX Export")
+        onnx_path = None
+        if export_options.get("path") or export_options.get("save_path"):
+            logger.warning(
+                "[NextRec CLI Warning] export_onnx.path/save_path is deprecated; "
+                "ONNX will be saved to best/checkpoint paths."
+            )
+        onnx_best_path = Path(model.best_path).with_suffix(".onnx")
+        onnx_ckpt_path = Path(model.checkpoint_path).with_suffix(".onnx")
+        onnx_batch_size = export_options.get("batch_size", 1)
+        onnx_opset = export_options.get("opset_version", 18)
+        log_kv_lines(
+            [
+                ("ONNX best path", onnx_best_path),
+                ("ONNX checkpoint path", onnx_ckpt_path),
+                ("Batch size", onnx_batch_size),
+                ("Opset", onnx_opset),
+                ("Dynamic batch", False),
+            ]
+        )
+        model.export_onnx(
+            save_path=onnx_best_path,
+            batch_size=onnx_batch_size,
+            opset_version=onnx_opset,
+        )
+        model.export_onnx(
+            save_path=onnx_ckpt_path,
+            batch_size=onnx_batch_size,
+            opset_version=onnx_opset,
+        )
+
 
 def predict_model(predict_config_path: str) -> None:
     """
@@ -492,12 +535,16 @@ def predict_model(predict_config_path: str) -> None:
     # Load checkpoint and ensure required parameters are passed
     checkpoint_base = Path(session_dir)
     if checkpoint_base.is_dir():
+        best_candidates = sorted(checkpoint_base.glob("*_best.pt"))
         candidates = sorted(checkpoint_base.glob("*.pt"))
-        if not candidates:
+        if best_candidates:
+            model_file = best_candidates[-1]
+        elif candidates:
+            model_file = candidates[-1]
+        else:
             raise FileNotFoundError(
                 f"[NextRec CLI Error]: Unable to find model checkpoint: {checkpoint_base}"
             )
-        model_file = candidates[-1]
         config_dir_for_features = checkpoint_base
     else:
         model_file = (
@@ -564,11 +611,32 @@ def predict_model(predict_config_path: str) -> None:
     )
 
     log_cli_section("Model")
+    use_onnx = bool(predict_cfg.get("use_onnx")) or bool(predict_cfg.get("onnx_path"))
+    onnx_path = predict_cfg.get("onnx_path") or cfg.get("onnx_path")
+    if onnx_path:
+        onnx_path = resolve_path(onnx_path, config_dir)
+    if use_onnx and onnx_path is None:
+        search_dir = (
+            checkpoint_base if checkpoint_base.is_dir() else checkpoint_base.parent
+        )
+        best_candidates = sorted(search_dir.glob("*_best.onnx"))
+        if best_candidates:
+            onnx_path = best_candidates[-1]
+        else:
+            candidates = sorted(search_dir.glob("*.onnx"))
+            if not candidates:
+                raise FileNotFoundError(
+                    f"[NextRec CLI Error]: Unable to find ONNX model in {search_dir}"
+                )
+            onnx_path = candidates[-1]
+
     log_kv_lines(
         [
             ("Model", model.__class__.__name__),
             ("Checkpoint", model_file),
             ("Device", predict_cfg.get("device", "cpu")),
+            ("Use ONNX", use_onnx),
+            ("ONNX path", onnx_path if use_onnx else "(disabled)"),
         ]
     )
 
@@ -582,7 +650,10 @@ def predict_model(predict_config_path: str) -> None:
     )
 
     data_path = resolve_path(predict_cfg["data_path"], config_dir)
-    batch_size = predict_cfg.get("batch_size", 512)
+    streaming = bool(predict_cfg.get("streaming", True))
+    chunk_size = int(predict_cfg.get("chunk_size", 20000))
+    batch_size = int(predict_cfg.get("batch_size", 512))
+    effective_batch_size = chunk_size if streaming else batch_size
 
     log_cli_section("Data")
     log_kv_lines(
@@ -594,18 +665,18 @@ def predict_model(predict_config_path: str) -> None:
                     "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)),
+            ("Batch size", effective_batch_size),
+            ("Chunk size", chunk_size),
+            ("Streaming", streaming),
         ]
     )
     logger.info("")
     pred_loader = rec_dataloader.create_dataloader(
         data=str(data_path),
-        batch_size=batch_size,
+        batch_size=1 if streaming else batch_size,
         shuffle=False,
-        streaming=predict_cfg.get("streaming", True),
-        chunk_size=predict_cfg.get("chunk_size", 20000),
+        streaming=streaming,
+        chunk_size=chunk_size,
         prefetch_factor=predict_cfg.get("prefetch_factor"),
     )
 
@@ -613,27 +684,43 @@ def predict_model(predict_config_path: str) -> None:
         "save_data_format", predict_cfg.get("save_format", "csv")
     )
     pred_name = predict_cfg.get("name", "pred")
-
-    save_path = checkpoint_base / "predictions" / f"{pred_name}.{save_format}"
+    pred_name_path = Path(pred_name)
+    if pred_name_path.is_absolute():
+        save_path = pred_name_path
+        if save_path.suffix == "":
+            save_path = save_path.with_suffix(f".{save_format}")
+    else:
+        save_path = checkpoint_base / "predictions" / f"{pred_name}.{save_format}"
 
     start = time.time()
     logger.info("")
-    result = model.predict(
-        data=pred_loader,
-        batch_size=batch_size,
-        include_ids=bool(id_columns),
-        return_dataframe=False,
-        save_path=str(save_path),
-        save_format=save_format,
-        num_workers=predict_cfg.get("num_workers", 0),
-    )
+    if use_onnx:
+        result = model.predict_onnx(
+            onnx_path=onnx_path,
+            data=pred_loader,
+            batch_size=effective_batch_size,
+            include_ids=bool(id_columns),
+            return_dataframe=False,
+            save_path=str(save_path),
+            save_format=save_format,
+            num_workers=predict_cfg.get("num_workers", 0),
+        )
+    else:
+        result = model.predict(
+            data=pred_loader,
+            batch_size=effective_batch_size,
+            include_ids=bool(id_columns),
+            return_dataframe=False,
+            save_path=str(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
-    )
+    if isinstance(result, (str, Path)):
+        output_path = Path(result)
+    else:
+        output_path = save_path
     logger.info(f"Prediction completed, results saved to: {output_path}")
     logger.info(f"Total time: {duration:.2f} seconds")
 

nextrec/data/data_processing.py

@@ -12,18 +12,21 @@ from typing import Any
 import numpy as np
 import pandas as pd
 import torch
+import polars as pl
+
 
 from nextrec.utils.torch_utils import to_numpy
 
 
-def get_column_data(data: dict | pd.DataFrame, name: str):
+def get_column_data(data: dict | pd.DataFrame | pl.DataFrame, name: str):
 
     if isinstance(data, dict):
         return data[name] if name in data else None
     elif isinstance(data, pd.DataFrame):
-        if name not in data.columns:
-            return None
         return data[name].values
+    elif isinstance(data, pl.DataFrame):
+        series = data.get_column(name)
+        return series.to_numpy()
     else:
         raise KeyError(f"Only dict or DataFrame supported, got {type(data)}")
 
@@ -33,6 +36,8 @@ def get_data_length(data: Any) -> int | None:
         return None
     if isinstance(data, pd.DataFrame):
         return len(data)
+    if isinstance(data, pl.DataFrame):
+        return data.height
     if isinstance(data, dict):
         if not data:
             return None
@@ -92,16 +97,6 @@ def split_dict_random(data_dict, test_size=0.2, random_state=None):
     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,