


'''simple docstring'''


import inspect
import unittest

from transformers import ViTConfig
from transformers.testing_utils import (
    require_accelerate,
    require_torch,
    require_torch_gpu,
    require_vision,
    slow,
    torch_device,
)
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
						import torch
						from torch import nn

						from transformers import ViTForImageClassification, ViTForMaskedImageModeling, ViTModel
						from transformers.models.vit.modeling_vit import VIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
						from PIL import Image

						from transformers import ViTImageProcessor






class 		lowerCAmelCase__	:




			def __init__(						self     :   int		,     lowerCamelCase__     :   Tuple		,     lowerCamelCase__     :   Tuple=13		,     lowerCamelCase__     :   Any=30		,     lowerCamelCase__     :   int=2		,     lowerCamelCase__     :   Dict=3		,     lowerCamelCase__     :   Tuple=True		,     lowerCamelCase__     :   Optional[Any]=True		,     lowerCamelCase__     :   Tuple=32		,     lowerCamelCase__     :   Union[str, Any]=5		,     lowerCamelCase__     :   Any=4		,     lowerCamelCase__     :   Union[str, Any]=37		,     lowerCamelCase__     :   Optional[Any]="gelu"		,     lowerCamelCase__     :   Any=0.1		,     lowerCamelCase__     :   Optional[int]=0.1		,     lowerCamelCase__     :   List[Any]=10		,     lowerCamelCase__     :   List[str]=0.0_2		,     lowerCamelCase__     :   int=None		,     lowerCamelCase__     :   Union[str, Any]=2		,     )      ->Any:



					'''simple docstring'''







					_UpperCAmelCase						:				List[str]	    =					parent
					_UpperCAmelCase						:				List[Any]	    =					batch_size
					_UpperCAmelCase						:				List[Any]	    =					image_size
					_UpperCAmelCase						:				Tuple	    =					patch_size
					_UpperCAmelCase						:				Dict	    =					num_channels
					_UpperCAmelCase						:				str	    =					is_training
					_UpperCAmelCase						:				Optional[Any]	    =					use_labels
					_UpperCAmelCase						:				str	    =					hidden_size
					_UpperCAmelCase						:				Any	    =					num_hidden_layers
					_UpperCAmelCase						:				Optional[int]	    =					num_attention_heads
					_UpperCAmelCase						:				Optional[Any]	    =					intermediate_size
					_UpperCAmelCase						:				Tuple	    =					hidden_act
					_UpperCAmelCase						:				Any	    =					hidden_dropout_prob
					_UpperCAmelCase						:				Tuple	    =					attention_probs_dropout_prob
					_UpperCAmelCase						:				Optional[Any]	    =					type_sequence_label_size
					_UpperCAmelCase						:				Union[str, Any]	    =					initializer_range
					_UpperCAmelCase						:				str	    =					scope
					_UpperCAmelCase						:				int	    =					encoder_stride

					# in ViT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
					_UpperCAmelCase						:				List[str]	    =					(image_size // patch_size) ** 2
					_UpperCAmelCase						:				str	    =					num_patches + 1




			def 							lowerCAmelCase__			(						self     :   str       )      ->List[Any]:



					'''simple docstring'''







					_UpperCAmelCase						:				Tuple	    =					floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]       )

					_UpperCAmelCase						:				Optional[int]	    =					None
					if self.use_labels:
							_UpperCAmelCase						:				Tuple	    =					ids_tensor([self.batch_size]		,     self.type_sequence_label_size       )

					_UpperCAmelCase						:				Union[str, Any]	    =					self.get_config()

					return config, pixel_values, labels




			def 							lowerCAmelCase__			(						self     :   Dict       )      ->str:



					'''simple docstring'''







					return ViTConfig(
					    image_size=self.image_size		,     patch_size=self.patch_size		,     num_channels=self.num_channels		,     hidden_size=self.hidden_size		,     num_hidden_layers=self.num_hidden_layers		,     num_attention_heads=self.num_attention_heads		,     intermediate_size=self.intermediate_size		,     hidden_act=self.hidden_act		,     hidden_dropout_prob=self.hidden_dropout_prob		,     attention_probs_dropout_prob=self.attention_probs_dropout_prob		,     is_decoder=lowerCamelCase__		,     initializer_range=self.initializer_range		,     encoder_stride=self.encoder_stride		,     )




			def 							lowerCAmelCase__			(						self     :   List[str]		,     lowerCamelCase__     :   Optional[Any]		,     lowerCamelCase__     :   Optional[int]		,     lowerCamelCase__     :   int       )      ->List[str]:



					'''simple docstring'''







					_UpperCAmelCase						:				Union[str, Any]	    =					ViTModel(config=lowerCamelCase__       )
					model.to(lowerCamelCase__       )
					model.eval()
					_UpperCAmelCase						:				List[str]	    =					model(lowerCamelCase__       )
					self.parent.assertEqual(result.last_hidden_state.shape		,     (self.batch_size, self.seq_length, self.hidden_size)       )




			def 							lowerCAmelCase__			(						self     :   Any		,     lowerCamelCase__     :   List[str]		,     lowerCamelCase__     :   Dict		,     lowerCamelCase__     :   Dict       )      ->int:



					'''simple docstring'''







					_UpperCAmelCase						:				int	    =					ViTForMaskedImageModeling(config=lowerCamelCase__       )
					model.to(lowerCamelCase__       )
					model.eval()
					_UpperCAmelCase						:				Optional[Any]	    =					model(lowerCamelCase__       )
					self.parent.assertEqual(
					    result.reconstruction.shape		,     (self.batch_size, self.num_channels, self.image_size, self.image_size)       )

					# test greyscale images
					_UpperCAmelCase						:				Optional[Any]	    =					1
					_UpperCAmelCase						:				int	    =					ViTForMaskedImageModeling(lowerCamelCase__       )
					model.to(lowerCamelCase__       )
					model.eval()

					_UpperCAmelCase						:				Optional[int]	    =					floats_tensor([self.batch_size, 1, self.image_size, self.image_size]       )
					_UpperCAmelCase						:				Dict	    =					model(lowerCamelCase__       )
					self.parent.assertEqual(result.reconstruction.shape		,     (self.batch_size, 1, self.image_size, self.image_size)       )




			def 							lowerCAmelCase__			(						self     :   str		,     lowerCamelCase__     :   List[str]		,     lowerCamelCase__     :   Any		,     lowerCamelCase__     :   str       )      ->Dict:



					'''simple docstring'''







					_UpperCAmelCase						:				List[Any]	    =					self.type_sequence_label_size
					_UpperCAmelCase						:				Tuple	    =					ViTForImageClassification(lowerCamelCase__       )
					model.to(lowerCamelCase__       )
					model.eval()
					_UpperCAmelCase						:				Optional[Any]	    =					model(lowerCamelCase__		,     labels=lowerCamelCase__       )
					self.parent.assertEqual(result.logits.shape		,     (self.batch_size, self.type_sequence_label_size)       )

					# test greyscale images
					_UpperCAmelCase						:				List[str]	    =					1
					_UpperCAmelCase						:				Any	    =					ViTForImageClassification(lowerCamelCase__       )
					model.to(lowerCamelCase__       )
					model.eval()

					_UpperCAmelCase						:				Dict	    =					floats_tensor([self.batch_size, 1, self.image_size, self.image_size]       )
					_UpperCAmelCase						:				str	    =					model(lowerCamelCase__       )
					self.parent.assertEqual(result.logits.shape		,     (self.batch_size, self.type_sequence_label_size)       )






			def 							lowerCAmelCase__			(						self     :   Optional[Any]       )      ->List[Any]:



					'''simple docstring'''







					_UpperCAmelCase						:				Dict	    =					self.prepare_config_and_inputs()
					(
					    (
					    _UpperCAmelCase
					)    ,			(
					    _UpperCAmelCase
					)    ,			(
					    _UpperCAmelCase
					)    ,			
					)						:				Dict	    =					config_and_inputs
					_UpperCAmelCase						:				List[str]	    =					{"pixel_values": pixel_values}
					return config, inputs_dict






@require_torch
class 		lowerCAmelCase__	(	UpperCAmelCase__			,				UpperCAmelCase__			,				unittest.TestCase		):
			lowerCAmelCase       :		List[str]				      =   (
			    (
			        ViTModel,
			        ViTForImageClassification,
			        ViTForMaskedImageModeling,
			    )
			    if is_torch_available()
			    else ()
			)
			lowerCAmelCase       :		int				      =   (
			    {"feature-extraction": ViTModel, "image-classification": ViTForImageClassification}
			    if is_torch_available()
			    else {}
			)
			lowerCAmelCase       :		Optional[Any]				      =   True

			lowerCAmelCase       :		Tuple				      =   False
			lowerCAmelCase       :		Any				      =   False
			lowerCAmelCase       :		str				      =   False




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



					'''simple docstring'''







					_UpperCAmelCase						:				Any	    =					ViTModelTester(self       )
					_UpperCAmelCase						:				Dict	    =					ConfigTester(self		,     config_class=lowerCamelCase__		,     has_text_modality=lowerCamelCase__		,     hidden_size=37       )




			def 							lowerCAmelCase__			(						self     :   str       )      ->int:



					'''simple docstring'''







					self.config_tester.run_common_tests()




			@unittest.skip(reason="ViT does not use inputs_embeds"       )
			def 							lowerCAmelCase__			(						self     :   List[Any]       )      ->int:



					'''simple docstring'''







					pass




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



					'''simple docstring'''







					_UpperCAmelCase    ,			_UpperCAmelCase						:				Dict	    =					self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes:
							_UpperCAmelCase						:				Union[str, Any]	    =					model_class(lowerCamelCase__       )
							self.assertIsInstance(model.get_input_embeddings()		,     (nn.Module)       )
							_UpperCAmelCase						:				List[str]	    =					model.get_output_embeddings()
							self.assertTrue(x is None or isinstance(lowerCamelCase__		,     nn.Linear       )       )




			def 							lowerCAmelCase__			(						self     :   Tuple       )      ->List[str]:



					'''simple docstring'''







					_UpperCAmelCase    ,			_UpperCAmelCase						:				Optional[int]	    =					self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes:
							_UpperCAmelCase						:				Optional[int]	    =					model_class(lowerCamelCase__       )
							_UpperCAmelCase						:				Union[str, Any]	    =					inspect.signature(model.forward       )
							# signature.parameters is an OrderedDict => so arg_names order is deterministic
							_UpperCAmelCase						:				Optional[int]	    =					[*signature.parameters.keys()]

							_UpperCAmelCase						:				Optional[int]	    =					["pixel_values"]
							self.assertListEqual(arg_names[:1]		,     lowerCamelCase__       )




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



					'''simple docstring'''







					_UpperCAmelCase						:				Tuple	    =					self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_model(*lowerCamelCase__       )




			def 							lowerCAmelCase__			(						self     :   int       )      ->Any:



					'''simple docstring'''







					_UpperCAmelCase						:				Union[str, Any]	    =					self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_masked_image_modeling(*lowerCamelCase__       )




			def 							lowerCAmelCase__			(						self     :   Dict       )      ->Tuple:



					'''simple docstring'''







					_UpperCAmelCase						:				Tuple	    =					self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_image_classification(*lowerCamelCase__       )






			@slow
			def 							lowerCAmelCase__			(						self     :   Optional[int]       )      ->Optional[Any]:



					'''simple docstring'''







					for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							_UpperCAmelCase						:				Any	    =					ViTModel.from_pretrained(lowerCamelCase__       )
							self.assertIsNotNone(lowerCamelCase__       )






def        __lowerCAmelCase       ():
		_UpperCAmelCase						:				Union[str, Any]	    =					Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"      )
		return image






@require_torch
@require_vision
class 		lowerCAmelCase__	(	unittest.TestCase		):




			@cached_property
			def 							lowerCAmelCase__			(						self     :   Tuple       )      ->Dict:



					'''simple docstring'''







					return ViTImageProcessor.from_pretrained("google/vit-base-patch16-224"       ) if is_vision_available() else None




			@slow
			def 							lowerCAmelCase__			(						self     :   Tuple       )      ->int:



					'''simple docstring'''







					_UpperCAmelCase						:				Optional[Any]	    =					ViTForImageClassification.from_pretrained("google/vit-base-patch16-224"       ).to(lowerCamelCase__       )

					_UpperCAmelCase						:				List[Any]	    =					self.default_image_processor
					_UpperCAmelCase						:				Optional[int]	    =					prepare_img()
					_UpperCAmelCase						:				Optional[int]	    =					image_processor(images=lowerCamelCase__		,     return_tensors="pt"       ).to(lowerCamelCase__       )

					# forward pass
					with torch.no_grad():
							_UpperCAmelCase						:				List[Any]	    =					model(**lowerCamelCase__       )

					# verify the logits
					_UpperCAmelCase						:				str	    =					torch.Size((1, 10_00)       )
					self.assertEqual(outputs.logits.shape		,     lowerCamelCase__       )

					_UpperCAmelCase						:				Any	    =					torch.tensor([-0.2_7_4_4, 0.8_2_1_5, -0.0_8_3_6]       ).to(lowerCamelCase__       )

					self.assertTrue(torch.allclose(outputs.logits[0, :3]		,     lowerCamelCase__		,     atol=1E-4       )       )




			@slow
			def 							lowerCAmelCase__			(						self     :   int       )      ->List[Any]:



					'''simple docstring'''







					_UpperCAmelCase						:				Optional[Any]	    =					ViTModel.from_pretrained("facebook/dino-vits8"       ).to(lowerCamelCase__       )

					_UpperCAmelCase						:				Dict	    =					ViTImageProcessor.from_pretrained("facebook/dino-vits8"		,     size=4_80       )
					_UpperCAmelCase						:				Union[str, Any]	    =					prepare_img()
					_UpperCAmelCase						:				Any	    =					image_processor(images=lowerCamelCase__		,     return_tensors="pt"       )
					_UpperCAmelCase						:				Dict	    =					inputs.pixel_values.to(lowerCamelCase__       )

					# forward pass
					with torch.no_grad():
							_UpperCAmelCase						:				Optional[int]	    =					model(lowerCamelCase__		,     interpolate_pos_encoding=lowerCamelCase__       )

					# verify the logits
					_UpperCAmelCase						:				Tuple	    =					torch.Size((1, 36_01, 3_84)       )
					self.assertEqual(outputs.last_hidden_state.shape		,     lowerCamelCase__       )

					_UpperCAmelCase						:				Optional[int]	    =					torch.tensor(
					    [[4.2_3_4_0, 4.3_9_0_6, -6.6_6_9_2], [4.5_4_6_3, 1.8_9_2_8, -6.7_2_5_7], [4.4_4_2_9, 0.8_4_9_6, -5.8_5_8_5]]       ).to(lowerCamelCase__       )

					self.assertTrue(torch.allclose(outputs.last_hidden_state[0, :3, :3]		,     lowerCamelCase__		,     atol=1E-4       )       )






			@slow
			@require_accelerate
			@require_torch_gpu
			def 							lowerCAmelCase__			(						self     :   Optional[int]       )      ->List[Any]:



					'''simple docstring'''







					_UpperCAmelCase						:				List[str]	    =					ViTModel.from_pretrained("facebook/dino-vits8"		,     torch_dtype=torch.floataa		,     device_map="auto"       )
					_UpperCAmelCase						:				Optional[Any]	    =					self.default_image_processor

					_UpperCAmelCase						:				Tuple	    =					prepare_img()
					_UpperCAmelCase						:				List[str]	    =					image_processor(images=lowerCamelCase__		,     return_tensors="pt"       )
					_UpperCAmelCase						:				List[str]	    =					inputs.pixel_values.to(lowerCamelCase__       )

					# forward pass to make sure inference works in fp16
					with torch.no_grad():
							_UpperCAmelCase						:				List[str]	    =					model(lowerCamelCase__       )







'''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      ):
		_UpperCAmelCase						:				List[str]	    =					[
		    "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      ):
		_UpperCAmelCase    ,			_UpperCAmelCase						:				Optional[Any]	    =					emb.weight.shape
		_UpperCAmelCase						:				str	    =					nn.Linear(__lowerCAmelCase					,       __lowerCAmelCase					,       bias=__lowerCAmelCase      )
		_UpperCAmelCase						:				Optional[int]	    =					emb.weight.data
		return lin_layer






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






def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase = WEIGHTS_NAME      ):
		_UpperCAmelCase						:				Optional[int]	    =					[]
		_UpperCAmelCase						:				Optional[Any]	    =					0
		os.makedirs(__lowerCAmelCase					,       exist_ok=__lowerCAmelCase      )

		for expert in range(__lowerCAmelCase      ):
				_UpperCAmelCase						:				Tuple	    =					switch_checkpoint_path + F"""-rank-{expert}.pt"""
				if os.path.isfile(__lowerCAmelCase      ):
						_UpperCAmelCase						:				Tuple	    =					torch.load(__lowerCAmelCase      )["model"]
						remove_ignore_keys_(__lowerCAmelCase      )
						_UpperCAmelCase						:				Dict	    =					rename_fairseq_keys(__lowerCAmelCase					,       __lowerCAmelCase      )
						_UpperCAmelCase						:				List[str]	    =					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						:				Tuple	    =					os.path.join(__lowerCAmelCase					,       weights_name.replace(".bin"					,       F"""-{len(__lowerCAmelCase      )+1:05d}-of-???.bin"""      )      )
		_UpperCAmelCase						:				Union[str, Any]	    =					torch.load(switch_checkpoint_path + "-shared.pt"      )["model"]
		remove_ignore_keys_(__lowerCAmelCase      )
		_UpperCAmelCase						:				Union[str, Any]	    =					rename_fairseq_keys(__lowerCAmelCase					,       __lowerCAmelCase      )
		_UpperCAmelCase						:				Any	    =					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						:				List[str]	    =					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						:				Union[str, Any]	    =					{}
		for idx, shard in enumerate(__lowerCAmelCase      ):
				_UpperCAmelCase						:				Tuple	    =					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						:				List[Any]	    =					shard_file

    # Add the metadata
		_UpperCAmelCase						:				Any	    =					{"total_size": total_size}
		_UpperCAmelCase						:				List[str]	    =					{"metadata": metadata, "weight_map": weight_map}

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

		return metadata, index


if __name__ == "__main__":
						lowerCamelCase__					=   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.',
						)
						lowerCamelCase__					=   parser.parse_args()
						lowerCamelCase__		,lowerCamelCase__					=   shard_on_the_fly(
						    args.nllb_moe_checkpoint_path,
						    args.pytorch_dump_folder_path,
						    128,
						    args.dtype,
						)

						lowerCamelCase__					=   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)
						lowerCamelCase__					=   NllbMoeModel.from_pretrained(args.pytorch_dump_folder_path)
						print('Done')
						model.save_pretrained(args.pytorch_dump_folder_path)



