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

nextrec/utils/file.py

@@ -2,11 +2,14 @@
 File I/O utilities for NextRec
 
 Date: create on 03/12/2025
+Checkpoint: edit on 06/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 """
 
+import yaml
 import pandas as pd
 import pyarrow.parquet as pq
+
 from pathlib import Path
 from typing import Generator
 
@@ -14,7 +17,7 @@ from typing import Generator
 def resolve_file_paths(path: str) -> tuple[list[str], str]:
     """
     Resolve file or directory path into a sorted list of files and file type.
-    
+
     Args: path: Path to a file or directory
     Returns: tuple: (list of file paths, file type)
     """
@@ -22,16 +25,23 @@ def resolve_file_paths(path: str) -> tuple[list[str], str]:
 
     if path_obj.is_file():
         file_type = path_obj.suffix.lower().lstrip(".")
-        assert file_type in ["csv", "parquet"], f"Unsupported file extension: {file_type}"
+        assert file_type in [
+            "csv",
+            "parquet",
+        ], f"Unsupported file extension: {file_type}"
         return [str(path_obj)], file_type
 
     if path_obj.is_dir():
         collected_files = [p for p in path_obj.iterdir() if p.is_file()]
         csv_files = [str(p) for p in collected_files if p.suffix.lower() == ".csv"]
-        parquet_files = [str(p) for p in collected_files if p.suffix.lower() == ".parquet"]
+        parquet_files = [
+            str(p) for p in collected_files if p.suffix.lower() == ".parquet"
+        ]
 
         if csv_files and parquet_files:
-            raise ValueError("Directory contains both CSV and Parquet files. Please keep a single format.")
+            raise ValueError(
+                "Directory contains both CSV and Parquet files. Please keep a single format."
+            )
         file_paths = csv_files if csv_files else parquet_files
         if not file_paths:
             raise ValueError(f"No CSV or Parquet files found in directory: {path}")
@@ -42,18 +52,24 @@ def resolve_file_paths(path: str) -> tuple[list[str], str]:
     raise ValueError(f"Invalid path: {path}")
 
 
-def read_table(file_path: str, file_type: str) -> pd.DataFrame:
-    if file_type == "csv":
-        return pd.read_csv(file_path)
-    return pd.read_parquet(file_path)
+def read_table(path: str | Path, data_format: str | None = None) -> pd.DataFrame:
+    data_path = Path(path)
+    fmt = data_format.lower() if data_format else data_path.suffix.lower().lstrip(".")
+    if data_path.is_dir() and not fmt:
+        fmt = "parquet"
+    if fmt in {"parquet", ""}:
+        return pd.read_parquet(data_path)
+    if fmt in {"csv", "txt"}:
+        return pd.read_csv(data_path)
+    raise ValueError(f"Unsupported data format: {data_path}")
+
 
 def load_dataframes(file_paths: list[str], file_type: str) -> list[pd.DataFrame]:
     return [read_table(fp, file_type) for fp in file_paths]
 
+
 def iter_file_chunks(
-    file_path: str,
-    file_type: str,
-    chunk_size: int
+    file_path: str, file_type: str, chunk_size: int
 ) -> Generator[pd.DataFrame, None, None]:
     if file_type == "csv":
         yield from pd.read_csv(file_path, chunksize=chunk_size)
@@ -68,3 +84,8 @@ def default_output_dir(path: str) -> Path:
     if path_obj.is_file():
         return path_obj.parent / f"{path_obj.stem}_preprocessed"
     return path_obj.with_name(f"{path_obj.name}_preprocessed")
+
+
+def read_yaml(path: str | Path):
+    with open(path, "r", encoding="utf-8") as file:
+        return yaml.safe_load(file) or {}

nextrec/utils/initializer.py

@@ -5,32 +5,77 @@ Date: create on 13/11/2025
 Author: Yang Zhou, zyaztec@gmail.com
 """
 
+from typing import Any, Dict, Set, cast
+
 import torch.nn as nn
+from torch.nn.init import _NonlinearityType
 
+KNOWN_NONLINEARITIES: Set[str] = {
+    "linear",
+    "conv1d",
+    "conv2d",
+    "conv3d",
+    "conv_transpose1d",
+    "conv_transpose2d",
+    "conv_transpose3d",
+    "sigmoid",
+    "tanh",
+    "relu",
+    "leaky_relu",
+    "selu",
+    "gelu",
+}
+
+
+def resolve_nonlinearity(activation: str | _NonlinearityType) -> _NonlinearityType:
+    if isinstance(activation, str):
+        if activation in KNOWN_NONLINEARITIES:
+            return cast(_NonlinearityType, activation)
+        # Fall back to linear for custom activations (gain handled separately).
+        return "linear"
+    return activation
 
-def get_initializer(init_type='normal', activation='linear', param=None):
-    param = param or {}
 
+def resolve_gain(activation: str | _NonlinearityType, param: Dict[str, Any]) -> float:
+    if "gain" in param:
+        return param["gain"]
+    nonlinearity = resolve_nonlinearity(activation)
     try:
-        gain = param.get('gain', nn.init.calculate_gain(activation, param.get('param', None)))
+        return nn.init.calculate_gain(nonlinearity, param.get("param"))
     except ValueError:
-        gain = 1.0  # for custom activations like 'dice'
-    
+        return 1.0  # custom activation with no gain estimate available
+
+
+def get_initializer(
+    init_type: str = "normal",
+    activation: str | _NonlinearityType = "linear",
+    param: Dict[str, Any] | None = None,
+):
+    param = param or {}
+    nonlinearity = resolve_nonlinearity(activation)
+    gain = resolve_gain(activation, param)
+
     def initializer_fn(tensor):
-        if init_type == 'xavier_uniform':
+        if init_type == "xavier_uniform":
             nn.init.xavier_uniform_(tensor, gain=gain)
-        elif init_type == 'xavier_normal':
+        elif init_type == "xavier_normal":
             nn.init.xavier_normal_(tensor, gain=gain)
-        elif init_type == 'kaiming_uniform':
-            nn.init.kaiming_uniform_(tensor, a=param.get('a', 0), nonlinearity=activation)
-        elif init_type == 'kaiming_normal':
-            nn.init.kaiming_normal_(tensor, a=param.get('a', 0), nonlinearity=activation)
-        elif init_type == 'orthogonal':
+        elif init_type == "kaiming_uniform":
+            nn.init.kaiming_uniform_(
+                tensor, a=param.get("a", 0), nonlinearity=nonlinearity
+            )
+        elif init_type == "kaiming_normal":
+            nn.init.kaiming_normal_(
+                tensor, a=param.get("a", 0), nonlinearity=nonlinearity
+            )
+        elif init_type == "orthogonal":
             nn.init.orthogonal_(tensor, gain=gain)
-        elif init_type == 'normal':
-            nn.init.normal_(tensor, mean=param.get('mean', 0.0), std=param.get('std', 0.0001))
-        elif init_type == 'uniform':
-            nn.init.uniform_(tensor, a=param.get('a', -0.05), b=param.get('b', 0.05))
+        elif init_type == "normal":
+            nn.init.normal_(
+                tensor, mean=param.get("mean", 0.0), std=param.get("std", 0.0001)
+            )
+        elif init_type == "uniform":
+            nn.init.uniform_(tensor, a=param.get("a", -0.05), b=param.get("b", 0.05))
         else:
             raise ValueError(f"Unknown init_type: {init_type}")
         return tensor

nextrec/utils/optimizer.py

@@ -8,16 +8,17 @@ Author: Yang Zhou, zyaztec@gmail.com
 import torch
 from typing import Iterable
 
+
 def get_optimizer(
     optimizer: str | torch.optim.Optimizer = "adam",
     params: Iterable[torch.nn.Parameter] | None = None,
-    **optimizer_params
+    **optimizer_params,
 ):
     if params is None:
         raise ValueError("params cannot be None. Please provide model parameters.")
 
-    if 'lr' not in optimizer_params:
-        optimizer_params['lr'] = 1e-3
+    if "lr" not in optimizer_params:
+        optimizer_params["lr"] = 1e-3
     if isinstance(optimizer, str):
         opt_name = optimizer.lower()
         if opt_name == "adam":
@@ -39,21 +40,36 @@ def get_optimizer(
         raise TypeError(f"Invalid optimizer type: {type(optimizer)}")
     return optimizer_fn
 
+
 def get_scheduler(
-    scheduler: str | torch.optim.lr_scheduler._LRScheduler | torch.optim.lr_scheduler.LRScheduler | type[torch.optim.lr_scheduler._LRScheduler] | type[torch.optim.lr_scheduler.LRScheduler] | None,
+    scheduler: (
+        str
+        | torch.optim.lr_scheduler._LRScheduler
+        | torch.optim.lr_scheduler.LRScheduler
+        | type[torch.optim.lr_scheduler._LRScheduler]
+        | type[torch.optim.lr_scheduler.LRScheduler]
+        | None
+    ),
     optimizer,
-    **scheduler_params
+    **scheduler_params,
 ):
     if isinstance(scheduler, str):
         if scheduler == "step":
-            scheduler_fn = torch.optim.lr_scheduler.StepLR(optimizer, **scheduler_params)
+            scheduler_fn = torch.optim.lr_scheduler.StepLR(
+                optimizer, **scheduler_params
+            )
         elif scheduler == "cosine":
-            scheduler_fn = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, **scheduler_params)
+            scheduler_fn = torch.optim.lr_scheduler.CosineAnnealingLR(
+                optimizer, **scheduler_params
+            )
         else:
             raise NotImplementedError(f"Unsupported scheduler: {scheduler}")
-    elif isinstance(scheduler, (torch.optim.lr_scheduler._LRScheduler, torch.optim.lr_scheduler.LRScheduler)):
+    elif isinstance(
+        scheduler,
+        (torch.optim.lr_scheduler._LRScheduler, torch.optim.lr_scheduler.LRScheduler),
+    ):
         scheduler_fn = scheduler
     else:
         raise TypeError(f"Invalid scheduler type: {type(scheduler)}")
-    
+
     return scheduler_fn