code: def train_model(
      trainer,
      model_file,
      checkpoint_file,
      args,
      train_set,
      dev_set,
      labels,
  ):
      tlogger.setLevel(logging.DEBUG)
  
      # TODO: use a (torch) dataloader to possibly speed up the GPU usage
      model = trainer.model
      optimizer = trainer.optimizer
  
      device = next(model.parameters()).device
      logger.info("Current device: %s" % device)
  
      label_map = {x: y for (y, x) in enumerate(labels)}
      label_tensors = {
          x: torch.tensor(y, requires_grad=False, device=device)
          for (y, x) in enumerate(labels)
      }
  
      process_outputs = lambda x: x
      if args.loss == Loss.CROSS:
          logger.info("Creating CrossEntropyLoss")
          loss_function = nn.CrossEntropyLoss()
      elif args.loss == Loss.WEIGHTED_CROSS:
          logger.info("Creating weighted cross entropy loss w/o log")
          loss_function = loss.weighted_cross_entropy_loss(
              [label_map[x[0]] for x in train_set], log_dampened=False
          )
      elif args.loss == Loss.LOG_CROSS:
          logger.info("Creating weighted cross entropy loss w/ log")
          loss_function = loss.weighted_cross_entropy_loss(
              [label_map[x[0]] for x in train_set], log_dampened=True
          )
      elif args.loss == Loss.FOCAL:
          try:
              from focal_loss.focal_loss import FocalLoss
          except ImportError:
              raise ImportError(
                  "focal_loss not installed.  Must `pip install focal_loss_torch` to use the --loss=focal feature"
              )
          logger.info(
              "Creating FocalLoss with loss %f", args.loss_focal_gamma
          )
          process_outputs = lambda x: torch.softmax(x, dim=1)
          loss_function = FocalLoss(gamma=args.loss_focal_gamma)
      else:
          raise ValueError("Unknown loss function {}".format(args.loss))
      loss_function.to(device)
  
      train_set_by_len = data.sort_dataset_by_len(train_set)
  
      if trainer.global_step > 0:
          # We reloaded the model, so let's report its current dev set score
          _ = score_dev_set(model, dev_set, args.dev_eval_scoring)
          logger.info("Reloaded model for continued training.")
          if trainer.best_score is not None:
              logger.info(
                  "Previous best score: %.5f", trainer.best_score
              )
  
      log_param_sizes(model)
  
      # https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html
      batch_starts = list(range(0, len(train_set), args.batch_size))
  
      if args.wandb:
          import wandb
  
          wandb_name = (
              args.wandb_name
              if args.wandb_name
              else "%s_classifier" % args.shorthand
          )
          wandb.init(name=wandb_name, config=args)
          wandb.run.define_metric("accuracy", summary="max")
          wandb.run.define_metric("macro_f1", summary="max")
          wandb.run.define_metric("epoch_loss", summary="min")
  
      for trainer.epochs_trained in range(
          trainer.epochs_trained, args.max_epochs
      ):
          running_loss = 0.0
          epoch_loss = 0.0
          shuffled = data.shuffle_dataset(train_set_by_len)
  
          model.train()
          logger.info("Starting epoch %d", trainer.epochs_trained)
          if args.log_norms:
              model.log_norms()
  
          random.shuffle(batch_starts)
          for batch_num, start_batch in enumerate(batch_starts):
              trainer.global_step += 1
              logger.debug(
                  "Starting batch: %d step %d",
                  start_batch,
                  trainer.global_step,
              )
  
              batch = shuffled[
                  start_batch : start_batch + args.batch_size
              ]
              batch_labels = torch.stack(
                  [label_tensors[x.sentiment] for x in batch]
              )
  
              # zero the parameter gradients
              optimizer.zero_grad()
  
              outputs = model(batch)
              outputs = process_outputs(outputs)
              batch_loss = loss_function(outputs, batch_labels)
              batch_loss.backward()
              optimizer.step()
  
              # print statistics
              running_loss += batch_loss.item()
              if (
                  (batch_num + 1) * args.batch_size
              ) % args.tick < args.batch_size:  # print every 2000 items
                  train_loss = running_loss / args.tick
                  logger.info(
                      "[%d, %5d] Average loss: %.3f",
                      trainer.epochs_trained + 1,
                      (batch_num + 1) * args.batch_size,
                      train_loss,
                  )
                  if args.wandb:
                      wandb.log(
                          {"train_loss": train_loss},
                          step=trainer.global_step,
                      )
                  if (
                      args.dev_eval_steps > 0
                      and ((batch_num + 1) * args.batch_size)
                      % args.dev_eval_steps
                      < args.batch_size
                  ):
                      logger.info("---- Interim analysis ----")
                      dev_score, accuracy, macro_f1 = score_dev_set(
                          model, dev_set, args.dev_eval_scoring
                      )
                      if args.wandb:
                          wandb.log(
                              {
                                  "accuracy": accuracy,
                                  "macro_f1": macro_f1,
                              },
                              step=trainer.global_step,
                          )
                      if (
                          trainer.best_score is None
                          or dev_score > trainer.best_score
                      ):
                          trainer.best_score = dev_score
                          trainer.save(model_file, save_optimizer=False)
                          logger.info(
                              "Saved new best score model!  Accuracy %.5f   Macro F1 %.5f   Epoch %5d   Batch %d"
                              % (
                                  accuracy,
                                  macro_f1,
                                  trainer.epochs_trained + 1,
                                  batch_num + 1,
                              )
                          )
                      model.train()
                  epoch_loss += running_loss
                  running_loss = 0.0
          # Add any leftover loss to the epoch_loss
          epoch_loss += running_loss
  
          logger.info(
              "Finished epoch %d  Total loss %.3f"
              % (trainer.epochs_trained + 1, epoch_loss)
          )
          dev_score, accuracy, macro_f1 = score_dev_set(
              model, dev_set, args.dev_eval_scoring
          )
          if args.wandb:
              wandb.log(
                  {
                      "accuracy": accuracy,
                      "macro_f1": macro_f1,
                      "epoch_loss": epoch_loss,
                  },
                  step=trainer.global_step,
              )
          if checkpoint_file:
              trainer.save(
                  checkpoint_file,
                  epochs_trained=trainer.epochs_trained + 1,
              )
          if args.save_intermediate_models:
              intermediate_file = intermediate_name(
                  model_file,
                  trainer.epochs_trained + 1,
                  args.dev_eval_scoring,
                  dev_score,
              )
              trainer.save(intermediate_file, save_optimizer=False)
          if (
              trainer.best_score is None
              or dev_score > trainer.best_score
          ):
              trainer.best_score = dev_score
              trainer.save(model_file, save_optimizer=False)
              logger.info(
                  "Saved new best score model!  Accuracy %.5f   Macro F1 %.5f   Epoch %5d"
                  % (accuracy, macro_f1, trainer.epochs_trained + 1)
              )
  
      if args.wandb:
          wandb.finish()
  
imports: import argparse
  import ast
  import logging
  import os
  import random
  import re
  from enum import Enum
  import torch
  import torch.nn
  from stanza.models.common import loss
  from stanza.models.common import utils
  from stanza.models.pos.vocab import CharVocab
  import stanza.models.classifiers.data
  from stanza.models.classifiers.trainer import Trainer
  from stanza.models.classifiers.utils import WVType, ExtraVectors, ModelType
  from stanza.utils.confusion import format_confusion, confusion_to_accuracy, confusion_to_macro_f1
label_desc: Use image node_id 10 for calling a global function with example usage: train_model(trainer, model_file, checkpoint_file, args, train_set, dev_set, labels) without return types