text: import argparse
  
  import torch
  import torch.nn.functional as F
  from torch.utils.data import DataLoader
  from torch_geometric.utils import negative_sampling
  
  import torch_geometric.transforms as T
  from torch_geometric.nn import GCNConv, SAGEConv
  
  from ogb.linkproppred import PygLinkPropPredDataset, Evaluator
  
  from logger import Logger
  
  
  class GCN(torch.nn.Module):
      def __init__(self, in_channels, hidden_channels, out_channels, num_layers, dropout):
          super(GCN, self).__init__()
  
          self.convs = torch.nn.ModuleList()
          self.convs.append(GCNConv(in_channels, hidden_channels, cached=True))
          for _ in range(num_layers - 2):
              self.convs.append(GCNConv(hidden_channels, hidden_channels, cached=True))
          self.convs.append(GCNConv(hidden_channels, out_channels, cached=True))
  
          self.dropout = dropout
  
      def reset_parameters(self):
          for conv in self.convs:
              conv.reset_parameters()
  
      def forward(self, x, adj_t):
          for conv in self.convs[:-1]:
              x = conv(x, adj_t)
              x = F.relu(x)
              x = F.dropout(x, p=self.dropout, training=self.training)
          x = self.convs[-1](x, adj_t)
          return x
  
  
  class SAGE(torch.nn.Module):
      def __init__(self, in_channels, hidden_channels, out_channels, num_layers, dropout):
          super(SAGE, self).__init__()
  
          self.convs = torch.nn.ModuleList()
          self.convs.append(SAGEConv(in_channels, hidden_channels))
          for _ in range(num_layers - 2):
              self.convs.append(SAGEConv(hidden_channels, hidden_channels))
          self.convs.append(SAGEConv(hidden_channels, out_channels))
  
          self.dropout = dropout
  
      def reset_parameters(self):
          for conv in self.convs:
              conv.reset_parameters()
  
      def forward(self, x, adj_t):
          for conv in self.convs[:-1]:
              x = conv(x, adj_t)
              x = F.relu(x)
              x = F.dropout(x, p=self.dropout, training=self.training)
          x = self.convs[-1](x, adj_t)
          return x
  
  
  class LinkPredictor(torch.nn.Module):
      def __init__(self, in_channels, hidden_channels, out_channels, num_layers, dropout):
          super(LinkPredictor, self).__init__()
  
          self.lins = torch.nn.ModuleList()
          self.lins.append(torch.nn.Linear(in_channels, hidden_channels))
          for _ in range(num_layers - 2):
              self.lins.append(torch.nn.Linear(hidden_channels, hidden_channels))
          self.lins.append(torch.nn.Linear(hidden_channels, out_channels))
  
          self.dropout = dropout
  
      def reset_parameters(self):
          for lin in self.lins:
              lin.reset_parameters()
  
      def forward(self, x_i, x_j):
          x = x_i * x_j
          for lin in self.lins[:-1]:
              x = lin(x)
              x = F.relu(x)
              x = F.dropout(x, p=self.dropout, training=self.training)
          x = self.lins[-1](x)
          return torch.sigmoid(x)
  
  
  def train(model, predictor, x, adj_t, split_edge, optimizer, batch_size):
  
      row, col, _ = adj_t.coo()
      edge_index = torch.stack([col, row], dim=0)
  
      model.train()
      predictor.train()
  
      pos_train_edge = split_edge["train"]["edge"].to(x.device)
  
      total_loss = total_examples = 0
      for perm in DataLoader(range(pos_train_edge.size(0)), batch_size, shuffle=True):
          optimizer.zero_grad()
  
          h = model(x, adj_t)
  
          edge = pos_train_edge[perm].t()
  
          pos_out = predictor(h[edge[0]], h[edge[1]])
          pos_loss = -torch.log(pos_out + 1e-15).mean()
  
          edge = negative_sampling(
              edge_index,
              num_nodes=x.size(0),
              num_neg_samples=perm.size(0),
              method="dense",
          )
  
          neg_out = predictor(h[edge[0]], h[edge[1]])
          neg_loss = -torch.log(1 - neg_out + 1e-15).mean()
  
          loss = pos_loss + neg_loss
          loss.backward()
  
          torch.nn.utils.clip_grad_norm_(x, 1.0)
          torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
          torch.nn.utils.clip_grad_norm_(predictor.parameters(), 1.0)
  
          optimizer.step()
  
          num_examples = pos_out.size(0)
          total_loss += loss.item() * num_examples
          total_examples += num_examples
  
      return total_loss / total_examples
  
  
  @torch.no_grad()
  def test(model, predictor, x, adj_t, split_edge, evaluator, batch_size):
      model.eval()
      predictor.eval()
  
      h = model(x, adj_t)
  
      pos_train_edge = split_edge["eval_train"]["edge"].to(x.device)
      pos_valid_edge = split_edge["valid"]["edge"].to(x.device)
      neg_valid_edge = split_edge["valid"]["edge_neg"].to(x.device)
      pos_test_edge = split_edge["test"]["edge"].to(x.device)
      neg_test_edge = split_edge["test"]["edge_neg"].to(x.device)
  
      pos_train_preds = []
      for perm in DataLoader(range(pos_train_edge.size(0)), batch_size):
          edge = pos_train_edge[perm].t()
          pos_train_preds += [predictor(h[edge[0]], h[edge[1]]).squeeze().cpu()]
      pos_train_pred = torch.cat(pos_train_preds, dim=0)
  
      pos_valid_preds = []
      for perm in DataLoader(range(pos_valid_edge.size(0)), batch_size):
          edge = pos_valid_edge[perm].t()
          pos_valid_preds += [predictor(h[edge[0]], h[edge[1]]).squeeze().cpu()]
      pos_valid_pred = torch.cat(pos_valid_preds, dim=0)
  
      neg_valid_preds = []
      for perm in DataLoader(range(neg_valid_edge.size(0)), batch_size):
          edge = neg_valid_edge[perm].t()
          neg_valid_preds += [predictor(h[edge[0]], h[edge[1]]).squeeze().cpu()]
      neg_valid_pred = torch.cat(neg_valid_preds, dim=0)
  
      pos_test_preds = []
      for perm in DataLoader(range(pos_test_edge.size(0)), batch_size):
          edge = pos_test_edge[perm].t()
          pos_test_preds += [predictor(h[edge[0]], h[edge[1]]).squeeze().cpu()]
      pos_test_pred = torch.cat(pos_test_preds, dim=0)
  
      neg_test_preds = []
      for perm in DataLoader(range(neg_test_edge.size(0)), batch_size):
          edge = neg_test_edge[perm].t()
          neg_test_preds += [predictor(h[edge[0]], h[edge[1]]).squeeze().cpu()]
      neg_test_pred = torch.cat(neg_test_preds, dim=0)
  
      results = {}
      for K in [10, 20, 30]:
          evaluator.K = K
          train_hits = evaluator.eval(
              {"y_pred_pos": pos_train_pred, "y_pred_neg": neg_valid_pred,}
          )[f"hits@{K}"]
          valid_hits = evaluator.eval(
              {"y_pred_pos": pos_valid_pred, "y_pred_neg": neg_valid_pred,}
          )[f"hits@{K}"]
          test_hits = evaluator.eval(
              {"y_pred_pos": pos_test_pred, "y_pred_neg": neg_test_pred,}
          )[f"hits@{K}"]
  
          results[f"Hits@{K}"] = (train_hits, valid_hits, test_hits)
  
      return results
  
  
  def main():
      parser = argparse.ArgumentParser(description="OGBL-DDI (GNN)")
      parser.add_argument("--device", type=int, default=0)
      parser.add_argument("--log_steps", type=int, default=1)
      parser.add_argument("--use_sage", action="store_true")
      parser.add_argument("--num_layers", type=int, default=2)
      parser.add_argument("--hidden_channels", type=int, default=256)
      parser.add_argument("--dropout", type=float, default=0.5)
      parser.add_argument("--batch_size", type=int, default=64 * 1024)
      parser.add_argument("--lr", type=float, default=0.005)
      parser.add_argument("--epochs", type=int, default=200)
      parser.add_argument("--eval_steps", type=int, default=5)
      parser.add_argument("--runs", type=int, default=10)
      args = parser.parse_args()
      print(args)
  
      device = f"cuda:{args.device}" if torch.cuda.is_available() else "cpu"
      device = torch.device(device)
  
      dataset = PygLinkPropPredDataset(name="ogbl-ddi", transform=T.ToSparseTensor())
      data = dataset[0]
      adj_t = data.adj_t.to(device)
  
      split_edge = dataset.get_edge_split()
  
      # We randomly pick some training samples that we want to evaluate on:
      torch.manual_seed(12345)
      idx = torch.randperm(split_edge["train"]["edge"].size(0))
      idx = idx[: split_edge["valid"]["edge"].size(0)]
      split_edge["eval_train"] = {"edge": split_edge["train"]["edge"][idx]}
  
      if args.use_sage:
          model = SAGE(
              args.hidden_channels,
              args.hidden_channels,
              args.hidden_channels,
              args.num_layers,
              args.dropout,
          ).to(device)
      else:
          model = GCN(
              args.hidden_channels,
              args.hidden_channels,
              args.hidden_channels,
              args.num_layers,
              args.dropout,
          ).to(device)
  
      emb = torch.nn.Embedding(data.num_nodes, args.hidden_channels).to(device)
      predictor = LinkPredictor(
          args.hidden_channels, args.hidden_channels, 1, args.num_layers, args.dropout
      ).to(device)
  
      evaluator = Evaluator(name="ogbl-ddi")
      loggers = {
          "Hits@10": Logger(args.runs, args),
          "Hits@20": Logger(args.runs, args),
          "Hits@30": Logger(args.runs, args),
      }
  
      for run in range(args.runs):
          torch.nn.init.xavier_uniform_(emb.weight)
          model.reset_parameters()
          predictor.reset_parameters()
          optimizer = torch.optim.Adam(
              list(model.parameters())
              + list(emb.parameters())
              + list(predictor.parameters()),
              lr=args.lr,
          )
  
          for epoch in range(1, 1 + args.epochs):
              loss = train(
                  model,
                  predictor,
                  emb.weight,
                  adj_t,
                  split_edge,
                  optimizer,
                  args.batch_size,
              )
  
              if epoch % args.eval_steps == 0:
                  results = test(
                      model,
                      predictor,
                      emb.weight,
                      adj_t,
                      split_edge,
                      evaluator,
                      args.batch_size,
                  )
                  for key, result in results.items():
                      loggers[key].add_result(run, result)
  
                  if epoch % args.log_steps == 0:
                      for key, result in results.items():
                          train_hits, valid_hits, test_hits = result
                          print(key)
                          print(
                              f"Run: {run + 1:02d}, "
                              f"Epoch: {epoch:02d}, "
                              f"Loss: {loss:.4f}, "
                              f"Train: {100 * train_hits:.2f}%, "
                              f"Valid: {100 * valid_hits:.2f}%, "
                              f"Test: {100 * test_hits:.2f}%"
                          )
                      print("---")
  
          for key in loggers.keys():
              print(key)
              loggers[key].print_statistics(run)
  
      for key in loggers.keys():
          print(key)
          loggers[key].print_statistics()
  
  
  if __name__ == "__main__":
      main()
  