code: def train_model_one_batch(
      epoch,
      batch_idx,
      model,
      training_batch,
      transition_tensors,
      process_outputs,
      model_loss_function,
      oracle,
      args,
  ):
      """
      Train the model for one batch
  
      The model itself will be updated, and a bunch of stats are returned
      It is unclear if this refactoring is useful in any way.  Might not be
  
      ... although the indentation does get pretty ridiculous if this is
      merged into train_model_one_epoch and then iterate_training
      """
      # now we add the state to the trees in the batch
      # the state is build as a bulk operation
      initial_states = model.initial_state_from_preterminals(
          [x.preterminals for x in training_batch],
          [x.tree for x in training_batch],
      )
      initial_states = [
          state._replace(gold_sequence=sequence)
          for (tree, sequence, _), state in zip(
              training_batch, initial_states
          )
      ]
      current_batch = initial_states
  
      transitions_correct = Counter()
      transitions_incorrect = Counter()
      repairs_used = Counter()
      fake_transitions_used = 0
  
      all_errors = []
      all_answers = []
  
      # we iterate through the batch in the following sequence:
      # predict the logits and the applied transition for each tree in the batch
      # collect errors
      #  - we always train to the desired one-hot vector
      #    this was a noticeable improvement over training just the
      #    incorrect transitions
      # determine whether the training can continue using the "student" transition
      #   or if we need to use teacher forcing
      # update all states using either the gold or predicted transition
      # any trees which are now finished are removed from the training cycle
      while len(current_batch) > 0:
          outputs, pred_transitions, _ = model.predict(
              current_batch, is_legal=False
          )
          gold_transitions = [
              x.gold_sequence[x.num_transitions()]
              for x in current_batch
          ]
          trans_tensor = [
              transition_tensors[gold_transition]
              for gold_transition in gold_transitions
          ]
          all_errors.append(outputs)
          all_answers.extend(trans_tensor)
  
          new_batch = []
          update_transitions = []
          for pred_transition, gold_transition, state in zip(
              pred_transitions, gold_transitions, current_batch
          ):
              # forget teacher forcing vs scheduled sampling
              # we're going with idiot forcing
              if pred_transition == gold_transition:
                  transitions_correct[gold_transition.short_name()] += 1
                  if state.num_transitions() + 1 < len(
                      state.gold_sequence
                  ):
                      if (
                          oracle is not None
                          and random.random()
                          < args["oracle_forced_errors"]
                      ):
                          # TODO: could randomly choose from the legal transitions
                          fake_transition = random.choice(
                              model.transitions
                          )
                          if fake_transition.is_legal(state, model):
                              _, new_sequence = oracle.fix_error(
                                  gold_transition,
                                  fake_transition,
                                  state.gold_sequence,
                                  state.num_transitions(),
                              )
                              if new_sequence is not None:
                                  new_batch.append(
                                      state._replace(
                                          gold_sequence=new_sequence
                                      )
                                  )
                                  update_transitions.append(
                                      fake_transition
                                  )
                                  fake_transitions_used = (
                                      fake_transitions_used + 1
                                  )
                                  continue
                      new_batch.append(state)
                      update_transitions.append(gold_transition)
                  continue
  
              transitions_incorrect[
                  gold_transition.short_name(),
                  pred_transition.short_name(),
              ] += 1
              # if we are on the final operation, there are two choices:
              #   - the parsing mode is IN_ORDER, and the final transition
              #     is the close to end the sequence, which has no alternatives
              #   - the parsing mode is something else, in which case
              #     we have no oracle anyway
              if state.num_transitions() + 1 >= len(
                  state.gold_sequence
              ):
                  continue
  
              if (
                  oracle is None
                  or epoch < args["oracle_initial_epoch"]
                  or not pred_transition.is_legal(state, model)
              ):
                  new_batch.append(state)
                  update_transitions.append(gold_transition)
                  continue
  
              repair_type, new_sequence = oracle.fix_error(
                  gold_transition,
                  pred_transition,
                  state.gold_sequence,
                  state.num_transitions(),
              )
              # we can only reach here on an error
              assert not repair_type.is_correct()
              repairs_used[repair_type] += 1
              if (
                  new_sequence is not None
                  and random.random() < args["oracle_frequency"]
              ):
                  new_batch.append(
                      state._replace(gold_sequence=new_sequence)
                  )
                  update_transitions.append(pred_transition)
              else:
                  new_batch.append(state)
                  update_transitions.append(gold_transition)
  
          if len(current_batch) > 0:
              # bulk update states - significantly faster
              current_batch = parse_transitions.bulk_apply(
                  model, new_batch, update_transitions, fail=True
              )
  
      errors = torch.cat(all_errors)
      answers = torch.cat(all_answers)
  
      errors = process_outputs(errors)
      tree_loss = model_loss_function(errors, answers)
      tree_loss.backward()
      if args["watch_regex"]:
          matched = False
          logger.info(
              "Watching %s   ... epoch %d batch %d",
              args["watch_regex"],
              epoch,
              batch_idx,
          )
          watch_regex = re.compile(args["watch_regex"])
          for n, p in model.named_parameters():
              if watch_regex.search(n):
                  matched = True
                  if p.requires_grad and p.grad is not None:
                      logger.info(
                          "  %s norm: %f grad: %f",
                          n,
                          torch.linalg.norm(p),
                          torch.linalg.norm(p.grad),
                      )
                  elif p.requires_grad:
                      logger.info(
                          "  %s norm: %f grad required, but is None!",
                          n,
                          torch.linalg.norm(p),
                      )
                  else:
                      logger.info(
                          "  %s norm: %f grad not required",
                          n,
                          torch.linalg.norm(p),
                      )
          if not matched:
              logger.info("  (none found!)")
      if torch.any(torch.isnan(tree_loss)):
          batch_loss = 0.0
          nans = 1
      else:
          batch_loss = tree_loss.item()
          nans = 0
  
      return EpochStats(
          batch_loss,
          transitions_correct,
          transitions_incorrect,
          repairs_used,
          fake_transitions_used,
          nans,
      )
  
imports: from collections import Counter
  from collections import namedtuple
  import copy
  from enum import Enum
  import logging
  from operator import itemgetter
  import os
  import random
  import re
  import sys
  import torch
  from torch import nn
  from stanza.models.common import pretrain
  from stanza.models.common import utils
  from stanza.models.common.foundation_cache import load_bert, load_charlm, load_pretrain, FoundationCache
  from stanza.models.common.large_margin_loss import LargeMarginInSoftmaxLoss
  from stanza.models.constituency import parse_transitions
  from stanza.models.constituency import parse_tree
  from stanza.models.constituency import transition_sequence
  from stanza.models.constituency import tree_reader
  from stanza.models.constituency.base_model import SimpleModel, UNARY_LIMIT
  from stanza.models.constituency.in_order_oracle import InOrderOracle
  from stanza.models.constituency.lstm_model import LSTMModel, StackHistory
  from stanza.models.constituency.parse_transitions import TransitionScheme
  from stanza.models.constituency.parse_tree import Tree
  from stanza.models.constituency.utils import retag_tags, retag_trees, build_optimizer, build_scheduler
  from stanza.models.constituency.utils import DEFAULT_LEARNING_EPS, DEFAULT_LEARNING_RATES, DEFAULT_LEARNING_RHO, DEFAULT_WEIGHT_DECAY
  from stanza.server.parser_eval import EvaluateParser, ParseResult
  from stanza.utils.get_tqdm import get_tqdm
  from stanza.utils.training.common import find_wordvec_pretrain, choose_charlm, find_charlm_file
  from stanza.resources.default_packages import default_charlms, default_pretrains
label_desc: Use image node_id 21 for calling a global function with example usage: train_model_one_batch(epoch, batch_idx, model, training_batch, transition_tensors, process_outputs, model_loss_function, oracle, args) and returns: EpochStats