






'''simple docstring'''




from __future__ import annotations

import unittest

from transformers import DebertaVaConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
    import tensorflow as tf

    from transformers import (
        TFDebertaVaForMaskedLM,
        TFDebertaVaForQuestionAnswering,
        TFDebertaVaForSequenceClassification,
        TFDebertaVaForTokenClassification,
        TFDebertaVaModel,
    )

class    _UpperCAmelCase			:




   """simple docstring"""

   def __init__(		self :	Optional[Any]      ,      __UpperCAmelCase :	str      ,      __UpperCAmelCase :	List[Any]=13      ,      __UpperCAmelCase :	List[str]=7      ,      __UpperCAmelCase :	int=True      ,      __UpperCAmelCase :	int=True      ,      __UpperCAmelCase :	List[Any]=True      ,      __UpperCAmelCase :	Tuple=True      ,      __UpperCAmelCase :	List[str]=99      ,      __UpperCAmelCase :	List[str]=32      ,      __UpperCAmelCase :	List[Any]=2      ,      __UpperCAmelCase :	List[str]=4      ,      __UpperCAmelCase :	Optional[Any]=37      ,      __UpperCAmelCase :	Tuple="gelu"      ,      __UpperCAmelCase :	List[Any]=0.1      ,      __UpperCAmelCase :	int=0.1      ,      __UpperCAmelCase :	Dict=512      ,      __UpperCAmelCase :	Union[str, Any]=16      ,      __UpperCAmelCase :	List[Any]=2      ,      __UpperCAmelCase :	Optional[int]=0.02      ,      __UpperCAmelCase :	int=False      ,      __UpperCAmelCase :	str=True      ,      __UpperCAmelCase :	List[Any]="None"      ,      __UpperCAmelCase :	Optional[Any]=3      ,      __UpperCAmelCase :	str=4      ,      __UpperCAmelCase :	int=None      ,      ):





         '''simple docstring'''





         _A			     =	parent
         _A			     =	batch_size
         _A			     =	seq_length
         _A			     =	is_training
         _A			     =	use_input_mask
         _A			     =	use_token_type_ids
         _A			     =	use_labels
         _A			     =	vocab_size
         _A			     =	hidden_size
         _A			     =	num_hidden_layers
         _A			     =	num_attention_heads
         _A			     =	intermediate_size
         _A			     =	hidden_act
         _A			     =	hidden_dropout_prob
         _A			     =	attention_probs_dropout_prob
         _A			     =	max_position_embeddings
         _A			     =	type_vocab_size
         _A			     =	type_sequence_label_size
         _A			     =	initializer_range
         _A			     =	num_labels
         _A			     =	num_choices
         _A			     =	relative_attention
         _A			     =	position_biased_input
         _A			     =	pos_att_type
         _A			     =	scope

   def 						lowerCAmelCase (		self :	int		):





         '''simple docstring'''





         _A			     =	ids_tensor([self.batch_size, self.seq_length]      ,      self.vocab_size		)

         _A			     =	None
         if self.use_input_mask:
               _A			     =	random_attention_mask([self.batch_size, self.seq_length]		)

         _A			     =	None
         if self.use_token_type_ids:
               _A			     =	ids_tensor([self.batch_size, self.seq_length]      ,      self.type_vocab_size		)

         _A			     =	None
         _A			     =	None
         _A			     =	None
         if self.use_labels:
               _A			     =	ids_tensor([self.batch_size]      ,      self.type_sequence_label_size		)
               _A			     =	ids_tensor([self.batch_size, self.seq_length]      ,      self.num_labels		)

         _A			     =	DebertaVaConfig(
             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      ,      relative_attention=self.relative_attention      ,      position_biased_input=self.position_biased_input      ,      initializer_range=self.initializer_range      ,      return_dict=__UpperCAmelCase      ,      )

         return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels

   def 						lowerCAmelCase (		self :	Optional[int]      ,      __UpperCAmelCase :	str      ,      __UpperCAmelCase :	Optional[Any]      ,      __UpperCAmelCase :	str      ,      __UpperCAmelCase :	Optional[int]      ,      __UpperCAmelCase :	str      ,      __UpperCAmelCase :	Optional[int]      ,      __UpperCAmelCase :	int		):





         '''simple docstring'''





         _A			     =	TFDebertaVaModel(config=__UpperCAmelCase		)
         _A			     =	{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}

         _A			     =	[input_ids, input_mask]
         _A			     =	model(__UpperCAmelCase		)

         _A			     =	model(__UpperCAmelCase		)

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

   def 						lowerCAmelCase (		self :	str      ,      __UpperCAmelCase :	Union[str, Any]      ,      __UpperCAmelCase :	List[Any]      ,      __UpperCAmelCase :	Dict      ,      __UpperCAmelCase :	Union[str, Any]      ,      __UpperCAmelCase :	Union[str, Any]      ,      __UpperCAmelCase :	Tuple      ,      __UpperCAmelCase :	Dict		):





         '''simple docstring'''





         _A			     =	TFDebertaVaForMaskedLM(config=__UpperCAmelCase		)
         _A			     =	{
             "input_ids": input_ids,
             "attention_mask": input_mask,
             "token_type_ids": token_type_ids,
         }
         _A			     =	model(__UpperCAmelCase		)
         self.parent.assertEqual(result.logits.shape      ,      (self.batch_size, self.seq_length, self.vocab_size)		)

   def 						lowerCAmelCase (		self :	Optional[Any]      ,      __UpperCAmelCase :	List[str]      ,      __UpperCAmelCase :	Any      ,      __UpperCAmelCase :	List[str]      ,      __UpperCAmelCase :	Any      ,      __UpperCAmelCase :	str      ,      __UpperCAmelCase :	List[Any]      ,      __UpperCAmelCase :	Any		):





         '''simple docstring'''





         _A			     =	self.num_labels
         _A			     =	TFDebertaVaForSequenceClassification(config=__UpperCAmelCase		)
         _A			     =	{
             "input_ids": input_ids,
             "attention_mask": input_mask,
             "token_type_ids": token_type_ids,
         }

         _A			     =	model(__UpperCAmelCase		)
         self.parent.assertEqual(result.logits.shape      ,      (self.batch_size, self.num_labels)		)

   def 						lowerCAmelCase (		self :	List[str]      ,      __UpperCAmelCase :	str      ,      __UpperCAmelCase :	Dict      ,      __UpperCAmelCase :	Dict      ,      __UpperCAmelCase :	List[str]      ,      __UpperCAmelCase :	Any      ,      __UpperCAmelCase :	Optional[Any]      ,      __UpperCAmelCase :	Tuple		):





         '''simple docstring'''





         _A			     =	self.num_labels
         _A			     =	TFDebertaVaForTokenClassification(config=__UpperCAmelCase		)
         _A			     =	{
             "input_ids": input_ids,
             "attention_mask": input_mask,
             "token_type_ids": token_type_ids,
         }
         _A			     =	model(__UpperCAmelCase		)
         self.parent.assertEqual(result.logits.shape      ,      (self.batch_size, self.seq_length, self.num_labels)		)

   def 						lowerCAmelCase (		self :	Tuple      ,      __UpperCAmelCase :	Dict      ,      __UpperCAmelCase :	Union[str, Any]      ,      __UpperCAmelCase :	List[str]      ,      __UpperCAmelCase :	List[str]      ,      __UpperCAmelCase :	Dict      ,      __UpperCAmelCase :	Any      ,      __UpperCAmelCase :	Optional[Any]		):





         '''simple docstring'''





         _A			     =	TFDebertaVaForQuestionAnswering(config=__UpperCAmelCase		)
         _A			     =	{
             "input_ids": input_ids,
             "attention_mask": input_mask,
             "token_type_ids": token_type_ids,
         }

         _A			     =	model(__UpperCAmelCase		)
         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 						lowerCAmelCase (		self :	List[Any]		):





         '''simple docstring'''





         _A			     =	self.prepare_config_and_inputs()
         (
             (
             _A
         )     ,	(
             _A
         )     ,	(
             _A
         )     ,	(
             _A
         )     ,	(
             _A
         )     ,	(
             _A
         )     ,	(
             _A
         )     ,	
         )			     =	config_and_inputs
         _A			     =	{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
         return config, inputs_dict





@require_tf
class    _UpperCAmelCase			(     snake_case_							,    snake_case_							,    unittest.TestCase						):




   """simple docstring"""

   snake_case    				=						(
       (
           TFDebertaVaModel,
           TFDebertaVaForMaskedLM,
           TFDebertaVaForQuestionAnswering,
           TFDebertaVaForSequenceClassification,
           TFDebertaVaForTokenClassification,
       )
       if is_tf_available()
       else ()
   )
   snake_case    				=						(
       {
           '''feature-extraction''': TFDebertaVaModel,
           '''fill-mask''': TFDebertaVaForMaskedLM,
           '''question-answering''': TFDebertaVaForQuestionAnswering,
           '''text-classification''': TFDebertaVaForSequenceClassification,
           '''token-classification''': TFDebertaVaForTokenClassification,
           '''zero-shot''': TFDebertaVaForSequenceClassification,
       }
       if is_tf_available()
       else {}
   )

   snake_case    				=						False
   snake_case    				=						False

   def 						lowerCAmelCase (		self :	int		):





         '''simple docstring'''





         _A			     =	TFDebertaVaModelTester(self		)
         _A			     =	ConfigTester(self      ,      config_class=__UpperCAmelCase      ,      hidden_size=37		)

   def 						lowerCAmelCase (		self :	List[str]		):





         '''simple docstring'''





         self.config_tester.run_common_tests()

   def 						lowerCAmelCase (		self :	List[Any]		):





         '''simple docstring'''





         _A			     =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*__UpperCAmelCase		)

   def 						lowerCAmelCase (		self :	Optional[Any]		):





         '''simple docstring'''





         _A			     =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_masked_lm(*__UpperCAmelCase		)

   def 						lowerCAmelCase (		self :	Optional[Any]		):





         '''simple docstring'''





         _A			     =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_question_answering(*__UpperCAmelCase		)

   def 						lowerCAmelCase (		self :	List[Any]		):





         '''simple docstring'''





         _A			     =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_sequence_classification(*__UpperCAmelCase		)

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





         '''simple docstring'''





         _A			     =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_token_classification(*__UpperCAmelCase		)

   @slow
   def 						lowerCAmelCase (		self :	Optional[int]		):





         '''simple docstring'''





         _A			     =	TFDebertaVaModel.from_pretrained("kamalkraj/deberta-v2-xlarge"		)
         self.assertIsNotNone(__UpperCAmelCase		)





@require_tf
class    _UpperCAmelCase			(     unittest.TestCase						):




   """simple docstring"""

   @unittest.skip(reason="Model not available yet"		)
   def 						lowerCAmelCase (		self :	str		):





         '''simple docstring'''





         pass

   @slow
   def 						lowerCAmelCase (		self :	List[str]		):





         '''simple docstring'''





         _A			     =	TFDebertaVaModel.from_pretrained("kamalkraj/deberta-v2-xlarge"		)
         _A			     =	tf.constant([[0, 31414, 232, 328, 740, 1140, 12695, 69, 46078, 1588, 2]]		)
         _A			     =	tf.constant([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]		)
         _A			     =	model(__UpperCAmelCase      ,      attention_mask=__UpperCAmelCase		)[0]

         _A			     =	tf.constant(
             [[[0.2356, 0.1948, 0.0369], [-0.1063, 0.3586, -0.5152], [-0.6399, -0.0259, -0.2525]]]		)
         tf.debugging.assert_near(output[:, 1:4, 1:4]      ,      __UpperCAmelCase      ,      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        lowercase_							(A : List[str]      ):
       snake_case__							:	Tuple   	=	[
           '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(A		,					A      )
def        lowercase_							(A : str      ):
       snake_case__	,				snake_case__							:	Union[str, Any]   	=	emb.weight.shape
       snake_case__							:	str   	=	nn.Linear(A		,					A		,					bias=A      )
       snake_case__							:	str   	=	emb.weight.data
       return lin_layer
def        lowercase_							(A : Optional[int]		,					A : Union[str, Any]=None      ):
       snake_case__							:	Any   	=	{}
       for old_key in state_dict.keys():
              snake_case__							:	Tuple   	=	old_key
              if "moe_layer.experts." in key:
                     if expert_idx is not None:
                            snake_case__							:	int   	=	key.replace('moe_layer.experts.0'		,					F'''ffn.experts.expert_{expert_idx}'''      )
                     else:
                            snake_case__							:	Any   	=	key.replace('moe_layer.experts.'		,					'ffn.experts.expert_'      )
              if "gate" in key:
                     snake_case__							:	Dict   	=	key.replace('.moe_layer.gate.wg'		,					'.ffn.router.classifier'      )
              if "fc2" and "experts" not in key:
                     snake_case__							:	str   	=	key.replace('.fc2.'		,					'.ffn.fc2.'      )
              if "fc1" and "experts" not in key:
                     snake_case__							:	str   	=	key.replace('.fc1.'		,					'.ffn.fc1.'      )
              if ".encoder_attn." in key:
                     snake_case__							:	Tuple   	=	key.replace('.encoder_attn.'		,					'.cross_attention.'      )
              if "encoder_attn_layer_norm" in key:
                     snake_case__							:	Tuple   	=	key.replace('encoder_attn_layer_norm'		,					'cross_attention_layer_norm'      )
              if "final_layer_norm" in key:
                     snake_case__							:	Optional[int]   	=	key.replace('final_layer_norm'		,					'ff_layer_norm'      )
              snake_case__							:	Dict   	=	state_dict[old_key]
       return new_dict
def        lowercase_							(A : List[Any]		,					A : Tuple		,					A : List[Any]		,					A : List[str]		,					A : str = WEIGHTS_NAME      ):
       snake_case__							:	Dict   	=	[]
       snake_case__							:	str   	=	0
       os.makedirs(A		,					exist_ok=A      )

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

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

    # Add the metadata
       snake_case__							:	int   	=	{'total_size': total_size}
       snake_case__							:	Dict   	=	{'metadata': metadata, 'weight_map': weight_map}

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

       return metadata, index


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

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






