nextrec 0.1.4__py3-none-any.whl → 0.1.7__py3-none-any.whl

nextrec/data/__init__.py

@@ -19,9 +19,9 @@ from nextrec.data.data_utils import (
 from nextrec.data import data_utils
 
 __all__ = [
-    "collate_fn",
-    "get_column_data",
-    "split_dict_random",
-    "build_eval_candidates",
-    "data_utils",
+    'collate_fn',
+    'get_column_data',
+    'split_dict_random',
+    'build_eval_candidates',
+    'data_utils',
 ]

nextrec/data/data_utils.py

@@ -10,7 +10,6 @@ import torch
 import numpy as np
 import pandas as pd
 
-
 def collate_fn(batch):
     """
     Custom collate function for batching tuples of tensors.
@@ -31,12 +30,12 @@ def collate_fn(batch):
 
     num_tensors = len(batch[0])
     result = []
-
+    
     for i in range(num_tensors):
         tensor_list = [item[i] for item in batch]
         stacked = torch.cat(tensor_list, dim=0)
         result.append(stacked)
-
+    
     return tuple(result)
 
 
@@ -54,9 +53,7 @@ def get_column_data(data: dict | pd.DataFrame, name: str):
         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
-):
+def split_dict_random(data_dict: dict, test_size: float=0.2, random_state:int|None=None):
     """Randomly split a dictionary of data into training and testing sets."""
     lengths = [len(v) for v in data_dict.values()]
     if len(set(lengths)) != 1:
@@ -72,13 +69,13 @@ def split_dict_random(
         if isinstance(v, np.ndarray):
             return v[idx]
         elif isinstance(v, pd.Series):
-            return v.iloc[idx].to_numpy()
+            return v.iloc[idx].to_numpy()  
         else:
-            v_arr = np.asarray(v, dtype=object)
+            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()}
+    test_dict  = {k: take(v, test_idx)  for k, v in data_dict.items()}
     return train_dict, test_dict
 
 
@@ -93,7 +90,7 @@ def build_eval_candidates(
     num_neg_per_pos: int = 50,
     random_seed: int = 2025,
 ) -> pd.DataFrame:
-    """Build evaluation candidates with positive and negative samples for each user."""
+    """Build evaluation candidates with positive and negative samples for each user.   """
     rng = np.random.default_rng(random_seed)
 
     users = df_all[user_col].unique()
@@ -102,7 +99,8 @@ def build_eval_candidates(
     rows = []
 
     user_hist_items = {
-        u: df_all[df_all[user_col] == u][item_col].unique() for u in users
+        u: df_all[df_all[user_col] == u][item_col].unique()
+        for u in users
     }
 
     for u in users:
@@ -114,9 +112,7 @@ def build_eval_candidates(
         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
 
@@ -131,6 +127,6 @@ def build_eval_candidates(
                 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