code: def relation_train(args):
      if "albert" in args.relation_model:
          RelationModel = AlbertForRelation
          args.relation_add_new_tokens = True
      else:
          RelationModel = BertForRelation
  
      device = torch.device(
          "cuda"
          if torch.cuda.is_available() and not args.no_cuda
          else "cpu"
      )
      n_gpu = torch.cuda.device_count()
  
      # train set
      if args.relation_do_train:
          (
              train_dataset,
              train_examples,
              train_nrel,
          ) = generate_relation_data(
              args.relation_train_file,
              use_gold=True,
              context_window=args.relation_context_window,
          )
      # dev set
      if (args.relation_do_eval and args.relation_do_train) or (
          args.relation_do_eval and not (args.relation_eval_test)
      ):
          (
              eval_dataset,
              eval_examples,
              eval_nrel,
          ) = generate_relation_data(
              os.path.join(
                  args.entity_output_dir,
                  args.relation_entity_predictions_dev,
              ),
              use_gold=args.relation_eval_with_gold,
              context_window=args.relation_context_window,
          )
      # test set
      if args.relation_eval_test:
          (
              test_dataset,
              test_examples,
              test_nrel,
          ) = generate_relation_data(
              os.path.join(
                  args.entity_output_dir,
                  args.relation_entity_predictions_test,
              ),
              use_gold=args.relation_eval_with_gold,
              context_window=args.relation_context_window,
          )
  
      relation_setseed(args, args.relation_seed)
  
      if not args.relation_do_train and not args.relation_do_eval:
          raise ValueError(
              "At least one of `do_train` or `do_eval` must be True."
          )
  
      if not os.path.exists(args.relation_output_dir):
          os.makedirs(args.relation_output_dir)
      if args.relation_do_train:
          logger.addHandler(
              logging.FileHandler(
                  os.path.join(args.relation_output_dir, "train.log"),
                  "w",
              )
          )
      else:
          logger.addHandler(
              logging.FileHandler(
                  os.path.join(args.relation_output_dir, "eval.log"),
                  "w",
              )
          )
      logger.info(sys.argv)
      logger.info(args)
      logger.info("device: {}, n_gpu: {}".format(device, n_gpu))
  
      # get label_list
      if os.path.exists(
          os.path.join(args.relation_output_dir, "label_list.json")
      ):
          with open(
              os.path.join(args.relation_output_dir, "label_list.json"),
              "r",
          ) as f:
              label_list = json.load(f)
      else:
          label_list = [args.relation_negative_label] + task_rel_labels[
              args.task
          ]
          with open(
              os.path.join(args.relation_output_dir, "label_list.json"),
              "w",
          ) as f:
              json.dump(label_list, f)
      label2id = {label: i for i, label in enumerate(label_list)}
      id2label = {i: label for i, label in enumerate(label_list)}
      num_labels = len(label_list)
  
      tokenizer = AutoTokenizer.from_pretrained(
          args.relation_model, do_lower_case=args.relation_do_lower_case
      )
      if args.relation_add_new_tokens:
          relation_add_marker_tokens(
              tokenizer, task_ner_labels[args.task]
          )
  
      if os.path.exists(
          os.path.join(args.relation_output_dir, "special_tokens.json")
      ):
          with open(
              os.path.join(
                  args.relation_output_dir, "special_tokens.json"
              ),
              "r",
          ) as f:
              special_tokens = json.load(f)
      else:
          special_tokens = {}
  
      if args.relation_do_eval and (
          args.relation_do_train or not (args.relation_eval_test)
      ):
          eval_features = relation_convert_examples_to_features(
              eval_examples,
              label2id,
              args.relation_max_seq_length,
              tokenizer,
              special_tokens,
              unused_tokens=not (args.relation_add_new_tokens),
          )
          logger.info("***** Dev *****")
          logger.info("  Num examples = %d", len(eval_examples))
          logger.info(
              "  Batch size = %d", args.relation_eval_batch_size
          )
          all_input_ids = torch.tensor(
              [f.input_ids for f in eval_features], dtype=torch.long
          )
          all_input_mask = torch.tensor(
              [f.input_mask for f in eval_features], dtype=torch.long
          )
          all_segment_ids = torch.tensor(
              [f.segment_ids for f in eval_features], dtype=torch.long
          )
          all_label_ids = torch.tensor(
              [f.label_id for f in eval_features], dtype=torch.long
          )
          all_sub_idx = torch.tensor(
              [f.sub_idx for f in eval_features], dtype=torch.long
          )
          all_obj_idx = torch.tensor(
              [f.obj_idx for f in eval_features], dtype=torch.long
          )
          eval_data = TensorDataset(
              all_input_ids,
              all_input_mask,
              all_segment_ids,
              all_label_ids,
              all_sub_idx,
              all_obj_idx,
          )
          eval_dataloader = DataLoader(
              eval_data, batch_size=args.relation_eval_batch_size
          )
          eval_label_ids = all_label_ids
      with open(
          os.path.join(args.relation_output_dir, "special_tokens.json"),
          "w",
      ) as f:
          json.dump(special_tokens, f)
  
      if args.relation_do_train:
          train_features = relation_convert_examples_to_features(
              train_examples,
              label2id,
              args.relation_max_seq_length,
              tokenizer,
              special_tokens,
              unused_tokens=not (args.relation_add_new_tokens),
          )
          if (
              args.relation_train_mode == "sorted"
              or args.relation_train_mode == "random_sorted"
          ):
              train_features = sorted(
                  train_features, key=lambda f: np.sum(f.input_mask)
              )
          else:
              random.shuffle(train_features)
          all_input_ids = torch.tensor(
              [f.input_ids for f in train_features], dtype=torch.long
          )
          all_input_mask = torch.tensor(
              [f.input_mask for f in train_features], dtype=torch.long
          )
          all_segment_ids = torch.tensor(
              [f.segment_ids for f in train_features], dtype=torch.long
          )
          all_label_ids = torch.tensor(
              [f.label_id for f in train_features], dtype=torch.long
          )
          all_sub_idx = torch.tensor(
              [f.sub_idx for f in train_features], dtype=torch.long
          )
          all_obj_idx = torch.tensor(
              [f.obj_idx for f in train_features], dtype=torch.long
          )
          train_data = TensorDataset(
              all_input_ids,
              all_input_mask,
              all_segment_ids,
              all_label_ids,
              all_sub_idx,
              all_obj_idx,
          )
          train_dataloader = DataLoader(
              train_data, batch_size=args.relation_train_batch_size
          )
          train_batches = [batch for batch in train_dataloader]
  
          num_train_optimization_steps = (
              len(train_dataloader) * args.relation_num_train_epochs
          )
  
          logger.info("***** Training *****")
          logger.info("  Num examples = %d", len(train_examples))
          logger.info(
              "  Batch size = %d", args.relation_train_batch_size
          )
          logger.info("  Num steps = %d", num_train_optimization_steps)
  
          best_result = None
          eval_step = max(
              1, len(train_batches) // args.relation_eval_per_epoch
          )
  
          lr = args.relation_learning_rate
          model = RelationModel.from_pretrained(
              args.relation_model,
              cache_dir=str(PYTORCH_PRETRAINED_BERT_CACHE),
              num_rel_labels=num_labels,
          )
          if hasattr(model, "bert"):
              model.bert.resize_token_embeddings(len(tokenizer))
          elif hasattr(model, "albert"):
              model.albert.resize_token_embeddings(len(tokenizer))
          else:
              raise TypeError("Unknown model class")
  
          model.to(device)
          if n_gpu > 1 and args.relation_single_card == False:
              model = torch.nn.DataParallel(model)
  
          param_optimizer = list(model.named_parameters())
          no_decay = ["bias", "LayerNorm.bias", "LayerNorm.weight"]
          optimizer_grouped_parameters = [
              {
                  "params": [
                      p
                      for n, p in param_optimizer
                      if not any(nd in n for nd in no_decay)
                  ],
                  "weight_decay": 0.01,
              },
              {
                  "params": [
                      p
                      for n, p in param_optimizer
                      if any(nd in n for nd in no_decay)
                  ],
                  "weight_decay": 0.0,
              },
          ]
          optimizer = AdamW(
              optimizer_grouped_parameters,
              lr=lr,
              correct_bias=not (args.relation_bertadam),
              no_deprecation_warning=True,
          )
          scheduler = get_linear_schedule_with_warmup(
              optimizer,
              int(
                  num_train_optimization_steps
                  * args.relation_warmup_proportion
              ),
              num_train_optimization_steps,
          )
  
          start_time = time.time()
          global_step = 0
          tr_loss = 0
          nb_tr_examples = 0
          nb_tr_steps = 0
          for epoch in range(int(args.relation_num_train_epochs)):
              model.train()
              logger.info(
                  "Start epoch #{} (lr = {})...".format(epoch, lr)
              )
              if (
                  args.relation_train_mode == "random"
                  or args.relation_train_mode == "random_sorted"
              ):
                  random.shuffle(train_batches)
              for step, batch in enumerate(train_batches):
                  batch = tuple(t.to(device) for t in batch)
                  (
                      input_ids,
                      input_mask,
                      segment_ids,
                      label_ids,
                      sub_idx,
                      obj_idx,
                  ) = batch
                  loss = model(
                      input_ids,
                      segment_ids,
                      input_mask,
                      label_ids,
                      sub_idx,
                      obj_idx,
                  )
                  if n_gpu > 1:
                      loss = loss.mean()
  
                  loss.backward()
  
                  tr_loss += loss.item()
                  nb_tr_examples += input_ids.size(0)
                  nb_tr_steps += 1
  
                  optimizer.step()
                  scheduler.step()
                  optimizer.zero_grad()
                  global_step += 1
  
                  if (step + 1) % eval_step == 0:
                      logger.info(
                          "Epoch: {}, Step: {} / {}, used_time = {:.2f}s, loss = {:.6f}".format(
                              epoch,
                              step + 1,
                              len(train_batches),
                              time.time() - start_time,
                              tr_loss / nb_tr_steps,
                          )
                      )
                      save_model = False
                      if args.relation_do_eval:
                          preds, result, logits = relation_evaluate(
                              model,
                              device,
                              eval_dataloader,
                              eval_label_ids,
                              num_labels,
                              e2e_ngold=eval_nrel,
                          )
                          model.train()
                          result["global_step"] = global_step
                          result["epoch"] = epoch
                          result["learning_rate"] = lr
                          result[
                              "batch_size"
                          ] = args.relation_train_batch_size
  
                          if (best_result is None) or (
                              result[args.eval_metric]
                              > best_result[args.eval_metric]
                          ):
                              best_result = result
                              logger.info(
                                  "!!! Best dev %s (lr=%s, epoch=%d): %.2f"
                                  % (
                                      args.eval_metric,
                                      str(lr),
                                      epoch,
                                      result[args.eval_metric] * 100.0,
                                  )
                              )
                              relation_save_trained_model(
                                  args.relation_output_dir,
                                  model,
                                  tokenizer,
                              )
  
      evaluation_results = {}
      if args.relation_do_eval:
          logger.info(special_tokens)
          if args.relation_eval_test:
              eval_dataset = test_dataset
              eval_examples = test_examples
              eval_features = relation_convert_examples_to_features(
                  test_examples,
                  label2id,
                  args.relation_max_seq_length,
                  tokenizer,
                  special_tokens,
                  unused_tokens=not (args.relation_add_new_tokens),
              )
              eval_nrel = test_nrel
              logger.info(special_tokens)
              logger.info("***** Test *****")
              logger.info("  Num examples = %d", len(test_examples))
              logger.info(
                  "  Batch size = %d", args.relation_eval_batch_size
              )
              all_input_ids = torch.tensor(
                  [f.input_ids for f in eval_features], dtype=torch.long
              )
              all_input_mask = torch.tensor(
                  [f.input_mask for f in eval_features],
                  dtype=torch.long,
              )
              all_segment_ids = torch.tensor(
                  [f.segment_ids for f in eval_features],
                  dtype=torch.long,
              )
              all_label_ids = torch.tensor(
                  [f.label_id for f in eval_features], dtype=torch.long
              )
              all_sub_idx = torch.tensor(
                  [f.sub_idx for f in eval_features], dtype=torch.long
              )
              all_obj_idx = torch.tensor(
                  [f.obj_idx for f in eval_features], dtype=torch.long
              )
              eval_data = TensorDataset(
                  all_input_ids,
                  all_input_mask,
                  all_segment_ids,
                  all_label_ids,
                  all_sub_idx,
                  all_obj_idx,
              )
              eval_dataloader = DataLoader(
                  eval_data, batch_size=args.relation_eval_batch_size
              )
              eval_label_ids = all_label_ids
          model = RelationModel.from_pretrained(
              args.relation_output_dir, num_rel_labels=num_labels
          )
          model.to(device)
          preds, result, logits = relation_evaluate(
              model,
              device,
              eval_dataloader,
              eval_label_ids,
              num_labels,
              e2e_ngold=eval_nrel,
          )
  
          logger.info("*** Evaluation Results ***")
          for key in sorted(result.keys()):
              logger.info("  %s = %s", key, str(result[key]))
  
          relation_print_pred_json(
              eval_dataset,
              eval_examples,
              preds,
              id2label,
              os.path.join(
                  args.relation_output_dir,
                  args.relation_prediction_file,
              ),
          )
  
file_path: DeepKE/example/triple/PURE/run.py
imports: import os
  import sys
  import argparse
  import hydra
  from hydra.utils import get_original_cwd
  import torch
  import logging
  import tqdm
  import wandb
  from models.const import *
  from models.data_structures import *
  from models.entityModels import *
  from models.entityUtils import *
  from models.relationModels import *
  from models.relationUtils import *
  from run_entity import *
  from run_relation import *
label_desc: Use image node_id 2 for calling a global function with example usage: relation_train(args) without return types