



"""simple docstring"""




import copy
import unittest

from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, 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, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
		import torch

		from transformers import (
		    MODEL_FOR_MULTIPLE_CHOICE_MAPPING,
		    MODEL_FOR_QUESTION_ANSWERING_MAPPING,
		    MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
		    MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
		    LayoutLMvaConfig,
		    LayoutLMvaForQuestionAnswering,
		    LayoutLMvaForSequenceClassification,
		    LayoutLMvaForTokenClassification,
		    LayoutLMvaModel,
		)
		from transformers.models.layoutlmva.modeling_layoutlmva import LAYOUTLMV3_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
		from PIL import Image

		from transformers import LayoutLMvaImageProcessor

class        snake_case__    :







						"""simple docstring"""




						def __init__(       self ,       __lowercase ,       __lowercase=2 ,       __lowercase=3 ,       __lowercase=4 ,       __lowercase=2 ,       __lowercase=7 ,       __lowercase=True ,       __lowercase=True ,       __lowercase=True ,       __lowercase=True ,       __lowercase=9_9 ,       __lowercase=3_6 ,       __lowercase=3 ,       __lowercase=4 ,       __lowercase=3_7 ,       __lowercase="gelu" ,       __lowercase=0.1 ,       __lowercase=0.1 ,       __lowercase=5_1_2 ,       __lowercase=1_6 ,       __lowercase=2 ,       __lowercase=0.0_2 ,       __lowercase=6 ,       __lowercase=6 ,       __lowercase=3 ,       __lowercase=4 ,       __lowercase=None ,       __lowercase=1_0_0_0 ,       ) ->				Dict:
									"""simple docstring"""




									a__       :						Optional[Any]							=  parent
									a__       :						Optional[int]							=  batch_size
									a__       :						List[str]							=  num_channels
									a__       :						Optional[Any]							=  image_size
									a__       :						int							=  patch_size
									a__       :						int							=  text_seq_length
									a__       :						Union[str, Any]							=  is_training
									a__       :						Any							=  use_input_mask
									a__       :						Optional[int]							=  use_token_type_ids
									a__       :						str							=  use_labels
									a__       :						Union[str, Any]							=  vocab_size
									a__       :						List[str]							=  hidden_size
									a__       :						int							=  num_hidden_layers
									a__       :						Optional[int]							=  num_attention_heads
									a__       :						Optional[Any]							=  intermediate_size
									a__       :						List[str]							=  hidden_act
									a__       :						str							=  hidden_dropout_prob
									a__       :						str							=  attention_probs_dropout_prob
									a__       :						Union[str, Any]							=  max_position_embeddings
									a__       :						List[str]							=  type_vocab_size
									a__       :						str							=  type_sequence_label_size
									a__       :						Optional[int]							=  initializer_range
									a__       :						int							=  coordinate_size
									a__       :						str							=  shape_size
									a__       :						Optional[Any]							=  num_labels
									a__       :						List[str]							=  num_choices
									a__       :						Optional[Any]							=  scope
									a__       :						str							=  range_bbox

									# LayoutLMv3's sequence length equals the number of text tokens + number of patches + 1 (we add 1 for the CLS token)
									a__       :						Optional[int]							=  text_seq_length
									a__       :						Dict							=  (image_size // patch_size) ** 2 + 1
									a__       :						str							=  self.text_seq_length + self.image_seq_length

						def 							SCREAMING_SNAKE_CASE__(       self  ) ->				Optional[int]:
									"""simple docstring"""




									a__       :						Any							=  ids_tensor([self.batch_size, self.text_seq_length] ,       self.vocab_size  )

									a__       :						Union[str, Any]							=  ids_tensor([self.batch_size, self.text_seq_length, 4] ,       self.range_bbox  )
									# Ensure that bbox is legal
									for i in range(bbox.shape[0]  ):
												for j in range(bbox.shape[1]  ):
															if bbox[i, j, 3] < bbox[i, j, 1]:
																		a__       :						Tuple							=  bbox[i, j, 3]
																		a__       :						List[Any]							=  bbox[i, j, 1]
																		a__       :						Optional[Any]							=  t
															if bbox[i, j, 2] < bbox[i, j, 0]:
																		a__       :						List[str]							=  bbox[i, j, 2]
																		a__       :						Union[str, Any]							=  bbox[i, j, 0]
																		a__       :						str							=  t

									a__       :						str							=  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]  )

									a__       :						List[str]							=  None
									if self.use_input_mask:
												a__       :						Any							=  random_attention_mask([self.batch_size, self.text_seq_length]  )

									a__       :						List[str]							=  None
									if self.use_token_type_ids:
												a__       :						Optional[Any]							=  ids_tensor([self.batch_size, self.text_seq_length] ,       self.type_vocab_size  )

									a__       :						int							=  None
									a__       :						Tuple							=  None
									if self.use_labels:
												a__       :						Union[str, Any]							=  ids_tensor([self.batch_size] ,       self.type_sequence_label_size  )
												a__       :						Optional[Any]							=  ids_tensor([self.batch_size, self.text_seq_length] ,       self.num_labels  )

									a__       :						List[str]							=  LayoutLMvaConfig(
									    vocab_size=self.vocab_size ,       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 ,       max_position_embeddings=self.max_position_embeddings ,       type_vocab_size=self.type_vocab_size ,       initializer_range=self.initializer_range ,       coordinate_size=self.coordinate_size ,       shape_size=self.shape_size ,       input_size=self.image_size ,       patch_size=self.patch_size ,       )

									return config, input_ids, bbox, pixel_values, token_type_ids, input_mask, sequence_labels, token_labels

						def 							SCREAMING_SNAKE_CASE__(       self ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase  ) ->				int:
									"""simple docstring"""




									a__       :						str							=  LayoutLMvaModel(config=__lowercase  )
									model.to(__lowercase  )
									model.eval()

									# text + image
									a__       :						List[str]							=  model(__lowercase ,       pixel_values=__lowercase  )
									a__       :						List[str]							=  model(
									    __lowercase ,       bbox=__lowercase ,       pixel_values=__lowercase ,       attention_mask=__lowercase ,       token_type_ids=__lowercase  )
									a__       :						Dict							=  model(__lowercase ,       bbox=__lowercase ,       pixel_values=__lowercase ,       token_type_ids=__lowercase  )
									a__       :						Union[str, Any]							=  model(__lowercase ,       bbox=__lowercase ,       pixel_values=__lowercase  )

									self.parent.assertEqual(result.last_hidden_state.shape ,       (self.batch_size, self.seq_length, self.hidden_size)  )

									# text only
									a__       :						Optional[int]							=  model(__lowercase  )
									self.parent.assertEqual(
									    result.last_hidden_state.shape ,       (self.batch_size, self.text_seq_length, self.hidden_size)  )

									# image only
									a__       :						Any							=  model(pixel_values=__lowercase  )
									self.parent.assertEqual(
									    result.last_hidden_state.shape ,       (self.batch_size, self.image_seq_length, self.hidden_size)  )

						def 							SCREAMING_SNAKE_CASE__(       self ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase  ) ->				Any:
									"""simple docstring"""




									a__       :						Union[str, Any]							=  self.num_labels
									a__       :						Optional[int]							=  LayoutLMvaForSequenceClassification(__lowercase  )
									model.to(__lowercase  )
									model.eval()
									a__       :						Any							=  model(
									    __lowercase ,       bbox=__lowercase ,       pixel_values=__lowercase ,       attention_mask=__lowercase ,       token_type_ids=__lowercase ,       labels=__lowercase ,       )
									self.parent.assertEqual(result.logits.shape ,       (self.batch_size, self.num_labels)  )

						def 							SCREAMING_SNAKE_CASE__(       self ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase  ) ->				Tuple:
									"""simple docstring"""




									a__       :						List[Any]							=  self.num_labels
									a__       :						List[Any]							=  LayoutLMvaForTokenClassification(config=__lowercase  )
									model.to(__lowercase  )
									model.eval()
									a__       :						List[Any]							=  model(
									    __lowercase ,       bbox=__lowercase ,       pixel_values=__lowercase ,       attention_mask=__lowercase ,       token_type_ids=__lowercase ,       labels=__lowercase ,       )
									self.parent.assertEqual(result.logits.shape ,       (self.batch_size, self.text_seq_length, self.num_labels)  )

						def 							SCREAMING_SNAKE_CASE__(       self ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase  ) ->				int:
									"""simple docstring"""




									a__       :						List[Any]							=  LayoutLMvaForQuestionAnswering(config=__lowercase  )
									model.to(__lowercase  )
									model.eval()
									a__       :						Dict							=  model(
									    __lowercase ,       bbox=__lowercase ,       pixel_values=__lowercase ,       attention_mask=__lowercase ,       token_type_ids=__lowercase ,       start_positions=__lowercase ,       end_positions=__lowercase ,       )
									self.parent.assertEqual(result.start_logits.shape ,       (self.batch_size, self.seq_length)  )
									self.parent.assertEqual(result.end_logits.shape ,       (self.batch_size, self.seq_length)  )

						def 							SCREAMING_SNAKE_CASE__(       self  ) ->				int:
									"""simple docstring"""




									a__       :						str							=  self.prepare_config_and_inputs()
									(
									    a__						
									)       :						Tuple							=  config_and_inputs
									a__       :						List[str]							=  {
									    """input_ids""": input_ids,
									    """bbox""": bbox,
									    """pixel_values""": pixel_values,
									    """token_type_ids""": token_type_ids,
									    """attention_mask""": input_mask,
									}
									return config, inputs_dict




@require_torch
class        snake_case__    (A__ ,     A__ ,     unittest.TestCase		):







						"""simple docstring"""




						__lowerCAmelCase					:List[Any]      =					False
						__lowerCAmelCase					:Dict      =					False
						__lowerCAmelCase					:Any      =					False

						__lowerCAmelCase					:Tuple      =					(
						    (
						        LayoutLMvaModel,
						        LayoutLMvaForSequenceClassification,
						        LayoutLMvaForTokenClassification,
						        LayoutLMvaForQuestionAnswering,
						    )
						    if is_torch_available()
						    else ()
						)
						__lowerCAmelCase					:str      =					(
						    {"document-question-answering": LayoutLMvaForQuestionAnswering, "feature-extraction": LayoutLMvaModel}
						    if is_torch_available()
						    else {}
						)

						def 							SCREAMING_SNAKE_CASE__(       self ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase  ) ->				Any:
									"""simple docstring"""




									return True

						def 							SCREAMING_SNAKE_CASE__(       self  ) ->				List[Any]:
									"""simple docstring"""




									a__       :						Union[str, Any]							=  LayoutLMvaModelTester(self  )
									a__       :						List[Any]							=  ConfigTester(self ,       config_class=__lowercase ,       hidden_size=3_7  )

						def 							SCREAMING_SNAKE_CASE__(       self ,       __lowercase ,       __lowercase ,       __lowercase=False  ) ->				List[str]:
									"""simple docstring"""




									a__       :						Optional[int]							=  copy.deepcopy(__lowercase  )
									if model_class in get_values(__lowercase  ):
												a__       :						Union[str, Any]							=  {
												    k: v.unsqueeze(1  ).expand(-1 ,       self.model_tester.num_choices ,       -1  ).contiguous()
												    if isinstance(__lowercase ,       torch.Tensor  ) and v.ndim > 1
												    else v
												    for k, v in inputs_dict.items()
												}
									if return_labels:
												if model_class in get_values(__lowercase  ):
															a__       :						Any							=  torch.ones(self.model_tester.batch_size ,       dtype=torch.long ,       device=__lowercase  )
												elif model_class in get_values(__lowercase  ):
															a__       :						Optional[Any]							=  torch.zeros(
															    self.model_tester.batch_size ,       dtype=torch.long ,       device=__lowercase  )
															a__       :						int							=  torch.zeros(
															    self.model_tester.batch_size ,       dtype=torch.long ,       device=__lowercase  )
												elif model_class in [
												    *get_values(__lowercase  ),
												]:
															a__       :						str							=  torch.zeros(
															    self.model_tester.batch_size ,       dtype=torch.long ,       device=__lowercase  )
												elif model_class in [
												    *get_values(__lowercase  ),
												]:
															a__       :						Any							=  torch.zeros(
															    (self.model_tester.batch_size, self.model_tester.text_seq_length) ,       dtype=torch.long ,       device=__lowercase ,       )

									return inputs_dict

						def 							SCREAMING_SNAKE_CASE__(       self  ) ->				List[str]:
									"""simple docstring"""




									self.config_tester.run_common_tests()

						def 							SCREAMING_SNAKE_CASE__(       self  ) ->				Tuple:
									"""simple docstring"""




									a__       :						List[str]							=  self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*__lowercase  )

						def 							SCREAMING_SNAKE_CASE__(       self  ) ->				Any:
									"""simple docstring"""




									a__       :						Optional[int]							=  self.model_tester.prepare_config_and_inputs()
									for type in ["absolute", "relative_key", "relative_key_query"]:
												a__       :						Optional[int]							=  type
												self.model_tester.create_and_check_model(*__lowercase  )

						def 							SCREAMING_SNAKE_CASE__(       self  ) ->				Tuple:
									"""simple docstring"""




									a__       :						Union[str, Any]							=  self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_sequence_classification(*__lowercase  )

						def 							SCREAMING_SNAKE_CASE__(       self  ) ->				Optional[int]:
									"""simple docstring"""




									a__       :						Optional[Any]							=  self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_token_classification(*__lowercase  )

						def 							SCREAMING_SNAKE_CASE__(       self  ) ->				List[Any]:
									"""simple docstring"""




									a__       :						str							=  self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_question_answering(*__lowercase  )

						@slow
						def 							SCREAMING_SNAKE_CASE__(       self  ) ->				Any:
									"""simple docstring"""




									for model_name in LAYOUTLMV3_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
												a__       :						Optional[int]							=  LayoutLMvaModel.from_pretrained(__lowercase  )
												self.assertIsNotNone(__lowercase  )




def      lowerCAmelCase_					(  )							->	str:







			"""simple docstring"""
			a__       :						List[Any]							=  Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png""")
			return image





@require_torch
class        snake_case__    (unittest.TestCase		):







						"""simple docstring"""




						@cached_property
						def 							SCREAMING_SNAKE_CASE__(       self  ) ->				Any:
									"""simple docstring"""




									return LayoutLMvaImageProcessor(apply_ocr=__lowercase  ) if is_vision_available() else None

						@slow
						def 							SCREAMING_SNAKE_CASE__(       self  ) ->				int:
									"""simple docstring"""




									a__       :						Any							=  LayoutLMvaModel.from_pretrained("""microsoft/layoutlmv3-base"""  ).to(__lowercase  )

									a__       :						List[Any]							=  self.default_image_processor
									a__       :						Union[str, Any]							=  prepare_img()
									a__       :						Any							=  image_processor(images=__lowercase ,       return_tensors="""pt"""  ).pixel_values.to(__lowercase  )

									a__       :						Dict							=  torch.tensor([[1, 2]]  )
									a__       :						Any							=  torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8]]  ).unsqueeze(0  )

									# forward pass
									a__       :						List[Any]							=  model(
									    input_ids=input_ids.to(__lowercase  ) ,       bbox=bbox.to(__lowercase  ) ,       pixel_values=pixel_values.to(__lowercase  ) ,       )

									# verify the logits
									a__       :						List[Any]							=  torch.Size((1, 1_9_9, 7_6_8)  )
									self.assertEqual(outputs.last_hidden_state.shape ,       __lowercase  )

									a__       :						int							=  torch.tensor(
									    [[-0.0_5_2_9, 0.3_6_1_8, 0.1_6_3_2], [-0.1_5_8_7, -0.1_6_6_7, -0.0_4_0_0], [-0.1_5_5_7, -0.1_6_7_1, -0.0_5_0_5]]  ).to(__lowercase  )

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









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_					(  _lowercase    :       List[str])							->	Union[str, Any]:







			"""simple docstring"""
			a__       :						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(_lowercase						,       _lowercase)

def      lowerCAmelCase_					(  _lowercase    :       List[Any])							->	Optional[Any]:







			"""simple docstring"""
			a__	,						a__       :						Union[str, Any]							=  emb.weight.shape
			a__       :						str							=  nn.Linear(_lowercase						,       _lowercase						,       bias=_lowercase)
			a__       :						Any							=  emb.weight.data
			return lin_layer

def      lowerCAmelCase_					(  _lowercase    :       int						,       _lowercase    :       int=None)							->	List[Any]:







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

def      lowerCAmelCase_					(  _lowercase    :       Tuple						,       _lowercase    :       Optional[int]						,       _lowercase    :       Any						,       _lowercase    :       Dict						,       _lowercase    :       str = WEIGHTS_NAME)							->	Tuple:







			"""simple docstring"""
			a__       :						Tuple							=  []
			a__       :						Optional[Any]							=  0
			os.makedirs(_lowercase						,       exist_ok=_lowercase)

			for expert in range(_lowercase):
						a__       :						str							=  switch_checkpoint_path + F'''-rank-{expert}.pt'''
						if os.path.isfile(_lowercase):
									a__       :						List[str]							=  torch.load(_lowercase)["""model"""]
									remove_ignore_keys_(_lowercase)
									a__       :						Tuple							=  rename_fairseq_keys(_lowercase						,       _lowercase)
									a__       :						str							=  os.path.join(
									    _lowercase						,       weights_name.replace(""".bin"""						,       F'''-{len(_lowercase)+1:05d}-of-???.bin'''))
									torch.save(_lowercase						,       _lowercase)
									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(_lowercase)[0]].dtype)

    # Add the last block
			a__       :						int							=  os.path.join(_lowercase						,       weights_name.replace(""".bin"""						,       F'''-{len(_lowercase)+1:05d}-of-???.bin'''))
			a__       :						Union[str, Any]							=  torch.load(switch_checkpoint_path + """-shared.pt""")["""model"""]
			remove_ignore_keys_(_lowercase)
			a__       :						List[str]							=  rename_fairseq_keys(_lowercase						,       _lowercase)
			a__       :						int							=  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(_lowercase) == 1:
						a__       :						Optional[int]							=  os.path.join(_lowercase						,       _lowercase)
						torch.save(_lowercase						,       _lowercase)
						return {weights_name: sharded_state_dicts[0]}, None
			else:
						torch.save(_lowercase						,       _lowercase)
			# Otherwise, let's build the index
			a__       :						List[str]							=  {}
			for idx, shard in enumerate(_lowercase):
						a__       :						Union[str, Any]							=  weights_name.replace(""".bin"""						,       F'''-{idx+1:05d}-of-{len(_lowercase):05d}.bin''')
						a__       :						List[str]							=  os.path.join(_lowercase						,       weights_name.replace(""".bin"""						,       F'''-{idx+1:05d}-of-???.bin'''))
						os.rename(_lowercase						,       os.path.join(_lowercase						,       _lowercase))
						for key in shard:
									a__       :						Tuple							=  shard_file

    # Add the metadata
			a__       :						Tuple							=  {"""total_size""": total_size}
			a__       :						Optional[Any]							=  {"""metadata""": metadata, """weight_map""": weight_map}

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

			return metadata, index


if __name__ == "__main__":
		_lowercase		:				Any      		=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		:				Tuple      		=parser.parse_args()
		_lowercase		,					_lowercase		:				List[Any]      		=shard_on_the_fly(
		    args.nllb_moe_checkpoint_path,
		    args.pytorch_dump_folder_path,
		    128,
		    args.dtype,
		)

		_lowercase		:				int      		=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		:				List[str]      		=NllbMoeModel.from_pretrained(args.pytorch_dump_folder_path)
		print("Done")
		model.save_pretrained(args.pytorch_dump_folder_path)




