code: def main(cfg):
      cwd = get_original_cwd()
      os.chdir(cwd)
  
      if not os.path.exists(
          os.path.join(cfg.data_dir, "train_distant.json")
      ):
          raise FileNotFoundError(
              "Sorry, the file: 'train_annotated.json' is too big to upload to github, \
              please manually download to 'data/' from DocRED GoogleDrive https://drive.google.com/drive/folders/1c5-0YwnoJx8NS6CV2f-NoTHR__BdkNqw"
          )
  
      device = torch.device(
          "cuda:0" if torch.cuda.is_available() else "cpu"
      )
  
      config = AutoConfig.from_pretrained(
          cfg.config_name
          if cfg.config_name
          else cfg.model_name_or_path,
          num_labels=cfg.num_class,
      )
      tokenizer = AutoTokenizer.from_pretrained(
          cfg.tokenizer_name
          if cfg.tokenizer_name
          else cfg.model_name_or_path,
      )
  
      Dataset = ReadDataset(
          cfg, cfg.dataset, tokenizer, cfg.max_seq_length
      )
  
      train_file = os.path.join(cfg.data_dir, cfg.train_file)
      dev_file = os.path.join(cfg.data_dir, cfg.dev_file)
      test_file = os.path.join(cfg.data_dir, cfg.test_file)
      train_features = Dataset.read(train_file)
      dev_features = Dataset.read(dev_file)
      test_features = Dataset.read(test_file)
  
      model = AutoModel.from_pretrained(
          cfg.model_name_or_path,
          from_tf=bool(".ckpt" in cfg.model_name_or_path),
          config=config,
      )
      wandb.watch(model, log="all")
  
      config.cls_token_id = tokenizer.cls_token_id
      config.sep_token_id = tokenizer.sep_token_id
      config.transformer_type = cfg.transformer_type
  
      set_seed(cfg)
      model = DocREModel(config, cfg, model, num_labels=cfg.num_labels)
      if cfg.train_from_saved_model != "":
          model.load_state_dict(
              torch.load(cfg.train_from_saved_model)["checkpoint"]
          )
          print(
              "load saved model from {}.".format(
                  cfg.train_from_saved_model
              )
          )
  
      # if torch.cuda.device_count() > 1:
      #    print("Let's use", torch.cuda.device_count(), "GPUs!")
      #    model = torch.nn.DataParallel(model, device_ids = list(range(torch.cuda.device_count())))
      model.to(device)
  
      train(cfg, model, train_features, dev_features, test_features)
  
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 3 for calling a global function with example usage: main(cfg) without return types