import os from pathlib import Path p = Path(__file__).parents[1] OUTPUT_DIR_TRAIN=os.path.abspath(os.path.join(p, '..', 'data/raw/ml-1m/train.dat')) OUTPUT_DIR_TEST=os.path.abspath(os.path.join(p, '..', 'data/raw/ml-1m/test.dat')) ROOT_DIR=os.path.abspath(os.path.join(p, '..', 'data/raw/ml-1m/ratings.dat')) NUM_USERS=6040 NUM_TEST_RATINGS=10 def count_rating_per_user(): rating_per_user={} for line in open(ROOT_DIR): line=line.split('::') user_nr=int(line[0]) if user_nr in rating_per_user: rating_per_user[user_nr]+=1 else: rating_per_user[user_nr]=1 return rating_per_user def train_test_split(): user_rating=count_rating_per_user() test_counter=0 next_user=1 train_writer=open(OUTPUT_DIR_TRAIN, 'w') test_writer=open(OUTPUT_DIR_TEST, 'w') for line in open(ROOT_DIR): splitted_line=line.split('::') user_nr=int(splitted_line[0]) if user_rating[user_nr]<=NUM_TEST_RATINGS*2: next_user+=1 continue try: if user_nr==next_user: write_test_samples=True next_user+=1 if write_test_samples==True: test_writer.write(line) test_counter+=1 if test_counter>=NUM_TEST_RATINGS: test_counter=0 write_test_samples=False else: train_writer.write(line) except KeyError: print('Key not found') continue if __name__ == "__main__": train_test_split()
这段代码用于将MovieLens 1M数据集的评分数据划分为训练集和测试集。
- 首先，使用Path库获取当前文件的父级目录，也就是项目根目录。
- 接着，定义输出训练集和测试集文件的路径。
- 然后，定义数据集中用户的数量和每个用户的测试评分数量。
- 下一步是定义一个函数count_rating_per_user()，该函数将读取ratings.dat文件的每一行，拆分为user_id、movie_id、rating和timestamp，并计算每个用户的评分数。该函数返回一个字典，其中每个键表示用户ID，每个值表示该用户的评分数。
- train_test_split()函数使用count_rating_per_user()函数获取每个用户的评分数量。该函数读取文件中的每一行，拆分为user_id、movie_id、rating和timestamp，并通过将其user_id与每个用户的评分数量进行比较来确定该行是否写入测试集。如果该行对应的user_id在当前的用户评分数量计数器中具有的评分数小于或等于 NUM_TEST_RATINGS * 2，则该行被跳过。否则，该行将被写入测试集中，并且当前用户的评分数量计数器将递增。如果计数器的值达到 NUM_TEST_RATINGS，则将write_test_samples标记设置为False，以停止将数据写入测试文件，并开始将数据写入训练文件。
这段代码是为了将ml-1m数据集的评分数据ratings.dat划分为训练集和测试集，并将它们保存为两个文件：train.dat和test.dat。具体来说，它完成以下几个任务：
- 计算每个用户给了多少次评分，以便在划分数据集时可以控制每个用户的评分数。
- 按照每个用户的评分数来划分训练集和测试集，每个用户的前
NUM_TEST_RATINGS个评分划分为测试集，其余评分划分为训练集。 - 将划分后的训练集和测试集保存为
train.dat和test.dat两个文件。
在这段代码中，首先使用了pathlib和os两个库，用于处理文件路径。接下来，使用count_rating_per_user()函数计算每个用户的评分数。这个函数使用了ratings.dat文件，并将其按行读入。每行的第一项是用户编号，所以可以统计每个用户给出了多少次评分。
接下来，train_test_split()函数读取 ratings.dat 文件，并且按照之前的规则将数据划分为训练集和测试集。对于每个用户，如果其评分数小于等于 NUM_TEST_RATINGS 的两倍，则将所有评分都划分到训练集。否则，将其前 NUM_TEST_RATINGS 个评分划分到测试集，其余评分划分到训练集。通过一个 test_counter 变量来跟踪每个用户的测试集评分数是否达到了 NUM_TEST_RATINGS 的限制，从而保证测试集的大小。在写入数据时，使用两个文件句柄 train_writer 和 test_writer，分别写入训练集和测试集。
最后，在 main 函数中调用 train_test_split() 函数，生成两个文件 train.dat 和 test.dat，它们位于 ml-1m 数据集的 raw 文件夹中。
