Sequential Recommendation for cold-start users with meta transitional learning
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

utils.py 4.9KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164
  1. import sys
  2. import copy
  3. import torch
  4. import random
  5. import numpy as np
  6. from collections import defaultdict, Counter
  7. from multiprocessing import Process, Queue
  8. # sampler for batch generation
  9. def random_neq(l, r, s,user_train):
  10. # t = np.random.randint(l, r)
  11. # while t in s:
  12. # t = np.random.randint(l, r)
  13. # return t
  14. user = random.choice(list(user_train.keys()))
  15. item = random.choice(user_train[user])
  16. while item in s:
  17. user = random.choice(list(user_train.keys()))
  18. item = random.choice(user_train[user])
  19. return item
  20. def trans_to_cuda(variable):
  21. if torch.cuda.is_available():
  22. return variable.cuda()
  23. else:
  24. return variable
  25. def trans_to_cpu(variable):
  26. if torch.cuda.is_available():
  27. return variable.cpu()
  28. else:
  29. return variable
  30. # train/val/test data generation
  31. def data_load(fname, num_sample):
  32. usernum = 0
  33. itemnum = 0
  34. user_train = defaultdict(list)
  35. # assume user/item index starting from 1
  36. f = open('/home/maheri/metaTL/data/%s/%s_train.csv' % (fname, fname), 'r')
  37. for line in f:
  38. u, i, t = line.rstrip().split('\t')
  39. u = int(u)
  40. i = int(i)
  41. usernum = max(u, usernum)
  42. itemnum = max(i, itemnum)
  43. user_train[u].append(i)
  44. f.close()
  45. # read in new users for testing
  46. user_input_test = {}
  47. user_input_valid = {}
  48. user_valid = {}
  49. user_test = {}
  50. User_test_new = defaultdict(list)
  51. f = open('/home/maheri/metaTL/data/%s/%s_test_new_user.csv' % (fname, fname), 'r')
  52. for line in f:
  53. u, i, t = line.rstrip().split('\t')
  54. u = int(u)
  55. i = int(i)
  56. User_test_new[u].append(i)
  57. f.close()
  58. for user in User_test_new:
  59. if len(User_test_new[user]) > num_sample:
  60. if random.random()<0.3:
  61. user_input_valid[user] = User_test_new[user][:num_sample]
  62. user_valid[user] = []
  63. user_valid[user].append(User_test_new[user][num_sample])
  64. else:
  65. user_input_test[user] = User_test_new[user][:num_sample]
  66. user_test[user] = []
  67. user_test[user].append(User_test_new[user][num_sample])
  68. return [user_train, usernum, itemnum, user_input_test, user_test, user_input_valid, user_valid]
  69. class DataLoader(object):
  70. def __init__(self, user_train, user_test, itemnum, parameter):
  71. self.curr_rel_idx = 0
  72. self.bs = parameter['batch_size']
  73. self.maxlen = parameter['K']
  74. self.valid_user = []
  75. for u in user_train:
  76. if len(user_train[u]) < self.maxlen or len(user_test[u]) < 1: continue
  77. self.valid_user.append(u)
  78. self.num_tris = len(self.valid_user)
  79. self.train = user_train
  80. self.test = user_test
  81. self.itemnum = itemnum
  82. def next_one_on_eval(self):
  83. if self.curr_tri_idx == self.num_tris:
  84. return "EOT", "EOT"
  85. u = self.valid_user[self.curr_tri_idx]
  86. self.curr_tri_idx += 1
  87. seq = np.zeros([self.maxlen], dtype=np.int32)
  88. pos = np.zeros([self.maxlen - 1], dtype=np.int32)
  89. # neg = np.zeros([self.maxlen*30 - 1], dtype=np.int32)
  90. neg = np.zeros([self.maxlen * 5], dtype=np.int32)
  91. idx = self.maxlen - 1
  92. ts = set(self.train[u])
  93. for i in reversed(self.train[u]):
  94. seq[idx] = i
  95. if idx > 0:
  96. pos[idx - 1] = i
  97. # if i != 0: neg[idx - 1] = random_neq(1, self.itemnum + 1, ts,self.train)
  98. idx -= 1
  99. if idx == -1: break
  100. for i in range(len(neg)):
  101. neg[i] = random_neq(1, self.itemnum + 1, ts,self.train)
  102. curr_rel = u
  103. support_triples, support_negative_triples, query_triples, negative_triples = [], [], [], []
  104. for idx in range(self.maxlen-1):
  105. support_triples.append([seq[idx],curr_rel,pos[idx]])
  106. # support_negative_triples.append([seq[idx],curr_rel,neg[idx]])
  107. # support_negative_triples.append([seq[-1],curr_rel,neg[idx]])
  108. # for idx in range(len(neg)):
  109. # support_negative_triples.append([seq[-1],curr_rel,neg[idx]])
  110. for j in range(5):
  111. for idx in range(self.maxlen-1):
  112. support_negative_triples.append([seq[idx], curr_rel, neg[j * self.maxlen + idx]])
  113. rated = ts
  114. rated.add(0)
  115. query_triples.append([seq[-1],curr_rel,self.test[u][0]])
  116. for _ in range(100):
  117. t = np.random.randint(1, self.itemnum + 1)
  118. while t in rated: t = np.random.randint(1, self.itemnum + 1)
  119. negative_triples.append([seq[-1],curr_rel,t])
  120. support_triples = [support_triples]
  121. support_negative_triples = [support_negative_triples]
  122. query_triples = [query_triples]
  123. negative_triples = [negative_triples]
  124. return [support_triples, support_negative_triples, query_triples, negative_triples], curr_rel