text: import argparse
  import copy
  
  import numpy as np
  import torch
  import torch.nn.functional as F
  from tqdm import tqdm
  
  from cogdl.data import Dataset
  from cogdl.models.supervised_model import (
      SupervisedHomogeneousNodeClassificationModel,
  )
  from cogdl.trainers.supervised_model_trainer import SupervisedHomogeneousNodeClassificationTrainer
  from cogdl.utils.self_auxiliary_task import (
      EdgeMask,
      PairwiseDistance,
      Distance2Clusters,
      PairwiseAttrSim,
      Distance2ClustersPP,
  )
  from . import register_universal_trainer
  
  
  @register_universal_trainer("self_auxiliary_task")
  class SelfAuxiliaryTaskTrainer(SupervisedHomogeneousNodeClassificationTrainer):
      @staticmethod
      def add_args(parser: argparse.ArgumentParser):
          """Add trainer-specific arguments to the parser."""
          # fmt: off
          parser.add_argument('--auxiliary-task', default="none", type=str)
          parser.add_argument('--alpha', default=10, type=float)
          parser.add_argument('--label-mask', default=0, type=float)
          # fmt: on
  
      @classmethod
      def build_trainer_from_args(cls, args):
          return cls(args)
  
      def __init__(self, args):
          self.device = args.device_id[0] if not args.cpu else "cpu"
          self.patience = args.patience
          self.max_epoch = args.max_epoch
          self.lr = args.lr
          self.weight_decay = args.weight_decay
          self.auxiliary_task = args.auxiliary_task
          self.hidden_size = args.hidden_size
          self.alpha = args.alpha
          self.label_mask = args.label_mask
  
      def resplit_data(self, data):
          trained = torch.where(data.train_mask)[0]
          perm = np.random.permutation(trained.shape[0])
          preserve_nnz = int(len(perm) * (1 - self.label_mask))
          preserved = trained[perm[:preserve_nnz]]
          masked = trained[perm[preserve_nnz:]]
          data.train_mask = torch.full((data.train_mask.shape[0],), False, dtype=torch.bool)
          data.train_mask[preserved] = True
          data.test_mask[masked] = True
  
      def fit(self, model: SupervisedHomogeneousNodeClassificationModel, dataset: Dataset):
          # self.resplit_data(dataset.data)
          self.data = dataset.data
          self.original_data = dataset.data
          self.data.apply(lambda x: x.to(self.device))
          self.original_data.apply(lambda x: x.to(self.device))
          if self.auxiliary_task == "edgemask":
              self.agent = EdgeMask(self.data.edge_index, self.data.x, self.hidden_size, self.device)
          elif self.auxiliary_task == "pairwise-distance":
              self.agent = PairwiseDistance(self.data.edge_index, self.data.x, self.hidden_size, 3, self.device)
          elif self.auxiliary_task == "distance2clusters":
              self.agent = Distance2Clusters(self.data.edge_index, self.data.x, self.hidden_size, 30, self.device)
          elif self.auxiliary_task == "pairwise-attr-sim":
              self.agent = PairwiseAttrSim(self.data.edge_index, self.data.x, self.hidden_size, 5, self.device)
          elif self.auxiliary_task == "distance2clusters++":
              self.agent = Distance2ClustersPP(
                  self.data.edge_index, self.data.x, self.data.y, self.hidden_size, 5, 1, self.device
              )
          else:
              raise Exception(
                  "auxiliary task must be edgemask, pairwise-distance, distance2clusters, distance2clusters++ or pairwise-attr-sim"
              )
          self.model = model
  
          self.optimizer = torch.optim.Adam(
              list(model.parameters()) + list(self.agent.linear.parameters()), lr=self.lr, weight_decay=self.weight_decay
          )
          self.model.to(self.device)
          epoch_iter = tqdm(range(self.max_epoch))
  
          best_score = 0
          best_loss = np.inf
          max_score = 0
          min_loss = np.inf
  
          for epoch in epoch_iter:
              if self.auxiliary_task == "distance2clusters++" and epoch % 40 == 0:
                  self.agent.update_cluster()
              self._train_step()
              train_acc, _ = self._test_step(split="train")
              val_acc, val_loss = self._test_step(split="val")
              epoch_iter.set_description(f"Epoch: {epoch:03d}, Train: {train_acc:.4f}, Val: {val_acc:.4f}")
              if val_loss <= min_loss or val_acc >= max_score:
                  if val_loss <= best_loss:  # and val_acc >= best_score:
                      best_loss = val_loss
                      best_score = val_acc
                      best_model = copy.deepcopy(self.model)
                  min_loss = np.min((min_loss, val_loss))
                  max_score = np.max((max_score, val_acc))
          print(f"Valid accurracy = {best_score}")
  
          return best_model
  
      def _train_step(self):
          self.data.edge_index, self.data.x = self.agent.transform_data()
          self.model.train()
          self.optimizer.zero_grad()
          embeddings = self.model.get_embeddings(self.data.x, self.data.edge_index)
          loss = self.model.node_classification_loss(self.data) + self.alpha * self.agent.make_loss(embeddings)
          loss.backward()
          self.optimizer.step()
  
      def _test_step(self, split="val"):
          self.data = self.original_data
          self.model.eval()
          if split == "train":
              mask = self.data.train_mask
          elif split == "val":
              mask = self.data.val_mask
          else:
              mask = self.data.test_mask
  
          logits = self.model.predict(self.data)
          loss = F.nll_loss(logits[mask], self.data.y[mask]).item()
  
          pred = logits[mask].max(1)[1]
          acc = pred.eq(self.data.y[mask]).sum().item() / mask.sum().item()
          return acc, loss
  