text: # Copyright (c) Facebook, Inc. and its affiliates.
  #
  # This source code is licensed under the MIT license found in the
  # LICENSE file in the root directory of this source tree.
  
  from . import LegacyFairseqLRScheduler, register_lr_scheduler
  import logging
  import ast
  
  logger = logging.getLogger(__name__)
  logger.setLevel(logging.WARNING)
  
  
  @register_lr_scheduler("manual")
  class ManualSchedule(LegacyFairseqLRScheduler):
      """Decay the LR on a manual schedule."""
  
      def __init__(self, args, optimizer):
          super().__init__(args, optimizer)
  
          self.epoch2lr = self.parse_manuallr_args(args.epoch2lr)
          self.update2lr = self.parse_manuallr_args(args.update2lr)
          logger.info("@@@ ManualSchedule epoch2lr={}".format(self.epoch2lr))
          logger.info("@@@ ManualSchedule update2lr={}".format(self.update2lr))
  
          if 1 in self.epoch2lr:
              self.lr = self.epoch2lr[1]
          elif 1 in self.update2lr:
              self.lr = self.update2lr[1]
          else:
              self.lr = args.lr[0]
          self.optimizer.set_lr(self.lr)  # Set the beginning of the epoch.
  
      def parse_manuallr_args(self, lr_args_str):
          lr_dict = ast.literal_eval(lr_args_str.replace(" ", ""))
          if not isinstance(lr_dict, dict):
              raise ValueError("epoch2lr/update2lr must be abel to evaluated to a dict")
  
          lr_args = {}
          logger.info("@@@ after parsing input dictionary lr_dict = {}".format(lr_dict))
          for key, val in lr_dict.items():
              if "," in key:
                  for k in key.split(","):
                      lr_args[int(k)] = float(val)
              elif "-" in key:
                  s = int(key.split("-")[0])
                  e = int(key.split("-")[1])
                  for k in range(s, e + 1, 1):
                      lr_args[k] = float(val)
              else:
                  lr_args[int(key)] = float(val)
  
          return lr_args
  
      @staticmethod
      def add_args(parser):
          """Add arguments to the parser for this LR scheduler."""
          # fmt: off
          parser.add_argument(
              "--epoch2lr",
              type=str,
              metavar="DICT",
              default="{}",
              help="a dictionary used to set lr for each epoch manually",
          )
          parser.add_argument(
              "--update2lr",
              type=str,
              metavar="DICT",
              default="{}",
              help="a dictionary used to set lr for each update manually",
          )
          # fmt: on
  
      def state_dict(self):
          return {"lr": self.lr}
  
      def load_state_dict(self, state_dict):
          if "lr" in state_dict:
              self.lr = state_dict["lr"]
  
      def get_next_lr(self, epoch):
          manual_keys = [k for k in self.epoch2lr if k <= epoch]
          if manual_keys:
              manual_lr = self.epoch2lr[max(manual_keys)]
          else:
              logger.warning(
                  "@@@ epoch={} does not exist in manual lr input. epoch2lr={}...".format(
                      epoch,
                      list(self.epoch2lr.items())[
                          : min(10, len(self.epoch2lr.keys()) - 1)
                      ],
                  )
              )
              manual_lr = self.optimizer.get_lr()
          return manual_lr
  
      def step_begin_epoch(self, epoch):
          """Update the learning rate at the beginning of the given epoch."""
          self.lr = self.get_next_lr(epoch)
          self.optimizer.set_lr(self.lr)
          return self.optimizer.get_lr()
  
      def step_update(self, num_updates):
          """Update the learning rate after each update."""
          manual_keys = [k for k in self.update2lr if k <= num_updates]
          if manual_keys:
              manual_lr = self.update2lr[max(manual_keys)]
          else:
              logger.warning(
                  "epoch={} does not exist in manual lr input update2lr={}...".format(
                      num_updates,
                      list(self.update2lr.items())[
                          : min(10, len(self.update2lr.keys()) - 1)
                      ],
                  )
              )
              manual_lr = self.optimizer.get_lr()
  
          self.optimizer.set_lr(manual_lr)
          return self.optimizer.get_lr()
  