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

nextrec/data/data_processing.py

@@ -8,10 +8,11 @@ Author: Yang Zhou, zyaztec@gmail.com
 import torch
 import numpy as np
 import pandas as pd
-from typing import Any, Mapping
+from typing import Any
 
 
 def get_column_data(data: dict | pd.DataFrame, name: str):
+
     if isinstance(data, dict):
         return data[name] if name in data else None
     elif isinstance(data, pd.DataFrame):
@@ -23,21 +24,21 @@ def get_column_data(data: dict | pd.DataFrame, name: str):
             return getattr(data, name)
         raise KeyError(f"Unsupported data type for extracting column {name}")
 
+
 def split_dict_random(
-    data_dict: dict,
-    test_size: float = 0.2,
-    random_state: int | None = None
+    data_dict: dict, test_size: float = 0.2, random_state: int | None = None
 ):
+
     lengths = [len(v) for v in data_dict.values()]
     if len(set(lengths)) != 1:
         raise ValueError(f"Length mismatch: {lengths}")
-    
+
     n = lengths[0]
     rng = np.random.default_rng(random_state)
     perm = rng.permutation(n)
     cut = int(round(n * (1 - test_size)))
     train_idx, test_idx = perm[:cut], perm[cut:]
-    
+
     def take(v, idx):
         if isinstance(v, np.ndarray):
             return v[idx]
@@ -46,12 +47,22 @@ def split_dict_random(
         else:
             v_arr = np.asarray(v, dtype=object)
             return v_arr[idx]
-    
+
     train_dict = {k: take(v, train_idx) for k, v in data_dict.items()}
     test_dict = {k: take(v, test_idx) for k, v in data_dict.items()}
     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,
@@ -65,7 +76,7 @@ def build_eval_candidates(
 ) -> pd.DataFrame:
     """
     Build evaluation candidates with positive and negative samples for each user.
-    
+
     Args:
         df_all: Full interaction DataFrame
         user_col: Name of the user ID column
@@ -76,7 +87,7 @@ def build_eval_candidates(
         num_pos_per_user: Number of positive samples per user (default: 5)
         num_neg_per_pos: Number of negative samples per positive (default: 50)
         random_seed: Random seed for reproducibility (default: 2025)
-        
+
     Returns:
         pd.DataFrame: Evaluation candidates with features
     """
@@ -85,8 +96,10 @@ def build_eval_candidates(
     users = df_all[user_col].unique()
     all_items = item_features[item_col].unique()
     rows = []
-    user_hist_items = {u: df_all[df_all[user_col] == u][item_col].unique() for u in users}
-    
+    user_hist_items = {
+        u: df_all[df_all[user_col] == u][item_col].unique() for u in users
+    }
+
     for u in users:
         df_user = df_all[df_all[user_col] == u]
         pos_items = df_user[df_user[label_col] == 1][item_col].unique()
@@ -94,7 +107,9 @@ def build_eval_candidates(
             continue
         pos_items = pos_items[:num_pos_per_user]
         seen_items = set(user_hist_items[u])
-        neg_pool = np.setdiff1d(all_items, np.fromiter(seen_items, dtype=all_items.dtype))
+        neg_pool = np.setdiff1d(
+            all_items, np.fromiter(seen_items, dtype=all_items.dtype)
+        )
         if len(neg_pool) == 0:
             continue
         for pos in pos_items:
@@ -105,31 +120,30 @@ def build_eval_candidates(
             rows.append((u, pos, 1))
             for ni in neg_items:
                 rows.append((u, ni, 0))
-    
+
     eval_df = pd.DataFrame(rows, columns=[user_col, item_col, label_col])
-    eval_df = eval_df.merge(user_features, on=user_col, how='left')
-    eval_df = eval_df.merge(item_features, on=item_col, how='left')
+    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 get_user_ids(
-    data: Any,
-    id_columns: list[str] | str | None = None
+    data: Any, id_columns: list[str] | str | None = None
 ) -> np.ndarray | None:
     """
     Extract user IDs from various data structures.
-    
+
     Args:
         data: Data source (DataFrame, dict, or batch dict)
         id_columns: List or single ID column name(s) (default: None)
-        
+
     Returns:
         np.ndarray | None: User IDs as numpy array, or None if not found
     """
     id_columns = (
-        id_columns if isinstance(id_columns, list)
-        else [id_columns] if isinstance(id_columns, str)
-        else []
+        id_columns
+        if isinstance(id_columns, list)
+        else [id_columns] if isinstance(id_columns, str) else []
     )
     if not id_columns:
         return None
@@ -138,12 +152,16 @@ def get_user_ids(
     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)
+            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])

nextrec/data/data_utils.py

@@ -13,23 +13,34 @@ 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.data_processing import get_column_data, split_dict_random, build_eval_candidates, get_user_ids
-from nextrec.utils.file import resolve_file_paths, iter_file_chunks, read_table, load_dataframes, default_output_dir
+from nextrec.data.data_processing import (
+    get_column_data,
+    split_dict_random,
+    build_eval_candidates,
+    get_user_ids,
+)
+from nextrec.utils.file import (
+    resolve_file_paths,
+    iter_file_chunks,
+    read_table,
+    load_dataframes,
+    default_output_dir,
+)
 
 __all__ = [
     # Batch utilities
-    'collate_fn',
-    'batch_to_dict',
-    'stack_section',
+    "collate_fn",
+    "batch_to_dict",
+    "stack_section",
     # Data processing
-    'get_column_data',
-    'split_dict_random',
-    'build_eval_candidates',
-    'get_user_ids',
+    "get_column_data",
+    "split_dict_random",
+    "build_eval_candidates",
+    "get_user_ids",
     # File utilities
-    'resolve_file_paths',
-    'iter_file_chunks',
-    'read_table',
-    'load_dataframes',
-    'default_output_dir',
-]
+    "resolve_file_paths",
+    "iter_file_chunks",
+    "read_table",
+    "load_dataframes",
+    "default_output_dir",
+]