text: # -*- coding: utf-8 -*- #
  """*********************************************************************************************"""
  #   FileName     [ utility/fix_ckpt.py ]
  #   Synopsis     [ scripts to fix older checkpoints ]
  #   Author       [ Andy T. Liu (Andi611) ]
  #   Copyright    [ Copyleft(c), Speech Lab, NTU, Taiwan ]
  """*********************************************************************************************"""
  
  
  """
  Usage:
  This .py helps fix the `torch serialization ModuleNotFoundError` issue, 
  which occurs when the model.py directory is changed.
  Make sure you understand exactly what this script does before proceeding.
  """
  
  
  ###############
  # IMPORTATION #
  ###############
  import os
  import sys
  import torch
  
  
  def check_model_equiv(model1, model2):
      for p1, p2 in zip(model1.parameters(), model2.parameters()):
          if p1.data.ne(p2.data).sum() > 0:
              return False
          if not torch.equal(p1[0], p2[0]):
              return False
          if not torch.equal(p1[1].data, p2[1].data):
              return False
      return True
  
  
  def copyParams(module_src, module_dest):
      params1 = module_src.named_parameters()
      params2 = module_dest.named_parameters()
      dict_params2 = dict(params2)
      for name1, param1 in params1:
          if name1 in dict_params2:
              dict_params2[name1].data.copy_(param1.data)
  
  
  def main():
  
      input_ckpt = sys.argv[1]
  
      # load model with old setting
      from transformer.nn_transformer import SPEC_TRANSFORMER
      options = {'ckpt_file'     : input_ckpt,
                  'load_pretrain' : 'True',
                  'no_grad'       : 'True',
                  'dropout'       : 'default',
                  'spec_aug'      : 'False',
                  'spec_aug_prev' : 'True',
                  'weighted_sum'  : 'False',
                  'select_layer'  : -1,
                  'permute_input' : 'False' }
      old_transformer = SPEC_TRANSFORMER(options, inp_dim=-1)
  
      # build model with new setting
      from upstream.mockingjay.model import TransformerForMaskedAcousticModel
      model = TransformerForMaskedAcousticModel(old_transformer.model_config, old_transformer.inp_dim, old_transformer.inp_dim).to(torch.device('cuda'))
  
      # load old to new
      assert not check_model_equiv(old_transformer.model, model.Transformer)
      copyParams(old_transformer.model, model.Transformer)
      assert check_model_equiv(old_transformer.model, model.Transformer)
  
      assert not check_model_equiv(old_transformer.SpecHead, model.SpecHead)
      copyParams(old_transformer.SpecHead, model.SpecHead)
      assert check_model_equiv(old_transformer.SpecHead, model.SpecHead)
  
      global_step = old_transformer.all_states['Global_step']
      settings = old_transformer.all_states['Settings']
  
      # save
      all_states = {
          'SpecHead': model.SpecHead.state_dict(),
          'Transformer': model.Transformer.state_dict(),
          'Global_step': global_step,
          'Settings': settings
      }
      new_ckpt_path = input_ckpt.replace('.ckpt', '-new.ckpt')
      torch.save(all_states, new_ckpt_path)
  
      print('Done fixing ckpt: ', input_ckpt, 'to: ', new_ckpt_path)
  
  
  if __name__ == '__main__':
      main()
  