content: 
  from textbrewer.distiller_utils import *
  from textbrewer.distiller_basic import BasicDistiller
  from pyemd import emd_with_flow
  from scipy.special import softmax
  
  class EMDDistiller(BasicDistiller):
      """
      BERT-EMD
  
      Args:
          train_config (:class:`TrainingConfig`): training configuration.
          distill_config (:class:`DistillationConfig`): distillation configuration.
          model_T (:class:`torch.nn.Module`): teacher model.
          model_S (:class:`torch.nn.Module`): student model.
          adaptor_T (Callable): teacher model's adaptor.
          adaptor_S (Callable): student model's adaptor.
          emd (dict): configuration for EMD
  
      The roles of `adaptor_T` and `adaptor_S` are explained in :py:func:`adaptor`.
  
      """
      def __init__(self, train_config,
                   distill_config,
                   model_T,
                   model_S,
                   adaptor_T,
                   adaptor_S,
                   emd: Optional[Dict] = None):
          # emd:{'layer_num_S' : 4+1,   
          #      'layer_num_T' : 12+1,
          #      'feature': 'hidden', 
          #      'loss': 'mse', 
          #      'weight' : 1.0, 
          # 'proj':['linear',312,768]}
  
          super(EMDDistiller, self).__init__(train_config, distill_config, model_T, model_S, adaptor_T, adaptor_S)
  
          self.projs = []
          self.projs_group = []
          # for im in self.d_config.intermediate_matches:
          #     if im.proj is not None:
          #         projection = im.proj[0]
          #         dim_in = im.proj[1]
          #         dim_out = im.proj[2]
          #         self.projs_group.append(im.proj[3])
          #         self.projs.append(PROJ_MAP[projection](dim_in,dim_out))
          #         self.projs[-1].to(self.t_config.device)
          #     else:
          #         self.projs.append(None)
          #         self.projs_group.append(None)
          
          self.d_config.is_caching_logits = False
  
          self.layer_num_S = emd['layer_num_S']
          self.layer_num_T = emd['layer_num_T']
          self.emd_feature = emd['feature']
          self.emd_loss_type = emd['loss']
          self.emd_loss_weight = emd['weight']
          if self.emd_feature != 'hidden' or self.emd_loss_type != 'hidden_mse':
              raise NotImplementedError
  
          self.feature_weight_S = np.ones(self.layer_num_S - 1) / (self.layer_num_S - 1) # excluding the embedding layer
          self.feature_weight_T = np.ones(self.layer_num_T - 1) / (self.layer_num_T - 1)
  
          if isinstance(emd['proj'],list) and len(emd['proj']) > 0:
              projection,dim_in,dim_out = emd['proj']
              for im in range(self.layer_num_S):
                  self.projs.append(PROJ_MAP[projection](dim_in,dim_out))
                  self.projs[-1].to(self.t_config.device)
                  self.projs_group.append(dict())
  
  
  
      def train_on_batch(self, batch, args):
  
          (teacher_batch, results_T), (student_batch, results_S) = get_outputs_from_batch(batch, self.t_config.device, self.model_T, self.model_S, args)
  
          results_T = post_adaptor(self.adaptor_T(teacher_batch,results_T))
          results_S = post_adaptor(self.adaptor_S(student_batch,results_S))
  
          total_loss, losses_dict = self.compute_loss(results_S, results_T)
  
          return total_loss, losses_dict
  
  
      def compute_loss(self,results_S,results_T):
  
          losses_dict = dict()
  
          total_loss  = 0
          if 'logits' in results_T and 'logits' in results_S:
              logits_list_T = results_T['logits']  # list of tensor
              logits_list_S = results_S['logits']  # list of tensor
              total_kd_loss = 0
              if 'logits_mask' in results_S:
                  masks_list_S = results_S['logits_mask']
                  logits_list_S = select_logits_with_mask(logits_list_S,masks_list_S)  #(mask_sum, num_of_class)
              if 'logits_mask' in results_T:
                  masks_list_T = results_T['logits_mask']
                  logits_list_T = select_logits_with_mask(logits_list_T,masks_list_T)  #(mask_sum, num_of_class)
  
              for l_T,l_S in zip(logits_list_T,logits_list_S):
                  if self.d_config.temperature_scheduler is not None:
                      temperature = self.d_config.temperature_scheduler(l_S, l_T, self.d_config.temperature)
                  else:
                      temperature = self.d_config.temperature
                  total_kd_loss += self.kd_loss(l_S, l_T, temperature)
  
              total_loss += total_kd_loss * self.d_config.kd_loss_weight
              losses_dict['unweighted_kd_loss'] = total_kd_loss
  
          inters_T = {feature: results_T.get(feature,[]) for feature in FEATURES}
          inters_S = {feature: results_S.get(feature,[]) for feature in FEATURES}
          inputs_mask_T = results_T.get('inputs_mask',None)
          inputs_mask_S = results_S.get('inputs_mask',None)
  
          #hidden states and embedding
          feature = self.emd_feature
          emd_loss_weight = self.emd_loss_weight
          loss_type = self.emd_loss_type
          match_loss = MATCH_LOSS_MAP[loss_type]
  
          feature_maps_S = inters_S[feature][1:] # list of features
          feature_maps_T = inters_T[feature][1:] # list of features
  
          embeddings_S = inters_S[feature][0]
          embeddings_T = inters_T[feature][0]
  
          assert isinstance(feature_maps_S, (tuple,list))
          assert isinstance(feature_maps_T, (tuple,list))
          assert isinstance(feature_maps_S[0], torch.Tensor)
          assert isinstance(feature_maps_T[0], torch.Tensor)
          assert len(feature_maps_S) == self.layer_num_S - 1
          assert len(feature_maps_T) == self.layer_num_T - 1
  
  
          if len(self.projs) > 0:
              assert len(self.projs) == self.layer_num_S
              embeddings_S = self.projs[0](embeddings_S)
              feature_maps_S = [proj(s) for proj,s in zip(self.projs[1:],feature_maps_S)]
  
  
          feature_num_S = len(feature_maps_S)
          feature_num_T = len(feature_maps_T)
          feature_num_A = feature_num_S + feature_num_T
  
          distance_matrix = torch.zeros([feature_num_A, feature_num_A]).to(feature_maps_S[0])
          for s in range(feature_num_S):
              f_S = feature_maps_S[s]
              for t in range(feature_num_T):
                  f_T = feature_maps_T[t]
                  distance_matrix[s][t+feature_num_S] = distance_matrix[t+feature_num_S][s] = match_loss(f_S, f_T, mask=inputs_mask_S) 
  
          feature_weight_S = np.concatenate([self.feature_weight_S, np.zeros(feature_num_T)])
          feature_weight_T = np.concatenate([np.zeros(feature_num_S), self.feature_weight_T])
  
  
          _, trans_matrix = emd_with_flow(feature_weight_S, feature_weight_T, distance_matrix.detach().cpu().numpy().astype('float64'))
          trans_matrix = torch.tensor(trans_matrix).to(distance_matrix)
  
  
          emd_loss = torch.sum(trans_matrix * distance_matrix)
  
          total_loss += emd_loss * emd_loss_weight
  
          losses_dict[f'unweighted_{feature}_{loss_type}_emd'] = emd_loss
  
          if (self.feature_weight_S<=0).any() or (self.feature_weight_T<=0).any():
              import sys
              logger.info(f"{self.feature_weight_S}")
              logger.info(f"{self.feature_weight_T}")
  
          if np.isnan(self.feature_weight_S).any() or np.isnan(self.feature_weight_T).any():
              import sys
              logger.info(f"{self.feature_weight_S}")
              logger.info(f"{self.feature_weight_T}")
              sys.exit()
  
          #feature_weight_S = np.copy(self.feature_weight_S)
          #feature_weight_T = np.copy(self.feature_weight_T)
          #self.feature_weight_S, self.feature_weight_T = get_new_feature_weight(
          #    trans_matrix, distance_matrix.detach(), feature_weight_S, feature_weight_T, self.d_config.temperature)
  
          #embedding matching
          embedding_loss = match_loss(embeddings_S, embeddings_T, mask=inputs_mask_S)
          total_loss += embedding_loss * emd_loss_weight #sharing the same weight
          losses_dict[f'unweighted_embedding_{loss_type}'] = embedding_loss
  
  
          if 'losses' in results_S:
              total_hl_loss = 0
              for loss in results_S['losses']:
                  # in case of multi-GPU
                  total_hl_loss += loss.mean() 
              total_loss += total_hl_loss * self.d_config.hard_label_weight
              losses_dict['unweighted_hard_label_loss'] = total_hl_loss
          return total_loss, losses_dict
  
  
  def get_new_feature_weight(trans_matrix, distance_matrix, feature_weight_S, feature_weight_T, temperature):
  
      #distance_matrix = distance_matrix.detach().cpu().numpy().astype('float64')
      trans_weight_S = torch.sum(trans_matrix     * distance_matrix, dim=-1).cpu().numpy()
      trans_weight_T = torch.sum(trans_matrix.t() * distance_matrix, dim=-1).cpu().numpy()
      # new_student_weight = torch.zeros(stu_layer_num).cuda()
  
      feature_num_S = len(feature_weight_S)
      feature_num_T = len(feature_weight_T)
  
      feature_weight_S = trans_weight_S[:feature_num_S] / feature_weight_S
      weight_sum = np.sum(feature_weight_S)
      for i in range(feature_num_S):
          if feature_weight_S[i] != 0:
              feature_weight_S[i] = weight_sum / feature_weight_S[i]
     #feature_weight_S = weight_sum / (feature_weight_S + 1e-8)
  
      feature_weight_T = trans_weight_T[feature_num_S:] / feature_weight_T
      weight_sum = np.sum(feature_weight_T)
      for j in range(feature_num_T):
          if feature_weight_T[j] != 0:
              feature_weight_T[j] = weight_sum / feature_weight_T[j]
      #feature_weight_T = weight_sum / (feature_weight_T + 1e-8)
  
      feature_weight_S = softmax(feature_weight_S / temperature)
      feature_weight_T = softmax(feature_weight_T / temperature)
  
      return feature_weight_S, feature_weight_T
  