content: import json
  import os
  from typing import Union, List, Optional, Dict
  from .presets import *
  
  class Config:
      """Base class for TrainingConfig and DistillationConfig."""
      def __init__(self,**kwargs):
          pass
  
      @classmethod
      def from_json_file(cls, json_filename):
          """Construct configurations from a json file."""
          with open(json_filename,'r') as f:
              json_data = json.load(f)
          return cls.from_dict(json_data)
  
      @classmethod
      def from_dict(cls, dict_object):
          """Construct configurations from a dict."""
          config = cls(**dict_object)
          return config
  
      def __str__(self):
          str = ""
          for k,v in self.__dict__.items():
              str += f"{k} : {v}\n"
          return str
  
      def __repr__(self):
          classname = self.__class__.__name__
          return classname +":\n"+self.__str__()
  
  
  class TrainingConfig(Config):
      """
      Configurations related to general model training.
  
      Args:
          gradient_accumulation_steps (int): accumulates gradients before optimization to reduce GPU memory usage. It calls ``optimizer.step()`` every `gradient_accumulation_steps` backward steps.
          ckpt_frequency (int): stores model weights `ckpt_frequency` times every epoch.
          ckpt_epoch_frequency (int): stores model weights every `ckpt_epoch_frequency` epochs.
          ckpt_steps (int):  if *num_steps* is passes to ``distiller.train()``, saves the model every **ckpt_steps**, meanwhile ignore `ckpt_frequency` and `ckpt_epoch_frequency` .
          log_dir (str): directory to save the tensorboard log file. Set it to ``None`` to disable tensorboard.
          output_dir (str): directory to save model weights.
          device (str or torch.device): training on CPU or GPU.
          fp16 (bool): if ``True``, enables mixed precision training using Apex.
          fp16_opt_level(str): Pure or mixed precision optimization level. Accepted values are "O0", "O1", "O2", and "O3". See Apex documenation for details.
          data_parallel (bool): If ``True``, wraps the models with ``torch.nn.DataParallel``.
          local_rank (int): the local rank of the current processes. A non-nagative value means that we are in the distributed training mode with ``DistributedDataParallel``.  
      Note:
          * To perform data parallel (DP) training, you could either wrap the models with ``torch.nn.DataParallel`` outside TextBrewer by yourself, or leave the work for TextBrewer by setting **data_parallel** to ``True``.
          * To enable both data parallel training and mixed precision training, you should set **data_parallel** to ``True``, and DO NOT wrap the models by yourself.
          * In some experiments, we have observed slowing down in the speed with ``torch.nn.DataParallel``.
          * To perform distributed data parallel (DDP) training, you should call ``torch.distributed.init_process_group`` before intializing a TrainingConfig; and pass the **raw** (unwrapped) model when initializing the distiller.
          * DP and DDP are mutual exclusive.
      Example::
  
          # Usually just need to set log_dir and output_dir and leave others default
          train_config = TrainingConfig(log_dir=my_log_dir, output_dir=my_output_dir)
          
          # Stores the model at the end of each epoch
          train_config = TrainingConfig(ckpt_frequency=1, ckpt_epoch_frequency=1)
          # Stores the model twice (at the middle and at the end) in each epoch
          train_config = TrainingConfig(ckpt_frequency=2, ckpt_epoch_frequency=1)
          # Stores the model once every two epochs
          train_config = TrainingConfig(ckpt_frequency=1, ckpt_epoch_frequency=2)
  
      """
      def __init__(self,gradient_accumulation_steps = 1,
                   ckpt_frequency = 1,
                   ckpt_epoch_frequency = 1,
                   ckpt_steps = None,
                   log_dir = None,
                   output_dir = './saved_models',
                   device = 'cuda',
                   fp16 = False,
                   fp16_opt_level = 'O1',
                   data_parallel = False,
                   local_rank = -1
                   ):
          super(TrainingConfig, self).__init__()
  
          self.gradient_accumulation_steps =gradient_accumulation_steps
          self.ckpt_frequency = ckpt_frequency
          self.ckpt_epoch_frequency = ckpt_epoch_frequency
          self.ckpt_steps = ckpt_steps
          self.log_dir = log_dir
          self.output_dir = output_dir
          self.device = device
          self.fp16 = fp16
          self.fp16_opt_level = fp16_opt_level
          self.data_parallel = data_parallel
  
          self.local_rank = local_rank
          if self.local_rank == -1 or torch.distributed.get_rank() == 0:
              if not os.path.exists(self.output_dir):
                  os.makedirs(self.output_dir)
  
  
  class IntermediateMatch:
      def __init__(self,layer_T: Union[int,List[int]], layer_S: Union[int,List[int]],
                   weight: float, loss: str, feature: str, proj: Optional[List] = None):
          self.layer_T = layer_T
          self.layer_S = layer_S
          self.feature = feature
          self.weight = weight
          self.loss = loss
          self.proj = proj
          assert feature in FEATURES
          if proj:
              assert proj[0] in PROJ_MAP.keys()
              assert type(proj[1]) is int and type(proj[2]) is int
              if len(proj)==3:
                  self.proj.append(dict())   # ['linear', dim_T, dim_S, {...}]
              else:
                  assert type(proj[3]) is dict
  
      def __str__(self):
          str = ""
          for k,v in self.__dict__.items():
              str += f"{k} : {v}, "
          return str[:-2]
  
      def __repr__(self):
          classname = self.__class__.__name__
          return '\n'+classname +": "+self.__str__()
  
      @classmethod
      def from_dict(cls,dict_object):
          if dict_object is None:
              return None
          else:
              return cls(**dict_object)
  
  
  class DistillationConfig(Config):
      """
      Configurations related to distillation methods. It defines the total loss to be optimized:
      
      .. math::
  
          \mathcal{L}_{total}=  \mathcal{L}_{KD} * w_{KD} + \mathcal{L}_{hl} * w_{hl} + sum(\\textrm{intermediate_losses})
      
      where
  
          * :math:`\mathcal{L}_{KD}` is the KD loss on logits, :math:`w_{KD}` is its weight;
          * :math:`\mathcal{L}_{hl}` is the sum of ``losses`` returned by the adaptor and :math:`w_{hl}` is its weight;
          * intermediate_losses is defined via `intermediate_matches`.
      
      Args:
          temperature (float) :temperature for the distillation. The teacher and student models' logits will be divided by the temperature in computing the KD loss. The temperature typicially ranges from 1 to 10. We found that temperature higher than 1 usually leads to better performance.
          temperature_scheduler: dynamically adjusts temperature. See :data:`~textbrewer.presets.TEMPERATURE_SCHEDULER` for all available options.
          kd_loss_type (str): KD loss function for the ``logits`` term returned by the adaptor, can be ``'ce'`` or ``'mse'``. See :data:`~textbrewer.presets.KD_LOSS_MAP`.
          kd_loss_weight (float): the weight for the KD loss.
          hard_label_weight (float): the weight for the sum of ``losses`` term returned by the adaptor. ``losses`` may include the losses on the ground-truth labels and other user-defined losses. 
          kd_loss_weight_scheduler: Dynamically adjusts KD loss weight. See :data:`~textbrewer.presets.WEIGHT_SCHEDULER` for all available options.
          hard_label_weight_scheduler: Dynamically adjusts the weight of the sum of ``losses``. See :data:`~textbrewer.presets.WEIGHT_SCHEDULER` for all available options.
          probability_shift (bool): if ``True``, switch the ground-truth label's logit and the largest logit predicted by the teacher, to make the ground-truth label's logit largest. Requires ``labels`` term returned by the adaptor.
          is_caching_logits (bool): if ``True``, caches the batches and the output logits of the teacher model in memory, so that those logits will only be calcuated once. It will speed up the distillation process. This feature is **only available** for :class:`~textbrewer.BasicDistiller` and :class:`~textbrewer.MultiTeacherDistiller`, and only when distillers' ``train()`` method is called with ``num_steps=None``. It is suitable for small and medium datasets.
          intermediate_matches (`List[Dict]`) : Configuration for intermediate feature matching. Each element in the list is a dict, representing a pair of matching config. 
      
      The dict in `intermediate_matches` contains the following keys:
  
          * '**layer_T**': `layer_T` (*int*): selects the layer_T-th layer of teacher model.
          * '**layer_S**': `layer_S` (*int*): selects the layer_S-th layer of student model.
  
          .. Note::
          
              1. `layer_T` and `layer_S` indicate layers in ``attention`` or ``hidden`` list in the returned dict of the adaptor, rather than the actual layers in the model. 
              2. If the loss is :func:`fst <textbrewer.losses.fsp_loss>` or :func:`nst <textbrewer.losses.mmd_loss>`, two layers have to be chosen from the teacher and the student respectively. In this case, `layer_T` and `layer_S` are lists of two ints. See the example below.
  
          * '**feature**': `feature` (*str*): features of intermediate layers. It can be:
  
              * '**attention**' : attention matrix, of the shape (*batch_size*, *num_heads*, *length*, *length*) or (*batch_size*, *length*, *length*)
              * '**hidden**'：hidden states, of the shape (*batch_size*, *length*, *hidden_dim*).
  
          * '**loss**' : `loss` (*str*) : loss function. See :data:`~textbrewer.presets.MATCH_LOSS_MAP` for available losses. Currently includes: ``'attention_mse'``, ``'attention_ce'``, ``'hidden_mse'``, ``'nst'``, etc.
          * '**weight**': `weight` (float) : weight for the loss.
          * '**proj**' : `proj` (*List*, optional) : if the teacher and the student have the same feature dimension, it is optional; otherwise it is required. It is the mapping function to match teacher and student intermediate feature dimension. It is a list, with these elements:
  
              * **proj[0]** (*str*): mapping function, can be ``'linear'``, ``'relu'``, ``'tanh'``. See :data:`~textbrewer.presets.PROJ_MAP`.
              * **proj[1]** (*int*): feature dimension of student model.
              * **proj[2]** (*int*): feature dimension of teacher model.
              * **proj[3]** (*dict*): optional, provides configurations such as learning rate. If not provided, the learning rate and optimizer configurations will follow the default config of the optimizer, otherwise it will use the ones specified here.
  
      Example::
  
          from textbrewer import DistillationConfig
  
          # simple configuration: use default values, or try different temperatures 
          distill_config = DistillationConfig(temperature=8)
  
          # adding intermediate feature matching
          # under this setting, the returned dict results_T/S of adaptor_T/S should contain 'hidden' key.
          # The mse loss between teacher's results_T['hidden'][10] and student's results_S['hidden'][3] will be computed
          distill_config = DistillationConfig(
              temperature=8,
              intermediate_matches = [{'layer_T':10, 'layer_S':3, 'feature':'hidden', 'loss':'hidden_mse', 'weight':1}]
          )
  
          # multiple inatermediate feature matching. The teacher and the student have a  hidden_dim of 768 and 384 respectively. 
          distill_config = DistillationConfig(
              temperature = 8, 
              intermediate_matches = [ \\
              {'layer_T':0,  'layer_S':0, 'feature':'hidden','loss': 'hidden_mse', 'weight' : 1,'proj':['linear',384,768]},
              {'layer_T':4,  'layer_S':1, 'feature':'hidden','loss': 'hidden_mse', 'weight' : 1,'proj':['linear',384,768]},
              {'layer_T':8,  'layer_S':2, 'feature':'hidden','loss': 'hidden_mse', 'weight' : 1,'proj':['linear',384,768]},
              {'layer_T':12, 'layer_S':3, 'feature':'hidden','loss': 'hidden_mse', 'weight' : 1,'proj':['linear',384,768]}]
          )
  
      """
      def __init__(self,temperature=4,
                        temperature_scheduler = 'none',
                        hard_label_weight=0,
                        hard_label_weight_scheduler = 'none',
                        kd_loss_type='ce',
                        kd_loss_weight=1,
                        kd_loss_weight_scheduler = 'none',
                        probability_shift = False,
                        intermediate_matches:Optional[List[Dict]]=None,
                        is_caching_logits = False):
          super(DistillationConfig, self).__init__()
  
          self.temperature = temperature
          self.temperature_scheduler = None
          if temperature_scheduler != 'none':
              assert temperature_scheduler in TEMPERATURE_SCHEDULER, \
                      f"Invalid temperature_scheduler {temperature_scheduler}"
              self.temperature_scheduler = TEMPERATURE_SCHEDULER[temperature_scheduler]
  
          self.hard_label_weight = hard_label_weight
          self.hard_label_weight_scheduler = None
          if hard_label_weight_scheduler != 'none':
              assert hard_label_weight_scheduler in WEIGHT_SCHEDULER, \
                      "Invalid hard_label_weight_scheduler"
              self.hard_label_weight_scheduler = WEIGHT_SCHEDULER[hard_label_weight_scheduler]
  
          self.kd_loss_type = kd_loss_type
          self.kd_loss_weight = kd_loss_weight
          self.kd_loss_weight_scheduler = None
          if kd_loss_weight_scheduler != 'none':
              assert kd_loss_weight_scheduler in WEIGHT_SCHEDULER, \
                      "Invalid kd_loss_weight_scheduler"
              self.kd_loss_weight_scheduler = WEIGHT_SCHEDULER[kd_loss_weight_scheduler]
  
          self.probability_shift = probability_shift
  
          self.intermediate_matches:[List[IntermediateMatch]] = []
          if intermediate_matches:
              self.intermediate_matches = [IntermediateMatch.from_dict(im) for im in intermediate_matches]
  
          self.is_caching_logits = is_caching_logits
  