Function main

datasets/Beauty/data_augment_candi_gen.py:492–680

(args)

def main(args):
    max_seq_length = args.max_seq_length
    max_predictions_per_seq = args.max_predictions_per_seq
    masked_lm_prob = args.masked_lm_prob
    mask_prob = args.mask_prob
    dupe_factor = args.dupe_factor
    prop_sliding_window = args.prop_sliding_window
    pool_size = args.pool_size

    output_dir = args.data_dir
    dataset_name = args.dataset_name

    if not os.path.isdir(output_dir):
        print(output_dir + ' is not exist')
        print(os.getcwd())
        exit(1)
    os.mkdir(output_dir + 'train')
    os.mkdir(output_dir + 'test')

    dataset = data_partition(output_dir + dataset_name + '.txt')
    [user_train, user_valid, user_test, usernum, itemnum] = dataset
    cc = 0.0
    max_len = 0
    min_len = 100000
    for u in user_train:
        cc += len(user_train[u])
        max_len = max(len(user_train[u]), max_len)
        min_len = min(len(user_train[u]), min_len)

    print('average sequence length: %.2f' % (cc / len(user_train)))
    print('max:{}, min:{}'.format(max_len, min_len))

    print('len_train:{}, len_valid:{}, len_test:{}, usernum:{}, itemnum:{}'.
          format(
              len(user_train),
              len(user_valid), len(user_test), usernum, itemnum))

    for idx, u in enumerate(user_train):
        if idx <= 1:
            print(user_train[u])
            print(user_valid[u])
            print(user_test[u])

    # put validate into train
    for u in user_train:
        if u in user_valid:
            user_train[u].extend(user_valid[u])

    # get the max index of the data
    user_train_data = {
        'user_' + str(k): ['item_' + str(item) for item in v]
        for k, v in user_train.items() if len(v) > 0
    }
    user_test_data = {
        'user_' + str(u):
        ['item_' + str(item) for item in (user_train[u] + user_test[u])]
        for u in user_train if len(user_train[u]) > 0 and len(user_test[u]) > 0
    }