nextrec 0.3.1__py3-none-any.whl → 0.3.3__py3-none-any.whl

nextrec/data/__init__.py

@@ -18,7 +18,7 @@ from nextrec.data.data_utils import (
     read_table,
     load_dataframes,
 )
-from nextrec.basic.features import FeatureSpecMixin
+from nextrec.basic.features import FeatureSet
 from nextrec.data import data_utils
 from nextrec.data.dataloader import (
     TensorDictDataset,
@@ -38,7 +38,7 @@ __all__ = [
     'iter_file_chunks',
     'read_table',
     'load_dataframes',
-    'FeatureSpecMixin',
+    'FeatureSet',
     'data_utils',
     'TensorDictDataset',
     'FileDataset',

nextrec/data/data_utils.py

@@ -5,8 +5,9 @@ import numpy as np
 import pandas as pd
 import pyarrow.parquet as pq
 from pathlib import Path
+from typing import Any, Mapping, Sequence
 
-def _stack_section(batch: list[dict], section: str):
+def stack_section(batch: list[dict], section: str):
     """Stack one section of the batch (features/labels/ids)."""
     entries = [item.get(section) for item in batch if item.get(section) is not None]
     if not entries:
@@ -39,9 +40,9 @@ def collate_fn(batch):
                 "ids": first.get("ids"),
             }
         return {
-            "features": _stack_section(batch, "features") or {},
-            "labels": _stack_section(batch, "labels"),
-            "ids": _stack_section(batch, "ids"),
+            "features": stack_section(batch, "features") or {},
+            "labels": stack_section(batch, "labels"),
+            "ids": stack_section(batch, "ids"),
         }
 
     # Fallback: stack tuples/lists of tensors
@@ -190,3 +191,78 @@ def build_eval_candidates(
     eval_df = eval_df.merge(user_features, on=user_col, how='left')
     eval_df = eval_df.merge(item_features, on=item_col, how='left')
     return eval_df
+
+def batch_to_dict(batch_data: Any, include_ids: bool = True) -> dict:
+    """Standardize a dataloader batch into a dict of features, labels, and ids."""
+    if not (isinstance(batch_data, Mapping) and "features" in batch_data):
+        raise TypeError(
+            "[BaseModel-batch_to_dict Error] Batch data must be a dict with 'features' produced by the current DataLoader."
+        )
+    return {
+        "features": batch_data.get("features", {}),
+        "labels": batch_data.get("labels"),
+        "ids": batch_data.get("ids") if include_ids else None,
+    }
+
+
+# def get_user_ids(
+#     data: dict | pd.DataFrame | None, user_id_column: str = "user_id"
+# ) -> np.ndarray | None:
+#     """Extract user IDs from a dataset dict or DataFrame."""
+#     if data is None:
+#         return None
+#     if isinstance(data, pd.DataFrame) and user_id_column in data.columns:
+#         return np.asarray(data[user_id_column].values)
+#     if isinstance(data, dict) and user_id_column in data:
+#         return np.asarray(data[user_id_column])
+#     return None
+
+
+# def get_user_ids_from_batch(
+#     batch_dict: Mapping[str, Any], id_columns: Sequence[str] | None = None
+# ) -> np.ndarray | None:
+#     """Extract the prioritized user id column from a batch dict."""
+#     ids_container = batch_dict.get("ids") if isinstance(batch_dict, Mapping) else None
+#     if not ids_container:
+#         return None
+
+#     batch_user_id = None
+#     if id_columns:
+#         for id_name in id_columns:
+#             if id_name in ids_container:
+#                 batch_user_id = ids_container[id_name]
+#                 break
+#     if batch_user_id is None:
+#         batch_user_id = next(iter(ids_container.values()), None)
+#     if batch_user_id is None:
+#         return None
+
+#     if isinstance(batch_user_id, torch.Tensor):
+#         ids_np = batch_user_id.detach().cpu().numpy()
+#     else:
+#         ids_np = np.asarray(batch_user_id)
+#     if ids_np.ndim == 0:
+#         ids_np = ids_np.reshape(1)
+#     return ids_np.reshape(ids_np.shape[0])
+
+
+def get_user_ids(data, id_columns: list[str] | str | None = None) -> np.ndarray | None:
+    id_columns = id_columns if isinstance(id_columns, list) else [id_columns] if isinstance(id_columns, str) else []
+    if not id_columns:
+        return None
+
+    main_id = id_columns[0]  
+    if isinstance(data, pd.DataFrame) and main_id in data.columns:
+        arr = np.asarray(data[main_id].values)
+        return arr.reshape(arr.shape[0])
+    if isinstance(data, dict):
+        ids_container = data.get("ids")
+        if isinstance(ids_container, dict) and main_id in ids_container:
+            val = ids_container[main_id]
+            val = val.detach().cpu().numpy() if isinstance(val, torch.Tensor) else np.asarray(val)
+            return val.reshape(val.shape[0])
+        if main_id in data:
+            arr = np.asarray(data[main_id])
+            return arr.reshape(arr.shape[0])
+
+    return None

nextrec/data/dataloader.py

@@ -2,11 +2,10 @@
 Dataloader definitions
 
 Date: create on 27/10/2025
-Checkpoint: edit on 29/11/2025
+Checkpoint: edit on 02/12/2025
 Author: Yang Zhou,zyaztec@gmail.com
 """
 import os
-import tqdm
 import torch
 import logging
 import numpy as np
@@ -18,15 +17,11 @@ from typing import cast
 
 from torch.utils.data import DataLoader, Dataset, IterableDataset
 from nextrec.data.preprocessor import DataProcessor
-from nextrec.basic.features import DenseFeature, SparseFeature, SequenceFeature, FeatureSpecMixin
+from nextrec.basic.features import DenseFeature, SparseFeature, SequenceFeature, FeatureSet
 
 from nextrec.basic.loggers import colorize
-from nextrec.data import (
-    get_column_data,
-    collate_fn,
-    resolve_file_paths,
-    read_table,
-)
+from nextrec.data import get_column_data, collate_fn, resolve_file_paths, read_table
+from nextrec.utils import to_tensor
 
 class TensorDictDataset(Dataset):
     """Dataset returning sample-level dicts matching the unified batch schema."""
@@ -52,7 +47,7 @@ class TensorDictDataset(Dataset):
         sample_ids = {name: tensor[idx] for name, tensor in self.ids.items()} if self.ids else None
         return {"features": sample_features, "labels": sample_labels, "ids": sample_ids}
 
-class FileDataset(FeatureSpecMixin, IterableDataset):
+class FileDataset(FeatureSet, IterableDataset):
     def __init__(self, 
                  file_paths: list[str],                      # file paths to read, containing CSV or Parquet files
                  dense_features: list[DenseFeature],         # dense feature definitions
@@ -67,44 +62,37 @@ class FileDataset(FeatureSpecMixin, IterableDataset):
         self.chunk_size = chunk_size
         self.file_type = file_type
         self.processor = processor
-        self._set_feature_config(dense_features, sparse_features, sequence_features, target_columns, id_columns)
+        self.set_all_features(dense_features, sparse_features, sequence_features, target_columns, id_columns)
         self.current_file_index = 0
         self.total_files = len(file_paths)
     
     def __iter__(self):
         self.current_file_index = 0
-        self._file_pbar = None
-        if self.total_files > 1:
-            self._file_pbar = tqdm.tqdm(total=self.total_files, desc="Files", unit="file", position=0, leave=True, bar_format='{desc}: {percentage:3.0f}%|{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}]')
         for file_path in self.file_paths:
             self.current_file_index += 1
-            if self._file_pbar is not None:
-                self._file_pbar.update(1)
-            elif self.total_files == 1:
+            if self.total_files == 1:
                 file_name = os.path.basename(file_path)
                 logging.info(f"Processing file: {file_name}")
             if self.file_type == 'csv':
-                yield from self._read_csv_chunks(file_path)
+                yield from self.read_csv_chunks(file_path)
             elif self.file_type == 'parquet':
-                yield from self._read_parquet_chunks(file_path)
-        if self._file_pbar is not None:
-            self._file_pbar.close()
+                yield from self.read_parquet_chunks(file_path)
     
-    def _read_csv_chunks(self, file_path: str):
+    def read_csv_chunks(self, file_path: str):
         chunk_iterator = pd.read_csv(file_path, chunksize=self.chunk_size)
         for chunk in chunk_iterator:
-            tensors = self._dataframe_to_tensors(chunk)
+            tensors = self.dataframeto_tensors(chunk)
             yield tensors
     
-    def _read_parquet_chunks(self, file_path: str):
+    def read_parquet_chunks(self, file_path: str):
         parquet_file = pq.ParquetFile(file_path)
         for batch in parquet_file.iter_batches(batch_size=self.chunk_size):
             chunk = batch.to_pandas()            
-            tensors = self._dataframe_to_tensors(chunk)
+            tensors = self.dataframeto_tensors(chunk)
             yield tensors
             del chunk
                 
-    def _dataframe_to_tensors(self, df: pd.DataFrame) -> dict | None:
+    def dataframeto_tensors(self, df: pd.DataFrame) -> dict | None:
         if self.processor is not None:
             if not self.processor.is_fitted:
                 raise ValueError("[DataLoader Error] DataProcessor must be fitted before using in streaming mode")
@@ -120,7 +108,7 @@ class FileDataset(FeatureSpecMixin, IterableDataset):
         return batch
 
 
-class RecDataLoader(FeatureSpecMixin):
+class RecDataLoader(FeatureSet):
     def __init__(self,
                  dense_features: list[DenseFeature] | None = None,
                  sparse_features: list[SparseFeature] | None = None,
@@ -129,7 +117,7 @@ class RecDataLoader(FeatureSpecMixin):
                  id_columns: str | list[str] | None = None,
                  processor: DataProcessor | None = None):
         self.processor = processor
-        self._set_feature_config(dense_features, sparse_features, sequence_features, target, id_columns)
+        self.set_all_features(dense_features, sparse_features, sequence_features, target, id_columns)
 
     def create_dataloader(self,
                          data: dict | pd.DataFrame | str | DataLoader,
@@ -140,13 +128,13 @@ class RecDataLoader(FeatureSpecMixin):
         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)
+            return self.create_from_path(path=data, batch_size=batch_size, shuffle=shuffle, load_full=load_full, chunk_size=chunk_size)
         elif isinstance(data, (dict, pd.DataFrame)):
-            return self._create_from_memory(data=data, batch_size=batch_size, shuffle=shuffle)
+            return self.create_from_memory(data=data, batch_size=batch_size, shuffle=shuffle)
         else:
             raise ValueError(f"[RecDataLoader Error] Unsupported data type: {type(data)}")
     
-    def _create_from_memory(self, 
+    def create_from_memory(self, 
                            data: dict | pd.DataFrame,
                            batch_size: int,
                            shuffle: bool) -> DataLoader:
@@ -162,7 +150,7 @@ class RecDataLoader(FeatureSpecMixin):
         dataset = TensorDictDataset(tensors)
         return DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, collate_fn=collate_fn)
     
-    def _create_from_path(self,
+    def create_from_path(self,
                          path: str,
                          batch_size: int,
                          shuffle: bool,
@@ -179,7 +167,6 @@ class RecDataLoader(FeatureSpecMixin):
                 except OSError:
                     pass
                 try:
-                    df = read_table(file_path, file_type)
                     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.") from exc
@@ -187,11 +174,11 @@ class RecDataLoader(FeatureSpecMixin):
                 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.") from exc
-            return self._create_from_memory(combined_df, batch_size, shuffle,)
+            return self.create_from_memory(combined_df, batch_size, shuffle,)
         else:
-            return self._load_files_streaming(file_paths, file_type, batch_size, chunk_size, shuffle)
+            return self.load_files_streaming(file_paths, file_type, batch_size, chunk_size, shuffle)
 
-    def _load_files_streaming(self,
+    def load_files_streaming(self,
                              file_paths: list[str],
                              file_type: str,
                              batch_size: int,
@@ -201,20 +188,10 @@ class RecDataLoader(FeatureSpecMixin):
             logging.warning("[RecDataLoader Warning] Shuffle is ignored in streaming mode (IterableDataset).")
         if batch_size != 1:
             logging.warning("[RecDataLoader Warning] Streaming mode enforces batch_size=1; tune chunk_size to control memory/throughput.")
-        dataset = FileDataset(
-            file_paths=file_paths,
-            dense_features=self.dense_features,
-            sparse_features=self.sparse_features,
-            sequence_features=self.sequence_features,
-            target_columns=self.target_columns,
-            id_columns=self.id_columns,
-            chunk_size=chunk_size,
-            file_type=file_type,
-            processor=self.processor
-        )
+        dataset = FileDataset(file_paths=file_paths, dense_features=self.dense_features, sparse_features=self.sparse_features, sequence_features=self.sequence_features, target_columns=self.target_columns, id_columns=self.id_columns, chunk_size=chunk_size, file_type=file_type, processor=self.processor)
         return DataLoader(dataset, batch_size=1, collate_fn=collate_fn)
 
-def _normalize_sequence_column(column, feature: SequenceFeature) -> np.ndarray:
+def normalize_sequence_column(column, feature: SequenceFeature) -> np.ndarray:
     if isinstance(column, pd.Series):
         column = column.tolist()
     if isinstance(column, (list, tuple)):
@@ -250,25 +227,27 @@ def _normalize_sequence_column(column, feature: SequenceFeature) -> np.ndarray:
         column = column.reshape(-1, 1)
     return np.asarray(column, dtype=np.int64)
 
-
-def build_tensors_from_data(  # noqa: C901
+def build_tensors_from_data( 
     data: dict | pd.DataFrame,
     raw_data: dict | pd.DataFrame,
     features: list,
     target_columns: list[str],
     id_columns: list[str]
 ) -> dict | None:
-    feature_tensors: dict[str, torch.Tensor] = {}
+    feature_tensors = {}
     for feature in features:
         column = get_column_data(data, feature.name)
         if column is None:
             raise ValueError(f"[RecDataLoader Error] Feature column '{feature.name}' not found in data")
-        if isinstance(feature, SequenceFeature):
-            tensor = torch.from_numpy(_normalize_sequence_column(column, feature))
+        if isinstance(feature, SequenceFeature): # sequence feature will do padding/truncation again to avoid the case when input data is not preprocessed
+            arr = normalize_sequence_column(column, feature)
+            tensor = to_tensor(arr, dtype=torch.long)
         elif isinstance(feature, DenseFeature):
-            tensor = torch.from_numpy(np.asarray(column, dtype=np.float32))
+            arr = np.asarray(column, dtype=np.float32)
+            tensor = to_tensor(arr, dtype=torch.float32)
         else:
-            tensor = torch.from_numpy(np.asarray(column, dtype=np.int64))
+            arr = np.asarray(column, dtype=np.int64)
+            tensor = to_tensor(arr, dtype=torch.long)
         feature_tensors[feature.name] = tensor
     label_tensors = None
     if target_columns:
@@ -277,7 +256,7 @@ def build_tensors_from_data(  # noqa: C901
             column = get_column_data(data, target_name)
             if column is None:
                 continue
-            label_tensor = torch.from_numpy(np.asarray(column, dtype=np.float32))
+            label_tensor = to_tensor(np.asarray(column, dtype=np.float32), dtype=torch.float32)
             if label_tensor.dim() == 2 and label_tensor.shape[0] == 1 and label_tensor.shape[1] > 1:
                 label_tensor = label_tensor.t()
             if label_tensor.shape[1:] == (1,):
@@ -298,7 +277,7 @@ def build_tensors_from_data(  # noqa: C901
                 id_arr = np.asarray(column, dtype=np.int64)
             except Exception as exc:
                 raise TypeError( f"[RecDataLoader Error] ID column '{id_col}' must contain numeric values. Received dtype={np.asarray(column).dtype}, error: {exc}") from exc
-            id_tensors[id_col] = torch.from_numpy(id_arr)
+            id_tensors[id_col] = to_tensor(id_arr, dtype=torch.long)
     if not feature_tensors:
         return None
     return {"features": feature_tensors, "labels": label_tensors, "ids": id_tensors}

nextrec/data/preprocessor.py

@@ -2,6 +2,7 @@
 DataProcessor for data preprocessing including numeric, sparse, sequence features and target processing.
 
 Date: create on 13/11/2025
+Checkpoint: edit on 02/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 """
 from __future__ import annotations
@@ -32,11 +33,11 @@ from nextrec.data.data_utils import (
     default_output_dir,
 )
 from nextrec.basic.session import resolve_save_path
-from nextrec.basic.features import FeatureSpecMixin
+from nextrec.basic.features import FeatureSet
 from nextrec.__version__ import __version__
 
 
-class DataProcessor(FeatureSpecMixin):
+class DataProcessor(FeatureSet):
     """DataProcessor for data preprocessing including numeric, sparse, sequence features and target processing.
     
     Examples:
@@ -70,7 +71,7 @@ class DataProcessor(FeatureSpecMixin):
         self.scalers: Dict[str, Any] = {}
         self.label_encoders: Dict[str, LabelEncoder] = {}
         self.target_encoders: Dict[str, Dict[str, int]] = {}
-        self._set_target_id_config([], [])
+        self.set_target_id([], [])
 
     def add_numeric_feature(
         self, 
@@ -129,7 +130,7 @@ class DataProcessor(FeatureSpecMixin):
             'target_type': target_type,
             'label_map': label_map
         }
-        self._set_target_id_config(list(self.target_features.keys()), [])
+        self.set_target_id(list(self.target_features.keys()), [])
         
     def _hash_string(self, s: str, hash_size: int) -> int:
         return int(hashlib.md5(str(s).encode()).hexdigest(), 16) % hash_size

nextrec/models/generative/__init__.py

@@ -0,0 +1,5 @@
+from .hstu import HSTU
+
+__all__ = [
+    "HSTU",
+]