code: def train(args, model, train_features, dev_features, test_features):
      def logging(s, print_=True, log_=True):
          if print_:
              print(s)
          if log_ and args.log_dir != "":
              with open(args.log_dir, "a+") as f_log:
                  f_log.write(s + "\n")
  
      def finetune(features, optimizer, num_epoch, num_steps, model):
          cur_model = (
              model.module if hasattr(model, "module") else model
          )
          device = torch.device(
              "cuda:0" if torch.cuda.is_available() else "cpu"
          )
          if args.train_from_saved_model != "":
              best_score = torch.load(args.train_from_saved_model)[
                  "best_f1"
              ]
              epoch_delta = (
                  torch.load(args.train_from_saved_model)["epoch"] + 1
              )
          else:
              epoch_delta = 0
              best_score = -1
          train_dataloader = DataLoader(
              features,
              batch_size=args.train_batch_size,
              shuffle=True,
              collate_fn=collate_fn,
              drop_last=True,
          )
          train_iterator = [
              epoch + epoch_delta for epoch in range(num_epoch)
          ]
          total_steps = int(
              len(train_dataloader)
              * num_epoch
              // args.gradient_accumulation_steps
          )
          warmup_steps = int(total_steps * args.warmup_ratio)
          scheduler = get_linear_schedule_with_warmup(
              optimizer,
              num_warmup_steps=warmup_steps,
              num_training_steps=total_steps,
          )
          print("Total steps: {}".format(total_steps))
          print("Warmup steps: {}".format(warmup_steps))
          global_step = 0
          log_step = 100
          total_loss = 0
  
          # scaler = GradScaler()
          for epoch in train_iterator:
              start_time = time.time()
              optimizer.zero_grad()
  
              for step, batch in enumerate(train_dataloader):
                  model.train()
  
                  inputs = {
                      "input_ids": batch[0].to(device),
                      "attention_mask": batch[1].to(device),
                      "labels": batch[2],
                      "entity_pos": batch[3],
                      "hts": batch[4],
                  }
                  # with autocast():
                  outputs = model(**inputs)
                  loss = outputs[0] / args.gradient_accumulation_steps
                  total_loss += loss.item()
                  #    scaler.scale(loss).backward()
  
                  loss.backward()
  
                  if step % args.gradient_accumulation_steps == 0:
                      # scaler.unscale_(optimizer)
                      if args.max_grad_norm > 0:
                          # torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), args.max_grad_norm)
                          torch.nn.utils.clip_grad_norm_(
                              cur_model.parameters(), args.max_grad_norm
                          )
                      # scaler.step(optimizer)
                      # scaler.update()
                      # scheduler.step()
                      optimizer.step()
                      scheduler.step()
                      optimizer.zero_grad()
                      global_step += 1
                      num_steps += 1
                      if global_step % log_step == 0:
                          cur_loss = total_loss / log_step
                          elapsed = time.time() - start_time
                          logging(
                              "| epoch {:2d} | step {:4d} | min/b {:5.2f} | lr {} | train loss {:5.3f}".format(
                                  epoch,
                                  global_step,
                                  elapsed / 60,
                                  scheduler.get_last_lr(),
                                  cur_loss * 1000,
                              )
                          )
                          total_loss = 0
                          start_time = time.time()
  
                          wandb.log({"train_loss": cur_loss})
  
                  if (step + 1) == len(train_dataloader) - 1 or (
                      args.evaluation_steps > 0
                      and num_steps % args.evaluation_steps == 0
                      and step % args.gradient_accumulation_steps == 0
                  ):
                      # if step ==0:
                      logging("-" * 89)
                      eval_start_time = time.time()
                      dev_score, dev_output = evaluate(
                          args, model, dev_features, tag="dev"
                      )
  
                      logging(
                          "| epoch {:3d} | time: {:5.2f}s | dev_result:{}".format(
                              epoch,
                              time.time() - eval_start_time,
                              dev_output,
                          )
                      )
  
                      wandb.log({"dev_result": dev_output})
  
                      logging("-" * 89)
                      if dev_score > best_score:
                          best_score = dev_score
                          logging(
                              "| epoch {:3d} | best_f1:{}".format(
                                  epoch, best_score
                              )
                          )
  
                          wandb.log({"best_f1": best_score})
  
                          if args.save_path != "":
                              torch.save(
                                  {
                                      "epoch": epoch,
                                      "checkpoint": cur_model.state_dict(),
                                      "best_f1": best_score,
                                      "optimizer": optimizer.state_dict(),
                                  },
                                  args.save_path,
                                  _use_new_zipfile_serialization=False,
                              )
                              logging(
                                  "| successfully save model at: {}".format(
                                      args.save_path
                                  )
                              )
                              logging("-" * 89)
          return num_steps
  
      cur_model = model.module if hasattr(model, "module") else model
      extract_layer = ["extractor", "bilinear"]
      bert_layer = ["bert_model"]
      optimizer_grouped_parameters = [
          {
              "params": [
                  p
                  for n, p in cur_model.named_parameters()
                  if any(nd in n for nd in bert_layer)
              ],
              "lr": args.bert_lr,
          },
          {
              "params": [
                  p
                  for n, p in cur_model.named_parameters()
                  if any(nd in n for nd in extract_layer)
              ],
              "lr": 1e-4,
          },
          {
              "params": [
                  p
                  for n, p in cur_model.named_parameters()
                  if not any(
                      nd in n for nd in extract_layer + bert_layer
                  )
              ]
          },
      ]
  
      optimizer = AdamW(
          optimizer_grouped_parameters,
          lr=args.learning_rate,
          eps=args.adam_epsilon,
      )
      if args.train_from_saved_model != "":
          optimizer.load_state_dict(
              torch.load(args.train_from_saved_model)["optimizer"]
          )
          print(
              "load saved optimizer from {}.".format(
                  args.train_from_saved_model
              )
          )
  
      num_steps = 0
      set_seed(args)
      model.zero_grad()
      finetune(
          train_features,
          optimizer,
          args.num_train_epochs,
          num_steps,
          model,
      )
  
file_path: DeepKE/example/re/document/run.py
imports: import os
  import time
  import hydra
  from hydra.utils import get_original_cwd
  import numpy
  import torch
  import ujson
  from torch.utils.data import DataLoader
  from transformers import AutoConfig, AutoModel, AutoTokenizer
  from transformers.optimization import AdamW, get_linear_schedule_with_warmup
  from deepke.relation_extraction.document import *
  import wandb
label_desc: Use image node_id 1 for calling a global function with example usage: train(args, model, train_features, dev_features, test_features) and returns: num_steps