text: import dgl
  import torch as th
  import torch.nn.functional as F
  from dgl.dataloading.negative_sampler import Uniform
  from . import BaseTask, register_task
  from ..dataset import build_dataset
  from ..utils import Evaluator
  
  
  @register_task("link_prediction")
  class LinkPrediction(BaseTask):
      r"""
      Link prediction tasks.
  
      Attributes
      -----------
      dataset : NodeClassificationDataset
          Task-related dataset
  
      evaluator : Evaluator
          offer evaluation metric
  
      Methods
      ---------
      get_graph :
          return a graph
      get_loss_fn :
          return a loss function
      """
  
      def __init__(self, args):
          super(LinkPrediction, self).__init__()
          self.n_dataset = args.dataset
          self.logger = args.logger
          self.dataset = build_dataset(args.dataset, 'link_prediction', logger=self.logger)
          # self.evaluator = Evaluator()
          self.train_hg, self.val_hg, self.test_hg = self.dataset.get_idx()
          if self.val_hg is None and self.test_hg is None:
              pass
          else:
              self.val_hg = self.val_hg.to(args.device)
              self.test_hg = self.test_hg.to(args.device)
          self.evaluator = Evaluator(args.seed)
          if not hasattr(args, 'score_fn'):
              self.ScorePredictor = HeteroDistMultPredictor()
              args.score_fn = 'distmult'
          elif args.score_fn == 'dot-product':
              self.ScorePredictor = HeteroDotProductPredictor()
          elif args.score_fn == 'distmult':
              self.ScorePredictor = HeteroDistMultPredictor()
          self.negative_sampler = Uniform(1)
  
          if args.dataset in ['wn18', 'FB15k', 'FB15k-237']:
              self.evaluation_metric = 'mrr'
          else:
              self.evaluation_metric = 'roc_auc'
  
          args.logger.info('[Init Task] The task: link prediction, the dataset: {}, the evaluation metric is {}, '
                           'the score function: {} '.format(self.n_dataset, self.evaluation_metric, args.score_fn))
  
      def get_out_ntype(self):
          ntype = []
          for l in self.dataset.target_link:
              ntype.append(l[0])
              ntype.append(l[2])
          return set(ntype)
  
      def get_graph(self):
          return self.dataset.g
  
      def get_loss_fn(self):
          return F.binary_cross_entropy_with_logits
  
      def get_evaluator(self, name):
          if name == 'acc':
              return self.evaluator.author_link_prediction
          elif name == 'mrr':
              return self.evaluator.mrr_
          elif name == 'academic_lp':
              return self.evaluator.author_link_prediction
          elif name == 'roc_auc':
              return self.evaluator.cal_roc_auc
  
      def evaluate(self, n_embedding, r_embedding=None, mode='test'):
          r"""
  
          Parameters
          ----------
          n_embedding: th.Tensor
              the embedding of nodes
          r_embedding: th.Tensor
              the embedding of relation types
          mode: str
              the evaluation mode, train/valid/test
          Returns
          -------
  
          """
          if self.evaluation_metric == 'acc':
              acc = self.evaluator.author_link_prediction
              return dict(Accuracy=acc)
          elif self.evaluation_metric == 'mrr':
              mrr = self.evaluator.mrr_(n_embedding['_N'], self.dict2emd(r_embedding), self.dataset.train_triplets,
                                        self.dataset.valid_triplets, self.dataset.test_triplets,
                                        hits=[1, 3, 10], eval_bz=100)
              return dict(MRR=mrr)
          elif self.evaluation_metric == 'roc_auc':
              if mode == 'test':
                  eval_hg = self.test_hg
              elif mode == 'valid':
                  eval_hg = self.val_hg
              else:
                  raise ValueError('Mode error, supported test and valid.')
              negative_graph = self.construct_negative_graph(eval_hg)
              p_score = th.sigmoid(self.ScorePredictor(eval_hg, n_embedding, r_embedding))
              n_score = th.sigmoid(self.ScorePredictor(negative_graph, n_embedding, r_embedding))
              p_label = th.ones(len(p_score), device=p_score.device)
              n_label = th.zeros(len(n_score), device=p_score.device)
              roc_auc = self.evaluator.cal_roc_auc(th.cat((p_label, n_label)).cpu(), th.cat((p_score, n_score)).cpu())
              loss = F.binary_cross_entropy_with_logits(th.cat((p_score, n_score)), th.cat((p_label, n_label)))
              return dict(roc_auc=roc_auc, loss=loss)
          else:
              return self.evaluator.link_prediction
  
      def downstream_evaluate(self, logits, evaluation_metric):
          if evaluation_metric == 'academic_lp':
              auc, macro_f1, micro_f1 = self.evaluator.author_link_prediction(logits, self.dataset.train_batch,
                                                                              self.dataset.test_batch)
              return dict(AUC=auc, Macro_f1=macro_f1, Mirco_f1=micro_f1)
  
      def get_batch(self):
          return self.dataset.train_batch, self.dataset.test_batch
  
      def get_train(self):
          return self.train_hg
  
      def get_labels(self):
          return self.dataset.get_labels()
  
      def dict2emd(self, r_embedding):
          r_emd = []
          for i in range(self.dataset.num_rels):
              r_emd.append(r_embedding[str(i)])
          return th.stack(r_emd).squeeze()
  
      def construct_negative_graph(self, hg):
          e_dict = {
              etype: hg.edges(etype=etype, form='eid')
              for etype in hg.canonical_etypes}
          neg_srcdst = self.negative_sampler(hg, e_dict)
          neg_pair_graph = dgl.heterograph(neg_srcdst,
                                           {ntype: hg.number_of_nodes(ntype) for ntype in hg.ntypes})
          return neg_pair_graph
  
  
  class HeteroDotProductPredictor(th.nn.Module):
      """
      References: `documentation of dgl <https://docs.dgl.ai/guide/training-link.html#heterogeneous-graphs>_`
      
      """
  
      def forward(self, edge_subgraph, x, *args, **kwargs):
          """
          Parameters
          ----------
          edge_subgraph: dgl.Heterograph
              the prediction graph only contains the edges of the target link
          x: dict[str: th.Tensor]
              the embedding dict. The key only contains the nodes involving with the target link.
      
          Returns
          -------
          score: th.Tensor
              the prediction of the edges in edge_subgraph
          """
  
          with edge_subgraph.local_scope():
              for ntype in edge_subgraph.ntypes:
                  edge_subgraph.nodes[ntype].data['x'] = x[ntype]
                  for etype in edge_subgraph.canonical_etypes:
                      edge_subgraph.apply_edges(
                          dgl.function.u_dot_v('x', 'x', 'score'), etype=etype)
              score = edge_subgraph.edata['score']
              if isinstance(score, dict):
                  result = []
                  for _, value in score.items():
                      result.append(value)
                  score = th.cat(result)
              return score.squeeze()
  
  
  class HeteroDistMultPredictor(th.nn.Module):
  
      def forward(self, edge_subgraph, x, r_embedding, *args, **kwargs):
          """
          DistMult factorization (Yang et al. 2014) as the scoring function,
          which is known to perform well on standard link prediction benchmarks when used on its own.
      
          In DistMult, every relation r is associated with a diagonal matrix :math:`R_{r} \in \mathbb{R}^{d \times d}`
          and a triple (s, r, o) is scored as
      
          .. math::
              f(s, r, o)=e_{s}^{T} R_{r} e_{o}
      
          Parameters
          ----------
          edge_subgraph: dgl.Heterograph
              the prediction graph only contains the edges of the target link
          x: dict[str: th.Tensor]
              the node embedding dict. The key only contains the nodes involving with the target link.
          r_embedding: th.Tensor
              the all relation types embedding
      
          Returns
          -------
          score: th.Tensor
              the prediction of the edges in edge_subgraph
          """
          with edge_subgraph.local_scope():
              for ntype in edge_subgraph.ntypes:
                  edge_subgraph.nodes[ntype].data['x'] = x[ntype]
              for etype in edge_subgraph.canonical_etypes:
                  e = r_embedding[etype[1]]
                  n = edge_subgraph.num_edges(etype)
                  if 1 == len(edge_subgraph.canonical_etypes):
                      edge_subgraph.edata['e'] = e.expand(n, -1)
                  else:
                      edge_subgraph.edata['e'] = {etype: e.expand(n, -1)}
                  edge_subgraph.apply_edges(
                      dgl.function.u_mul_e('x', 'e', 's'), etype=etype)
                  edge_subgraph.apply_edges(
                      dgl.function.e_mul_v('s', 'x', 'score'), etype=etype)
  
              score = edge_subgraph.edata['score']
              if isinstance(score, dict):
                  result = []
                  for _, value in score.items():
                      result.append(th.sum(value, dim=1))
                  score = th.cat(result)
              else:
                  score = th.sum(score, dim=1)
              return score
  