推荐算法领域的数据集总是从1开始,或是一串数字,每次处理的时候都要多一个user2id
的操作,实在是麻烦
干脆在使用数据集前就处理好,并保存下user2id
字典,方便后续查询
注意一下:
- sep要改成当前数据集的分隔符(’ ‘,’\t’)
- names改成当前数据集的列名
代码如下:
import numpy as np
import pandas as pd
from sklearn.preprocessing import LabelEncoder
def load_mat():
data_path = '../dataset/ml-100k/u.data'
df_data = pd.read_csv(data_path, header = None, sep='\t', names =['user_id', 'item_id', 'rating','time'])
lbe_user = LabelEncoder()
lbe_user.fit(df_data['user_id'].unique())
converted_user = lbe_user.transform(df_data['user_id'])
lbe_item = LabelEncoder() # 弄成离散的
lbe_item.fit(df_data['item_id'].unique())
converted_item = lbe_item.transform(df_data['item_id'])
converted_data = pd.DataFrame()
converted_data['user_id'] = converted_user
converted_data['item_id'] = converted_item
converted_data['rating'] = df_data['rating']
# 对应关系
user2id = {}
for user in lbe_user.classes_:
user2id.update({user: lbe_user.transform([user])[0]})
item2id = {}
for item in lbe_item.classes_:
item2id.update({item: lbe_item.transform([item])[0]})
return converted_data,user2id,item2id
def save(converted_data,user2id,item2id):
sort = converted_data.sort_values(by=['user_id'])
sort.to_csv('../dataset/ml-100k/data_converted', header=None, index=False)
np.save('../dataset/ml-100k/user2id.npy', user2id)
np.save('../dataset/ml-100k/item2id.npy', item2id)
print('successfully saved')
if __name__ == '__main__':
converted_data,user2id,item2id = load_mat()
save(converted_data,user2id,item2id)