#!/usr/bin/env python
# coding: utf-8

# 
# # 基于 BipartiteGraphSage 的二部图无监督学习
# 
# 
# 二部图是电子商务推荐场景中很常见的一种图，GraphScope提供了针对二部图处理学习任务的模型。本次教程，我们将会展示GraphScope如何使用BipartiteGraphSage算法在二部图上训练一个无监督学习模型。
# 
# 本次教程的学习任务是链接预测，通过计算在图中用户顶点和商品顶点之间存在边的概率来预测链接。
# 
# 在这一任务中，我们使用GraphScope内置的BipartiteGraphSage算法在 [U2I](http://graph-learn-dataset.oss-cn-zhangjiakou.aliyuncs.com/u2i.zip) 数据集上训练一个模型，这一训练模型可以用来预测用户顶点和商品顶点之间的链接。这一任务可以被看作在一个异构链接网络上的无监督训练任务。
# 
# 在这一任务中，BipartiteGraphSage算法会将图中的结构信息和属性信息压缩为每个节点上的低维嵌入向量，这些嵌入和表征可以进一步用来预测节点间的链接。
# 
# 这一教程将会分为以下几个步骤：
# 
# - 启动GraphScope的学习引擎，并将图关联到引擎上
# - 使用内置的GCN模型定义训练过程，并定义相关的超参
# - 开始训练

# In[ ]:


# Install graphscope package if you are NOT in the Playground

get_ipython().system('pip3 install graphscope')
get_ipython().system('pip3 uninstall -y importlib_metadata  # Address an module conflict issue on colab.google. Remove this line if you are not on colab.')


# In[ ]:


# Import the graphscope module.

import graphscope

graphscope.set_option(show_log=False)  # enable logging


# In[ ]:


# Load u2i dataset

from graphscope.dataset import load_u2i

graph = load_u2i()


# 
# ## Launch learning engine
# 
# 然后，我们需要定义一个特征列表用于图的训练。训练特征集合必须从点的属性集合中选取。在这个例子中，我们选择了 "feature" 属性作为训练特征集，这一特征集也是 U2I 数据中用户顶点和商品顶点的特征集。
# 
# 借助定义的特征列表，接下来，我们使用 [graphlearn](https://graphscope.io/docs/reference/session.html#graphscope.Session.graphlearn) 方法来开启一个学习引擎。
# 
# 在这个例子中，我们在 "graphlearn" 方法中，指定在数据中 "u" 类型的顶点和 "i" 类型顶点和 "u-i" 类型边上进行模型训练。
# 

# In[ ]:


# launch a learning engine.

lg = graphscope.graphlearn(
    graph,
    nodes=[("u", ["feature"]), ("i", ["feature"])],
    edges=[(("u", "u-i", "i"), ["weight"]), (("i", "u-i_reverse", "u"), ["weight"])],
)


# 
# 
# 这里我们使用内置的`BipartiteGraphSage`模型定义训练过程。你可以在 [Graph Learning Model](https://graphscope.io/docs/learning_engine.html#data-model) 获取更多内置学习模型的信息。
# 
# 
# 在本次示例中，我们使用 tensorflow 作为神经网络后端训练器。

# In[ ]:


import numpy as np
import tensorflow as tf
from graphscope.learning.examples import BipartiteGraphSage
from graphscope.learning.graphlearn.python.model.tf.optimizer import get_tf_optimizer
from graphscope.learning.graphlearn.python.model.tf.trainer import LocalTFTrainer


# Unsupervised GraphSage.
def train(config, graph):
    def model_fn():
        return BipartiteGraphSage(
            graph,
            config["batch_size"],
            config["hidden_dim"],
            config["output_dim"],
            config["hops_num"],
            config["u_neighs_num"],
            config["i_neighs_num"],
            u_features_num=config["u_features_num"],
            u_categorical_attrs_desc=config["u_categorical_attrs_desc"],
            i_features_num=config["i_features_num"],
            i_categorical_attrs_desc=config["i_categorical_attrs_desc"],
            neg_num=config["neg_num"],
            use_input_bn=config["use_input_bn"],
            act=config["act"],
            agg_type=config["agg_type"],
            need_dense=config["need_dense"],
            in_drop_rate=config["drop_out"],
            ps_hosts=config["ps_hosts"],
        )

    trainer = LocalTFTrainer(
        model_fn,
        epoch=config["epoch"],
        optimizer=get_tf_optimizer(
            config["learning_algo"], config["learning_rate"], config["weight_decay"]
        ),
    )

    trainer.train()
    u_embs = trainer.get_node_embedding("u")
    np.save("u_emb", u_embs)
    i_embs = trainer.get_node_embedding("i")
    np.save("i_emb", i_embs)


# Define hyperparameters
config = {
    "batch_size": 128,
    "hidden_dim": 128,
    "output_dim": 128,
    "u_features_num": 1,
    "u_categorical_attrs_desc": {"0": ["u_id", 10000, 64]},
    "i_features_num": 1,
    "i_categorical_attrs_desc": {"0": ["i_id", 10000, 64]},
    "hops_num": 1,
    "u_neighs_num": [10],
    "i_neighs_num": [10],
    "neg_num": 10,
    "learning_algo": "adam",
    "learning_rate": 0.001,
    "weight_decay": 0.0005,
    "epoch": 5,
    "use_input_bn": True,
    "act": tf.nn.leaky_relu,
    "agg_type": "gcn",
    "need_dense": True,
    "drop_out": 0.0,
    "ps_hosts": None,
}


# ## 执行训练过程
# 
# 
# 在定义完训练过程和超参后，现在我们可以使用学习引擎和定义的超参开始训练过程。

# In[ ]:


train(config, lg)


# In[ ]: