code: def train(args):
      # load data
      logger.debug("max_dec_len: %d" % args["max_dec_len"])
      logger.debug(
          "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.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"])
      save_name = (
          args["save_name"]
          if args["save_name"]
          else "{}_mwt_expander.pt".format(args["shorthand"])
      )
      model_file = os.path.join(args["save_dir"], save_name)
  
      # 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:
          logger.warning("Skip training because no data available...")
          return
      dev_mwt = dev_doc.get_mwt_expansions(False)
      if len(dev_batch) == 0 and args.get("dict_only", False):
          logger.warning(
              "Training data available, but dev data has no MWTs.  Only training a dict based MWT"
          )
          args["dict_only"] = True
  
      # train a dictionary-based MWT expander
      trainer = Trainer(args=args, vocab=vocab, device=args["device"])
      logger.info("Training dictionary-based MWT expander...")
      trainer.train_dict(
          train_batch.doc.get_mwt_expansions(evaluation=False)
      )
      logger.info("Evaluating on dev set...")
      dev_preds = trainer.predict_dict(
          dev_batch.doc.get_mwt_expansions(evaluation=True)
      )
      doc = copy.deepcopy(dev_batch.doc)
      doc.set_mwt_expansions(dev_preds, fake_dependencies=True)
      CoNLL.write_doc2conll(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:
          # train a seq2seq model
          logger.info("Training seq2seq-based MWT expander...")
          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}"
  
          if args["wandb"]:
              import wandb
  
              wandb_name = (
                  args["wandb_name"]
                  if args["wandb_name"]
                  else "%s_mwt" % 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")
  
          # 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 = []
              for i, batch in enumerate(dev_batch):
                  preds = trainer.predict(batch)
                  dev_preds += preds
              if args.get("ensemble_dict", False) and args.get(
                  "ensemble_early_stop", False
              ):
                  logger.info("[Ensembling dict with seq2seq model...]")
                  dev_preds = trainer.ensemble(
                      dev_batch.doc.get_mwt_expansions(evaluation=True),
                      dev_preds,
                  )
              doc = copy.deepcopy(dev_batch.doc)
              doc.set_mwt_expansions(dev_preds, fake_dependencies=True)
              CoNLL.write_doc2conll(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]
              ):
                  current_lr *= args["lr_decay"]
                  trainer.change_lr(current_lr)
  
              dev_score_history += [dev_score]
  
          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
              )
          )
  
          # 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_mwt_expansions(evaluation=True),
                  best_dev_preds,
              )
              doc = copy.deepcopy(dev_batch.doc)
              doc.set_mwt_expansions(dev_preds, fake_dependencies=True)
              CoNLL.write_doc2conll(doc, system_pred_file)
              _, _, dev_score = scorer.score(
                  system_pred_file, gold_file
              )
              logger.info(
                  "Ensemble dev F1 = {:.2f}".format(dev_score * 100)
              )
              best_f = max(best_f, dev_score)
  
imports: import sys
  import os
  import shutil
  import time
  from datetime import datetime
  import argparse
  import logging
  import numpy
  import random
  import torch
  from torch import nn, optim
  import copy
  from stanza.models.mwt.data import DataLoader
  from stanza.models.mwt.vocab import Vocab
  from stanza.models.mwt.trainer import Trainer
  from stanza.models.mwt import scorer
  from stanza.models.common import utils
  import stanza.models.common.seq2seq_constant
  from stanza.models.common.doc import Document
  from stanza.utils.conll import CoNLL
  from stanza.models import _training_logging
label_desc: Use image node_id 4 for calling a global function with example usage: train(args) without return types