code: def entity_train(args):
      if "albert" in args.entity_model:
          logger.info("Use Albert: %s" % args.entity_model)
          args.entity_use_albert = True
  
      entity_setseed(args.entity_seed)
  
      if not os.path.exists(args.entity_output_dir):
          os.makedirs(args.entity_output_dir)
  
      if args.entity_do_train:
          logger.addHandler(
              logging.FileHandler(
                  os.path.join(args.entity_output_dir, "train.log"), "w"
              )
          )
      else:
          logger.addHandler(
              logging.FileHandler(
                  os.path.join(args.entity_output_dir, "eval.log"), "w"
              )
          )
  
      logger.info(sys.argv)
      logger.info(args)
  
      ner_label2id, ner_id2label = get_labelmap(
          task_ner_labels[args.task]
      )
  
      num_ner_labels = len(task_ner_labels[args.task]) + 1
      model = EntityModel(args, num_ner_labels=num_ner_labels)
  
      # print(args.dev_data)
      dev_data = Dataset(args.dev_data)
      dev_samples, dev_ner = convert_dataset_to_samples(
          dev_data,
          args.entity_max_span_length,
          ner_label2id=ner_label2id,
          context_window=args.entity_context_window,
      )
      dev_batches = batchify(dev_samples, args.entity_eval_batch_size)
  
      if args.entity_do_train:
          train_data = Dataset(args.train_data)
          train_samples, train_ner = convert_dataset_to_samples(
              train_data,
              args.entity_max_span_length,
              ner_label2id=ner_label2id,
              context_window=args.entity_context_window,
          )
          train_batches = batchify(
              train_samples, args.entity_train_batch_size
          )
          best_result = 0.0
  
          param_optimizer = list(model.bert_model.named_parameters())
          optimizer_grouped_parameters = [
              {
                  "params": [
                      p for n, p in param_optimizer if "bert" in n
                  ]
              },
              {
                  "params": [
                      p for n, p in param_optimizer if "bert" not in n
                  ],
                  "lr": args.entity_task_learning_rate,
              },
          ]
          optimizer = AdamW(
              optimizer_grouped_parameters,
              lr=args.entity_learning_rate,
              correct_bias=not (args.entity_bertadam),
          )
          t_total = len(train_batches) * args.entity_num_epoch
          scheduler = get_linear_schedule_with_warmup(
              optimizer,
              int(t_total * args.entity_warmup_proportion),
              t_total,
          )
  
          tr_loss = 0
          tr_examples = 0
          global_step = 0
          eval_step = len(train_batches) // args.entity_eval_per_epoch
          for _ in tqdm(range(args.entity_num_epoch)):
              if args.entity_train_shuffle:
                  random.shuffle(train_batches)
              for i in tqdm(range(len(train_batches))):
                  output_dict = model.run_batch(
                      train_batches[i], training=True
                  )
                  loss = output_dict["ner_loss"]
                  loss.backward()
  
                  tr_loss += loss.item()
                  tr_examples += len(train_batches[i])
                  global_step += 1
  
                  optimizer.step()
                  scheduler.step()
                  optimizer.zero_grad()
  
                  if global_step % args.entity_print_loss_step == 0:
                      logger.info(
                          "Epoch=%d, iter=%d, loss=%.5f"
                          % (_, i, tr_loss / tr_examples)
                      )
                      tr_loss = 0
                      tr_examples = 0
  
                  if global_step % eval_step == 0:
                      f1 = entity_evaluate(model, dev_batches, dev_ner)
                      if f1 > best_result:
                          best_result = f1
                          logger.info(
                              "!!! Best valid (epoch=%d): %.2f"
                              % (_, f1 * 100)
                          )
                          entity_save_model(model, args)
  
      if args.entity_do_eval:
          args.entity_bert_model_dir = args.entity_output_dir
          model = EntityModel(args, num_ner_labels=num_ner_labels)
  
          test_data = Dataset(args.dev_data)
          prediction_file = os.path.join(
              args.entity_output_dir, args.entity_dev_pred_filename
          )
          test_samples, test_ner = convert_dataset_to_samples(
              test_data,
              args.entity_max_span_length,
              ner_label2id=ner_label2id,
              context_window=args.entity_context_window,
          )
          test_batches = batchify(
              test_samples, args.entity_eval_batch_size
          )
          entity_evaluate(model, test_batches, test_ner)
          entity_output_ner_predictions(
              model,
              test_batches,
              test_data,
              output_file=prediction_file,
              ner_id2label=ner_id2label,
          )
  
          if args.entity_eval_test:
              test_data = Dataset(args.test_data)
              prediction_file = os.path.join(
                  args.entity_output_dir, args.entity_test_pred_filename
              )
              test_samples, test_ner = convert_dataset_to_samples(
                  test_data,
                  args.entity_max_span_length,
                  ner_label2id=ner_label2id,
                  context_window=args.entity_context_window,
              )
              test_batches = batchify(
                  test_samples, args.entity_eval_batch_size
              )
              entity_evaluate(model, test_batches, test_ner)
              entity_output_ner_predictions(
                  model,
                  test_batches,
                  test_data,
                  output_file=prediction_file,
                  ner_id2label=ner_id2label,
              )
  
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 1 for calling a global function with example usage: entity_train(args) without return types