Update utils/data_processing.py

utils/data_processing.py

@@ -56,7 +56,7 @@ def get_data(dataset_name, different_new_nodes_between_val_and_test=False, rando
   node_features = np.load('./data/ml_{}_node.npy'.format(dataset_name)) 
 
 
-  #ASHUTOSH
+  ############################ ASHUTOSH #####################################
   graph_df = graph_df.rename(columns={
     'user_id': 'u',
     'item_id': 'i',
@@ -64,7 +64,9 @@ def get_data(dataset_name, different_new_nodes_between_val_and_test=False, rando
     'state_label': 'label',
     # You can drop or handle the features column separately if needed
   })
-  ###################
+  # Add edge index column
+  graph_df['idx'] = range(len(graph_df))
+  ############################################################################
     
   if randomize_features:
     node_features = np.random.rand(node_features.shape[0], node_features.shape[1])
@@ -77,6 +79,7 @@ def get_data(dataset_name, different_new_nodes_between_val_and_test=False, rando
   labels = graph_df.label.values
   timestamps = graph_df.ts.values
 
+  
   full_data = Data(sources, destinations, timestamps, edge_idxs, labels)
 
   random.seed(2020)