| 5 | import scipy.sparse as sp |
| 6 | |
| 7 | class Data(object): |
| 8 | def __init__(self, args): |
| 9 | |
| 10 | self.path = args.data_path + args.dataset |
| 11 | self.n_batch = args.n_batch |
| 12 | self.batch_size = args.batch_size |
| 13 | self.train_num = args.train_num |
| 14 | self.sample_num = args.sample_num |
| 15 | |
| 16 | try: |
| 17 | train_file = self.path + '/train.pkl' |
| 18 | test_file = self.path + '/test.pkl' |
| 19 | with open(train_file, 'rb') as f: |
| 20 | train_mat = pickle.load(f) |
| 21 | with open(test_file, 'rb') as f: |
| 22 | test_mat = pickle.load(f) |
| 23 | except Exception as e: |
| 24 | print("Try an alternative way of reading the data.") |
| 25 | train_file = self.path + '/train_index.pkl' |
| 26 | test_file = self.path + '/test_index.pkl' |
| 27 | with open(train_file, 'rb') as f: |
| 28 | train_index = pickle.load(f) |
| 29 | with open(test_file, 'rb') as f: |
| 30 | test_index = pickle.load(f) |
| 31 | train_row, train_col = train_index[0], train_index[1] |
| 32 | n_user = max(train_row) + 1 |
| 33 | n_item = max(train_col) + 1 |
| 34 | train_mat = sp.coo_matrix((np.ones(len(train_row)), (train_row, train_col)), shape=[n_user, n_item]) |
| 35 | test_row, test_col = test_index[0], test_index[1] |
| 36 | test_mat = sp.coo_matrix((np.ones(len(test_row)), (test_row, test_col)), shape=[n_user, n_item]) |
| 37 | |
| 38 | # get number of users and items |
| 39 | self.n_users, self.n_items = train_mat.shape[0], train_mat.shape[1] |
| 40 | self.n_train, self.n_test = len(train_mat.row), len(test_mat.row) |
| 41 | |
| 42 | self.print_statistics() |
| 43 | |
| 44 | self.R = train_mat.todok() |
| 45 | self.train_items, self.test_set = {}, {} |
| 46 | train_uid, train_iid = train_mat.row, train_mat.col |
| 47 | for i in range(len(train_uid)): |
| 48 | uid = train_uid[i] |
| 49 | iid = train_iid[i] |
| 50 | if uid not in self.train_items: |
| 51 | self.train_items[uid] = [iid] |
| 52 | else: |
| 53 | self.train_items[uid].append(iid) |
| 54 | test_uid, test_iid = test_mat.row, test_mat.col |
| 55 | for i in range(len(test_uid)): |
| 56 | uid = test_uid[i] |
| 57 | iid = test_iid[i] |
| 58 | if uid not in self.test_set: |
| 59 | self.test_set[uid] = [iid] |
| 60 | else: |
| 61 | self.test_set[uid].append(iid) |
| 62 | |
| 63 | def get_adj_mat(self): |
| 64 | adj_mat = self.create_adj_mat() |