nextrec 0.1.1__py3-none-any.whl → 0.1.2__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,6 +10,7 @@ import torch
 import numpy as np
 import pandas as pd
 
+
 def collate_fn(batch):
     """
     Custom collate function for batching tuples of tensors.
@@ -30,12 +31,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)
 
 
@@ -53,7 +54,9 @@ 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:
@@ -69,13 +72,13 @@ def split_dict_random(data_dict: dict, test_size: float=0.2, random_state:int|No
         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
 
 
@@ -90,7 +93,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()
@@ -99,8 +102,7 @@ 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:
@@ -112,7 +114,9 @@ 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
 
@@ -127,6 +131,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