code: def train(args):
      # load data
      logger.info(
          "[Loading data with batch size {}...]".format(
              args["batch_size"]
          )
      )
      train_doc = CoNLL.conll2doc(input_file=args["train_file"])
      train_batch = DataLoader(
          train_doc, args["batch_size"], args, evaluation=False
      )
      vocab = train_batch.vocab
      args["vocab_size"] = vocab["char"].size
      args["pos_vocab_size"] = vocab["pos"].size
      dev_doc = CoNLL.conll2doc(input_file=args["eval_file"])
      dev_batch = DataLoader(
          dev_doc,
          args["batch_size"],
          args,
          vocab=vocab,
          evaluation=True,
      )
  
      utils.ensure_dir(args["save_dir"])
      model_file = build_model_filename(args)
      logger.info("Using full savename: %s", model_file)
  
      # pred and gold path
      system_pred_file = args["output_file"]
      gold_file = args["gold_file"]
  
      utils.print_config(args)
  
      # skip training if the language does not have training or dev data
      if len(train_batch) == 0 or len(dev_batch) == 0:
          logger.warning(
              "[Skip training because no training data available...]"
          )
          return
  
      # start training
      # train a dictionary-based lemmatizer
      logger.info("Building lemmatizer in %s", model_file)
      trainer = Trainer(args=args, vocab=vocab, device=args["device"])
      logger.info("[Training dictionary-based lemmatizer...]")
      trainer.train_dict(train_batch.doc.get([TEXT, UPOS, LEMMA]))
      logger.info("Evaluating on dev set...")
      dev_preds = trainer.predict_dict(dev_batch.doc.get([TEXT, UPOS]))
      dev_batch.doc.set([LEMMA], dev_preds)
      CoNLL.write_doc2conll(dev_batch.doc, system_pred_file)
      _, _, dev_f = scorer.score(system_pred_file, gold_file)
      logger.info("Dev F1 = {:.2f}".format(dev_f * 100))
  
      if args.get("dict_only", False):
          # save dictionaries
          trainer.save(model_file)
      else:
          if args["wandb"]:
              import wandb
  
              wandb_name = (
                  args["wandb_name"]
                  if args["wandb_name"]
                  else "%s_lemmatizer" % 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")
  
          # train a seq2seq model
          logger.info("[Training seq2seq-based lemmatizer...]")
          global_step = 0
          max_steps = len(train_batch) * args["num_epoch"]
          dev_score_history = []
          best_dev_preds = []
          current_lr = args["lr"]
          global_start_time = time.time()
          format_str = "{}: step {}/{} (epoch {}/{}), loss = {:.6f} ({:.3f} sec/batch), lr: {:.6f}"
  
          # start training
          for epoch in range(1, args["num_epoch"] + 1):
              train_loss = 0
              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(
                              datetime.now().strftime(
                                  "%Y-%m-%d %H:%M:%S"
                              ),
                              global_step,
                              max_steps,
                              epoch,
                              args["num_epoch"],
                              loss,
                              duration,
                              current_lr,
                          )
                      )
  
              # eval on dev
              logger.info("Evaluating on dev set...")
              dev_preds = []
              dev_edits = []
              for i, batch in enumerate(dev_batch):
                  preds, edits = trainer.predict(
                      batch, args["beam_size"]
                  )
                  dev_preds += preds
                  if edits is not None:
                      dev_edits += edits
              dev_preds = trainer.postprocess(
                  dev_batch.doc.get([TEXT]), dev_preds, edits=dev_edits
              )
  
              # try ensembling with dict if necessary
              if args.get("ensemble_dict", False):
                  logger.info("[Ensembling dict with seq2seq model...]")
                  dev_preds = trainer.ensemble(
                      dev_batch.doc.get([TEXT, UPOS]), dev_preds
                  )
              dev_batch.doc.set([LEMMA], dev_preds)
              CoNLL.write_doc2conll(dev_batch.doc, system_pred_file)
              _, _, dev_score = scorer.score(
                  system_pred_file, gold_file
              )
  
              train_loss = (
                  train_loss
                  / train_batch.num_examples
                  * args["batch_size"]
              )  # avg loss per batch
              logger.info(
                  "epoch {}: train_loss = {:.6f}, dev_score = {:.4f}".format(
                      epoch, train_loss, dev_score
                  )
              )
  
              if args["wandb"]:
                  wandb.log(
                      {"train_loss": train_loss, "dev_score": dev_score}
                  )
  
              # save best model
              if epoch == 1 or dev_score > max(dev_score_history):
                  trainer.save(model_file)
                  logger.info("new best model saved.")
                  best_dev_preds = dev_preds
  
              # lr schedule
              if (
                  epoch > args["decay_epoch"]
                  and dev_score <= dev_score_history[-1]
                  and args["optim"] in ["sgd", "adagrad"]
              ):
                  current_lr *= args["lr_decay"]
                  trainer.update_lr(current_lr)
  
              dev_score_history += [dev_score]
              logger.info("")
  
          logger.info("Training ended with {} epochs.".format(epoch))
  
          if args["wandb"]:
              wandb.finish()
  
          best_f, best_epoch = (
              max(dev_score_history) * 100,
              np.argmax(dev_score_history) + 1,
          )
          logger.info(
              "Best dev F1 = {:.2f}, at epoch = {}".format(
                  best_f, best_epoch
              )
          )
  
imports: import logging
  import sys
  import os
  import shutil
  import time
  from datetime import datetime
  import argparse
  import numpy
  import random
  import torch
  from torch import nn, optim
  from stanza.models.lemma.data import DataLoader
  from stanza.models.lemma.vocab import Vocab
  from stanza.models.lemma.trainer import Trainer
  from stanza.models.lemma import scorer, edit
  from stanza.models.common import utils
  import stanza.models.common.seq2seq_constant
  from stanza.models.common.doc import *
  from stanza.utils.conll import CoNLL
  from stanza.models import _training_logging
label_desc: Use image node_id 5 for calling a global function with example usage: train(args) without return types