code: def train(args):
      model_file = model_file_name(args)
      utils.ensure_dir(os.path.split(model_file)[0])
  
      # load pretrained vectors if needed
      pretrain = load_pretrain(args)
  
      # TODO: refactor.  the exact same thing is done in the tagger
      if args["charlm"]:
          if args["charlm_shorthand"] is None:
              raise ValueError(
                  "CharLM Shorthand is required for loading pretrained CharLM model..."
              )
          logger.info("Using pretrained contextualized char embedding")
          if not args["charlm_forward_file"]:
              args[
                  "charlm_forward_file"
              ] = "{}/{}_forward_charlm.pt".format(
                  args["charlm_save_dir"], args["charlm_shorthand"]
              )
          if not args["charlm_backward_file"]:
              args[
                  "charlm_backward_file"
              ] = "{}/{}_backward_charlm.pt".format(
                  args["charlm_save_dir"], args["charlm_shorthand"]
              )
  
      # load data
      logger.info(
          "Loading data with batch size {}...".format(
              args["batch_size"]
          )
      )
      train_data, _, _ = CoNLL.conll2dict(input_file=args["train_file"])
      # possibly augment the training data with some amount of fake data
      # based on the options chosen
      logger.info("Original data size: {}".format(len(train_data)))
      train_data.extend(
          augment_punct(
              train_data,
              args["augment_nopunct"],
              keep_original_sentences=False,
          )
      )
      logger.info("Augmented data size: {}".format(len(train_data)))
      train_doc = Document(train_data)
      train_batch = DataLoader(
          train_doc,
          args["batch_size"],
          args,
          pretrain,
          evaluation=False,
      )
      vocab = train_batch.vocab
      dev_doc = CoNLL.conll2doc(input_file=args["eval_file"])
      dev_batch = DataLoader(
          dev_doc,
          args["batch_size"],
          args,
          pretrain,
          vocab=vocab,
          evaluation=True,
          sort_during_eval=True,
      )
  
      # pred and gold path
      system_pred_file = args["output_file"]
      gold_file = args["gold_file"]
  
      # skip training if the language does not have training or dev data
      if len(train_batch) == 0 or len(dev_batch) == 0:
          logger.info("Skip training because no data available...")
          sys.exit(0)
  
      if args["wandb"]:
          import wandb
  
          wandb_name = (
              args["wandb_name"]
              if args["wandb_name"]
              else "%s_depparse" % args["shorthand"]
          )
          wandb.init(name=wandb_name, config=args)
          wandb.run.define_metric("train_loss", summary="min")
          wandb.run.define_metric("dev_score", summary="max")
  
      logger.info("Training parser...")
      trainer = Trainer(
          args=args,
          vocab=vocab,
          pretrain=pretrain,
          device=args["device"],
      )
  
      global_step = 0
      max_steps = args["max_steps"]
      dev_score_history = []
      best_dev_preds = []
      current_lr = args["lr"]
      global_start_time = time.time()
      format_str = "Finished STEP {}/{}, loss = {:.6f} ({:.3f} sec/batch), lr: {:.6f}"
  
      using_amsgrad = False
      last_best_step = 0
      # start training
      train_loss = 0
      while True:
          do_break = False
          for i, batch in enumerate(train_batch):
              start_time = time.time()
              global_step += 1
              loss = trainer.update(batch, eval=False)  # update step
              train_loss += loss
              if global_step % args["log_step"] == 0:
                  duration = time.time() - start_time
                  logger.info(
                      format_str.format(
                          global_step,
                          max_steps,
                          loss,
                          duration,
                          current_lr,
                      )
                  )
  
              if global_step % args["eval_interval"] == 0:
                  # eval on dev
                  logger.info("Evaluating on dev set...")
                  dev_preds = []
                  for batch in dev_batch:
                      preds = trainer.predict(batch)
                      dev_preds += preds
                  dev_preds = utils.unsort(
                      dev_preds, dev_batch.data_orig_idx
                  )
  
                  dev_batch.doc.set(
                      [HEAD, DEPREL], [y for x in dev_preds for y in x]
                  )
                  CoNLL.write_doc2conll(dev_batch.doc, system_pred_file)
                  _, _, dev_score = scorer.score(
                      system_pred_file, gold_file
                  )
  
                  train_loss = (
                      train_loss / args["eval_interval"]
                  )  # avg loss per batch
                  logger.info(
                      "step {}: train_loss = {:.6f}, dev_score = {:.4f}".format(
                          global_step, train_loss, dev_score
                      )
                  )
  
                  if args["wandb"]:
                      wandb.log(
                          {
                              "train_loss": train_loss,
                              "dev_score": dev_score,
                          }
                      )
  
                  train_loss = 0
  
                  # save best model
                  if len(dev_score_history) == 0 or dev_score > max(
                      dev_score_history
                  ):
                      last_best_step = global_step
                      trainer.save(model_file)
                      logger.info("new best model saved.")
                      best_dev_preds = dev_preds
  
                  dev_score_history += [dev_score]
  
              if (
                  global_step - last_best_step
                  >= args["max_steps_before_stop"]
              ):
                  if not using_amsgrad:
                      logger.info("Switching to AMSGrad")
                      last_best_step = global_step
                      using_amsgrad = True
                      trainer.optimizer = optim.Adam(
                          trainer.model.parameters(),
                          amsgrad=True,
                          lr=args["lr"],
                          betas=(0.9, args["beta2"]),
                          eps=1e-6,
                      )
                  else:
                      do_break = True
                      break
  
              if global_step >= args["max_steps"]:
                  do_break = True
                  break
  
          if do_break:
              break
  
          train_batch.reshuffle()
  
      logger.info("Training ended with {} steps.".format(global_step))
  
      if args["wandb"]:
          wandb.finish()
  
      if len(dev_score_history) > 0:
          best_f, best_eval = (
              max(dev_score_history) * 100,
              np.argmax(dev_score_history) + 1,
          )
          logger.info(
              "Best dev F1 = {:.2f}, at iteration = {}".format(
                  best_f, best_eval * args["eval_interval"]
              )
          )
      else:
          logger.info("Dev set never evaluated.  Saving final model.")
          trainer.save(model_file)
  
imports: import sys
  import os
  import shutil
  import time
  import argparse
  import logging
  import numpy
  import random
  import torch
  from torch import nn, optim
  import stanza.models.depparse.data
  from stanza.models.depparse.data import DataLoader
  from stanza.models.depparse.trainer import Trainer
  from stanza.models.depparse import scorer
  from stanza.models.common import utils
  from stanza.models.common import pretrain
  from stanza.models.common.data import augment_punct
  from stanza.models.common.doc import *
  from stanza.utils.conll import CoNLL
  from stanza.models import _training_logging
label_desc: Use image node_id 6 for calling a global function with example usage: train(args) without return types