


from __future__ import annotations

import unittest

from transformers import FunnelConfig, is_tf_available
from transformers.testing_utils import require_tf

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 (
	    TFFunnelBaseModel,
	    TFFunnelForMaskedLM,
	    TFFunnelForMultipleChoice,
	    TFFunnelForPreTraining,
	    TFFunnelForQuestionAnswering,
	    TFFunnelForSequenceClassification,
	    TFFunnelForTokenClassification,
	    TFFunnelModel,
	)





class  A       :


							def __init__(self ,							lowerCAmelCase ,							lowerCAmelCase=1_3 ,							lowerCAmelCase=7 ,							lowerCAmelCase=True ,							lowerCAmelCase=True ,							lowerCAmelCase=True ,							lowerCAmelCase=True ,							lowerCAmelCase=9_9 ,							lowerCAmelCase=[1, 1, 2] ,							lowerCAmelCase=1 ,							lowerCAmelCase=3_2 ,							lowerCAmelCase=4 ,							lowerCAmelCase=8 ,							lowerCAmelCase=3_7 ,							lowerCAmelCase="gelu_new" ,							lowerCAmelCase=0.1 ,							lowerCAmelCase=0.1 ,							lowerCAmelCase=0.0 ,							lowerCAmelCase=5_1_2 ,							lowerCAmelCase=3 ,							lowerCAmelCase=0.02 ,							lowerCAmelCase=3 ,							lowerCAmelCase=4 ,							lowerCAmelCase=None ,							lowerCAmelCase=False ,							):
									__lowercase=     parent
									__lowercase=     batch_size
									__lowercase=     seq_length
									__lowercase=     is_training
									__lowercase=     use_input_mask
									__lowercase=     use_token_type_ids
									__lowercase=     use_labels
									__lowercase=     vocab_size
									__lowercase=     block_sizes
									__lowercase=     num_decoder_layers
									__lowercase=     d_model
									__lowercase=     n_head
									__lowercase=     d_head
									__lowercase=     d_inner
									__lowercase=     hidden_act
									__lowercase=     hidden_dropout
									__lowercase=     attention_dropout
									__lowercase=     activation_dropout
									__lowercase=     max_position_embeddings
									__lowercase=     type_vocab_size
									__lowercase=     2
									__lowercase=     num_labels
									__lowercase=     num_choices
									__lowercase=     scope
									__lowercase=     initializer_std

									# Used in the tests to check the size of the first attention layer
									__lowercase=     n_head
									# Used in the tests to check the size of the first hidden state
									__lowercase=     self.d_model
									# Used in the tests to check the number of output hidden states/attentions
									__lowercase=     sum(self.block_sizes   ) + (0 if base else self.num_decoder_layers)
									# FunnelModel adds two hidden layers: input embeddings and the sum of the upsampled encoder hidden state with
									# the last hidden state of the first block (which is the first hidden state of the decoder).
									if not base:
											__lowercase=     self.num_hidden_layers + 2


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

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

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

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

									__lowercase=     FunnelConfig(
									    vocab_size=self.vocab_size ,							block_sizes=self.block_sizes ,							num_decoder_layers=self.num_decoder_layers ,							d_model=self.d_model ,							n_head=self.n_head ,							d_head=self.d_head ,							d_inner=self.d_inner ,							hidden_act=self.hidden_act ,							hidden_dropout=self.hidden_dropout ,							attention_dropout=self.attention_dropout ,							activation_dropout=self.activation_dropout ,							max_position_embeddings=self.max_position_embeddings ,							type_vocab_size=self.type_vocab_size ,							initializer_std=self.initializer_std ,							)

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


							def      _A			(self ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							):
									__lowercase=     TFFunnelModel(config=__lowerCAmelCase   )
									__lowercase=     {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
									__lowercase=     model(__lowerCAmelCase   )

									__lowercase=     [input_ids, input_mask]
									__lowercase=     model(__lowerCAmelCase   )

									__lowercase=     model(__lowerCAmelCase   )
									self.parent.assertEqual(result.last_hidden_state.shape ,							(self.batch_size, self.seq_length, self.d_model)   )

									__lowercase=     False
									__lowercase=     TFFunnelModel(config=__lowerCAmelCase   )
									__lowercase=     model(__lowerCAmelCase   )
									self.parent.assertEqual(result.last_hidden_state.shape ,							(self.batch_size, self.seq_length, self.d_model)   )

									__lowercase=     False
									__lowercase=     TFFunnelModel(config=__lowerCAmelCase   )
									__lowercase=     model(__lowerCAmelCase   )
									self.parent.assertEqual(result.last_hidden_state.shape ,							(self.batch_size, self.seq_length, self.d_model)   )


							def      _A			(self ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							):
									__lowercase=     TFFunnelBaseModel(config=__lowerCAmelCase   )
									__lowercase=     {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
									__lowercase=     model(__lowerCAmelCase   )

									__lowercase=     [input_ids, input_mask]
									__lowercase=     model(__lowerCAmelCase   )

									__lowercase=     model(__lowerCAmelCase   )
									self.parent.assertEqual(result.last_hidden_state.shape ,							(self.batch_size, 2, self.d_model)   )

									__lowercase=     False
									__lowercase=     TFFunnelBaseModel(config=__lowerCAmelCase   )
									__lowercase=     model(__lowerCAmelCase   )
									self.parent.assertEqual(result.last_hidden_state.shape ,							(self.batch_size, 3, self.d_model)   )

									__lowercase=     False
									__lowercase=     TFFunnelBaseModel(config=__lowerCAmelCase   )
									__lowercase=     model(__lowerCAmelCase   )
									self.parent.assertEqual(result.last_hidden_state.shape ,							(self.batch_size, 2, self.d_model)   )


							def      _A			(self ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							):
									__lowercase=     TFFunnelForPreTraining(config=__lowerCAmelCase   )
									__lowercase=     {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
									__lowercase=     model(__lowerCAmelCase   )
									self.parent.assertEqual(result.logits.shape ,							(self.batch_size, self.seq_length)   )


							def      _A			(self ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							):
									__lowercase=     TFFunnelForMaskedLM(config=__lowerCAmelCase   )
									__lowercase=     {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
									__lowercase=     model(__lowerCAmelCase   )
									self.parent.assertEqual(result.logits.shape ,							(self.batch_size, self.seq_length, self.vocab_size)   )


							def      _A			(self ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							):
									__lowercase=     self.num_labels
									__lowercase=     TFFunnelForSequenceClassification(config=__lowerCAmelCase   )
									__lowercase=     {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
									__lowercase=     model(__lowerCAmelCase   )
									self.parent.assertEqual(result.logits.shape ,							(self.batch_size, self.num_labels)   )


							def      _A			(self ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							):
									__lowercase=     self.num_choices
									__lowercase=     TFFunnelForMultipleChoice(config=__lowerCAmelCase   )
									__lowercase=     tf.tile(tf.expand_dims(__lowerCAmelCase ,							1   ) ,							(1, self.num_choices, 1)   )
									__lowercase=     tf.tile(tf.expand_dims(__lowerCAmelCase ,							1   ) ,							(1, self.num_choices, 1)   )
									__lowercase=     tf.tile(tf.expand_dims(__lowerCAmelCase ,							1   ) ,							(1, self.num_choices, 1)   )
									__lowercase=     {
									    '''input_ids''': multiple_choice_inputs_ids,
									    '''attention_mask''': multiple_choice_input_mask,
									    '''token_type_ids''': multiple_choice_token_type_ids,
									}
									__lowercase=     model(__lowerCAmelCase   )
									self.parent.assertEqual(result.logits.shape ,							(self.batch_size, self.num_choices)   )


							def      _A			(self ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							):
									__lowercase=     self.num_labels
									__lowercase=     TFFunnelForTokenClassification(config=__lowerCAmelCase   )
									__lowercase=     {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
									__lowercase=     model(__lowerCAmelCase   )
									self.parent.assertEqual(result.logits.shape ,							(self.batch_size, self.seq_length, self.num_labels)   )


							def      _A			(self ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase ,							):
									__lowercase=     TFFunnelForQuestionAnswering(config=__lowerCAmelCase   )
									__lowercase=     {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
									__lowercase=     model(__lowerCAmelCase   )
									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      _A			(self   ):
									__lowercase=     self.prepare_config_and_inputs()
									(
									    __lowercase					
									)=     config_and_inputs
									__lowercase=     {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
									return config, inputs_dict





@require_tf
class  A       (  A_   ,							A_   ,							unittest.TestCase					):
							UpperCamelCase_  :				Dict			=(
							    (
							        TFFunnelModel,
							        TFFunnelForMaskedLM,
							        TFFunnelForPreTraining,
							        TFFunnelForQuestionAnswering,
							        TFFunnelForTokenClassification,
							    )
							    if is_tf_available()
							    else ()
							)
							UpperCamelCase_  :				str			=(
							    {
							        '''feature-extraction''': (TFFunnelBaseModel, TFFunnelModel),
							        '''fill-mask''': TFFunnelForMaskedLM,
							        '''question-answering''': TFFunnelForQuestionAnswering,
							        '''text-classification''': TFFunnelForSequenceClassification,
							        '''token-classification''': TFFunnelForTokenClassification,
							        '''zero-shot''': TFFunnelForSequenceClassification,
							    }
							    if is_tf_available()
							    else {}
							)
							UpperCamelCase_  :				List[Any]			=False
							UpperCamelCase_  :				Optional[Any]			=False


							def      _A			(self   ):
									__lowercase=     TFFunnelModelTester(self   )
									__lowercase=     ConfigTester(self ,							config_class=__lowerCAmelCase   )


							def      _A			(self   ):
									self.config_tester.run_common_tests()


							def      _A			(self   ):
									__lowercase=     self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*__lowerCAmelCase   )


							def      _A			(self   ):
									__lowercase=     self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_pretraining(*__lowerCAmelCase   )


							def      _A			(self   ):
									__lowercase=     self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_masked_lm(*__lowerCAmelCase   )


							def      _A			(self   ):
									__lowercase=     self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_token_classification(*__lowerCAmelCase   )







							def      _A			(self   ):
									__lowercase=     self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_question_answering(*__lowerCAmelCase   )





@require_tf
class  A       (  A_   ,							unittest.TestCase					):
							UpperCamelCase_  :				Dict			=(
							    (TFFunnelBaseModel, TFFunnelForMultipleChoice, TFFunnelForSequenceClassification) if is_tf_available() else ()
							)
							UpperCamelCase_  :				Any			=False
							UpperCamelCase_  :				Optional[Any]			=False


							def      _A			(self   ):
									__lowercase=     TFFunnelModelTester(self ,							base=__lowerCAmelCase   )
									__lowercase=     ConfigTester(self ,							config_class=__lowerCAmelCase   )


							def      _A			(self   ):
									self.config_tester.run_common_tests()


							def      _A			(self   ):
									__lowercase=     self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_base_model(*__lowerCAmelCase   )


							def      _A			(self   ):
									__lowercase=     self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_sequence_classification(*__lowerCAmelCase   )







							def      _A			(self   ):
									__lowercase=     self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_multiple_choice(*__lowerCAmelCase   )














import argparse

import torch
from torch import nn

from transformers import MaMaaaConfig, MaMaaaForConditionalGeneration

def   __lowercase       (			snake_case			):







				"""simple docstring"""






				__magic_name__		:Optional[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(snake_case,	snake_case			)

def   __lowercase       (			snake_case			):







				"""simple docstring"""






				__magic_name__  ,					__magic_name__		:Tuple   =   emb.weight.shape
				__magic_name__		:int   =   nn.Linear(snake_case,	snake_case,	bias=snake_case			)
				__magic_name__		:str   =   emb.weight.data
				return lin_layer

def   __lowercase       (			snake_case			):







				"""simple docstring"""






				__magic_name__		:int   =   torch.load(snake_case,	map_location='''cpu'''			)
				__magic_name__		:Optional[Any]   =   mam_aaa['''args'''] or mam_aaa['''cfg''']['''model''']
				__magic_name__		:List[Any]   =   mam_aaa['''model''']
				remove_ignore_keys_(snake_case			)
				__magic_name__		:Tuple   =   state_dict['''encoder.embed_tokens.weight'''].shape[0]

				__magic_name__		:List[str]   =   MaMaaaConfig(
				    vocab_size=snake_case,	max_position_embeddings=1_0_2_4,	encoder_layers=args.encoder_layers,	decoder_layers=args.decoder_layers,	encoder_attention_heads=args.encoder_attention_heads,	decoder_attention_heads=args.decoder_attention_heads,	encoder_ffn_dim=args.encoder_ffn_embed_dim,	decoder_ffn_dim=args.decoder_ffn_embed_dim,	d_model=args.encoder_embed_dim,	encoder_layerdrop=args.encoder_layerdrop,	decoder_layerdrop=args.decoder_layerdrop,	dropout=args.dropout,	attention_dropout=args.attention_dropout,	activation_dropout=args.activation_dropout,	activation_function='''relu''',	)

				__magic_name__		:int   =   state_dict['''decoder.embed_tokens.weight''']
				__magic_name__		:List[str]   =   MaMaaaForConditionalGeneration(snake_case			)
				model.model.load_state_dict(snake_case,	strict=snake_case			)
				__magic_name__		:List[str]   =   make_linear_from_emb(model.model.shared			)

				return model


if __name__ == "__main__":
					SCREAMING_SNAKE_CASE__   :      Optional[Any]				=		argparse.ArgumentParser()
					# Required parameters
					parser.add_argument("""fairseq_path""", type=str, help="""path to a model.pt on local filesystem.""")
					parser.add_argument("""pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model.""")
					SCREAMING_SNAKE_CASE__   :      int				=		parser.parse_args()
					SCREAMING_SNAKE_CASE__   :      Any				=		convert_fairseq_mamaaa_checkpoint_from_disk(args.fairseq_pathß)
					model.save_pretrained(args.pytorch_dump_folder_path)






