RippleNet代码精读

main.py代码

整体来说，这段代码的功能是：

• 解析命令行参数，其中包括数据集名称、模型参数等。
• 加载数据。
• 训练模型。

import argparse
import numpy as np
from data_loader import load_data  # 导入数据加载函数
from train import train  # 导入训练函数

np.random.seed(555)  # 设置随机种子，确保结果可重复

parser = argparse.ArgumentParser()  # 创建参数解析器

# 添加命令行参数
parser.add_argument('--dataset', type=str, default='movie', help='which dataset to use')  # 数据集名称，默认为'movie'
parser.add_argument('--dim', type=int, default=16, help='dimension of entity and relation embeddings')  # 实体和关系嵌入的维度，默认为16
parser.add_argument('--n_hop', type=int, default=2, help='maximum hops')  # 最大传播步数，默认为2
parser.add_argument('--kge_weight', type=float, default=0.01, help='weight of the KGE term')  # KGE项的权重，默认为0.01
parser.add_argument('--l2_weight', type=float, default=1e-7, help='weight of the l2 regularization term')  # L2正则化项的权重，默认为1e-7
parser.add_argument('--lr', type=float, default=0.02, help='learning rate')  # 学习率，默认为0.02
parser.add_argument('--batch_size', type=int, default=1024, help='batch size')  # 批大小，默认为1024
parser.add_argument('--n_epoch', type=int, default=10, help='the number of epochs')  # 训练的轮数，默认为10
parser.add_argument('--n_memory', type=int, default=32, help='size of ripple set for each hop')  # 每个传播步的ripple set大小，默认为32
parser.add_argument('--item_update_mode', type=str, default='plus_transform',
                    help='how to update item at the end of each hop')  # 每步传播结束后如何更新项目，默认为'plus_transform'
parser.add_argument('--using_all_hops', type=bool, default=True,
                    help='whether using outputs of all hops or just the last hop when making prediction')  # 是否使用所有传播步的输出，或仅使用最后一步，默认为True

# 默认设置为Book-Crossing的参数
'''
parser.add_argument('--dataset', type=str, default='book', help='which dataset to use')
parser.add_argument('--dim', type=int, default=4, help='dimension of entity and relation embeddings')
parser.add_argument('--n_hop', type=int, default=2, help='maximum hops')
parser.add_argument('--kge_weight', type=float, default=1e-2, help='weight of the KGE term')
parser.add_argument('--l2_weight', type=float, default=1e-5, help='weight of the l2 regularization term')
parser.add_argument('--lr', type=float, default=1e-3, help='learning rate')
parser.add_argument('--batch_size', type=int, default=1024, help='batch size')
parser.add_argument('--n_epoch', type=int, default=10, help='the number of epochs')
parser.add_argument('--n_memory', type=int, default=32, help='size of ripple set for each hop')
parser.add_argument('--item_update_mode', type=str, default='plus_transform',
                    help='how to update item at the end of each hop')
parser.add_argument('--using_all_hops', type=bool, default=True,
                    help='whether using outputs of all hops or just the last hop when making prediction')
'''

parser.add_argument('--use_cuda', type=bool, default=True, help='whether to use gpu')  # 是否使用GPU，默认为True
args = parser.parse_args()  # 解析命令行参数，并将其存储在args对象中

show_loss = False  # 是否显示损失，默认为False
data_info = load_data(args)  # 载入数据
train(args, data_info, show_loss)  # 训练模型

 train.py

这段代码实现了一个基于知识图谱的推荐系统模型的训练和评估过程。

1. import numpy as np: 导入NumPy库，用于数值计算。
2. import torch: 导入PyTorch库，用于构建和训练神经网络。
3. from model import RippleNet: 从自定义的模型文件中导入RippleNet模型。
4. def train(args, data_info, show_loss):: 定义了一个名为train的函数，用于模型训练。
   • 参数args包含了训练过程中的各种设置。
   • data_info包含了训练数据、验证数据、测试数据、实体数量、关系数量以及ripple set等信息。
   • show_loss用于控制是否显示损失值。
5. 在train函数内部：
   • 将数据信息解包赋值给各个变量。
   • 创建了RippleNet模型实例。
   • 如果args.use_cuda为True，则将模型转移到GPU上。
   • 使用Adam优化器来更新模型参数。
   • 在每个epoch中，对训练数据进行随机打乱，并进行训练。
   • 训练过程中，计算损失并进行反向传播优化。
   • 完成一个epoch的训练后，对模型进行评估，包括计算训练集、验证集和测试集的AUC和准确率，并输出评估结果。
6. def get_feed_dict(args, model, data, ripple_set, start, end):: 定义了一个函数用于生成喂入模型的数据字典。
   • 将训练数据切片，并提取所需的项目、标签、以及ripple set。
   • 如果使用GPU，则将数据转移到GPU上。
   • 返回喂入模型的数据字典。
7. def evaluation(args, model, data, ripple_set, batch_size):: 定义了一个评估函数，用于评估模型性能。
   • 对给定数据集进行评估，计算AUC和准确率。
   • 返回评估结果的均值。
8. 整体上，这段代码实现了模型的训练过程，并在每个epoch结束后评估模型性能。

import numpy as np
import torch

from model import RippleNet  # 导入RippleNet模型


def train(args, data_info, show_loss):
    # 解包data_info
    train_data = data_info[0]
    eval_data = data_info[1]
    test_data = data_info[2]
    n_entity = data_info[3]
    n_relation = data_info[4]
    ripple_set = data_info[5]

    # 初始化RippleNet模型
    model = RippleNet(args, n_entity, n_relation)
    if args.use_cuda:
        model.cuda()  # 如果可用，将模型移动到GPU上
    optimizer = torch.optim.Adam(
        filter(lambda p: p.requires_grad, model.parameters()),  # 过滤参数以进行优化
        args.lr,  # 学习率
    )

    for step in range(args.n_epoch):
        # 训练
        np.random.shuffle(train_data)  # 打乱训练数据
        start = 0
        while start < train_data.shape[0]:
            # 从模型获取返回字典
            return_dict = model(*get_feed_dict(args, model, train_data, ripple_set, start, start + args.batch_size))
            loss = return_dict["loss"]  # 从返回字典中提取损失

            optimizer.zero_grad()  # 清除梯度
            loss.backward()  # 反向传播
            optimizer.step()  # 优化器步骤

            start += args.batch_size  # 移动到下一批次
            if show_loss:
                print('%.1f%% %.4f' % (start / train_data.shape[0] * 100, loss.item()))  # 打印损失进度

        # 评估
        train_auc, train_acc = evaluation(args, model, train_data, ripple_set, args.batch_size)
        eval_auc, eval_acc = evaluation(args, model, eval_data, ripple_set, args.batch_size)
        test_auc, test_acc = evaluation(args, model, test_data, ripple_set, args.batch_size)

        # 打印评估指标
        print('epoch %d    train auc: %.4f  acc: %.4f    eval auc: %.4f  acc: %.4f    test auc: %.4f  acc: %.4f'
              % (step, train_auc, train_acc, eval_auc, eval_acc, test_auc, test_acc))


def get_feed_dict(args, model, data, ripple_set, start, end):
    # 为模型输入准备feed字典
    items = torch.LongTensor(data[start:end, 1])  # 物品ID
    labels = torch.LongTensor(data[start:end, 2])  # 标签
    memories_h, memories_r, memories_t = [], [], []  # 初始化记忆张量列表
    for i in range(args.n_hop):
        # 提取每个跳数的记忆
        memories_h.append(torch.LongTensor([ripple_set[user][i][0] for user in data[start:end, 0]]))
        memories_r.append(torch.LongTensor([ripple_set[user][i][1] for user in data[start:end, 0]]))
        memories_t.append(torch.LongTensor([ripple_set[user][i][2] for user in data[start:end, 0]]))
    if args.use_cuda:
        # 如果可用，将张量移到GPU
        items = items.cuda()
        labels = labels.cuda()
        memories_h = list(map(lambda x: x.cuda(), memories_h))
        memories_r = list(map(lambda x: x.cuda(), memories_r))
        memories_t = list(map(lambda x: x.cuda(), memories_t))
    return items, labels, memories_h, memories_r, memories_t


def evaluation(args, model, data, ripple_set, batch_size):
    start = 0
    auc_list = []
    acc_list = []
    model.eval()  # 将模型设置为评估模式
    while start < data.shape[0]:
        # 对每个批次进行评估
        auc, acc = model.evaluate(*get_feed_dict(args, model, data, ripple_set, start, start + batch_size))
        auc_list.append(auc)
        acc_list.append(acc)
        start += batch_size  # 移动到下一批次
    model.train()  # 将模型设置回训练模式
    return float(np.mean(auc_list)), float(np.mean(acc_list))

 model.py

这段代码定义了一个名为 RippleNet 的 PyTorch 模型类。

1. 初始化方法 (__init__):

   • 这个方法在创建类的实例时调用。它初始化了模型的各个组件，如实体嵌入 (entity_emb)、关系嵌入 (relation_emb) 和变换矩阵 (transform_matrix)。
   • 这里还定义了损失函数 (criterion)，使用了二元交叉熵损失 (BCELoss)。

2. 前向传播方法 (forward):

   • 这个方法定义了模型的前向传播过程。给定输入数据 (items, labels, memories_h, memories_r, memories_t)，它计算了模型的输出。
   • 首先，将输入的物品 (items) 通过实体嵌入层得到物品的嵌入表示。
   • 接着，对于每个阶段 (hop)，计算输出 (o_list)。
   • 最后，根据使用所有阶段还是最后一个阶段，预测了物品的评分 (scores)。

3. 损失计算方法 (_compute_loss):

   • 这个方法计算了模型的总损失，包括基本损失 (base_loss)、知识图嵌入损失 (kge_loss) 和 L2 正则化损失 (l2_loss)。

4. 键-地址机制 (_key_addressing):

   • 这个方法实现了模型中的键-地址机制，用于根据记忆和物品的当前嵌入来更新物品的嵌入。

5. 物品嵌入更新 (_update_item_embedding):

   • 这个方法根据不同的物品更新模式来更新物品的嵌入。

6. 评估方法 (evaluate):

   • 这个方法用于评估模型的性能，计算了给定输入数据的 AUC 和准确率 (accuracy)。

总体来说，这个代码定义了一个推荐模型 RippleNet，它利用知识图嵌入来增强推荐的效果，并可以通过评估方法来评估模型的性能。

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from sklearn.metrics import roc_auc_score


class RippleNet(nn.Module):
    def __init__(self, args, n_entity, n_relation):
        super(RippleNet, self).__init__()

        self._parse_args(args, n_entity, n_relation)

        # 实体嵌入层
        self.entity_emb = nn.Embedding(self.n_entity, self.dim)
        # 关系嵌入层
        self.relation_emb = nn.Embedding(self.n_relation, self.dim * self.dim)
        # 变换矩阵
        self.transform_matrix = nn.Linear(self.dim, self.dim, bias=False)
        # 损失函数
        self.criterion = nn.BCELoss()

    def _parse_args(self, args, n_entity, n_relation):
        self.n_entity = n_entity
        self.n_relation = n_relation
        self.dim = args.dim
        self.n_hop = args.n_hop
        self.kge_weight = args.kge_weight
        self.l2_weight = args.l2_weight
        self.lr = args.lr
        self.n_memory = args.n_memory
        self.item_update_mode = args.item_update_mode
        self.using_all_hops = args.using_all_hops

    def forward(
        self,
        items: torch.LongTensor,
        labels: torch.LongTensor,
        memories_h: list,
        memories_r: list,
        memories_t: list,
    ):
        # 物品嵌入
        item_embeddings = self.entity_emb(items)
        h_emb_list = []
        r_emb_list = []
        t_emb_list = []
        for i in range(self.n_hop):
            # 每个跳数的头实体嵌入
            h_emb_list.append(self.entity_emb(memories_h[i]))
            # 每个跳数的关系嵌入
            r_emb_list.append(
                self.relation_emb(memories_r[i]).view(
                    -1, self.n_memory, self.dim, self.dim
                )
            )
            # 每个跳数的尾实体嵌入
            t_emb_list.append(self.entity_emb(memories_t[i]))

        # 计算键-地址机制输出
        o_list, item_embeddings = self._key_addressing(
            h_emb_list, r_emb_list, t_emb_list, item_embeddings
        )
        # 预测物品评分
        scores = self.predict(item_embeddings, o_list)

        # 计算损失
        return_dict = self._compute_loss(
            scores, labels, h_emb_list, t_emb_list, r_emb_list
        )
        return_dict["scores"] = scores

        return return_dict

    def _compute_loss(self, scores, labels, h_emb_list, t_emb_list, r_emb_list):
        # 基本损失
        base_loss = self.criterion(scores, labels.float())

        # 知识图嵌入损失
        kge_loss = 0
        for hop in range(self.n_hop):
            # 计算知识图嵌入损失
            h_expanded = torch.unsqueeze(h_emb_list[hop], dim=2)
            t_expanded = torch.unsqueeze(t_emb_list[hop], dim=3)
            hRt = torch.squeeze(
                torch.matmul(torch.matmul(h_expanded, r_emb_list[hop]), t_expanded)
            )
            kge_loss += torch.sigmoid(hRt).mean()
        kge_loss = -self.kge_weight * kge_loss

        # L2 正则化损失
        l2_loss = 0
        for hop in range(self.n_hop):
            l2_loss += (h_emb_list[hop] * h_emb_list[hop]).sum()
            l2_loss += (t_emb_list[hop] * t_emb_list[hop]).sum()
            l2_loss += (r_emb_list[hop] * r_emb_list[hop]).sum()
        l2_loss = self.l2_weight * l2_loss

        # 总损失
        loss = base_loss + kge_loss + l2_loss
        return dict(base_loss=base_loss, kge_loss=kge_loss, l2_loss=l2_loss, loss=loss)

    def _key_addressing(self, h_emb_list, r_emb_list, t_emb_list, item_embeddings):
        o_list = []
        for hop in range(self.n_hop):
            # 头实体嵌入扩展维度
            h_expanded = torch.unsqueeze(h_emb_list[hop], dim=3)
            Rh = torch.squeeze(torch.matmul(r_emb_list[hop], h_expanded))
            # 物品嵌入扩展维度
            v = torch.unsqueeze(item_embeddings, dim=2)
            probs = torch.squeeze(torch.matmul(Rh, v))
            probs_normalized = F.softmax(probs, dim=1)
            probs_expanded = torch.unsqueeze(probs_normalized, dim=2)
            # 计算输出
            o = (t_emb_list[hop] * probs_expanded).sum(dim=1)

            # 更新物品嵌入
            item_embeddings = self._update_item_embedding(item_embeddings, o)
            o_list.append(o)
        return o_list, item_embeddings

    def _update_item_embedding(self, item_embeddings, o):
        # 根据不同的物品更新模式更新物品嵌入
        if self.item_update_mode == "replace":
            item_embeddings = o
        elif self.item_update_mode == "plus":
            item_embeddings = item_e

preprocess.py

这段代码是用于数据预处理的 Python 脚本。它的功能如下：

1. 定义了一些全局变量，包括文件名、分隔符和评分阈值。
2. 包含了三个函数：
   • read_item_index_to_entity_id_file(): 从文件中读取物品索引到实体 ID 的映射关系。
   • convert_rating(): 将原始评分数据转换为模型所需的格式，并保存为 ratings_final.txt 文件。
   • convert_kg(): 将知识图谱数据转换为模型所需的格式，并保存为 kg_final.txt 文件。
3. 在脚本的末尾，通过命令行参数指定数据集，然后依次调用上述函数进行数据预处理。

整体来说，这段代码的主要目的是将原始的评分数据和知识图谱数据转换为模型所需的格式，并保存到文件中，以便后续训练模型时使用。

import argparse
import numpy as np

# 定义不同数据集的文件名和分隔符
RATING_FILE_NAME = dict({'movie': 'ratings.dat', 'book': 'BX-Book-Ratings.csv', 'news': 'ratings.txt'})
SEP = dict({'movie': '::', 'book': ';', 'news': '\t'})
THRESHOLD = dict({'movie': 4, 'book': 0, 'news': 0})

# 读取物品索引到实体ID映射的文件
def read_item_index_to_entity_id_file():
    file = '../data/' + DATASET + '/item_index2entity_id_rehashed.txt'
    print('读取物品索引到实体ID文件: ' + file + ' ...')
    i = 0
    # 逐行读取文件内容
    for line in open(file, encoding='utf-8').readlines():
        # 获取物品索引和实体ID
        item_index = line.strip().split('\t')[0]
        satori_id = line.strip().split('\t')[1]
        # 将旧物品索引映射到新的连续索引
        item_index_old2new[item_index] = i
        # 将实体ID映射到新的连续索引
        entity_id2index[satori_id] = i
        i += 1

# 转换评分数据
def convert_rating():
    file = '../data/' + DATASET + '/' + RATING_FILE_NAME[DATASET]

    print('读取评分文件 ...')
    item_set = set(item_index_old2new.values())
    user_pos_ratings = dict()
    user_neg_ratings = dict()

    # 逐行读取评分文件，跳过表头
    for line in open(file, encoding='utf-8').readlines()[1:]:
        array = line.strip().split(SEP[DATASET])

        # 对书籍数据集，去除前后引号
        if DATASET == 'book':
            array = list(map(lambda x: x[1:-1], array))

        item_index_old = array[1]
        if item_index_old not in item_index_old2new:  # 跳过不在最终物品集中的物品
            continue
        item_index = item_index_old2new[item_index_old]

        user_index_old = int(array[0])
        rating = float(array[2])

        # 根据阈值将评分分类为正面或负面
        if rating >= THRESHOLD[DATASET]:
            if user_index_old not in user_pos_ratings:
                user_pos_ratings[user_index_old] = set()
            user_pos_ratings[user_index_old].add(item_index)
        else:
            if user_index_old not in user_neg_ratings:
                user_neg_ratings[user_index_old] = set()
            user_neg_ratings[user_index_old].add(item_index)

    print('转换评分文件 ...')
    writer = open('../data/' + DATASET + '/ratings_final.txt', 'w', encoding='utf-8')
    user_cnt = 0
    user_index_old2new = dict()
    for user_index_old, pos_item_set in user_pos_ratings.items():
        if user_index_old not in user_index_old2new:
            user_index_old2new[user_index_old] = user_cnt
            user_cnt += 1
        user_index = user_index_old2new[user_index_old]

        # 写入正面评分
        for item in pos_item_set:
            writer.write('%d\t%d\t1\n' % (user_index, item))
        
        # 从未观看物品集合中选择一些写入负面评分
        unwatched_set = item_set - pos_item_set
        if user_index_old in user_neg_ratings:
            unwatched_set -= user_neg_ratings[user_index_old]
        # 随机选择和正面评分数量相同的未观看物品
        for item in np.random.choice(list(unwatched_set), size=len(pos_item_set), replace=False):
            writer.write('%d\t%d\t0\n' % (user_index, item))
    writer.close()
    print('用户数量: %d' % user_cnt)
    print('物品数量: %d' % len(item_set))

# 转换知识图谱数据
def convert_kg():
    print('转换知识图谱文件 ...')
    entity_cnt = len(entity_id2index)
    relation_cnt = 0

    writer = open('../data/' + DATASET + '/kg_final.txt', 'w', encoding='utf-8')

    files = []
    if DATASET == 'movie':
        files.append(open('../data/' + DATASET + '/kg_part1_rehashed.txt', encoding='utf-8'))
        files.append(open('../data/' + DATASET + '/kg_part2_rehashed.txt', encoding='utf-8'))
    else:
        files.append(open('../data/' + DATASET + '/kg_rehashed.txt', encoding='utf-8'))

    # 逐行读取知识图谱文件
    for file in files:
        for line in file:
            array = line.strip().split('\t')
            head_old = array[0]
            relation_old = array[1]
            tail_old = array[2]

            # 将旧的头实体ID映射到新的连续索引
            if head_old not in entity_id2index:
                entity_id2index[head_old] = entity_cnt
                entity_cnt += 1
            head = entity_id2index[head_old]

            # 将旧的尾实体ID映射到新的连续索引
            if tail_old not in entity_id2index:
                entity_id2index[tail_old] = entity_cnt
                entity_cnt += 1
            tail = entity_id2index[tail_old]

            # 将旧的关系ID映射到新的连续索引
            if relation_old not in relation_id2index:
                relation_id2index[relation_old] = relation_cnt
                relation_cnt += 1
            relation = relation_id2index[relation_old]

            # 写入新的头实体ID、关系ID、尾实体ID
            writer.write('%d\t%d\t%d\n' % (head, relation, tail))

    writer.close()
    print('实体数量（包含物品）: %d' % entity_cnt)
    print('关系数量: %d' % relation_cnt)

if __name__ == '__main__':
    np.random.seed(555)

    parser = argparse.ArgumentParser()
    parser.add_argument('-d', '--dataset', type=str, default='movie', help='选择要预处理的数据集')
    args = parser.parse_args()
    DATASET = args.dataset

    entity_id2index = dict()
    relation_id2index = dict()
    item_index_old2new = dict()

    read_item_index_to_entity_id_file()
    convert_rating()
    convert_kg()

    print('完成')

data_loader.py

这段代码是用于处理和预处理推荐系统中的用户评分数据和知识图谱数据，以便构建用于训练推荐系统模型的数据集。以下是每个函数的详细功能解释：

load_data(args)

• 功能：加载并处理数据，包括用户评分数据和知识图谱数据。
• 参数：args，包含数据集和其他参数的命令行参数对象。
• 返回值：训练数据、验证数据、测试数据、实体数量、关系数量、Ripple Set。

load_rating(args)

• 功能：读取评分文件，并拆分成训练集、验证集和测试集。
• 参数：args，包含数据集和其他参数的命令行参数对象。
• 返回值：训练数据、验证数据、测试数据、用户历史记录字典。

dataset_split(rating_np)

• 功能：将评分数据拆分为训练集、验证集和测试集，并构建用户历史记录字典。
• 参数：rating_np，包含评分数据的NumPy数组。
• 返回值：训练数据、验证数据、测试数据、用户历史记录字典。

load_kg(args)

• 功能：读取知识图谱文件，并构建知识图谱。
• 参数：args，包含数据集和其他参数的命令行参数对象。
• 返回值：实体数量、关系数量、知识图谱字典。

construct_kg(kg_np)

• 功能：构建知识图谱，将知识图谱数据转换为字典格式。
• 参数：kg_np，包含知识图谱数据的NumPy数组。
• 返回值：知识图谱字典。

get_ripple_set(args, kg, user_history_dict)

• 功能：根据用户历史记录构建Ripple Set，这是用于推荐系统模型的一种扩展用户历史的方法。
• 参数：
  • args，包含数据集和其他参数的命令行参数对象。
  • kg，知识图谱字典。
  • user_history_dict，用户历史记录字典。
• 返回值：Ripple Set字典。

函数的工作流程

1. 加载数据：

   • 调用load_data函数，该函数依次调用load_rating和load_kg来加载评分数据和知识图谱数据。
   • 读取评分文件并进行数据拆分。
   • 读取知识图谱文件并构建知识图谱。
   • 根据用户历史记录和知识图谱构建Ripple Set。

2. 读取和拆分评分数据：

   • load_rating函数读取评分文件，并将其拆分为训练集、验证集和测试集。
   • dataset_split函数根据一定比例（6:2:2）拆分数据，并生成用户历史记录字典。

3. 读取和构建知识图谱：

   • load_kg函数读取知识图谱文件，并调用construct_kg函数将其构建为字典格式。

4. 构建Ripple Set：

   • get_ripple_set函数根据用户的历史记录和知识图谱，为每个用户构建多跳Ripple Set，用于扩展用户的历史记录。

这个过程旨在为推荐系统模型提供预处理后的数据，便于模型的训练和评估。

import argparse
import collections
import os
import numpy as np

# 加载并处理数据，包括用户评分数据和知识图谱数据
def load_data(args):
    train_data, eval_data, test_data, user_history_dict = load_rating(args)
    n_entity, n_relation, kg = load_kg(args)
    ripple_set = get_ripple_set(args, kg, user_history_dict)
    return train_data, eval_data, test_data, n_entity, n_relation, ripple_set

# 读取评分文件，并拆分成训练集、验证集和测试集
def load_rating(args):
    print('读取评分文件 ...')

    # 定义评分文件路径
    rating_file = '../data/' + args.dataset + '/ratings_final'
    # 如果有.npy文件，则直接加载，否则加载.txt文件并保存为.npy文件
    if os.path.exists(rating_file + '.npy'):
        rating_np = np.load(rating_file + '.npy')
    else:
        rating_np = np.loadtxt(rating_file + '.txt', dtype=np.int32)
        np.save(rating_file + '.npy', rating_np)

    return dataset_split(rating_np)

# 将评分数据拆分为训练集、验证集和测试集，并构建用户历史记录字典
def dataset_split(rating_np):
    print('拆分数据集 ...')

    # 设定验证集和测试集比例
    eval_ratio = 0.2
    test_ratio = 0.2
    n_ratings = rating_np.shape[0]

    # 随机选择验证集和测试集的索引
    eval_indices = np.random.choice(n_ratings, size=int(n_ratings * eval_ratio), replace=False)
    left = set(range(n_ratings)) - set(eval_indices)
    test_indices = np.random.choice(list(left), size=int(n_ratings * test_ratio), replace=False)
    train_indices = list(left - set(test_indices))

    # 遍历训练数据，仅保留有正面评分的用户
    user_history_dict = dict()
    for i in train_indices:
        user = rating_np[i][0]
        item = rating_np[i][1]
        rating = rating_np[i][2]
        if rating == 1:
            if user not in user_history_dict:
                user_history_dict[user] = []
            user_history_dict[user].append(item)

    # 仅保留有历史记录的用户的数据
    train_indices = [i for i in train_indices if rating_np[i][0] in user_history_dict]
    eval_indices = [i for i in eval_indices if rating_np[i][0] in user_history_dict]
    test_indices = [i for i in test_indices if rating_np[i][0] in user_history_dict]

    train_data = rating_np[train_indices]
    eval_data = rating_np[eval_indices]
    test_data = rating_np[test_indices]

    return train_data, eval_data, test_data, user_history_dict

# 读取知识图谱文件，并构建知识图谱
def load_kg(args):
    print('读取知识图谱文件 ...')

    # 定义知识图谱文件路径
    kg_file = '../data/' + args.dataset + '/kg_final'
    # 如果有.npy文件，则直接加载，否则加载.txt文件并保存为.npy文件
    if os.path.exists(kg_file + '.npy'):
        kg_np = np.load(kg_file + '.npy')
    else:
        kg_np = np.loadtxt(kg_file + '.txt', dtype=np.int32)
        np.save(kg_file + '.npy', kg_np)

    # 计算实体和关系的数量
    n_entity = len(set(kg_np[:, 0]) | set(kg_np[:, 2]))
    n_relation = len(set(kg_np[:, 1]))

    # 构建知识图谱
    kg = construct_kg(kg_np)

    return n_entity, n_relation, kg

# 构建知识图谱，将知识图谱数据转换为字典格式
def construct_kg(kg_np):
    print('构建知识图谱 ...')
    kg = collections.defaultdict(list)
    for head, relation, tail in kg_np:
        kg[head].append((tail, relation))
    return kg

# 根据用户历史记录构建Ripple Set
def get_ripple_set(args, kg, user_history_dict):
    print('构建Ripple Set ...')

    # 用户 -> [(第0跳的头实体, 第0跳的关系, 第0跳的尾实体), (第1跳的头实体, 第1跳的关系, 第1跳的尾实体), ...]
    ripple_set = collections.defaultdict(list)

    for user in user_history_dict:
        for h in range(args.n_hop):
            memories_h = []
            memories_r = []
            memories_t = []

            # 第0跳：初始实体为用户历史记录中的物品
            if h == 0:
                tails_of_last_hop = user_history_dict[user]
            else:
                tails_of_last_hop = ripple_set[user][-1][2]

            # 根据上一跳的尾实体，查找当前跳的关系和尾实体
            for entity in tails_of_last_hop:
                for tail_and_relation in kg[entity]:
                    memories_h.append(entity)
                    memories_r.append(tail_and_relation[1])
                    memories_t.append(tail_and_relation[0])

            # 如果当前Ripple Set为空，则复制上一跳的Ripple Set
            if len(memories_h) == 0:
                ripple_set[user].append(ripple_set[user][-1])
            else:
                # 为每个用户采样固定大小的1跳记忆
                replace = len(memories_h) < args.n_memory
                indices = np.random.choice(len(memories_h), size=args.n_memory, replace=replace)
                memories_h = [memories_h[i] for i in indices]
                memories_r = [memories_r[i] for i in indices]
                memories_t = [memories_t[i] for i in indices]
                ripple_set[user].append((memories_h, memories_r, memories_t))

    return ripple_set