KGCN---pytorch代码（1）---data_loader
代码：
import pandas as pd
import numpy as np
from sklearn.preprocessing import LabelEncoder
from sklearn.model_selection import train_test_split
import random
class DataLoader:
'''
Data Loader class which makes dataset for training / knowledge graph dictionary
'''
def __init__(self, data):
self.cfg = {
'movie': {
'item2id_path': 'data/movie/item_index2entity_id.txt',
'kg_path': 'data/movie/kg.txt',
'rating_path': 'data/movie/ratings.csv',
'rating_sep': ',',
'threshold': 4.0
},
'music': {
'item2id_path': 'data/music/item_index2entity_id.txt',
'kg_path': 'data/music/kg.txt',
'rating_path': 'data/music/user_artists.dat',
'rating_sep': '\t',
'threshold': 0.0
}
}
self.data = data
df_item2id = pd.read_csv(self.cfg[data]['item2id_path'], sep='\t', header=None, names=['item','id'])
df_kg = pd.read_csv(self.cfg[data]['kg_path'], sep='\t', header=None, names=['head','relation','tail'])
df_rating = pd.read_csv(self.cfg[data]['rating_path'], sep=self.cfg[data]['rating_sep'], names=['userID', 'itemID', 'rating'], skiprows=1)
# df_rating['itemID'] and df_item2id['item'] both represents old entity ID
df_rating = df_rating[df_rating['itemID'].isin(df_item2id['item'])] #确保处理的评分数据只涉及已知的项目，即那些可以在知识图谱中找到对应实体的项目
df_rating.reset_index(inplace=True, drop=True) #在过滤评分数据之后，df_rating 的索引可能会变得不连续，即某些行被移除后，剩余行的索引不再是一个连续的整数序列。
self.df_item2id = df_item2id
self.df_kg = df_kg
self.df_rating = df_rating
self.user_encoder = LabelEncoder()
self.entity_encoder = LabelEncoder()
self.relation_encoder = LabelEncoder()
self._encoding()
def _encoding(self):
'''
Fit each label encoder and encode knowledge graph
'''
self.user_encoder.fit(self.df_rating['userID'])
# df_item2id['id'] and df_kg[['head', 'tail']] represents new entity ID
self.entity_encoder.fit(pd.concat([self.df_item2id['id'], self.df_kg['head'], self.df_kg['tail']])) #来自 df_item2id 的 id 列和知识图谱 (df_kg) 中的 head 和 tail 列中的实体ID
self.relation_encoder.fit(self.df_kg['relation']) #学习 df_kg 数据帧中 relation 列的所有唯一关系，并为它们分配一个唯一的整数标签
# encode df_kg #transform 方法将知识图谱中的每个原始标签替换为对应的整数编码
self.df_kg['head'] = self.entity_encoder.transform(self.df_kg['head'])
self.df_kg['tail'] = self.entity_encoder.transform(self.df_kg['tail'])
self.df_kg['relation'] = self.relation_encoder.transform(self.df_kg['relation'])
def _build_dataset(self): #准备训练集
'''
Build dataset for training (rating data)
It contains negative sampling process
'''
print('Build dataset dataframe ...', end=' ')
# df_rating update
df_dataset = pd.DataFrame()
df_dataset['userID'] = self.user_encoder.transform(self.df_rating['userID']) #将 df_rating 中的 userID 转换为数值编码
# update to new id
item2id_dict = dict(zip(self.df_item2id['item'], self.df_item2id['id']))
self.df_rating['itemID'] = self.df_rating['itemID'].apply(lambda x: item2id_dict[x]) #更新 df_rating 中的 itemID，使其从旧ID映射到新的实体ID。
df_dataset['itemID'] = self.entity_encoder.transform(self.df_rating['itemID']) #对更新后的 itemID 进行数值编码，将结果存储到 df_dataset['itemID']
df_dataset['label'] = self.df_rating['rating'].apply(lambda x: 0 if x < self.cfg[self.data]['threshold'] else 1) #将 df_rating['rating'] 转换为二分类标签
# negative sampling
df_dataset = df_dataset[df_dataset['label']==1] #为了平衡正负样本的数量，从数据集中只选择正样本（label==1），然后对每个用户进行负采样。负采样的目标是为每个正样本随机选择一个未与该用户互动的项目作为负样本。
# df_dataset requires columns to have new entity ID
full_item_set = set(range(len(self.entity_encoder.classes_)))
user_list = []
item_list = []
label_list = []
for user, group in df_dataset.groupby(['userID']): #按照 userID 对数据集进行分组
item_set = set(group['itemID']) #该用户已经有正样本的项目ID集合
negative_set = full_item_set - item_set #得到该用户未互动的项目ID集合
negative_sampled = random.sample(negative_set, len(item_set)) #随机选择与正样本数量相等的项目ID作为负样本
user_list.extend([user] * len(negative_sampled))
item_list.extend(negative_sampled)
label_list.extend([0] * len(negative_sampled))
negative = pd.DataFrame({'userID': user_list, 'itemID': item_list, 'label': label_list})
df_dataset = pd.concat([df_dataset, negative])
#随机打乱数据集+重置索引
df_dataset = df_dataset.sample(frac=1, replace=False, random_state=999)
df_dataset.reset_index(inplace=True, drop=True)
print('Done')
return df_dataset
def _construct_kg(self):
'''
Construct knowledge graph
Knowledge graph is dictionary form
'head': [(relation, tail), ...]
'''
print('Construct knowledge graph ...', end=' ')
kg = dict()
for i in range(len(self.df_kg)):
head = self.df_kg.iloc[i]['head']
relation = self.df_kg.iloc[i]['relation']
tail = self.df_kg.iloc[i]['tail']
if head in kg:
kg[head].append((relation, tail))
else:
kg[head] = [(relation, tail)]
if tail in kg:
kg[tail].append((relation, head))
else:
kg[tail] = [(relation, head)]
print('Done')
return kg
def load_dataset(self):
return self._build_dataset()
def load_kg(self):
return self._construct_kg()
def get_encoders(self): #获取在数据预处理过程中创建的三个主要编码器
return (self.user_encoder, self.entity_encoder, self.relation_encoder)
def get_num(self): #获取每个编码器编码类别数量的方式
return (len(self.user_encoder.classes_), len(self.entity_encoder.classes_), len(self.relation_encoder.classes_))
DataLoader类：
__init__：
1.self.cfg
两种数据集的路径（movie/music）
2.self.data
3.df_item2id / df_kg / df_rating
- df_item2id：把item_index2entity_id.txt读入，['item','id']
- df_kg：把data/music/kg.txt读入，['head','relation','tail']
- df_rating：把data/music/user_artists.dat读入，['userID', 'itemID', 'rating']，只保留df_item2id中存在的itemID行，并重置索引。
4.对知识图谱中的head,relation,tail进行编码
_build_dataset：（准备训练集）
1.df_dataset：['userID','itemID','label']
将 df_rating 中的 userID 转换为数值编码
更新 df_rating 中的 itemID，使其从旧ID映射到新的实体ID
将 df_rating['rating'] 根据阈值转换为二分类标签
2.df_dataset负采样
为了平衡正负样本的数量，从数据集中只选择正样本（label==1），然后对每个用户进行负采样。负采样的目标是为每个正样本随机选择一个未与该用户互动的项目作为负样本。
3.随机打乱数据集+重置索引
_construct_kg：
将知识图谱转变为字典形式：'head': [(relation1, tail1),(relation2,tail2), ...]
load_dataset：
调用_build_dataset函数，形成训练集
load_kg：
调用_construct_kg函数，重整知识图谱的格式
get_encoders：
返回三个主要编码器：user_encoder，entity_encoder，relation_encoder
get_num：
获取每个编码器编码类别数量