






'''simple docstring'''






import argparse
import json
import os

import torch
from torch import nn

from transformers import NllbMoeConfig, NllbMoeModel
from transformers.modeling_utils import dtype_byte_size
from transformers.utils import WEIGHTS_INDEX_NAME, WEIGHTS_NAME





def lowerCamelCase__      (						__lowerCamelCase	:							List[Any]							):







   '''simple docstring'''




   _UpperCAmelCase       :       List[Any]							   =[
       'encoder.version',
       'decoder.version',
       'model.encoder.version',
       'model.decoder.version',
       'decoder.output_projection.weight',
       '_float_tensor',
       'encoder.embed_positions._float_tensor',
       'decoder.embed_positions._float_tensor',
   ]
   for k in ignore_keys:
      state_dict.pop(__lowerCamelCase	,   __lowerCamelCase							)





def lowerCamelCase__      (						__lowerCamelCase	:							Union[str, Any]							):







   '''simple docstring'''




   _UpperCAmelCase       :       List[str]							   =emb.weight.shape
   _UpperCAmelCase       :       Any							   =nn.Linear(__lowerCamelCase	,   __lowerCamelCase	,   bias=__lowerCamelCase							)
   _UpperCAmelCase       :       int							   =emb.weight.data
   return lin_layer





def lowerCamelCase__      (						__lowerCamelCase	:							str	,   __lowerCamelCase	:							int=None							):







   '''simple docstring'''




   _UpperCAmelCase       :       Optional[Any]							   ={}
   for old_key in state_dict.keys():
      _UpperCAmelCase       :       str							   =old_key
      if "moe_layer.experts." in key:
         if expert_idx is not None:
            _UpperCAmelCase       :       str							   =key.replace('moe_layer.experts.0'	,   f"ffn.experts.expert_{expert_idx}"							)
         else:
            _UpperCAmelCase       :       List[Any]							   =key.replace('moe_layer.experts.'	,   'ffn.experts.expert_'							)
      if "gate" in key:
         _UpperCAmelCase       :       Optional[int]							   =key.replace('.moe_layer.gate.wg'	,   '.ffn.router.classifier'							)
      if "fc2" and "experts" not in key:
         _UpperCAmelCase       :       Any							   =key.replace('.fc2.'	,   '.ffn.fc2.'							)
      if "fc1" and "experts" not in key:
         _UpperCAmelCase       :       str							   =key.replace('.fc1.'	,   '.ffn.fc1.'							)
      if ".encoder_attn." in key:
         _UpperCAmelCase       :       Optional[int]							   =key.replace('.encoder_attn.'	,   '.cross_attention.'							)
      if "encoder_attn_layer_norm" in key:
         _UpperCAmelCase       :       str							   =key.replace('encoder_attn_layer_norm'	,   'cross_attention_layer_norm'							)
      if "final_layer_norm" in key:
         _UpperCAmelCase       :       Union[str, Any]							   =key.replace('final_layer_norm'	,   'ff_layer_norm'							)
      _UpperCAmelCase       :       Any							   =state_dict[old_key]
   return new_dict





def lowerCamelCase__      (						__lowerCamelCase	:							Union[str, Any]	,   __lowerCamelCase	:							Dict	,   __lowerCamelCase	:							int	,   __lowerCamelCase	:							int	,   __lowerCamelCase	:							str = WEIGHTS_NAME							):







   '''simple docstring'''




   _UpperCAmelCase       :       Any							   =[]
   _UpperCAmelCase       :       List[str]							   =0
   os.makedirs(__lowerCamelCase	,   exist_ok=__lowerCamelCase							)

   for expert in range(__lowerCamelCase							):
      _UpperCAmelCase       :       Union[str, Any]							   =switch_checkpoint_path + f"-rank-{expert}.pt"
      if os.path.isfile(__lowerCamelCase							):
         _UpperCAmelCase       :       Union[str, Any]							   =torch.load(__lowerCamelCase							)['model']
         remove_ignore_keys_(__lowerCamelCase							)
         _UpperCAmelCase       :       Optional[Any]							   =rename_fairseq_keys(__lowerCamelCase	,   __lowerCamelCase							)
         _UpperCAmelCase       :       Optional[int]							   =os.path.join(
             __lowerCamelCase	,   weights_name.replace('.bin'	,   f"-{len(__lowerCamelCase							)+1:05d}-of-???.bin"							)							)
         torch.save(__lowerCamelCase	,   __lowerCamelCase							)
         sharded_state_dicts.append(expert_state.keys()							)
         total_size += sum([value.numel() for key, value in expert_state.items()]							) * dtype_byte_size(
             expert_state[list(__lowerCamelCase							)[0]].dtype							)

    # Add the last block
   _UpperCAmelCase       :       List[str]							   =os.path.join(__lowerCamelCase	,   weights_name.replace('.bin'	,   f"-{len(__lowerCamelCase							)+1:05d}-of-???.bin"							)							)
   _UpperCAmelCase       :       Any							   =torch.load(switch_checkpoint_path + '-shared.pt'							)['model']
   remove_ignore_keys_(__lowerCamelCase							)
   _UpperCAmelCase       :       Union[str, Any]							   =rename_fairseq_keys(__lowerCamelCase	,   __lowerCamelCase							)
   _UpperCAmelCase       :       str							   =shared_weights['decoder.embed_tokens.weight']
   sharded_state_dicts.append(shared_weights.keys()							)

   # If we only have the shared weights (dummy model/experts saved on the same file)
   if len(__lowerCamelCase							) == 1:
      _UpperCAmelCase       :       Union[str, Any]							   =os.path.join(__lowerCamelCase	,   __lowerCamelCase							)
      torch.save(__lowerCamelCase	,   __lowerCamelCase							)
      return {weights_name: sharded_state_dicts[0]}, None
   else:
      torch.save(__lowerCamelCase	,   __lowerCamelCase							)
   # Otherwise, let's build the index
   _UpperCAmelCase       :       Optional[int]							   ={}
   for idx, shard in enumerate(__lowerCamelCase							):
      _UpperCAmelCase       :       Optional[Any]							   =weights_name.replace('.bin'	,   f"-{idx+1:05d}-of-{len(__lowerCamelCase							):05d}.bin"							)
      _UpperCAmelCase       :       List[Any]							   =os.path.join(__lowerCamelCase	,   weights_name.replace('.bin'	,   f"-{idx+1:05d}-of-???.bin"							)							)
      os.rename(__lowerCamelCase	,   os.path.join(__lowerCamelCase	,   __lowerCamelCase							)							)
      for key in shard:
         _UpperCAmelCase       :       Any							   =shard_file

    # Add the metadata
   _UpperCAmelCase       :       List[Any]							   ={'total_size': total_size}
   _UpperCAmelCase       :       Optional[int]							   ={'metadata': metadata, 'weight_map': weight_map}

   with open(os.path.join(__lowerCamelCase	,   __lowerCamelCase							)	,   'w'	,   encoding='utf-8'							) as f:
      _UpperCAmelCase       :       int							   =json.dumps(__lowerCamelCase	,   indent=2	,   sort_keys=__lowerCamelCase							) + '\n'
      f.write(__lowerCamelCase							)

   return metadata, index


if __name__ == "__main__":
 lowercase							=argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     '--nllb_moe_checkpoint_path',
     default='/home/arthur_huggingface_co/fairseq/weights/checkpoints/model_moe_54b/checkpoint_2_300000',
     type=str,
     required=False,
     help='Path to a directory containing a folder per layer. Follows the original Google format.',
 )
 parser.add_argument('--dtype', default='float32', type=str, required=False, help='dtype of the saved model')
 parser.add_argument(
     '--pytorch_dump_folder_path',
     default='/home/arthur_huggingface_co/fairseq/weights/checkpoints/hf-converted-moe-54b',
     type=str,
     required=False,
     help='Path to the output pytorch model.',
 )
 lowercase							=parser.parse_args()
 lowercase,       lowercase							=shard_on_the_fly(
     args.nllb_moe_checkpoint_path,
     args.pytorch_dump_folder_path,
     128,
     args.dtype,
 )

 lowercase							=NllbMoeConfig.from_pretrained(
     'facebook/nllb-200-3.3B', encoder_sparse_step=4, decoder_sparse_step=4, num_experts=128
 )
 config.save_pretrained(args.pytorch_dump_folder_path)
 lowercase							=NllbMoeModel.from_pretrained(args.pytorch_dump_folder_path)
 print('Done')
 model.save_pretrained(args.pytorch_dump_folder_path)





'''simple docstring'''






import shutil
import tempfile
import unittest

import numpy as np
import pytest

from transformers import is_speech_available, is_vision_available
from transformers.testing_utils import require_torch


if is_vision_available():
 from transformers import TvltImageProcessor

if is_speech_available():
 from transformers import TvltFeatureExtractor

from transformers import TvltProcessor






@require_torch
class  __magic_name__							(		unittest.TestCase    ):






       def       lowerCAmelCase		(      self)						->  str:




          '''simple docstring'''
          _UpperCAmelCase       :       Dict							   ='ZinengTang/tvlt-base'
          _UpperCAmelCase       :       Dict							   =tempfile.mkdtemp()






       def       lowerCAmelCase		(      self  ,							**snake_case)						->  Union[str, Any]:




          '''simple docstring'''
          return TvltImageProcessor.from_pretrained(self.checkpoint  ,							**snake_case)






       def       lowerCAmelCase		(      self  ,							**snake_case)						->  Dict:




          '''simple docstring'''
          return TvltFeatureExtractor.from_pretrained(self.checkpoint  ,							**snake_case)






       def       lowerCAmelCase		(      self)						->  Any:




          '''simple docstring'''
          shutil.rmtree(self.tmpdirname)






       def       lowerCAmelCase		(      self)						->  Optional[int]:




          '''simple docstring'''
          _UpperCAmelCase       :       Any							   =self.get_image_processor()
          _UpperCAmelCase       :       Optional[Any]							   =self.get_feature_extractor()

          _UpperCAmelCase       :       str							   =TvltProcessor(image_processor=snake_case  ,							feature_extractor=snake_case)
          processor.save_pretrained(self.tmpdirname)
          _UpperCAmelCase       :       str							   =TvltProcessor.from_pretrained(self.tmpdirname)

          self.assertIsInstance(processor.feature_extractor  ,							snake_case)
          self.assertIsInstance(processor.image_processor  ,							snake_case)






       def       lowerCAmelCase		(      self)						->  Union[str, Any]:




          '''simple docstring'''
          _UpperCAmelCase       :       str							   =self.get_image_processor()
          _UpperCAmelCase       :       List[Any]							   =self.get_feature_extractor()

          _UpperCAmelCase       :       str							   =TvltProcessor(image_processor=snake_case  ,							feature_extractor=snake_case)

          _UpperCAmelCase       :       Optional[int]							   =np.ones([1_2_0_0_0])

          _UpperCAmelCase       :       str							   =feature_extractor(snake_case  ,							return_tensors='np')
          _UpperCAmelCase       :       Dict							   =processor(audio=snake_case  ,							return_tensors='np')

          for key in audio_dict.keys():
             self.assertAlmostEqual(audio_dict[key].sum()  ,							input_processor[key].sum()  ,							delta=1E-2)






       def       lowerCAmelCase		(      self)						->  int:




          '''simple docstring'''
          _UpperCAmelCase       :       Dict							   =self.get_image_processor()
          _UpperCAmelCase       :       int							   =self.get_feature_extractor()

          _UpperCAmelCase       :       int							   =TvltProcessor(image_processor=snake_case  ,							feature_extractor=snake_case)

          _UpperCAmelCase       :       Union[str, Any]							   =np.ones([3, 2_2_4, 2_2_4])

          _UpperCAmelCase       :       Optional[Any]							   =image_processor(snake_case  ,							return_tensors='np')
          _UpperCAmelCase       :       List[Any]							   =processor(images=snake_case  ,							return_tensors='np')

          for key in image_dict.keys():
             self.assertAlmostEqual(image_dict[key].sum()  ,							input_processor[key].sum()  ,							delta=1E-2)






       def       lowerCAmelCase		(      self)						->  List[str]:




          '''simple docstring'''
          _UpperCAmelCase       :       List[str]							   =self.get_image_processor()
          _UpperCAmelCase       :       Dict							   =self.get_feature_extractor()

          _UpperCAmelCase       :       Optional[int]							   =TvltProcessor(image_processor=snake_case  ,							feature_extractor=snake_case)

          _UpperCAmelCase       :       Optional[int]							   =np.ones([1_2_0_0_0])
          _UpperCAmelCase       :       str							   =np.ones([3, 2_2_4, 2_2_4])

          _UpperCAmelCase       :       Optional[int]							   =processor(audio=snake_case  ,							images=snake_case)

          self.assertListEqual(list(inputs.keys())  ,							['audio_values', 'audio_mask', 'pixel_values', 'pixel_mask'])

          # test if it raises when no input is passed
          with pytest.raises(snake_case):
             processor()






       def       lowerCAmelCase		(      self)						->  Optional[int]:




          '''simple docstring'''
          _UpperCAmelCase       :       Optional[int]							   =self.get_image_processor()
          _UpperCAmelCase       :       Tuple							   =self.get_feature_extractor()

          _UpperCAmelCase       :       Dict							   =TvltProcessor(image_processor=snake_case  ,							feature_extractor=snake_case)

          self.assertListEqual(
              processor.model_input_names  ,							image_processor.model_input_names + feature_extractor.model_input_names  ,							msg='`processor` and `image_processor`+`feature_extractor` model input names do not match'  ,							)

