code: def update_bert_learning_rate(args, optimizer, epochs_trained):
      """
      Update the learning rate for the bert finetuning, if applicable
      """
      # would be nice to have a parameter group specific scheduler
      # however, there is an issue with the optimizer we had the most success with, madgrad
      # when the learning rate is 0 for a group, it still learns by some
      # small amount because of the eps parameter
      # in fact, that is enough to make the learning for the bert in the
      # second half broken
      for base_param_group in optimizer.param_groups:
          if base_param_group["param_group_name"] == "base":
              break
      else:
          raise AssertionError(
              "There should always be a base parameter group"
          )
      for param_group in optimizer.param_groups:
          if param_group["param_group_name"] == "bert":
              # Occasionally a model goes haywire and forgets how to use the transformer
              # So far we have only seen this happen with Electra on the non-NML version of PTB
              # We tried fixing that with an increasing transformer learning rate, but that
              # didn't fully resolve the problem
              # Switching to starting the finetuning after a few epochs seems to help a lot, though
              old_lr = param_group["lr"]
              if (
                  args["bert_finetune_begin_epoch"] is not None
                  and epochs_trained < args["bert_finetune_begin_epoch"]
              ):
                  param_group["lr"] = 0.0
              elif (
                  args["bert_finetune_end_epoch"] is not None
                  and epochs_trained >= args["bert_finetune_end_epoch"]
              ):
                  param_group["lr"] = 0.0
              elif (
                  args["multistage"]
                  and epochs_trained < args["epochs"] // 2
              ):
                  param_group["lr"] = (
                      base_param_group["lr"]
                      * args["stage1_bert_learning_rate"]
                  )
              else:
                  param_group["lr"] = (
                      base_param_group["lr"]
                      * args["bert_learning_rate"]
                  )
              if param_group["lr"] != old_lr:
                  logger.info(
                      "Setting %s finetuning rate from %f to %f",
                      param_group["param_group_name"],
                      old_lr,
                      param_group["lr"],
                  )
  
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 18 for calling a global function with example usage: update_bert_learning_rate(args, optimizer, epochs_trained) without return types